Skip to content

Commit 1d77714

Browse files
zzp1221zzp1221
committed
test: cover tool safety patch lines
1 parent 26cab0a commit 1d77714

4 files changed

Lines changed: 120 additions & 34 deletions

File tree

‎tests/tools/safety/test_tool_safety_guard.py‎

Lines changed: 83 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -6,21 +6,25 @@
66

77
from __future__ import annotations
88

9+
import ast
910
import json
1011
import time
1112
from pathlib import Path
1213
from unittest.mock import AsyncMock
1314

1415
import pytest
1516

17+
from trpc_agent_sdk.code_executors import CodeBlock
18+
from trpc_agent_sdk.code_executors import CodeExecutionInput
19+
from trpc_agent_sdk.code_executors import BaseCodeExecutor
20+
from trpc_agent_sdk.code_executors import create_code_execution_result
1621
from trpc_agent_sdk.context import create_agent_context
17-
from trpc_agent_sdk.tools import FunctionTool
22+
from trpc_agent_sdk.filter import FilterResult
1823
from trpc_agent_sdk.tools.safety import SafetyDecision
1924
from trpc_agent_sdk.tools.safety import SafetyGuardedCodeExecutor
20-
from trpc_agent_sdk.tools.safety import ToolSafetyBlockedError
21-
from trpc_agent_sdk.tools.safety import ToolSafetyFinding
2225
from trpc_agent_sdk.tools.safety import SafetyRiskLevel
2326
from trpc_agent_sdk.tools.safety import ToolSafetyAuditLogger
27+
from trpc_agent_sdk.tools.safety import ToolSafetyBlockedError
2428
from trpc_agent_sdk.tools.safety import ToolSafetyFilter
2529
from trpc_agent_sdk.tools.safety import ToolSafetyGuard
2630
from trpc_agent_sdk.tools.safety import ToolSafetyPolicy
@@ -30,11 +34,8 @@
3034
from trpc_agent_sdk.tools.safety import apply_tool_safety_span_attributes
3135
from trpc_agent_sdk.tools.safety import extract_script_from_tool_args
3236
from trpc_agent_sdk.tools.safety import load_tool_safety_policy
33-
from trpc_agent_sdk.code_executors import CodeBlock
34-
from trpc_agent_sdk.code_executors import CodeExecutionInput
35-
from trpc_agent_sdk.code_executors import BaseCodeExecutor
36-
from trpc_agent_sdk.code_executors import create_code_execution_result
37-
from trpc_agent_sdk.filter import FilterResult
37+
from trpc_agent_sdk.tools.safety import _scanner as scanner_module
38+
from trpc_agent_sdk.tools.safety import _types as types_module
3839
from trpc_agent_sdk.tools.safety._guard import command_to_args
3940
from trpc_agent_sdk.tools.safety._guard import infer_language_from_key
4041
from trpc_agent_sdk.tools.safety._types import max_risk_level
@@ -310,6 +311,71 @@ def test_scanner_command_arg_prefixes_and_duration_helpers():
310311
]
311312

312313

314+
def test_scanner_edge_branches_from_direct_scans(monkeypatch):
315+
scanner = ToolSafetyScanner(policy(max_parallel_tasks=2, max_loop_iterations=3))
316+
317+
assert scanner.scan(ToolSafetyScanRequest(script="", language="python", env=[])).decision == SafetyDecision.ALLOW
318+
assert scanner._scan_bash_line("", 1) == [] # pylint: disable=protected-access
319+
320+
monkeypatch.setattr(scanner_module, "split_shell_commands", lambda tokens: [[]])
321+
assert scanner._scan_bash_line("echo ok", 1) == [] # pylint: disable=protected-access
322+
323+
scripts = [
324+
("funcs = [print]\nfuncs[0]('x')", None),
325+
("import os\nos.remove('.env')", "TSG-FILE-DENIED-PATH"),
326+
("open('/etc/passwd', mode='w')", "TSG-FILE-SYSTEM-WRITE"),
327+
("f = open('out.txt')", None),
328+
("import concurrent.futures\nconcurrent.futures.ThreadPoolExecutor()", "TSG-RESOURCE-PARALLELISM"),
329+
("for i in range(limit):\n pass", None),
330+
("for i in range(2, 10, 2):\n pass", "TSG-RESOURCE-LARGE-LOOP"),
331+
]
332+
for script, rule_id in scripts:
333+
report = scanner.scan(ToolSafetyScanRequest(script=script, language="python"))
334+
if rule_id:
335+
assert any(finding.rule_id == rule_id for finding in report.findings)
336+
337+
338+
def test_python_analyzer_helper_edge_branches():
339+
analyzer = scanner_module._PythonAnalyzer("pass\n", policy()) # pylint: disable=protected-access
340+
341+
assert analyzer.network_target_from_arg(ast.parse("url", mode="eval").body) == ""
342+
assert analyzer.node_line(ast.Pass(lineno=1, col_offset=0)) == "pass"
343+
assert analyzer.node_line(ast.Pass(lineno=99, col_offset=0)) == ""
344+
assert analyzer.call_name(ast.Attribute(value=ast.Constant(value=1), attr="field", ctx=ast.Load())) == "field"
345+
assert analyzer.call_name(ast.Constant(value=None)) == ""
346+
assert analyzer.name_of(ast.Attribute(value=ast.Name(id="obj", ctx=ast.Load()), attr="token", ctx=ast.Load())) == "token"
347+
assert analyzer.name_of(ast.Constant(value=None)) == ""
348+
assert analyzer.constant_string(None) == ""
349+
assert analyzer.constant_string(ast.parse("f'key={value}'", mode="eval").body) == "key={}"
350+
assert analyzer.constant_string(ast.parse("['a', value]", mode="eval").body) == ""
351+
assert analyzer.constant_string(ast.parse("('a', 'b')", mode="eval").body) == "a b"
352+
assert analyzer.path_literal_from_node(None) == ""
353+
assert analyzer.path_literal_from_node(ast.parse("Path('.env')", mode="eval").body) == ".env"
354+
assert analyzer.path_literal_from_node(ast.parse("Path(value)", mode="eval").body) == ""
355+
assert analyzer.path_literal_from_node(ast.parse("Path('.env').expanduser().read_text", mode="eval").body) == ".env"
356+
assert analyzer.path_literal_from_node(ast.parse("config.read_text", mode="eval").body) == ""
357+
assert analyzer.range_count(ast.parse("range(limit)", mode="eval").body) is None
358+
assert analyzer.range_count(ast.parse("range(2, 10, 2)", mode="eval").body) == 4
359+
360+
361+
def test_scanner_module_helper_edge_branches():
362+
assert scanner_module.is_secret_name("") is False
363+
assert scanner_module.numeric_constant(ast.parse("value", mode="eval").body) is None
364+
assert scanner_module.is_install_invocation(["python", "-m", "pip", "install", "demo"]) is False
365+
assert scanner_module.is_recursive_delete(["echo", "-rf"]) is False
366+
assert scanner_module.redirection_targets(["echo", "x", ">/etc/app.conf"]) == ["/etc/app.conf"]
367+
assert scanner_module.timeout_seconds_from_tokens(["timeout", "--kill-after=1s", "10m"]) == 600
368+
assert scanner_module.timeout_seconds_from_tokens(["timeout", "--preserve-status", "10m"]) == 600
369+
assert scanner_module.timeout_seconds_from_tokens(["timeout", "--preserve-status"]) is None
370+
assert scanner_module.duration_token_seconds("never") is None
371+
assert scanner_module.parallel_tasks_from_tokens(["parallel", "-j", "4"]) == 4
372+
assert scanner_module.parallel_tasks_from_tokens(["parallel", "--jobs=7"]) == 7
373+
assert scanner_module.parallel_tasks_from_tokens(["parallel", "echo"]) == 0
374+
assert scanner_module.parallel_tasks_from_tokens(["xargs", "-P3"]) == 3
375+
assert scanner_module.int_token("many") is None
376+
assert scanner_module.line_has_allowlisted_url("echo no-url", policy()) is False
377+
378+
313379
def test_policy_loader_and_matching_branches(tmp_path):
314380
policy_file = tmp_path / "policy.yaml"
315381
policy_file.write_text(
@@ -340,8 +406,11 @@ def test_policy_loader_and_matching_branches(tmp_path):
340406
assert loaded.is_command_allowed("custom-tool --flag") is True
341407
assert loaded.is_command_allowed(None) is False
342408
assert loaded.is_denied_path(None) is False
409+
assert loaded.is_denied_path("config.yaml") is False
343410
assert loaded.is_system_write_path("/") is True
344411
assert loaded.is_system_write_path(None) is False
412+
assert ToolSafetyPolicy(allowed_domains=[""]).is_domain_allowed("api.example.com") is False
413+
assert ToolSafetyPolicy(system_write_paths=[""]).is_system_write_path("/etc/passwd") is False
345414
assert load_tool_safety_policy(None).name == "default"
346415

347416

@@ -370,6 +439,12 @@ def test_report_and_type_helpers_cover_empty_and_fallback_paths():
370439
assert json.loads(report.to_json(indent=None))["risk_count"] == 0
371440

372441

442+
def test_max_risk_level_fallback_when_order_has_no_match(monkeypatch):
443+
monkeypatch.setattr(types_module, "RISK_LEVEL_ORDER", {})
444+
445+
assert types_module.max_risk_level(["unexpected"]) == SafetyRiskLevel.NONE
446+
447+
373448
def test_audit_event_contains_required_fields(tmp_path):
374449
report = ToolSafetyScanner(policy()).scan(
375450
ToolSafetyScanRequest(script="rm -rf /tmp/demo", language="bash", tool_metadata={"name": "Bash"}))

‎trpc_agent_sdk/tools/safety/_guard.py‎

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -121,23 +121,31 @@ async def execute_code(
121121
script=block.code,
122122
language=block.language,
123123
cwd=self._work_dir(),
124-
tool_metadata={"name": "CodeExecutor", "executor": self.delegate.__class__.__name__},
124+
tool_metadata={
125+
"name": "CodeExecutor",
126+
"executor": self.delegate.__class__.__name__,
127+
},
125128
)
126129
try:
127130
self.guard.check(request)
128131
except ToolSafetyBlockedError as exc:
129-
return create_code_execution_result(stderr=f"Tool safety guard blocked code execution: {exc.report.to_json(indent=None)}")
132+
return create_code_execution_result(
133+
stderr=f"Tool safety guard blocked code execution: {exc.report.to_json(indent=None)}")
130134
if code_execution_input.code and not code_execution_input.code_blocks:
131135
request = ToolSafetyScanRequest(
132136
script=code_execution_input.code,
133137
language="python",
134138
cwd=self._work_dir(),
135-
tool_metadata={"name": "CodeExecutor", "executor": self.delegate.__class__.__name__},
139+
tool_metadata={
140+
"name": "CodeExecutor",
141+
"executor": self.delegate.__class__.__name__,
142+
},
136143
)
137144
try:
138145
self.guard.check(request)
139146
except ToolSafetyBlockedError as exc:
140-
return create_code_execution_result(stderr=f"Tool safety guard blocked code execution: {exc.report.to_json(indent=None)}")
147+
return create_code_execution_result(
148+
stderr=f"Tool safety guard blocked code execution: {exc.report.to_json(indent=None)}")
141149
return await self.delegate.execute_code(invocation_context, code_execution_input)
142150

143151
def _work_dir(self) -> str | None:

‎trpc_agent_sdk/tools/safety/_policy.py‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@
1414

1515
import yaml
1616

17-
1817
DEFAULT_ALLOWED_COMMANDS = [
1918
"awk",
2019
"cat",

‎trpc_agent_sdk/tools/safety/_scanner.py‎

Lines changed: 25 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -27,12 +27,10 @@
2727
from ._types import max_risk_level
2828
from ._types import risk_level_value
2929

30-
3130
SECRET_VALUE_RE = re.compile(
3231
r"(?i)(api[_-]?key|access[_-]?token|secret|password|private[_-]?key|authorization)\s*[:=]\s*"
33-
r"['\"]?([A-Za-z0-9_./+=:-]{12,})"
34-
)
35-
URL_RE = re.compile(r"(?i)\bhttps?://[^\s'\"<>)]+" )
32+
r"['\"]?([A-Za-z0-9_./+=:-]{12,})")
33+
URL_RE = re.compile(r"(?i)\bhttps?://[^\s'\"<>)]+")
3634
ENV_VAR_RE = re.compile(r"\$(?:\{)?([A-Za-z_][A-Za-z0-9_]*)")
3735
SHELL_CONTROL_RE = re.compile(r"(\|\||&&|;|\||`|\$\(|>|<)")
3836
FORK_BOMB_RE = re.compile(r":\s*\(\s*\)\s*\{\s*:\s*\|\s*:\s*&\s*\}\s*;\s*:")
@@ -79,15 +77,21 @@
7977
"SECRET",
8078
"TOKEN",
8179
}
82-
SENSITIVE_SINK_NAMES = {"print", "logging.info", "logging.warning", "logging.error", "logger.info", "logger.warning",
83-
"logger.error", "sys.stdout.write", "sys.stderr.write"}
80+
SENSITIVE_SINK_NAMES = {
81+
"print", "logging.info", "logging.warning", "logging.error", "logger.info", "logger.warning", "logger.error",
82+
"sys.stdout.write", "sys.stderr.write"
83+
}
8484
NETWORK_PY_MODULES = {"requests", "aiohttp", "httpx", "urllib", "socket", "websocket"}
85-
SUBPROCESS_CALLS = {"subprocess.run", "subprocess.call", "subprocess.Popen", "subprocess.check_call",
86-
"subprocess.check_output", "os.system", "os.popen", "os.spawnl", "os.spawnle", "os.spawnlp",
87-
"os.spawnlpe", "os.spawnv", "os.spawnve", "os.spawnvp", "os.spawnvpe", "pty.spawn"}
85+
SUBPROCESS_CALLS = {
86+
"subprocess.run", "subprocess.call", "subprocess.Popen", "subprocess.check_call", "subprocess.check_output",
87+
"os.system", "os.popen", "os.spawnl", "os.spawnle", "os.spawnlp", "os.spawnlpe", "os.spawnv", "os.spawnve",
88+
"os.spawnvp", "os.spawnvpe", "pty.spawn"
89+
}
8890
FILE_OPEN_CALLS = {"open", "Path.open", "pathlib.Path.open"}
89-
DANGEROUS_FILE_CALLS = {"os.remove", "os.unlink", "os.rmdir", "shutil.rmtree", "Path.unlink", "Path.rmdir",
90-
"pathlib.Path.unlink", "pathlib.Path.rmdir"}
91+
DANGEROUS_FILE_CALLS = {
92+
"os.remove", "os.unlink", "os.rmdir", "shutil.rmtree", "Path.unlink", "Path.rmdir", "pathlib.Path.unlink",
93+
"pathlib.Path.rmdir"
94+
}
9195
INSTALL_PY_MODULES = {"pip", "ensurepip"}
9296
SHELL_WRITE_COMMANDS = {"cp", "install", "mkdir", "mv", "tee", "touch"}
9397
PARALLEL_PY_CALLS = {
@@ -137,8 +141,8 @@ def scan(self, request: ToolSafetyScanRequest) -> ToolSafetyReport:
137141
redacted = any(item.redacted for item in findings)
138142
rule_ids = [item.rule_id for item in findings]
139143
primary_rule_id = primary_rule(findings)
140-
blocked = decision == SafetyDecision.DENY or (
141-
decision == SafetyDecision.NEEDS_HUMAN_REVIEW and self.policy.block_on_review)
144+
blocked = decision == SafetyDecision.DENY or (decision == SafetyDecision.NEEDS_HUMAN_REVIEW
145+
and self.policy.block_on_review)
142146
report = ToolSafetyReport(
143147
decision=decision,
144148
risk_level=risk_level,
@@ -383,8 +387,7 @@ def _scan_bash_line(self, line: str, line_no: int, *, from_args: bool = False) -
383387
line_no=line_no,
384388
))
385389
parallel_tasks = parallel_tasks_from_tokens(command_tokens)
386-
if parallel_tasks is not None and (
387-
parallel_tasks == 0 or parallel_tasks > self.policy.max_parallel_tasks):
390+
if parallel_tasks is not None and (parallel_tasks == 0 or parallel_tasks > self.policy.max_parallel_tasks):
388391
findings.append(
389392
finding(
390393
"TSG-RESOURCE-PARALLELISM",
@@ -407,7 +410,8 @@ def _scan_bash_line(self, line: str, line_no: int, *, from_args: bool = False) -
407410
line_no=line_no,
408411
))
409412
if command == "rm" and is_recursive_delete(command_tokens):
410-
risk = SafetyRiskLevel.CRITICAL if targets_system_or_denied_path(command_tokens, self.policy) else SafetyRiskLevel.HIGH
413+
risk = SafetyRiskLevel.CRITICAL if targets_system_or_denied_path(command_tokens,
414+
self.policy) else SafetyRiskLevel.HIGH
411415
findings.append(
412416
finding(
413417
"TSG-FILE-RECURSIVE-DELETE",
@@ -461,7 +465,8 @@ def _scan_bash_line(self, line: str, line_no: int, *, from_args: bool = False) -
461465
SafetyRiskLevel.HIGH,
462466
"Shell redirection references a denied sensitive path.",
463467
evidence,
464-
"Remove direct access to credential paths such as .env, ~/.ssh, and cloud credential files.",
468+
"Remove direct access to credential paths such as .env, ~/.ssh, "
469+
"and cloud credential files.",
465470
line_no=line_no,
466471
))
467472
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:
700705
)
701706
if keyword_bool(node, "capture_output") or any(
702707
self.call_name(keyword.value) == "subprocess.PIPE"
703-
for keyword in node.keywords
704-
if keyword.arg in {"stdout", "stderr"}):
708+
for keyword in node.keywords if keyword.arg in {"stdout", "stderr"}):
705709
self.add(
706710
"TSG-RESOURCE-OUTPUT-CAPTURE",
707711
"resource_abuse",
@@ -736,7 +740,8 @@ def handle_open_call(self, node: ast.Call, call_name: str) -> None:
736740
if kw.arg == "mode":
737741
mode = self.constant_string(kw.value) or mode
738742
if path_text and self.policy.is_denied_path(path_text):
739-
risk_type = "sensitive_information_leak" if "r" in mode and not any(flag in mode for flag in "wa+") else "dangerous_file_operation"
743+
risk_type = "sensitive_information_leak" if "r" in mode and not any(
744+
flag in mode for flag in "wa+") else "dangerous_file_operation"
740745
self.add(
741746
"TSG-FILE-DENIED-PATH",
742747
risk_type,
@@ -1023,7 +1028,6 @@ def range_count(self, node: ast.Call) -> int | None:
10231028
start, stop = numbers[:2]
10241029
step = numbers[2] if len(numbers) >= 3 and numbers[2] != 0 else 1
10251030
return max((stop - start + (step - 1)) // step, 0)
1026-
return None
10271031

10281032

10291033
def normalize_language(language: str | None) -> str:

0 commit comments

Comments
 (0)