tinyagent-py 0.0.9__py3-none-any.whl → 0.0.12__py3-none-any.whl

tinyagent/code_agent/modal_sandbox.py

@@ -63,6 +63,7 @@ def create_sandbox(
     pip_install: Sequence[str] | None = None,
     image_name: str = "tinyagent-sandbox-image",
     app_name: str = "persistent-code-session",
+    force_build: bool = False,
     **sandbox_kwargs,
 ) -> Tuple[modal.Sandbox, modal.App]:
     """Create (or lookup) a `modal.Sandbox` pre-configured for code execution.
@@ -99,7 +100,7 @@ def create_sandbox(
 
     # Build image -----------------------------------------------------------
     agent_image = (
-        modal.Image.debian_slim(python_version=python_version)
+        modal.Image.debian_slim(python_version=python_version,force_build=force_build)
         .apt_install(*apt_packages)
         .pip_install(*full_pip_list)
     )
@@ -196,6 +197,7 @@ class SandboxSession:
         modal_secrets: modal.Secret,
         *,
         timeout: int = 5 * 60,
+        
         **create_kwargs,
     ) -> None:
         self.modal_secrets = modal_secrets

tinyagent/code_agent/providers/modal_provider.py

@@ -26,6 +26,7 @@ class ModalProvider(CodeExecutionProvider):
         default_packages: Optional[List[str]] = None,
         apt_packages: Optional[List[str]] = None,
         python_version: Optional[str] = None,
+        authorized_imports: list[str] | None = None,
         modal_secrets: Dict[str, Union[str, None]] | None = None,
         lazy_init: bool = True,
         sandbox_name: str = "tinycodeagent-sandbox",
@@ -48,6 +49,7 @@ class ModalProvider(CodeExecutionProvider):
                 (git, curl, …) so you only need to specify the extras.
             python_version: Python version used for the sandbox image. If
                 ``None`` the current interpreter version is used.
+            authorized_imports: Optional allow-list of modules the user code is permitted to import. Supports wildcard patterns (e.g. "pandas.*"). If ``None`` the safety layer blocks only the predefined dangerous modules.
         """
 
         # Resolve default values ------------------------------------------------
@@ -70,6 +72,7 @@ class ModalProvider(CodeExecutionProvider):
         self.default_packages: List[str] = default_packages
         self.apt_packages: List[str] = apt_packages
         self.python_version: str = python_version
+        self.authorized_imports = authorized_imports
 
         # ----------------------------------------------------------------------
         final_packages = list(set(self.default_packages + (pip_packages or [])))
@@ -89,6 +92,7 @@ class ModalProvider(CodeExecutionProvider):
         self.modal_secrets = modal.Secret.from_dict(self.secrets)
         self.app = None
         self._app_run_python = None
+        self.is_trusted_code = kwargs.get("trust_code", False)
         
         self._setup_modal_app()
         
@@ -164,18 +168,30 @@ class ModalProvider(CodeExecutionProvider):
         if self.executed_default_codes:
             print("✔️ default codes already executed")
             full_code = "\n".join(self.code_tools_definitions) +"\n\n"+code
+            # Code tools and default code are trusted, user code is not
         else:
             full_code = "\n".join(self.code_tools_definitions) +"\n\n"+ "\n".join(self.default_python_codes) + "\n\n" + code
             self.executed_default_codes = True
+            # First execution includes framework code which is trusted
         
         # Use Modal's native execution methods
         if self.local_execution:
-            # Use Modal's .local() method for local execution
-            return self._app_run_python.local(full_code, globals_dict or {}, locals_dict or {})
+            return self._app_run_python.local(
+                full_code,
+                globals_dict or {},
+                locals_dict or {},
+                self.authorized_imports,
+                self.is_trusted_code,
+            )
         else:
-            # Use Modal's .remote() method for remote execution
             with self.app.run():
-                return self._app_run_python.remote(full_code, globals_dict or {}, locals_dict or {})
+                return self._app_run_python.remote(
+                    full_code,
+                    globals_dict or {},
+                    locals_dict or {},
+                    self.authorized_imports,
+                    self.is_trusted_code,
+                )
     
     def _log_response(self, response: Dict[str, Any]):
         """Log the response from code execution."""
@@ -184,15 +200,31 @@ class ModalProvider(CodeExecutionProvider):
         print("#########################<printed_output>#########################")
         print(response["printed_output"])
         print("#########################</printed_output>#########################")
-        print("#########################<return_value>#########################")
-        print(response["return_value"])
-        print("#########################</return_value>#########################")
-        print("#########################<stderr>#########################")
-        print(response["stderr"])
-        print("#########################</stderr>#########################")
-        print("#########################<traceback>#########################")
-        print(response["error_traceback"])
-        print("#########################</traceback>#########################")
+        if response.get("return_value",None) not in [None,""]:
+            print("#########################<return_value>#########################")
+            print(response["return_value"])
+            print("#########################</return_value>#########################")
+        if response.get("stderr",None) not in [None,""]:
+            print("#########################<stderr>#########################")
+            print(response["stderr"])
+            print("#########################</stderr>#########################")
+        if response.get("error_traceback",None) not in [None,""]:
+            print("#########################<traceback>#########################")
+            # Check if this is a security exception and highlight it in red if so
+            error_text = response["error_traceback"]
+            if "SECURITY" in error_text:
+                try:
+                    from ..modal_sandbox import COLOR
+                except ImportError:
+                    # Fallback colors if modal_sandbox is not available
+                    COLOR = {
+                        "RED": "\033[91m",
+                        "ENDC": "\033[0m",
+                    }
+                print(f"{COLOR['RED']}{error_text}{COLOR['ENDC']}")
+            else:
+                print(error_text)
+            print("#########################</traceback>#########################")
     
     async def cleanup(self):
         """Clean up Modal resources."""

tinyagent/code_agent/safety.py

@@ -0,0 +1,546 @@
+# TinyAgent code execution safety utilities
+# -----------------------------------------
+#
+# This helper module defines *very* lightweight safeguards that are applied
+# before running any user-supplied Python code inside the Modal sandbox.
+# The goal is **not** to build a full blown secure interpreter (this would
+# require a much more sophisticated setup à la Pyodide or the `python-secure`
+# project).  Instead we implement the following pragmatic defence layers:
+#
+# 1.  Static AST inspection of the submitted code to detect direct `import` or
+#     `from … import …` statements that reference known dangerous modules
+#     (e.g. `os`, `subprocess`, …).  This prevents the *most common* attack
+#     vector where an LLM attempts to read or modify the host file-system or
+#     spawn sub-processes.
+# 2.  Runtime patching of the built-in `__import__` hook so that *dynamic*
+#     imports carried out via `importlib` or `__import__(…)` are blocked at
+#     execution time as well.
+# 3.  Static AST inspection to detect calls to dangerous functions like `exec`,
+#     `eval`, `compile`, etc. that could be used to bypass security measures.
+# 4.  Runtime patching of built-in dangerous functions to prevent their use
+#     at execution time.
+#
+# The chosen approach keeps the TinyAgent runtime *fast* and *lean* while
+# still providing a reasonable first line of defence against obviously
+# malicious code.
+
+from __future__ import annotations
+
+import ast
+import builtins
+import warnings
+from typing import Iterable, List, Set, Sequence, Any, Callable
+import contextlib
+
+__all__ = [
+    "DANGEROUS_MODULES",
+    "DANGEROUS_FUNCTIONS",
+    "RUNTIME_BLOCKED_FUNCTIONS",
+    "validate_code_safety",
+    "install_import_hook",
+    "function_safety_context",
+]
+
+# ---------------------------------------------------------------------------
+# Threat model / deny-list
+# ---------------------------------------------------------------------------
+
+# Non-exhaustive list of modules that grant (direct or indirect) access to the
+# underlying operating system, spawn sub-processes, perform unrestricted I/O,
+# or allow the user to circumvent the static import analysis performed below.
+DANGEROUS_MODULES: Set[str] = {
+    "builtins",  # Gives access to exec/eval etc.
+    "ctypes",
+    "importlib",
+    "io",
+    "multiprocessing",
+    "os",
+    "pathlib",
+    "pty",
+    "shlex",
+    "shutil",
+    "signal",
+    "socket",
+    "subprocess",
+    "sys",
+    "tempfile",
+    "threading",
+    "webbrowser",
+}
+
+# List of dangerous built-in functions that could be used to bypass security
+# measures or execute arbitrary code
+DANGEROUS_FUNCTIONS: Set[str] = {
+    "exec",
+    "eval",
+    "compile",
+    "__import__",
+    "open",
+    "input",
+    "breakpoint",
+}
+
+# Functions that should be blocked at runtime (a subset of DANGEROUS_FUNCTIONS)
+RUNTIME_BLOCKED_FUNCTIONS: Set[str] = {
+    "exec",
+    "eval",
+}
+
+# Essential modules that are always allowed, even in untrusted code
+# These are needed for the framework to function properly
+ESSENTIAL_MODULES: Set[str] = {
+    "cloudpickle",
+    "tinyagent",
+    "json",
+    "time",
+    "datetime",
+    "requests",
+    
+}
+
+
+# ---------------------------------------------------------------------------
+# Utility helpers
+# ---------------------------------------------------------------------------
+
+def _is_allowed(module_root: str, allowed: Sequence[str] | None) -> bool:
+    """Return ``True`` if *module_root* is within *allowed* specification."""
+
+    if allowed is None:
+        # No explicit allow-list means everything that is **not** in the
+        # dangerous list is considered fine.
+        return True
+
+    # Fast path – wildcard allows everything.
+    if "*" in allowed:
+        return True
+
+    for pattern in allowed:
+        if pattern.endswith(".*"):
+            if module_root == pattern[:-2]:
+                return True
+        elif module_root == pattern:
+            return True
+    return False
+
+
+# ---------------------------------------------------------------------------
+# Static analysis helpers
+# ---------------------------------------------------------------------------
+
+def _iter_import_nodes(tree: ast.AST) -> Iterable[ast.AST]:
+    """Yield all *import* related nodes from *tree*."""
+    for node in ast.walk(tree):
+        if isinstance(node, (ast.Import, ast.ImportFrom)):
+            yield node
+
+
+def _extract_module_roots(node: ast.AST) -> List[str]:
+    """Return the *top-level* module names referenced in an import node."""
+    roots: list[str] = []
+    if isinstance(node, ast.Import):
+        for alias in node.names:
+            roots.append(alias.name.split(".")[0])
+    elif isinstance(node, ast.ImportFrom):
+        if node.module is not None:
+            roots.append(node.module.split(".")[0])
+    return roots
+
+
+def _check_for_dangerous_function_calls(tree: ast.AST, authorized_functions: Sequence[str] | None = None) -> Set[str]:
+    """
+    Check for calls to dangerous functions in the AST.
+    
+    Parameters
+    ----------
+    tree
+        The AST to check
+    authorized_functions
+        Optional white-list of dangerous functions that are allowed
+        
+    Returns
+    -------
+    Set[str]
+        Set of dangerous function names found in the code
+    """
+    dangerous_calls = set()
+    
+    # Convert authorized_functions to a set if it's not None and not a boolean
+    authorized_set = set(authorized_functions) if authorized_functions is not None and not isinstance(authorized_functions, bool) else set()
+    
+    for node in ast.walk(tree):
+        # Check for direct function calls: func()
+        if isinstance(node, ast.Call) and isinstance(node.func, ast.Name):
+            func_name = node.func.id
+            # Only flag if it's a known dangerous function
+            if func_name in DANGEROUS_FUNCTIONS and func_name not in authorized_set:
+                dangerous_calls.add(func_name)
+        
+        # Check for calls via string literals in exec/eval: exec("import os")
+        if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id in ["exec", "eval"]:
+            # Only check if the function itself is not authorized
+            if node.func.id not in authorized_set:
+                if node.args and isinstance(node.args[0], ast.Constant) and isinstance(node.args[0].value, str):
+                    dangerous_calls.add(f"{node.func.id} with string literal")
+        
+        # Check for attribute access: builtins.exec()
+        if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute):
+            if node.func.attr in DANGEROUS_FUNCTIONS:
+                if isinstance(node.func.value, ast.Name):
+                    module_name = node.func.value.id
+                    # Focus on builtins module access
+                    if module_name == "builtins" and node.func.attr not in authorized_set:
+                        func_name = f"{module_name}.{node.func.attr}"
+                        dangerous_calls.add(func_name)
+        
+        # Check for string manipulation that could be used to bypass security
+        # For example: e = "e" + "x" + "e" + "c"; e("import os")
+        if isinstance(node, ast.Assign):
+            for target in node.targets:
+                if isinstance(target, ast.Name) and isinstance(node.value, ast.BinOp):
+                    # Check if we're building a string that could be a dangerous function name
+                    potential_name = _extract_string_from_binop(node.value)
+                    if potential_name in DANGEROUS_FUNCTIONS and potential_name not in authorized_set:
+                        dangerous_calls.add(f"string manipulation to create '{potential_name}'")
+        
+        # Check for getattr(builtins, "exec") pattern
+        if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "getattr":
+            if len(node.args) >= 2 and isinstance(node.args[1], ast.Constant) and isinstance(node.args[1].value, str):
+                attr_name = node.args[1].value
+                if attr_name in DANGEROUS_FUNCTIONS and attr_name not in authorized_set:
+                    if isinstance(node.args[0], ast.Name) and node.args[0].id == "builtins":
+                        module_name = node.args[0].id
+                        dangerous_calls.add(f"getattr({module_name}, '{attr_name}')")
+    
+    return dangerous_calls
+
+
+def _extract_string_from_binop(node: ast.BinOp) -> str:
+    """
+    Attempt to extract a string from a binary operation node.
+    This helps detect string concatenation that builds dangerous function names.
+    
+    For example: "e" + "x" + "e" + "c" -> "exec"
+    
+    Parameters
+    ----------
+    node
+        The binary operation node
+        
+    Returns
+    -------
+    str
+        The extracted string, or empty string if not extractable
+    """
+    if isinstance(node, ast.Constant) and isinstance(node.value, str):
+        return node.value
+    
+    if not isinstance(node, ast.BinOp):
+        return ""
+    
+    # Handle string concatenation
+    if isinstance(node.op, ast.Add):
+        left_str = _extract_string_from_binop(node.left)
+        right_str = _extract_string_from_binop(node.right)
+        return left_str + right_str
+    
+    return ""
+
+
+def _detect_string_obfuscation(tree: ast.AST) -> bool:
+    """
+    Detect common string obfuscation techniques that might be used to bypass security.
+    
+    Parameters
+    ----------
+    tree
+        The AST to check
+        
+    Returns
+    -------
+    bool
+        True if suspicious string manipulation is detected
+    """
+    suspicious_patterns = False
+    
+    for node in ast.walk(tree):
+        # Check for chr() usage to build strings
+        if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "chr":
+            suspicious_patterns = True
+            break
+            
+        # Check for ord() usage in combination with string operations
+        if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "ord":
+            suspicious_patterns = True
+            break
+            
+        # Check for suspicious string joins with list comprehensions
+        if (isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) and 
+            node.func.attr == "join" and isinstance(node.args[0], ast.ListComp)):
+            suspicious_patterns = True
+            break
+            
+        # Check for base64 decoding
+        if (isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) and 
+            node.func.attr in ["b64decode", "b64encode", "b32decode", "b32encode", "b16decode", "b16encode"]):
+            suspicious_patterns = True
+            break
+            
+        # Check for string formatting that might be used to build dangerous code
+        if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) and node.func.attr == "format":
+            suspicious_patterns = True
+            break
+            
+    return suspicious_patterns
+
+
+def validate_code_safety(code: str, *, authorized_imports: Sequence[str] | None = None, 
+                        authorized_functions: Sequence[str] | None = None, trusted_code: bool = False) -> None:
+    """Static validation of user code.
+
+    Parameters
+    ----------
+    code
+        The user supplied source code (single string or multi-line).
+    authorized_imports
+        Optional white-list restricting which modules may be imported.  If
+        *None* every module that is not in :pydata:`DANGEROUS_MODULES` is
+        allowed.  Wildcards are supported – e.g. ``["numpy.*"]`` allows any
+        sub-package of *numpy*.
+    authorized_functions
+        Optional white-list of dangerous functions that are allowed.
+    trusted_code
+        If True, skip security checks. This should only be used for code that is part of the
+        framework, developer-provided tools, or default executed code.
+    """
+    # Skip security checks for trusted code
+    if trusted_code:
+        return
+
+    try:
+        tree = ast.parse(code, mode="exec")
+    except SyntaxError:
+        # If the code does not even parse we leave the error handling to the
+        # caller (who will attempt to compile / execute the code later on).
+        return
+
+    blocked: set[str] = set()
+    
+    # Convert authorized_imports to a set if it's not None and not a boolean
+    combined_allowed = None
+    if authorized_imports is not None and not isinstance(authorized_imports, bool):
+        combined_allowed = set(list(authorized_imports) + list(ESSENTIAL_MODULES))
+    
+    
+    for node in _iter_import_nodes(tree):
+        for root in _extract_module_roots(node):
+
+            # Check if module is explicitly allowed
+            if combined_allowed is not None:
+                allowed = _is_allowed(root, combined_allowed)
+            else:
+                # If no explicit allow-list, only allow if not in DANGEROUS_MODULES
+                allowed = root not in DANGEROUS_MODULES
+
+            if root in DANGEROUS_MODULES and allowed and combined_allowed is not None:
+                warnings.warn(
+                    f"⚠️  Importing dangerous module '{root}' was allowed due to authorized_imports configuration.",
+                    stacklevel=2,
+                )
+            
+            # Block dangerous modules unless explicitly allowed
+            if root in DANGEROUS_MODULES and not allowed:
+                blocked.add(root)
+            # If there is an explicit allow-list, block everything not on it
+            elif authorized_imports is not None and not isinstance(authorized_imports, bool) and not allowed and root not in ESSENTIAL_MODULES:
+                blocked.add(root)
+
+    # ------------------------------------------------------------------
+    # Detect direct calls to __import__ (e.g.  __import__("os")) in *untrusted* code
+    # ------------------------------------------------------------------
+    for _node in ast.walk(tree):
+        if isinstance(_node, ast.Call):
+            # Pattern: __import__(...)
+            if isinstance(_node.func, ast.Name) and _node.func.id == "__import__":
+                # Check if it's a direct call to the built-in __import__
+                raise ValueError("SECURITY VIOLATION: Usage of __import__ is not allowed in untrusted code.")
+            # Pattern: builtins.__import__(...)
+            if (
+                isinstance(_node.func, ast.Attribute)
+                and isinstance(_node.func.value, ast.Name)
+                and _node.func.attr == "__import__"
+                and _node.func.value.id == "builtins"
+            ):
+                raise ValueError("SECURITY VIOLATION: Usage of builtins.__import__ is not allowed in untrusted code.")
+
+    # ------------------------------------------------------------------
+    # Detect calls to dangerous functions (e.g. exec, eval) in *untrusted* code
+    # ------------------------------------------------------------------
+    dangerous_calls = _check_for_dangerous_function_calls(tree, authorized_functions)
+    if dangerous_calls:
+        offenders = ", ".join(sorted(dangerous_calls))
+        raise ValueError(f"SECURITY VIOLATION: Usage of dangerous function(s) {offenders} is not allowed in untrusted code.")
+
+    # ------------------------------------------------------------------
+    # Detect string obfuscation techniques that might be used to bypass security
+    # ------------------------------------------------------------------
+    if _detect_string_obfuscation(tree):
+        raise ValueError("SECURITY VIOLATION: Suspicious string manipulation detected that could be used to bypass security.")
+
+    if blocked:
+        offenders = ", ".join(sorted(blocked))
+        msg = f"SECURITY VIOLATION: Importing module(s) {offenders} is not allowed."
+        if authorized_imports is not None and not isinstance(authorized_imports, bool):
+            msg += " Allowed imports are: " + ", ".join(sorted(authorized_imports))
+        raise ValueError(msg)
+
+
+# ---------------------------------------------------------------------------
+# Runtime import hook
+# ---------------------------------------------------------------------------
+
+def install_import_hook(
+    *,
+    blocked_modules: Set[str] | None = None,
+    authorized_imports: Sequence[str] | None = None,
+    trusted_code: bool = False,
+) -> None:
+    """Monkey-patch the built-in ``__import__`` to deny run-time imports.
+
+    The hook is *process-wide* but extremely cheap to install.  It simply
+    checks the *root* package name against the provided *blocked_modules*
+    (defaults to :pydata:`DANGEROUS_MODULES`) and raises ``ImportError`` if the
+    import should be denied.
+
+    Calling this function **multiple times** is safe – only the first call
+    installs the wrapper, subsequent calls are ignored.
+    
+    Parameters
+    ----------
+    blocked_modules
+        Set of module names to block. Defaults to DANGEROUS_MODULES.
+    authorized_imports
+        Optional white-list restricting which modules may be imported.
+    trusted_code
+        If True, skip security checks. This should only be used for code that is part of the
+        framework, developer-provided tools, or default executed code.
+    """
+    # Skip security checks for trusted code
+    if trusted_code:
+        return
+
+    blocked_modules = blocked_modules or DANGEROUS_MODULES
+    
+    # Convert authorized_imports to a set if it's not None and not a boolean
+    authorized_set = set(authorized_imports) if authorized_imports is not None and not isinstance(authorized_imports, bool) else None
+    
+    # Create a combined set for allowed modules (essential + authorized)
+    combined_allowed = None
+    if authorized_set is not None:
+        combined_allowed = set(list(authorized_set) + list(ESSENTIAL_MODULES))
+
+    # Check if we have already installed the hook to avoid double-wrapping.
+    if getattr(builtins, "__tinyagent_import_hook_installed", False):
+        return
+
+    original_import = builtins.__import__
+
+    def _safe_import(
+        name: str,
+        globals=None,
+        locals=None,
+        fromlist=(),
+        level: int = 0,
+    ):  # type: ignore[override]
+        root = name.split(".")[0]
+        
+        # Check if module is explicitly allowed
+        if combined_allowed is not None:
+            allowed = _is_allowed(root, combined_allowed)
+        else:
+            # If no explicit allow-list, only allow if not in blocked_modules
+            allowed = root not in blocked_modules
+
+        if root in blocked_modules and allowed and authorized_set is not None:
+            warnings.warn(
+                f"⚠️  Importing dangerous module '{root}' was allowed due to authorized_imports configuration.",
+                stacklevel=2,
+            )
+        elif root in blocked_modules and not allowed:
+            error_msg = f"SECURITY VIOLATION: Import of module '{name}' is blocked by TinyAgent security policy"
+            if authorized_set is not None:
+                error_msg += f". Allowed imports are: {', '.join(sorted(authorized_set))}"
+            raise ImportError(error_msg)
+        elif authorized_set is not None and not allowed and root not in ESSENTIAL_MODULES:
+            error_msg = f"SECURITY VIOLATION: Import of module '{name}' is not in the authorized imports list"
+            if authorized_set:
+                error_msg += f": {', '.join(sorted(authorized_set))}"
+            raise ImportError(error_msg)
+
+        return original_import(name, globals, locals, fromlist, level)
+
+    builtins.__import__ = _safe_import  # type: ignore[assignment]
+    setattr(builtins, "__tinyagent_import_hook_installed", True)
+
+
+# ---------------------------------------------------------------------------
+# Runtime function hook
+# ---------------------------------------------------------------------------
+
+@contextlib.contextmanager
+def function_safety_context(
+    *,
+    blocked_functions: Set[str] | None = None,
+    authorized_functions: Sequence[str] | None = None,
+    trusted_code: bool = False,
+):
+    """
+    Context manager for safely executing code with dangerous functions blocked.
+    
+    Parameters
+    ----------
+    blocked_functions
+        Set of function names to block. Defaults to RUNTIME_BLOCKED_FUNCTIONS.
+    authorized_functions
+        Optional white-list of dangerous functions that are allowed.
+    trusted_code
+        If True, skip security checks. This should only be used for code that is part of the
+        framework, developer-provided tools, or default executed code.
+    """
+    if trusted_code:
+        yield
+        return
+        
+    # Install the function hook
+    blocked_functions = blocked_functions or RUNTIME_BLOCKED_FUNCTIONS
+    
+    # Convert authorized_functions to a set if it's not None and not a boolean
+    authorized_set = set(authorized_functions) if authorized_functions is not None and not isinstance(authorized_functions, bool) else set()
+    
+    # Store original functions
+    original_functions = {}
+    
+    # Replace dangerous functions with safe versions
+    for func_name in blocked_functions:
+        # Only block functions that exist in builtins
+        if hasattr(builtins, func_name) and func_name not in authorized_set:
+            original_functions[func_name] = getattr(builtins, func_name)
+            
+            # Create a closure to capture the function name
+            def make_safe_function(name):
+                def safe_function(*args, **kwargs):
+                    error_msg = f"SECURITY VIOLATION: Function '{name}' is blocked by TinyAgent security policy"
+                    if authorized_functions and not isinstance(authorized_functions, bool) and isinstance(authorized_functions, (list, tuple, set)):
+                        error_msg += f". Allowed functions are: {', '.join(sorted(authorized_set))}"
+                    raise RuntimeError(error_msg)
+                return safe_function
+            
+            # Replace the function
+            setattr(builtins, func_name, make_safe_function(func_name))
+    
+    try:
+        yield
+    finally:
+        # Restore original functions
+        for func_name, original_func in original_functions.items():
+            setattr(builtins, func_name, original_func)