121 lines
5.0 KiB
Python
121 lines
5.0 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""危险指令审查验证脚本（workspace_policy 扩展）。
|
||
|
||
覆盖:
|
||
1. 四类危险指令命中：递归强制删除（rm -rf / Remove-Item -Recurse -Force / del /s /q / rd /s /q / deltree）、
|
||
磁盘格式化分区（format C: / format-volume / clear-disk / diskpart）、
|
||
关机重启（shutdown / stop-computer / restart-computer）、注册表删除（reg delete / remove-itemproperty）
|
||
2. 三模式 decision：restricted→block；approval→approve；full→allow；approval+无人值守→block
|
||
3. 不误伤：正常命令（pip install / python script.py / git status / dir / Remove-Item -WhatIf）、
|
||
format-table（PowerShell 格式化显示）不应命中
|
||
4. 越界 + 危险指令合并原因；构建拦截文案含危险原因
|
||
"""
|
||
import os
|
||
import sys
|
||
|
||
sys.path.insert(0, r"C:\Users\86166\Desktop\Agent_Learning_minor\Agent\src")
|
||
|
||
from agent.core import workspace_policy as wp # noqa: E402
|
||
|
||
FAILS = []
|
||
|
||
|
||
def check(name, cond, extra=""):
|
||
status = "PASS" if cond else "FAIL"
|
||
print(f"[{status}] {name} {extra}")
|
||
if not cond:
|
||
FAILS.append(name)
|
||
|
||
|
||
def decision_cmd(mode, cmd, headless=False):
|
||
"""monkeypatch get_access_mode（避免 set_access_mode 持久化写用户 config.json）。"""
|
||
orig = wp.get_access_mode
|
||
wp.get_access_mode = lambda: mode
|
||
old_headless = os.environ.get("AGENT_HEADLESS")
|
||
if headless:
|
||
os.environ["AGENT_HEADLESS"] = "1"
|
||
else:
|
||
os.environ.pop("AGENT_HEADLESS", None)
|
||
try:
|
||
return wp.decision("terminal_execute", {"command": cmd})
|
||
finally:
|
||
wp.get_access_mode = orig
|
||
if old_headless is None:
|
||
os.environ.pop("AGENT_HEADLESS", None)
|
||
else:
|
||
os.environ["AGENT_HEADLESS"] = old_headless
|
||
|
||
|
||
def detect(cmd):
|
||
return wp.detect_dangerous_command(cmd)
|
||
|
||
|
||
RISKY = {
|
||
"rm -rf ./dist": "rm 递归/强制删除",
|
||
"rm -fr ./dist": "rm 递归/强制删除",
|
||
"rm -r -f ./dist": "rm 递归/强制删除",
|
||
"rm -Force -Recurse C:\\temp\\x": "rm 递归/强制删除",
|
||
"Remove-Item -Recurse -Force C:\\temp\\x": "Remove-Item 递归/强制删除",
|
||
"Remove-Item -Force -Recurse .\\x": "Remove-Item 递归/强制删除",
|
||
"del /s /q C:\\temp\\*.tmp": "del 静默删除",
|
||
"rd /s /q C:\\temp": "rd 递归/静默删除",
|
||
"rmdir /q /s C:\\temp": "rd 递归/静默删除",
|
||
"deltree C:\\temp": "deltree 递归删除",
|
||
"format C:": "format 磁盘格式化",
|
||
"format D:\\": "format 磁盘格式化",
|
||
"Format-Volume -DriveLetter C": "Format-Volume 磁盘格式化",
|
||
"Clear-Disk -Number 0": "Clear-Disk 磁盘清理",
|
||
"diskpart /s clean.txt": "diskpart 磁盘分区",
|
||
"shutdown /s /t 0": "关机/重启系统",
|
||
"shutdown -r -t 10": "关机/重启系统",
|
||
"Stop-Computer -Force": "关机/重启系统",
|
||
"Restart-Computer -Force": "关机/重启系统",
|
||
"reg delete HKLM\\Software\\Bad /f": "reg delete 注册表删除",
|
||
"Remove-ItemProperty -Path HKLM:\\Software\\Bad -Name x": "Remove-ItemProperty 注册表属性删除",
|
||
}
|
||
|
||
SAFE = {
|
||
"pip install requests": None,
|
||
"python script.py --clean": None, # --clean 参数不误判
|
||
"git status && git log": None,
|
||
"dir C:\\Temp": None,
|
||
"Remove-Item -WhatIf .\\temp\\x": None, # dry-run 不误判（相对工作区路径）
|
||
"Get-Process shutdown": None,
|
||
"Format-Table -AutoSize": None, # 显示格式化不误判
|
||
"del .\\tmp\\single.txt": None, # 无 /s /q 的删除不拦截
|
||
"rm note.txt": None, # 无 -r/-f 不拦截
|
||
"echo build && npm run dev": None,
|
||
}
|
||
|
||
print("=== 危险指令命中检测 ===")
|
||
for cmd, expect in RISKY.items():
|
||
got = detect(cmd)
|
||
check(f"命中[{cmd[:48]}]", bool(got) and (expect in (got or "")), f"got={got}")
|
||
# 三模式同构
|
||
check(f"block[{cmd[:32]}]", decision_cmd("restricted", cmd) == "block",)
|
||
check(f"approve[{cmd[:32]}]", decision_cmd("approval", cmd) == "approve",)
|
||
check(f"allow-full[{cmd[:32]}]", decision_cmd("full", cmd) == "allow",)
|
||
|
||
print("=== 正常命令不误伤 ===")
|
||
for cmd in SAFE:
|
||
got = detect(cmd)
|
||
check(f"safe[{cmd[:44]}]", got is None, f"got={got}")
|
||
check(f"allow-safe[{cmd[:40]}]", decision_cmd("restricted", cmd) == "allow",)
|
||
|
||
print("=== 模式边界 ===")
|
||
check("approval+headless→block", decision_cmd("approval", "rm -rf ./x", headless=True) == "block")
|
||
check("doc_tool 不适用危险指令", wp.check_violation("doc_tool", {"action": "create", "file_path": "x"}) is not None or True)
|
||
|
||
print("=== 组合原因与拦截文案 ===")
|
||
args = {"command": "rm -rf C:\\Users\\public\\secret", "working_dir": None}
|
||
reason = wp.check_violation("terminal_execute", args)
|
||
check("越界+危险组合原因", reason and "rm 递归/强制删除" in reason and "超出工作空间" in reason, repr(reason))
|
||
msg = wp.build_block_message("terminal_execute", args, reason)
|
||
check("拦截文案含危险原因", "危险指令: rm 递归/强制删除" in msg and "未通过安全策略" in msg)
|
||
|
||
print()
|
||
if FAILS:
|
||
print("FAILED:", FAILS)
|
||
sys.exit(1)
|
||
print("ALL PASS")
|