# -*- 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")