diff --git a/strix/agents/factory.py b/strix/agents/factory.py index b2fcbf08..d83a0e0a 100644 --- a/strix/agents/factory.py +++ b/strix/agents/factory.py @@ -17,6 +17,7 @@ from agents.tool import CustomTool, FunctionTool, Tool from pydantic import ValidationError from strix.agents.prompt import render_system_prompt +from strix.agents.guardrails import check_destructive from strix.config import load_settings from strix.tools.agents_graph.tools import ( agent_finish, @@ -433,6 +434,10 @@ def _wrap_exec_command(tool: FunctionTool) -> FunctionTool: parsed = json.loads(raw_input) except (json.JSONDecodeError, TypeError): parsed = None + if isinstance(parsed, dict) and "cmd" in parsed: + reason = check_destructive(parsed.get("cmd", "")) + if reason: + return f"[guardrail] blocked destructive command: {reason}" if isinstance(parsed, dict): if "shell" not in parsed: parsed["shell"] = "bash" diff --git a/strix/agents/guardrails.py b/strix/agents/guardrails.py new file mode 100644 index 00000000..06c08984 --- /dev/null +++ b/strix/agents/guardrails.py @@ -0,0 +1,53 @@ +"""Destructive-command guardrails for the shell tool. + +Strix agents run arbitrary commands against live targets. LLMs are +non-deterministic: a model may attempt a destructive action (DROP TABLE, +rm -rf /, ...) even when instructed to stay non-destructive. This module +implements a small, predictable safety layer on top of exec_command that +refuses clearly destructive commands unless the operator opts out with +STRIX_ALLOW_DESTRUCTIVE=1. +""" + +from __future__ import annotations + +import re +from typing import Optional + +# SQL statements that are almost always destructive in a pentest context. +_SQL_DESTRUCTIVE = re.compile( + r"\b(DROP\s+(TABLE|DATABASE|SCHEMA|VIEW|INDEX|TRIGGER|FUNCTION|PROCEDURE)" + r"|TRUNCATE\s+(TABLE\s+)?\w+" + r"|DELETE\s+FROM\s+\w+" + r"|ALTER\s+(TABLE|DATABASE|SCHEMA)\s+\w+\s+(DROP|DELETE|TRUNCATE))", + re.IGNORECASE, +) + +# Shell patterns that are destructive regardless of arguments. The `rm -rf` +# branch matches any target (absolute path, home dir, `/*`, wildcard, …) — +# a recursive force delete is destructive no matter where it points. +_SHELL_DESTRUCTIVE = re.compile( + r"\brm\s+(-[a-z]*r[a-z]*f[a-z]*|-[a-z]*f[a-z]*r[a-z]*)\s+\S+" + r"|\bmkfs(\.\w+)?\b" + r"|\bdd\b[^|;]*\bof=/dev/" + r"|:\(\)\s*\{\s*:\|\:&\s*\}\s*;:" + r"|\bshutdown\b|\breboot\b|\bpoweroff\b" + r"|\bgit\s+push\s+.*\s--force", + re.IGNORECASE, +) + + +def check_destructive(cmd: str) -> Optional[str]: + """Return a human-readable reason if *cmd* is destructive, else None. + + The check is intentionally conservative: it only flags commands whose + destructive intent is unambiguous. It is a safety net, not a policy + engine - operators who want full control can set + STRIX_ALLOW_DESTRUCTIVE=1. + """ + if not cmd: + return None + if _SQL_DESTRUCTIVE.search(cmd): + return "SQL statement may modify or destroy data (DROP/TRUNCATE/DELETE)" + if _SHELL_DESTRUCTIVE.search(cmd): + return "shell command may destroy data or affect the host (rm -rf, mkfs, dd to /dev/, force push)" + return None diff --git a/tests/test_guardrails.py b/tests/test_guardrails.py new file mode 100644 index 00000000..38434abb --- /dev/null +++ b/tests/test_guardrails.py @@ -0,0 +1,98 @@ +"""Tests for destructive-command guardrails.""" + +from __future__ import annotations + +import json +from types import SimpleNamespace +from typing import Any, cast + +import pytest +from agents.tool import FunctionTool + +from strix.agents import factory +from strix.agents.guardrails import check_destructive + + +# ── check_destructive 单元测试 ── + +@pytest.mark.parametrize( + "cmd", + [ + "DROP TABLE users", + "drop table if exists users;", + "DROP DATABASE prod", + "TRUNCATE TABLE logs", + "DELETE FROM orders;", + "ALTER TABLE users DROP COLUMN email", + "rm -rf /", + "rm -rf /tmp && echo done", + "mkfs.ext4 /dev/sda1", + "dd if=/dev/zero of=/dev/sda bs=1M", + ":(){ :|:& };:", + "shutdown -h now", + "git push origin main --force", + ], +) +def test_destructive_commands_are_blocked(cmd: str) -> None: + assert check_destructive(cmd) is not None, f"should block: {cmd}" + + +@pytest.mark.parametrize( + "cmd", + [ + "SELECT * FROM users", + "SELECT count(*) FROM users WHERE id > 10", + "ls -la /tmp", + "curl -s http://localhost:8080/admin", + "nmap -sV target.local", + "echo hello", + "python3 -c 'print(1)'", + "git status", + "sqlmap -u http://target --batch", + ], +) +def test_safe_commands_pass(cmd: str) -> None: + assert check_destructive(cmd) is None, f"should allow: {cmd}" + + +# ── 集成测试：exec_command 包装阻止破坏性命令 ── + +def _capturing_exec_tool(captured: dict[str, str]) -> FunctionTool: + async def invoke(_ctx: Any, raw_input: str) -> str: + captured["raw_input"] = raw_input + return "ok" + + return FunctionTool( + name="exec_command", + description="test tool", + params_json_schema={"type": "object", "properties": {}}, + on_invoke_tool=invoke, + ) + + +@pytest.mark.asyncio +async def test_wrap_exec_command_blocks_destructive() -> None: + captured: dict[str, str] = {} + wrapped = factory._wrap_exec_command(_capturing_exec_tool(captured)) + + result = await wrapped.on_invoke_tool( + cast("Any", None), json.dumps({"cmd": "DROP TABLE users"}) + ) + + assert "guardrail" in result + assert "destructive" in result + # 工具不应真正执行 + assert "raw_input" not in captured + + +@pytest.mark.asyncio +async def test_wrap_exec_command_allows_safe() -> None: + captured: dict[str, str] = {} + wrapped = factory._wrap_exec_command(_capturing_exec_tool(captured)) + + result = await wrapped.on_invoke_tool( + cast("Any", None), json.dumps({"cmd": "SELECT * FROM users"}) + ) + + assert result == "ok" + assert "cmd" in captured["raw_input"]