From 6b7086b8667d3662db36b5ab47caaca79ce4d267 Mon Sep 17 00:00:00 2001 From: Ashutosh0x Date: Sat, 23 May 2026 19:09:37 +0530 Subject: [PATCH] security: add module blocklist for YAML agent config code references Add a _BLOCKED_MODULES set and _validate_module_reference() function to prevent importing dangerous standard library modules (os, subprocess, builtins, importlib, pickle, etc.) when resolving code references from YAML agent configurations. The existing CVE-2026-4810 fix blocks the 'args' key in YAML configs to prevent passing arguments to constructors. However, the resolve_code_reference(), resolve_fully_qualified_name(), and _resolve_tools() functions still call importlib.import_module() with no restriction on which modules can be imported. This allows an attacker to reference dangerous callables like os.system or subprocess.call in callback, tool, schema, or model code-reference fields. This commit adds validation at all three import points: - resolve_code_reference() (used for callbacks, schemas, model_code) - resolve_fully_qualified_name() (used for agent class resolution) - _resolve_tools() (used for user-defined tool resolution) The blocklist is only enforced when _ENFORCE_DENYLIST is True (set by the web dev server), matching the existing denylist behavior. Includes comprehensive tests verifying: - 11 different blocked modules are rejected in callback fields - 3 blocked modules are rejected in tool fields - Direct resolve_code_reference() calls are blocked - Direct resolve_fully_qualified_name() calls are blocked - google.adk.* modules continue to work (allowlist behavior) --- src/google/adk/agents/config_agent_utils.py | 60 ++++++++++++ tests/unittests/agents/test_agent_config.py | 103 ++++++++++++++++++++ 2 files changed, 163 insertions(+) diff --git a/src/google/adk/agents/config_agent_utils.py b/src/google/adk/agents/config_agent_utils.py index 53a0736231..9b2d65c68e 100644 --- a/src/google/adk/agents/config_agent_utils.py +++ b/src/google/adk/agents/config_agent_utils.py @@ -158,6 +158,7 @@ def _resolve_tools(self, tool_configs: list[ToolConfig]) -> list[Any]: obj = getattr(module, tool_config.name) else: # User-defined tools + _validate_module_reference(tool_config.name) module_path, obj_name = tool_config.name.rsplit(".", 1) module = importlib.import_module(module_path) obj = getattr(module, obj_name) @@ -490,6 +491,63 @@ def _resolve_agent_class(agent_class: str) -> type[Any]: _BLOCKED_YAML_KEYS = frozenset({"args"}) _ENFORCE_DENYLIST = False +# Modules that must never be imported via YAML agent configuration. +# These provide direct access to the operating system, process execution, +# or dynamic code evaluation and could be abused to achieve arbitrary +# code execution when referenced in callback, tool, schema, or model +# code-reference fields. +_BLOCKED_MODULES = frozenset({ + "os", + "subprocess", + "sys", + "builtins", + "importlib", + "shutil", + "socket", + "http", + "urllib", + "ctypes", + "multiprocessing", + "threading", + "signal", + "code", + "codeop", + "compileall", + "runpy", + "webbrowser", + "antigravity", + "pty", + "commands", + "pdb", + "profile", + "tempfile", + "shelve", + "pickle", + "marshal", +}) + + +def _validate_module_reference(fully_qualified_name: str) -> None: + """Validate that a module reference does not target a blocked module. + + Args: + fully_qualified_name: The fully-qualified Python name to validate + (e.g. ``"my_package.my_module.my_func"``). + + Raises: + ValueError: If the top-level module is in ``_BLOCKED_MODULES``. + """ + if not _ENFORCE_DENYLIST: + return + # Extract the top-level package from the fully-qualified name. + top_module = fully_qualified_name.split(".")[0] + if top_module in _BLOCKED_MODULES: + raise ValueError( + f"Blocked module reference: {fully_qualified_name!r}. " + f"Importing from the '{top_module}' module is not allowed in " + "agent configurations because it can execute arbitrary code." + ) + def _set_enforce_denylist(value: bool) -> None: global _ENFORCE_DENYLIST @@ -516,6 +574,7 @@ def _check_config_for_blocked_keys(node: Any, filename: str) -> None: def resolve_fully_qualified_name(name: str) -> Any: try: module_path, obj_name = name.rsplit(".", 1) + _validate_module_reference(name) module = importlib.import_module(module_path) return getattr(module, obj_name) except Exception as e: @@ -596,6 +655,7 @@ def resolve_code_reference(code_config: CodeConfig) -> Any: if not code_config or not code_config.name: raise ValueError("Invalid CodeConfig.") + _validate_module_reference(code_config.name) module_path, obj_name = code_config.name.rsplit(".", 1) module = importlib.import_module(module_path) obj = getattr(module, obj_name) diff --git a/tests/unittests/agents/test_agent_config.py b/tests/unittests/agents/test_agent_config.py index 1dc3cae225..86dc4ce88c 100644 --- a/tests/unittests/agents/test_agent_config.py +++ b/tests/unittests/agents/test_agent_config.py @@ -373,6 +373,109 @@ def test_from_config_blocks_args_when_enforced(tmp_path): config_agent_utils._set_enforce_denylist(False) +@pytest.mark.parametrize( + "blocked_module", + [ + "os.system", + "subprocess.call", + "subprocess.Popen", + "builtins.exec", + "builtins.eval", + "importlib.import_module", + "shutil.rmtree", + "socket.socket", + "ctypes.cdll", + "pickle.loads", + "marshal.loads", + ], +) +def test_blocked_module_in_callback_raises_when_enforced( + blocked_module: str, tmp_path: Path +): + """Verify that referencing blocked stdlib modules in callbacks is rejected.""" + config_file = tmp_path / "malicious.yaml" + config_file.write_text(f"""\ +name: evil_agent +model: gemini-2.0-flash +instruction: "harmless" +before_agent_callbacks: + - name: "{blocked_module}" +""") + + config_agent_utils._set_enforce_denylist(True) + try: + with pytest.raises(ValueError, match="Blocked module reference"): + config_agent_utils.from_config(str(config_file)) + finally: + config_agent_utils._set_enforce_denylist(False) + + +@pytest.mark.parametrize( + "blocked_module", + [ + "os.system", + "subprocess.run", + "builtins.exec", + ], +) +def test_blocked_module_in_tools_raises_when_enforced( + blocked_module: str, tmp_path: Path +): + """Verify that referencing blocked stdlib modules in tools is rejected.""" + config_file = tmp_path / "malicious.yaml" + config_file.write_text(f"""\ +name: evil_agent +model: gemini-2.0-flash +instruction: "harmless" +tools: + - name: "{blocked_module}" +""") + + config_agent_utils._set_enforce_denylist(True) + try: + with pytest.raises(ValueError, match="Blocked module reference"): + config_agent_utils.from_config(str(config_file)) + finally: + config_agent_utils._set_enforce_denylist(False) + + +def test_resolve_code_reference_blocks_os_when_enforced(): + """Verify resolve_code_reference blocks os module directly.""" + from google.adk.agents.common_configs import CodeConfig + + config_agent_utils._set_enforce_denylist(True) + try: + with pytest.raises(ValueError, match="Blocked module reference"): + config_agent_utils.resolve_code_reference( + CodeConfig(name="os.system") + ) + finally: + config_agent_utils._set_enforce_denylist(False) + + +def test_resolve_fully_qualified_name_blocks_subprocess_when_enforced(): + """Verify resolve_fully_qualified_name blocks subprocess module.""" + config_agent_utils._set_enforce_denylist(True) + try: + with pytest.raises(ValueError, match="Blocked module reference"): + config_agent_utils.resolve_fully_qualified_name("subprocess.Popen") + finally: + config_agent_utils._set_enforce_denylist(False) + + +def test_allowed_module_passes_when_enforced(tmp_path: Path): + """Verify that google.adk modules are NOT blocked by the module denylist.""" + config_agent_utils._set_enforce_denylist(True) + try: + # This should NOT raise — google.adk modules must remain allowed + result = config_agent_utils.resolve_fully_qualified_name( + "google.adk.agents.llm_agent.LlmAgent" + ) + assert result is LlmAgent + finally: + config_agent_utils._set_enforce_denylist(False) + + def test_create_workflow_from_yaml(tmp_path: Path): """Test creating a Workflow from a YAML file.""" yaml_content = """