test: cover tool safety patch lines

Author: zzp1221

tests/tools/safety/test_tool_safety_guard.py

@@ -6,21 +6,25 @@
 
 from __future__ import annotations
 
+import ast
 import json
 import time
 from pathlib import Path
 from unittest.mock import AsyncMock
 
 import pytest
 
+from trpc_agent_sdk.code_executors import CodeBlock
+from trpc_agent_sdk.code_executors import CodeExecutionInput
+from trpc_agent_sdk.code_executors import BaseCodeExecutor
+from trpc_agent_sdk.code_executors import create_code_execution_result
 from trpc_agent_sdk.context import create_agent_context
-from trpc_agent_sdk.tools import FunctionTool
+from trpc_agent_sdk.filter import FilterResult
 from trpc_agent_sdk.tools.safety import SafetyDecision
 from trpc_agent_sdk.tools.safety import SafetyGuardedCodeExecutor
-from trpc_agent_sdk.tools.safety import ToolSafetyBlockedError
-from trpc_agent_sdk.tools.safety import ToolSafetyFinding
 from trpc_agent_sdk.tools.safety import SafetyRiskLevel
 from trpc_agent_sdk.tools.safety import ToolSafetyAuditLogger
+from trpc_agent_sdk.tools.safety import ToolSafetyBlockedError
 from trpc_agent_sdk.tools.safety import ToolSafetyFilter
 from trpc_agent_sdk.tools.safety import ToolSafetyGuard
 from trpc_agent_sdk.tools.safety import ToolSafetyPolicy
@@ -30,11 +34,8 @@
 from trpc_agent_sdk.tools.safety import apply_tool_safety_span_attributes
 from trpc_agent_sdk.tools.safety import extract_script_from_tool_args
 from trpc_agent_sdk.tools.safety import load_tool_safety_policy
-from trpc_agent_sdk.code_executors import CodeBlock
-from trpc_agent_sdk.code_executors import CodeExecutionInput
-from trpc_agent_sdk.code_executors import BaseCodeExecutor
-from trpc_agent_sdk.code_executors import create_code_execution_result
-from trpc_agent_sdk.filter import FilterResult
+from trpc_agent_sdk.tools.safety import _scanner as scanner_module
+from trpc_agent_sdk.tools.safety import _types as types_module
 from trpc_agent_sdk.tools.safety._guard import command_to_args
 from trpc_agent_sdk.tools.safety._guard import infer_language_from_key
 from trpc_agent_sdk.tools.safety._types import max_risk_level
@@ -310,6 +311,71 @@ def test_scanner_command_arg_prefixes_and_duration_helpers():
         ]
 
 
+def test_scanner_edge_branches_from_direct_scans(monkeypatch):
+    scanner = ToolSafetyScanner(policy(max_parallel_tasks=2, max_loop_iterations=3))
+
+    assert scanner.scan(ToolSafetyScanRequest(script="", language="python", env=[])).decision == SafetyDecision.ALLOW
+    assert scanner._scan_bash_line("", 1) == []  # pylint: disable=protected-access
+
+    monkeypatch.setattr(scanner_module, "split_shell_commands", lambda tokens: [[]])
+    assert scanner._scan_bash_line("echo ok", 1) == []  # pylint: disable=protected-access
+
+    scripts = [
+        ("funcs = [print]\nfuncs[0]('x')", None),
+        ("import os\nos.remove('.env')", "TSG-FILE-DENIED-PATH"),
+        ("open('/etc/passwd', mode='w')", "TSG-FILE-SYSTEM-WRITE"),
+        ("f = open('out.txt')", None),
+        ("import concurrent.futures\nconcurrent.futures.ThreadPoolExecutor()", "TSG-RESOURCE-PARALLELISM"),
+        ("for i in range(limit):\n    pass", None),
+        ("for i in range(2, 10, 2):\n    pass", "TSG-RESOURCE-LARGE-LOOP"),
+    ]
+    for script, rule_id in scripts:
+        report = scanner.scan(ToolSafetyScanRequest(script=script, language="python"))
+        if rule_id:
+            assert any(finding.rule_id == rule_id for finding in report.findings)
+
+
+def test_python_analyzer_helper_edge_branches():
+    analyzer = scanner_module._PythonAnalyzer("pass\n", policy())  # pylint: disable=protected-access
+
+    assert analyzer.network_target_from_arg(ast.parse("url", mode="eval").body) == ""
+    assert analyzer.node_line(ast.Pass(lineno=1, col_offset=0)) == "pass"
+    assert analyzer.node_line(ast.Pass(lineno=99, col_offset=0)) == ""
+    assert analyzer.call_name(ast.Attribute(value=ast.Constant(value=1), attr="field", ctx=ast.Load())) == "field"
+    assert analyzer.call_name(ast.Constant(value=None)) == ""
+    assert analyzer.name_of(ast.Attribute(value=ast.Name(id="obj", ctx=ast.Load()), attr="token", ctx=ast.Load())) == "token"
+    assert analyzer.name_of(ast.Constant(value=None)) == ""
+    assert analyzer.constant_string(None) == ""
+    assert analyzer.constant_string(ast.parse("f'key={value}'", mode="eval").body) == "key={}"
+    assert analyzer.constant_string(ast.parse("['a', value]", mode="eval").body) == ""
+    assert analyzer.constant_string(ast.parse("('a', 'b')", mode="eval").body) == "a b"
+    assert analyzer.path_literal_from_node(None) == ""
+    assert analyzer.path_literal_from_node(ast.parse("Path('.env')", mode="eval").body) == ".env"
+    assert analyzer.path_literal_from_node(ast.parse("Path(value)", mode="eval").body) == ""
+    assert analyzer.path_literal_from_node(ast.parse("Path('.env').expanduser().read_text", mode="eval").body) == ".env"
+    assert analyzer.path_literal_from_node(ast.parse("config.read_text", mode="eval").body) == ""
+    assert analyzer.range_count(ast.parse("range(limit)", mode="eval").body) is None
+    assert analyzer.range_count(ast.parse("range(2, 10, 2)", mode="eval").body) == 4
+
+
+def test_scanner_module_helper_edge_branches():
+    assert scanner_module.is_secret_name("") is False
+    assert scanner_module.numeric_constant(ast.parse("value", mode="eval").body) is None
+    assert scanner_module.is_install_invocation(["python", "-m", "pip", "install", "demo"]) is False
+    assert scanner_module.is_recursive_delete(["echo", "-rf"]) is False
+    assert scanner_module.redirection_targets(["echo", "x", ">/etc/app.conf"]) == ["/etc/app.conf"]
+    assert scanner_module.timeout_seconds_from_tokens(["timeout", "--kill-after=1s", "10m"]) == 600
+    assert scanner_module.timeout_seconds_from_tokens(["timeout", "--preserve-status", "10m"]) == 600
+    assert scanner_module.timeout_seconds_from_tokens(["timeout", "--preserve-status"]) is None
+    assert scanner_module.duration_token_seconds("never") is None
+    assert scanner_module.parallel_tasks_from_tokens(["parallel", "-j", "4"]) == 4
+    assert scanner_module.parallel_tasks_from_tokens(["parallel", "--jobs=7"]) == 7
+    assert scanner_module.parallel_tasks_from_tokens(["parallel", "echo"]) == 0
+    assert scanner_module.parallel_tasks_from_tokens(["xargs", "-P3"]) == 3
+    assert scanner_module.int_token("many") is None
+    assert scanner_module.line_has_allowlisted_url("echo no-url", policy()) is False
+
+
 def test_policy_loader_and_matching_branches(tmp_path):
     policy_file = tmp_path / "policy.yaml"
     policy_file.write_text(
@@ -340,8 +406,11 @@ def test_policy_loader_and_matching_branches(tmp_path):
     assert loaded.is_command_allowed("custom-tool --flag") is True
     assert loaded.is_command_allowed(None) is False
     assert loaded.is_denied_path(None) is False
+    assert loaded.is_denied_path("config.yaml") is False
     assert loaded.is_system_write_path("/") is True
     assert loaded.is_system_write_path(None) is False
+    assert ToolSafetyPolicy(allowed_domains=[""]).is_domain_allowed("api.example.com") is False
+    assert ToolSafetyPolicy(system_write_paths=[""]).is_system_write_path("/etc/passwd") is False
     assert load_tool_safety_policy(None).name == "default"
 
 
@@ -370,6 +439,12 @@ def test_report_and_type_helpers_cover_empty_and_fallback_paths():
     assert json.loads(report.to_json(indent=None))["risk_count"] == 0
 
 
+def test_max_risk_level_fallback_when_order_has_no_match(monkeypatch):
+    monkeypatch.setattr(types_module, "RISK_LEVEL_ORDER", {})
+
+    assert types_module.max_risk_level(["unexpected"]) == SafetyRiskLevel.NONE
+
+
 def test_audit_event_contains_required_fields(tmp_path):
     report = ToolSafetyScanner(policy()).scan(
         ToolSafetyScanRequest(script="rm -rf /tmp/demo", language="bash", tool_metadata={"name": "Bash"}))

trpc_agent_sdk/tools/safety/_policy.py

@@ -14,7 +14,6 @@
 
 import yaml
 
-
 DEFAULT_ALLOWED_COMMANDS = [
     "awk",
     "cat",

trpc_agent_sdk/tools/safety/_scanner.py

@@ -27,12 +27,10 @@
 from ._types import max_risk_level
 from ._types import risk_level_value
 
-
 SECRET_VALUE_RE = re.compile(
     r"(?i)(api[_-]?key|access[_-]?token|secret|password|private[_-]?key|authorization)\s*[:=]\s*"
-    r"['\"]?([A-Za-z0-9_./+=:-]{12,})"
-)
-URL_RE = re.compile(r"(?i)\bhttps?://[^\s'\"<>)]+" )
+    r"['\"]?([A-Za-z0-9_./+=:-]{12,})")
+URL_RE = re.compile(r"(?i)\bhttps?://[^\s'\"<>)]+")
 ENV_VAR_RE = re.compile(r"\$(?:\{)?([A-Za-z_][A-Za-z0-9_]*)")
 SHELL_CONTROL_RE = re.compile(r"(\|\||&&|;|\||`|\$\(|>|<)")
 FORK_BOMB_RE = re.compile(r":\s*\(\s*\)\s*\{\s*:\s*\|\s*:\s*&\s*\}\s*;\s*:")
@@ -79,15 +77,21 @@
     "SECRET",
     "TOKEN",
 }
-SENSITIVE_SINK_NAMES = {"print", "logging.info", "logging.warning", "logging.error", "logger.info", "logger.warning",
-                        "logger.error", "sys.stdout.write", "sys.stderr.write"}
+SENSITIVE_SINK_NAMES = {
+    "print", "logging.info", "logging.warning", "logging.error", "logger.info", "logger.warning", "logger.error",
+    "sys.stdout.write", "sys.stderr.write"
+}
 NETWORK_PY_MODULES = {"requests", "aiohttp", "httpx", "urllib", "socket", "websocket"}
-SUBPROCESS_CALLS = {"subprocess.run", "subprocess.call", "subprocess.Popen", "subprocess.check_call",
-                    "subprocess.check_output", "os.system", "os.popen", "os.spawnl", "os.spawnle", "os.spawnlp",
-                    "os.spawnlpe", "os.spawnv", "os.spawnve", "os.spawnvp", "os.spawnvpe", "pty.spawn"}
+SUBPROCESS_CALLS = {
+    "subprocess.run", "subprocess.call", "subprocess.Popen", "subprocess.check_call", "subprocess.check_output",
+    "os.system", "os.popen", "os.spawnl", "os.spawnle", "os.spawnlp", "os.spawnlpe", "os.spawnv", "os.spawnve",
+    "os.spawnvp", "os.spawnvpe", "pty.spawn"
+}
 FILE_OPEN_CALLS = {"open", "Path.open", "pathlib.Path.open"}
-DANGEROUS_FILE_CALLS = {"os.remove", "os.unlink", "os.rmdir", "shutil.rmtree", "Path.unlink", "Path.rmdir",
-                        "pathlib.Path.unlink", "pathlib.Path.rmdir"}
+DANGEROUS_FILE_CALLS = {
+    "os.remove", "os.unlink", "os.rmdir", "shutil.rmtree", "Path.unlink", "Path.rmdir", "pathlib.Path.unlink",
+    "pathlib.Path.rmdir"
+}
 INSTALL_PY_MODULES = {"pip", "ensurepip"}
 SHELL_WRITE_COMMANDS = {"cp", "install", "mkdir", "mv", "tee", "touch"}
 PARALLEL_PY_CALLS = {
@@ -137,8 +141,8 @@ def scan(self, request: ToolSafetyScanRequest) -> ToolSafetyReport:
         redacted = any(item.redacted for item in findings)
         rule_ids = [item.rule_id for item in findings]
         primary_rule_id = primary_rule(findings)
-        blocked = decision == SafetyDecision.DENY or (
-            decision == SafetyDecision.NEEDS_HUMAN_REVIEW and self.policy.block_on_review)
+        blocked = decision == SafetyDecision.DENY or (decision == SafetyDecision.NEEDS_HUMAN_REVIEW
+                                                      and self.policy.block_on_review)
         report = ToolSafetyReport(
             decision=decision,
             risk_level=risk_level,
@@ -383,8 +387,7 @@ def _scan_bash_line(self, line: str, line_no: int, *, from_args: bool = False) -
                         line_no=line_no,
                     ))
             parallel_tasks = parallel_tasks_from_tokens(command_tokens)
-            if parallel_tasks is not None and (
-                    parallel_tasks == 0 or parallel_tasks > self.policy.max_parallel_tasks):
+            if parallel_tasks is not None and (parallel_tasks == 0 or parallel_tasks > self.policy.max_parallel_tasks):
                 findings.append(
                     finding(
                         "TSG-RESOURCE-PARALLELISM",
@@ -407,7 +410,8 @@ def _scan_bash_line(self, line: str, line_no: int, *, from_args: bool = False) -
                         line_no=line_no,
                     ))
             if command == "rm" and is_recursive_delete(command_tokens):
-                risk = SafetyRiskLevel.CRITICAL if targets_system_or_denied_path(command_tokens, self.policy) else SafetyRiskLevel.HIGH
+                risk = SafetyRiskLevel.CRITICAL if targets_system_or_denied_path(command_tokens,
+                                                                                 self.policy) else SafetyRiskLevel.HIGH
                 findings.append(
                     finding(
                         "TSG-FILE-RECURSIVE-DELETE",
@@ -461,7 +465,8 @@ def _scan_bash_line(self, line: str, line_no: int, *, from_args: bool = False) -
                             SafetyRiskLevel.HIGH,
                             "Shell redirection references a denied sensitive path.",
                             evidence,
-                            "Remove direct access to credential paths such as .env, ~/.ssh, and cloud credential files.",
+                            "Remove direct access to credential paths such as .env, ~/.ssh, "
+                            "and cloud credential files.",
                             line_no=line_no,
                         ))
             if command in {"cat", "grep", "awk", "sed", "tail", "head"} and any(
@@ -700,8 +705,7 @@ def handle_process_call(self, node: ast.Call, call_name: str) -> None:
             )
         if keyword_bool(node, "capture_output") or any(
                 self.call_name(keyword.value) == "subprocess.PIPE"
-                for keyword in node.keywords
-                if keyword.arg in {"stdout", "stderr"}):
+                for keyword in node.keywords if keyword.arg in {"stdout", "stderr"}):
             self.add(
                 "TSG-RESOURCE-OUTPUT-CAPTURE",
                 "resource_abuse",
@@ -736,7 +740,8 @@ def handle_open_call(self, node: ast.Call, call_name: str) -> None:
             if kw.arg == "mode":
                 mode = self.constant_string(kw.value) or mode
         if path_text and self.policy.is_denied_path(path_text):
-            risk_type = "sensitive_information_leak" if "r" in mode and not any(flag in mode for flag in "wa+") else "dangerous_file_operation"
+            risk_type = "sensitive_information_leak" if "r" in mode and not any(
+                flag in mode for flag in "wa+") else "dangerous_file_operation"
             self.add(
                 "TSG-FILE-DENIED-PATH",
                 risk_type,
@@ -1023,7 +1028,6 @@ def range_count(self, node: ast.Call) -> int | None:
             start, stop = numbers[:2]
             step = numbers[2] if len(numbers) >= 3 and numbers[2] != 0 else 1
             return max((stop - start + (step - 1)) // step, 0)
-        return None
 
 
 def normalize_language(language: str | None) -> str: