Forbid all modules by default except whitelist authorized_imports (huggingface#935)

Author: albertvillanova

docs/source/en/guided_tour.mdx

@@ -181,6 +181,9 @@ agent = CodeAgent(tools=[], model=model, additional_authorized_imports=['request
 agent.run("Could you get me the title of the page at url 'https://huggingface.co/blog'?")
 ```
 
+Additionally, as an extra security layer, access to submodule is forbidden by default, unless explicitly authorized within the import list.
+For instance, to access the `numpy.random` submodule, you need to add `'numpy.random'` to the `additional_authorized_imports` list.
+
 > [!WARNING]
 > The LLM can generate arbitrary code that will then be executed: do not add any unsafe imports!
 

docs/source/en/tutorials/secure_code_execution.mdx

@@ -52,7 +52,8 @@ We have re-built a more secure `LocalPythonExecutor` from the ground up.
 
 To be precise, this interpreter works by loading the Abstract Syntax Tree (AST) from your Code and executes it operation by operation, making sure to always follow certain rules:
 - By default, imports are disallowed unless they have been explicitly added to an authorization list by the user.
-   - Even so, because some innocuous packages like `re` can give access to potentially harmful packages as in `re.subprocess`, subpackages that match a list of dangerous patterns are not imported.
+- Furthermore, access to submodules is disabled by default, and each must be explicitly authorized in the import list as well.
+   - Note that some seemingly innocuous packages like `random` can give access to potentially harmful submodules, as in `random._os`.
  - The total count of elementary operations processed is capped to prevent infinite loops and resource bloating.
  - Any operation that has not been explicitly defined in our custom interpreter will raise an error.
 

src/smolagents/local_python_executor.py

@@ -24,7 +24,6 @@
 from collections.abc import Mapping
 from functools import wraps
 from importlib import import_module
-from importlib.util import find_spec
 from types import BuiltinFunctionType, FunctionType, ModuleType
 from typing import Any, Callable, Dict, List, Optional, Set, Tuple
 
@@ -114,19 +113,6 @@ def custom_print(*args):
     "complex": complex,
 }
 
-DANGEROUS_MODULES = [
-    "builtins",
-    "io",
-    "multiprocessing",
-    "os",
-    "pathlib",
-    "pty",
-    "shutil",
-    "socket",
-    "subprocess",
-    "sys",
-]
-
 DANGEROUS_FUNCTIONS = [
     "builtins.compile",
     "builtins.eval",
@@ -238,25 +224,11 @@ def _check_return(
         result = func(expression, state, static_tools, custom_tools, authorized_imports=authorized_imports)
         if "*" not in authorized_imports:
             if isinstance(result, ModuleType):
-                for module_name in DANGEROUS_MODULES:
-                    if (
-                        module_name not in authorized_imports
-                        and result.__name__ == module_name
-                        # builtins has no __file__ attribute
-                        and getattr(result, "__file__", "")
-                        == (getattr(import_module(module_name), "__file__", "") if find_spec(module_name) else "")
-                    ):
-                        raise InterpreterError(f"Forbidden access to module: {module_name}")
-            elif isinstance(result, dict) and result.get("__name__"):
-                for module_name in DANGEROUS_MODULES:
-                    if (
-                        module_name not in authorized_imports
-                        and result["__name__"] == module_name
-                        # builtins has no __file__ attribute
-                        and result.get("__file__", "")
-                        == (getattr(import_module(module_name), "__file__", "") if find_spec(module_name) else "")
-                    ):
-                        raise InterpreterError(f"Forbidden access to module: {module_name}")
+                if result.__name__ not in authorized_imports:
+                    raise InterpreterError(f"Forbidden access to module: {result.__name__}")
+            elif isinstance(result, dict) and result.get("__spec__"):
+                if result["__name__"] not in authorized_imports:
+                    raise InterpreterError(f"Forbidden access to module: {result['__name__']}")
             elif isinstance(result, (FunctionType, BuiltinFunctionType)):
                 for qualified_function_name in DANGEROUS_FUNCTIONS:
                     module_name, function_name = qualified_function_name.rsplit(".", 1)

tests/test_local_python_executor.py

@@ -27,7 +27,6 @@
 from smolagents.default_tools import BASE_PYTHON_TOOLS
 from smolagents.local_python_executor import (
     DANGEROUS_FUNCTIONS,
-    DANGEROUS_MODULES,
     InterpreterError,
     LocalPythonExecutor,
     PrintContainer,
@@ -41,6 +40,21 @@
 )
 
 
+# Non-exhaustive list of dangerous modules that should not be imported
+DANGEROUS_MODULES = [
+    "builtins",
+    "io",
+    "multiprocessing",
+    "os",
+    "pathlib",
+    "pty",
+    "shutil",
+    "socket",
+    "subprocess",
+    "sys",
+]
+
+
 # Fake function we will use as tool
 def add_two(x):
     return x + 2
@@ -504,10 +518,10 @@ def test_imports(self):
 
         # Test submodules are handled properly, thus not raising error
         code = "import numpy.random as rd\nrng = rd.default_rng(12345)\nrng.random()"
-        result, _ = evaluate_python_code(code, BASE_PYTHON_TOOLS, state={}, authorized_imports=["numpy"])
+        result, _ = evaluate_python_code(code, BASE_PYTHON_TOOLS, state={}, authorized_imports=["numpy.random"])
 
         code = "from numpy.random import default_rng as d_rng\nrng = d_rng(12345)\nrng.random()"
-        result, _ = evaluate_python_code(code, BASE_PYTHON_TOOLS, state={}, authorized_imports=["numpy"])
+        result, _ = evaluate_python_code(code, BASE_PYTHON_TOOLS, state={}, authorized_imports=["numpy.random"])
 
     def test_additional_imports(self):
         code = "import numpy as np"
@@ -1773,7 +1787,16 @@ def test_vulnerability_via_importlib(self, additional_authorized_imports, expect
         "code, additional_authorized_imports, expected_error",
         [
             # os submodule
-            ("import queue; queue.threading._os.system(':')", [], InterpreterError("Forbidden access to module: os")),
+            (
+                "import queue; queue.threading._os.system(':')",
+                [],
+                InterpreterError("Forbidden access to module: threading"),
+            ),
+            (
+                "import queue; queue.threading._os.system(':')",
+                ["threading"],
+                InterpreterError("Forbidden access to module: os"),
+            ),
             ("import random; random._os.system(':')", [], InterpreterError("Forbidden access to module: os")),
             (
                 "import random; random.__dict__['_os'].system(':')",
@@ -1783,22 +1806,42 @@ def test_vulnerability_via_importlib(self, additional_authorized_imports, expect
             (
                 "import doctest; doctest.inspect.os.system(':')",
                 ["doctest"],
+                InterpreterError("Forbidden access to module: inspect"),
+            ),
+            (
+                "import doctest; doctest.inspect.os.system(':')",
+                ["doctest", "inspect"],
                 InterpreterError("Forbidden access to module: os"),
             ),
             # subprocess submodule
             (
                 "import asyncio; asyncio.base_events.events.subprocess",
                 ["asyncio"],
+                InterpreterError("Forbidden access to module: asyncio.base_events"),
+            ),
+            (
+                "import asyncio; asyncio.base_events.events.subprocess",
+                ["asyncio", "asyncio.base_events"],
+                InterpreterError("Forbidden access to module: asyncio.events"),
+            ),
+            (
+                "import asyncio; asyncio.base_events.events.subprocess",
+                ["asyncio", "asyncio.base_events", "asyncio.events"],
                 InterpreterError("Forbidden access to module: subprocess"),
             ),
             # sys submodule
             (
                 "import queue; queue.threading._sys.modules['os'].system(':')",
                 [],
+                InterpreterError("Forbidden access to module: threading"),
+            ),
+            (
+                "import queue; queue.threading._sys.modules['os'].system(':')",
+                ["threading"],
                 InterpreterError("Forbidden access to module: sys"),
             ),
             # Allowed
-            ("import pandas; pandas.io", ["pandas"], None),
+            ("import pandas; pandas.io", ["pandas", "pandas.io"], None),
         ],
     )
     def test_vulnerability_via_submodules(self, code, additional_authorized_imports, expected_error):
@@ -1851,8 +1894,12 @@ def test_vulnerability_builtins_via_sys(self, additional_authorized_imports, add
     @pytest.mark.parametrize(
         "additional_authorized_imports, additional_tools, expected_error",
         [
-            ([], [], InterpreterError("Forbidden access to module: builtins")),
-            (["builtins", "os"], ["__import__"], None),
+            ([], [], InterpreterError("Forbidden access to module: smolagents.local_python_executor")),
+            (
+                ["builtins", "os"],
+                ["__import__"],
+                InterpreterError("Forbidden access to module: smolagents.local_python_executor"),
+            ),
         ],
     )
     def test_vulnerability_builtins_via_traceback(
@@ -1888,8 +1935,18 @@ def test_vulnerability_builtins_via_traceback(
     @pytest.mark.parametrize(
         "additional_authorized_imports, additional_tools, expected_error",
         [
-            ([], [], InterpreterError("Forbidden access to module: builtins")),
-            (["builtins", "os"], ["__import__"], None),
+            ([], [], InterpreterError("Forbidden access to module: warnings")),
+            (["warnings"], [], InterpreterError("Forbidden access to module: builtins")),
+            (
+                ["warnings", "builtins"],
+                [],
+                (
+                    InterpreterError("Forbidden access to function: __import__"),
+                    InterpreterError("Forbidden access to module: os"),
+                ),
+            ),
+            (["warnings", "builtins", "os"], [], (InterpreterError("Forbidden access to function: __import__"), None)),
+            (["warnings", "builtins", "os"], ["__import__"], None),
         ],
     )
     def test_vulnerability_builtins_via_class_catch_warnings(
@@ -1902,20 +1959,22 @@ def test_vulnerability_builtins_via_class_catch_warnings(
             from builtins import __import__
 
             executor.send_tools({"__import__": __import__})
-        with (
-            pytest.raises(type(expected_error), match=f".*{expected_error}")
-            if isinstance(expected_error, Exception)
-            else does_not_raise()
-        ):
+        if isinstance(expected_error, tuple):  # different error depending on patch status
+            expected_error = expected_error[patch_builtin_import_module]
+        if isinstance(expected_error, Exception):
+            expectation = pytest.raises(type(expected_error), match=f".*{expected_error}")
+        elif expected_error is None:
+            expectation = does_not_raise()
+        with expectation:
             executor(
                 dedent(
                     """
                     classes = {}.__class__.__base__.__subclasses__()
                     for cls in classes:
                         if cls.__name__ == "catch_warnings":
-                            builtins = cls()._module.__builtins__
-                            builtins_import = builtins["__import__"]
                             break
+                    builtins = cls()._module.__builtins__
+                    builtins_import = builtins["__import__"]
                     os_module = builtins_import('os')
                     os_module.system(":")
                     """