Merge pull request #10 from basnijholt/allow_all_imports

add configuration option "allow_all_imports"

Author: craigbarratt

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)
@@ -44,6 +57,9 @@ def check_isdir(path):
         _LOGGER.error("Folder %s not found in configuration folder", FOLDER)
         return False
 
+    hass.data.setdefault(DOMAIN, {})
+    hass.data[DOMAIN]["allow_all_imports"] = config[DOMAIN].get(CONF_ALLOW_ALL_IMPORTS)
+
     await compile_scripts(  # pylint: disable=unused-variable
         hass,
         event_func=event_func,

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,11 @@ 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"]
+            if global_ctx.hass is not None
+            else False
+        )
 
     async def ast_not_implemented(self, arg, *args):
         """Raise NotImplementedError exception for unimplemented AST types."""
@@ -399,7 +404,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 +414,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)