fix: add agent name validation to prevent arbitrary module imports

Co-authored-by: Sasha Sobran <asobran@google.com>
PiperOrigin-RevId: 889890969

Author: sasha-gitg

src/google/adk/cli/utils/agent_loader.py

@@ -19,6 +19,7 @@
 importlogging
 importos
 frompathlibimportPath
+importre
 importsys
 fromtypingimportAny
 fromtypingimportLiteral
@@ -187,8 +188,36 @@ def _load_from_yaml_config(
  ) +e.args[1:]
 raisee
 
+_VALID_AGENT_NAME_RE=re.compile(r"^[a-zA-Z0-9_]+$")
+
+def_validate_agent_name(self, agent_name: str) ->None:
+"""Validate agent name to prevent arbitrary module imports."""
+# Strip the special agent prefix for validation
+ifagent_name.startswith("__"):
+name_to_check=agent_name[2:]
+check_dir=os.path.abspath(SPECIAL_AGENTS_DIR)
+else:
+name_to_check=agent_name
+check_dir=self.agents_dir
+
+ifnotself._VALID_AGENT_NAME_RE.match(name_to_check):
+raiseValueError(
+f"Invalid agent name: {agent_name!r}. Agent names must be valid"
+" Python identifiers (letters, digits, and underscores only)."
+ )
+
+# Verify the agent exists on disk before allowing import
+agent_path=Path(check_dir) /name_to_check
+agent_file=Path(check_dir) /f"{name_to_check}.py"
+ifnot (agent_path.is_dir() oragent_file.is_file()):
+raiseValueError(
+f"Agent not found: {agent_name!r}. No matching directory or module"
+f" exists in '{os.path.join(check_dir, name_to_check)}'."
+ )
+
 def_perform_load(self, agent_name: str) ->Union[BaseAgent, App]:
 """Internal logic to load an agent"""
+self._validate_agent_name(agent_name)
 # Determine the directory to use for loading
 ifagent_name.startswith("__"):
 # Special agent: use special agents directory

tests/unittests/cli/utils/test_agent_loader.py

@@ -299,8 +299,6 @@ def test_error_messages_use_os_sep_consistently(self):
 loader.load_agent(agent_name)
 
 exc_info.match(re.escape(expected_path))
-exc_info.match(re.escape(f"{agent_name}{os.sep}root_agent.yaml"))
-exc_info.match(re.escape(f"<agents_dir>{os.sep}"))
 
 deftest_agent_loader_with_mocked_windows_path(self, monkeypatch):
 """Mock Path() to simulate Windows behavior and catch regressions.
@@ -333,19 +331,10 @@ def test_agent_not_found_error(self):
 withpytest.raises(ValueError) asexc_info:
 loader.load_agent("nonexistent_agent")
 
-expected_msg_part_1="No root_agent found for 'nonexistent_agent'."
-expected_msg_part_2= (
-"Searched in 'nonexistent_agent.agent.root_agent',"
-" 'nonexistent_agent.root_agent' and"
-" 'nonexistent_agent/root_agent.yaml'."
+assert"Agent not found: 'nonexistent_agent'"instr(exc_info.value)
+assertos.path.join(agents_dir, "nonexistent_agent") instr(
+exc_info.value
  )
-expected_msg_part_3= (
-f"Ensure '{agents_dir}/nonexistent_agent' is structured correctly"
- )
-
-assertexpected_msg_part_1instr(exc_info.value)
-assertexpected_msg_part_2instr(exc_info.value)
-assertexpected_msg_part_3instr(exc_info.value)
 
 deftest_agent_without_root_agent_error(self):
 """Test that appropriate error is raised when agent has no root_agent."""
@@ -567,26 +556,16 @@ def test_yaml_agent_not_found_error(self):
 """Test that appropriate error is raised when YAML agent is not found."""
 withtempfile.TemporaryDirectory() astemp_dir:
 loader=AgentLoader(temp_dir)
-agents_dir=temp_dir# For use in the expected message string
+agents_dir=temp_dir
 
 # Try to load nonexistent YAML agent
 withpytest.raises(ValueError) asexc_info:
 loader.load_agent("nonexistent_yaml_agent")
 
-expected_msg_part_1="No root_agent found for 'nonexistent_yaml_agent'."
-expected_msg_part_2= (
-"Searched in 'nonexistent_yaml_agent.agent.root_agent',"
-" 'nonexistent_yaml_agent.root_agent' and"
-" 'nonexistent_yaml_agent/root_agent.yaml'."
+assert"Agent not found: 'nonexistent_yaml_agent'"instr(exc_info.value)
+assertos.path.join(agents_dir, "nonexistent_yaml_agent") instr(
+exc_info.value
  )
-expected_msg_part_3= (
-f"Ensure '{agents_dir}/nonexistent_yaml_agent' is structured"
-" correctly"
- )
-
-assertexpected_msg_part_1instr(exc_info.value)
-assertexpected_msg_part_2instr(exc_info.value)
-assertexpected_msg_part_3instr(exc_info.value)
 
 deftest_yaml_agent_invalid_yaml_error(self):
 """Test that appropriate error is raised when YAML is invalid."""
@@ -778,20 +757,7 @@ def test_special_agent_not_found_error(self):
 withpytest.raises(ValueError) asexc_info:
 loader.load_agent("__nonexistent_special")
 
-expected_msg_part_1="No root_agent found for '__nonexistent_special'."
-expected_msg_part_2= (
-"Searched in 'nonexistent_special.agent.root_agent',"
-" 'nonexistent_special.root_agent' and"
-" 'nonexistent_special/root_agent.yaml'."
- )
-expected_msg_part_3= (
-f"Ensure '{special_agents_dir}/nonexistent_special' is structured"
-" correctly"
- )
-
-assertexpected_msg_part_1instr(exc_info.value)
-assertexpected_msg_part_2instr(exc_info.value)
-assertexpected_msg_part_3instr(exc_info.value)
+assert"Agent not found: '__nonexistent_special'"instr(exc_info.value)
 
 finally:
 # Restore original SPECIAL_AGENTS_DIR
@@ -1038,3 +1004,50 @@ def __init__(self):
 assert"random_folder"notinagents
 assert"data"notinagents
 assert"tmp"notinagents
+
+deftest_validate_agent_name_rejects_dotted_paths(self):
+"""Agent names with dots are rejected to prevent arbitrary module imports."""
+withtempfile.TemporaryDirectory() astemp_dir:
+loader=AgentLoader(temp_dir)
+fornamein ["os.path", "sys.modules", "subprocess.call"]:
+withpytest.raises(ValueError, match="Invalid agent name"):
+loader.load_agent(name)
+
+deftest_validate_agent_name_rejects_relative_imports(self):
+"""Agent names starting with dots are rejected."""
+withtempfile.TemporaryDirectory() astemp_dir:
+loader=AgentLoader(temp_dir)
+fornamein ["..foo", ".bar", "...baz"]:
+withpytest.raises(ValueError, match="Invalid agent name"):
+loader.load_agent(name)
+
+deftest_validate_agent_name_rejects_path_separators(self):
+"""Agent names with slashes or special characters are rejected."""
+withtempfile.TemporaryDirectory() astemp_dir:
+loader=AgentLoader(temp_dir)
+fornamein ["foo/bar", "foo\\bar", "foo-bar", "foo bar"]:
+withpytest.raises(ValueError, match="Invalid agent name"):
+loader.load_agent(name)
+
+deftest_validate_agent_name_allows_valid_names(self):
+"""Valid Python identifiers that exist on disk pass validation."""
+withtempfile.TemporaryDirectory() astemp_dir:
+temp_path=Path(temp_dir)
+fornamein ["my_agent", "Agent1", "_private"]:
+ (temp_path/name).mkdir(exist_ok=True)
+loader=AgentLoader(temp_dir)
+fornamein ["my_agent", "Agent1", "_private"]:
+# Should not raise ValueError for name validation;
+# may raise other errors because the agent has no root_agent
+withpytest.raises(Exception) asexc_info:
+loader.load_agent(name)
+assert"Invalid agent name"notinstr(exc_info.value)
+assert"Agent not found"notinstr(exc_info.value)
+
+deftest_validate_agent_name_rejects_nonexistent_agent(self):
+"""Valid identifiers that don't exist on disk are rejected before import."""
+withtempfile.TemporaryDirectory() astemp_dir:
+loader=AgentLoader(temp_dir)
+# 'subprocess' is a valid identifier but shouldn't be importable as an agent
+withpytest.raises(ValueError, match="Agent not found"):
+loader.load_agent("subprocess")