add configuration option "allow_all_imports"

Author: basnijholt

custom_components/pyscript/__init__.py

@@ -5,12 +5,15 @@
 import logging
 import os
 
+import voluptuous as vol
+
 from homeassistant.const import (
     EVENT_HOMEASSISTANT_STARTED,
     EVENT_HOMEASSISTANT_STOP,
     EVENT_STATE_CHANGED,
     SERVICE_RELOAD,
 )
+import homeassistant.helpers.config_validation as cv
 from homeassistant.loader import bind_hass
 
 from .const import DOMAIN, FOLDER, LOGGER_PATH, SERVICE_JUPYTER_KERNEL_START
@@ -24,10 +27,20 @@
 
 _LOGGER = logging.getLogger(LOGGER_PATH)
 
+CONF_ALLOW_ALL_IMPORTS = "allow_all_imports"
+
+CONFIG_SCHEMA = vol.Schema(
+    {
+        DOMAIN: vol.Schema(
+            {vol.Optional(CONF_ALLOW_ALL_IMPORTS, default=False): cv.boolean}
+        ),
+    },
+    extra=vol.ALLOW_EXTRA,
+)
+
 
 async def async_setup(hass, config):
     """Initialize the pyscript component."""
-
     handler_func = Handler(hass)
     event_func = Event(hass)
     trig_time_func = TrigTime(hass, handler_func)
@@ -163,6 +176,8 @@ async def stop_triggers(event):
     hass.bus.async_listen(EVENT_HOMEASSISTANT_STARTED, start_triggers)
     hass.bus.async_listen(EVENT_HOMEASSISTANT_STOP, stop_triggers)
 
+    hass.data.setdefault(DOMAIN, {})
+    hass.data[DOMAIN]["allow_all_imports"] = config[DOMAIN].get(CONF_ALLOW_ALL_IMPORTS)
     return True
 
 

custom_components/pyscript/eval.py

@@ -8,7 +8,7 @@
 import logging
 import sys
 
-from .const import ALLOWED_IMPORTS, LOGGER_PATH
+from .const import ALLOWED_IMPORTS, DOMAIN, LOGGER_PATH
 
 _LOGGER = logging.getLogger(LOGGER_PATH + ".eval")
 
@@ -361,6 +361,7 @@ def __init__(
         self.logger_handlers = set()
         self.logger = None
         self.set_logger_name(logger_name if logger_name is not None else self.name)
+        self.allow_all_imports = global_ctx.hass.data[DOMAIN]["allow_all_imports"]
 
     async def ast_not_implemented(self, arg, *args):
         """Raise NotImplementedError exception for unimplemented AST types."""
@@ -399,7 +400,7 @@ async def ast_module(self, arg):
     async def ast_import(self, arg):
         """Execute import."""
         for imp in arg.names:
-            if imp.name not in ALLOWED_IMPORTS:
+            if not self.allow_all_imports and imp.name not in ALLOWED_IMPORTS:
                 raise ModuleNotFoundError(f"import of {imp.name} not allowed")
             if imp.name not in sys.modules:
                 mod = importlib.import_module(imp.name)
@@ -409,7 +410,7 @@ async def ast_import(self, arg):
 
     async def ast_importfrom(self, arg):
         """Execute from X import Y."""
-        if arg.module not in ALLOWED_IMPORTS:
+        if not self.allow_all_imports and arg.module not in ALLOWED_IMPORTS:
             raise ModuleNotFoundError(f"import from {arg.module} not allowed")
         if arg.module not in sys.modules:
             mod = importlib.import_module(arg.module)