#!/usr/bin/env python3
"""实战：自动化扫描框架 CLI（Day 161 — 渗透测试框架）。

这是本课的"成品"：把范围门禁、只读探测、插件检测、阶段编排、
审计留痕、报告输出 全部串起来的一个命令行工具。

子命令：

    python3 04-pentest-framework.py plugins
        列出已注册插件（含 priority 与所属层）

    python3 04-pentest-framework.py authorize --hosts 127.0.0.1 --ports 8000-8002
        只做范围校验，不发起任何连接（签发前的"预演"）

    python3 04-pentest-framework.py lab
        启动本机实验服务（8000 伪装 SSH / 8080 伪装 HTTP），
        在回环范围内跑完整流水线，输出 JSON + Markdown 报告

    python3 04-pentest-framework.py scan --hosts <IP/CIDR> --ports <表达式> [选项]
        对**已授权**目标执行阶段 1–2；越界立即以退出码 2 终止

    python3 04-pentest-framework.py --self-test
        自检（全程回环，无外部网络）

退出码：2=越界/参数错误   4=结果不完整(有 filtered/error)   3=有 P0/P1   0=干净

⚠️ 本脚本**不包含任何利用能力**：
   - 探测只有 TCP connect + 读 banner（不发送任何字节）；
   - `exploit` 阶段是空实现，只输出"待人工验证"清单；
   - 不做爆破、不提交表单、不改动目标任何数据。
"""

from __future__ import annotations

import argparse
import json
import re
import sys
from pathlib import Path

from pentest_core import (
    LAB_SERVICES,
    AuthorizationError,
    Finding,
    Framework,
    LabServer,
    Phase,
    PhaseResult,
    PluginContext,
    Scope,
    ScopeError,
    build_report,
    default_lab_scope,
    render_markdown,
    tcp_probe,
    Target,
)

DAY_DIR = Path(__file__).resolve().parent.parent
DEFAULT_OUT = DAY_DIR / "out"


# ── 插件（事实层 → 推断层 → 结论层） ──────────────────────────────────

def sanitize(text: str, limit: int = 60) -> str:
    """banner 是目标控制的不可信文本：净化控制字符并截断。"""
    cleaned = re.sub(r"[\x00-\x1f\x7f]", " ", text)
    cleaned = re.sub(r"\s+", " ", cleaned).strip()
    return (cleaned[:limit] + "…") if len(cleaned) > limit else cleaned


def register_plugins(fw: Framework) -> None:
    @fw.plugin("banner_evidence", priority=20)
    def banner_evidence(ctx: PluginContext) -> list[Finding]:
        asset = ctx.asset
        if asset.state != "open" or not asset.banner:
            return []
        return [Finding(
            key="banner_evidence", title="服务问候语（事实记录）",
            severity="info", confidence="high", target=asset.target.key,
            evidence=f"service={asset.service}",
            detail={"banner": sanitize(asset.banner)},
        )]

    @fw.plugin("version_disclosure", priority=30)
    def version_disclosure(ctx: PluginContext) -> list[Finding]:
        asset = ctx.asset
        if asset.state != "open":
            return []
        tokens = [t for t in re.split(r"[\s/]+", asset.banner)
                  if any(ch.isdigit() for ch in t) and len(t) <= 24]
        if not tokens:
            return []
        return [Finding(
            key="version_disclosure", title=f"版本信息泄露: {' '.join(tokens[:2])}",
            severity="low", confidence="high", target=asset.target.key,
            evidence=f"banner={sanitize(asset.banner, 40)!r}",
        )]

    @fw.plugin("cleartext_protocol", priority=40)
    def cleartext_protocol(ctx: PluginContext) -> list[Finding]:
        asset = ctx.asset
        if asset.state == "open" and asset.service in {"telnet", "ftp", "http", "pop3", "imap"}:
            return [Finding(
                key="cleartext_protocol", title=f"明文协议暴露: {asset.service}",
                severity="high", confidence="medium", target=asset.target.key,
                evidence=f"port={asset.target.port} service={asset.service}",
            )]
        return []

    @fw.plugin("high_port_service", priority=50)
    def high_port_service(ctx: PluginContext) -> list[Finding]:
        """高位端口上的服务常常是临时/调试服务，缺少常规运维流程的保护。"""
        asset = ctx.asset
        if asset.state == "open" and asset.target.port > 1024:
            return [Finding(
                key="high_port_service", title=f"高位端口服务: {asset.target.port}",
                severity="info", confidence="medium", target=asset.target.key,
                evidence=f"service={asset.service}",
            )]
        return []


# ── 阶段定义 ───────────────────────────────────────────────────────────

def phase_collect(fw: Framework, assets) -> PhaseResult:
    """阶段 1 由 probe_many 完成，这里只做汇总（保持阶段签名一致）。"""
    stats = Framework.state_counts(assets)
    if stats["open"] == 0 and (stats["filtered"] or stats["error"]):
        return PhaseResult("collect", "failed",
                           f"全部不可判定: {stats}")
    return PhaseResult("collect", "ok")


def phase_detect(fw: Framework, assets) -> PhaseResult:
    findings = fw.run_plugins(assets)
    return PhaseResult("detect", "ok", findings=findings)


def phase_exploit_placeholder(fw: Framework, assets) -> PhaseResult:
    """★ 空实现 ★ 不产生任何网络行为，只登记需要人工验证的点位。"""
    pending = [Finding(key="needs_manual_verification",
                       title="待人工在授权窗口内验证",
                       severity="info", confidence="info",
                       target=a.target.key,
                       evidence=f"service={a.service}")
               for a in assets if a.state == "open"]
    return PhaseResult("exploit", "skipped",
                       "本课不提供利用能力；输出人工验证清单，无网络行为",
                       findings=pending)


def phase_report(fw: Framework, assets) -> PhaseResult:
    return PhaseResult("report", "ok")


PHASES = [
    Phase("detect", phase_detect, critical=True),
    Phase("exploit", phase_exploit_placeholder, optional=True),
    Phase("report", phase_report),
]


# ── 执行一轮 ───────────────────────────────────────────────────────────

def run_once(
    hosts: list[str],
    ports: list[str],
    scope: Scope,
    timeout: float,
    workers: int,
    rate: float,
    out_dir: Path,
    tag: str,
    quiet: bool = False,
) -> int:
    out_dir.mkdir(parents=True, exist_ok=True)
    audit_path = out_dir / f"{tag}-audit.jsonl"
    fw = Framework(scope, timeout=timeout, workers=workers, rate=rate, audit_path=audit_path)
    register_plugins(fw)

    if not quiet:
        print(f"[scope] networks={scope.networks} ports={scope.ports} ticket={scope.ticket}")

    try:
        targets = fw.build_targets(hosts, ports)
        fw.authorize(targets)                     # ★ 门禁：越界即中止
    except AuthorizationError as exc:
        print(f"⛔ AuthorizationError: {exc}")
        print("   （已记入审计日志：scope_check denied + run_abort exit=2）")
        return 2
    except ScopeError as exc:
        print(f"⛔ ScopeError: {exc}")
        return 2

    results = fw.probe_many(targets)
    stats = Framework.state_counts(results)
    if not quiet:
        print(f"[collect] 目标 {len(results)} 个 → open {stats['open']} / closed {stats['closed']}"
              f" / filtered {stats['filtered']} / error {stats['error']}"
              f"  (throttle_wait={fw.bucket.waited_s:.3f}s)")
    if stats["open"] == 0 and (stats["filtered"] or stats["error"]):
        # 让编排器把 detect 标成 skipped，而不是"0 个问题"
        phases = [Phase("detect", phase_detect, critical=True),
                  Phase("exploit", phase_exploit_placeholder, optional=True)]
        phase_results, _, findings = fw.run_phases(phases, [])
    else:
        phase_results, _, findings = fw.run_phases(PHASES, results)

    if not quiet:
        print(f"[detect]  插件 {len(fw.registry.names())} 个 → findings {len(findings)}")
        for p in phase_results:
            if p.status != "ok":
                print(f"[{p.name}] {p.status}：{p.reason}")

    params = {"timeout_s": timeout, "workers": workers, "rate_per_s": rate,
              "plugins": fw.registry.names(), "hosts_spec": hosts, "ports_spec": ports}
    report = build_report(scope, results, phase_results, findings, params, fw.audit)
    report["mode"] = tag
    json_path = out_dir / f"{tag}-report.json"
    md_path = out_dir / f"{tag}-report.md"
    json_path.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
    md_path.write_text(render_markdown(report), encoding="utf-8")
    if not quiet:
        print(f"[report]  已写入 {json_path.relative_to(DAY_DIR.parent.parent)} / "
              f"{md_path.relative_to(DAY_DIR.parent.parent)}")
        print(f"[exit]    {report['exit_code']}")
    return report["exit_code"]


# ── 子命令 ─────────────────────────────────────────────────────────────

def cmd_plugins() -> int:
    fw = Framework(default_lab_scope())
    register_plugins(fw)
    print(f"已注册插件 {len(fw.registry.names())} 个（按 priority 升序执行）：")
    for name, prio, fn in fw.registry.ordered():
        layer = "事实层" if prio < 30 else "推断层" if prio < 40 else "结论层"
        doc = (fn.__doc__ or "").strip().splitlines()[0] if fn.__doc__ else ""
        print(f"  [{prio:>3}] {name:<20} {layer:<6} {doc}")
    print("\n分层规则：事实层(10-20) → 推断层(30) → 结论层(40+)；同 priority 按注册顺序。")
    return 0


def cmd_authorize(hosts: list[str], ports: list[str], scope_path: str | None) -> int:
    scope = Scope.load(scope_path) if scope_path else default_lab_scope()
    fw = Framework(scope)
    try:
        targets = fw.build_targets(hosts, ports)
        approved = fw.authorize(targets)
    except AuthorizationError as exc:
        print(f"⛔ 拒绝：{exc}")
        return 2
    except ScopeError as exc:
        print(f"⛔ 范围/参数错误：{exc}")
        return 2
    print(f"✅ {len(approved)} 个目标通过范围校验（**未发起任何连接**）：")
    for t in approved[:20]:
        print(f"   - {t.key}")
    if len(approved) > 20:
        print(f"   … 其余 {len(approved) - 20} 个略")
    print("\n提示：authorize 是『预演』，用它在正式执行前确认范围理解与授权一致。")
    return 0


def cmd_lab(args) -> int:
    lab_ports = [8000, 8080]
    print(f"[lab]   启动本机实验服务: {lab_ports}（仅回环）")
    with LabServer({p: LAB_SERVICES[p] for p in lab_ports}) as _lab:
        return run_once(
            hosts=["127.0.0.1"],
            ports=[args.ports],
            scope=default_lab_scope(),
            timeout=args.timeout,
            workers=args.workers,
            rate=args.rate,
            out_dir=Path(args.out),
            tag="day161-lab",
        )


def cmd_scan(args) -> int:
    scope = Scope.load(args.scope) if args.scope else default_lab_scope()
    if args.scope is None:
        print("⚠️ 未指定 --scope，使用默认实验范围（仅回环 8000-8099）。")
        print("   真实项目请用 --scope 指向记录了授权单号的范围文件。")
    return run_once(
        hosts=args.hosts,
        ports=args.ports,
        scope=scope,
        timeout=args.timeout,
        workers=args.workers,
        rate=args.rate,
        out_dir=Path(args.out),
        tag=args.tag,
    )


def build_parser() -> argparse.ArgumentParser:
    ap = argparse.ArgumentParser(
        description="Day 161 自动化扫描框架（仅限授权目标）",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    ap.add_argument("--self-test", action="store_true")
    ap.add_argument("--out", default=str(DEFAULT_OUT), help="报告输出目录")
    sub = ap.add_subparsers(dest="cmd")

    sub.add_parser("plugins", help="列出已注册插件")

    p_auth = sub.add_parser("authorize", help="只做范围校验，不发起连接")
    p_auth.add_argument("--hosts", nargs="+", required=True)
    p_auth.add_argument("--ports", nargs="+", required=True)
    p_auth.add_argument("--scope", default=None)

    p_lab = sub.add_parser("lab", help="在本机实验服务上跑完整流水线")
    p_lab.add_argument("--ports", default="8000-8080")
    p_lab.add_argument("--timeout", type=float, default=0.5)
    p_lab.add_argument("--workers", type=int, default=8)
    p_lab.add_argument("--rate", type=float, default=50.0)

    p_scan = sub.add_parser("scan", help="对已授权目标执行阶段 1-2")
    p_scan.add_argument("--hosts", nargs="+", required=True, help="只接受 IP / CIDR")
    p_scan.add_argument("--ports", nargs="+", required=True, help="如 8000-8010 或 80,443")
    p_scan.add_argument("--scope", default=None, help="范围 JSON 文件路径")
    p_scan.add_argument("--timeout", type=float, default=0.5)
    p_scan.add_argument("--workers", type=int, default=8)
    p_scan.add_argument("--rate", type=float, default=50.0)
    p_scan.add_argument("--tag", default="day161-scan")
    return ap


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)

    if args.self_test:
        return _self_test()

    if args.cmd == "plugins":
        return cmd_plugins()
    if args.cmd == "authorize":
        return cmd_authorize(args.hosts, args.ports, args.scope)
    if args.cmd == "lab":
        return cmd_lab(args)
    if args.cmd == "scan":
        return cmd_scan(args)
    parser.print_help()
    return 0


def _self_test() -> int:
    import contextlib
    import io

    # ① 完整流水线（实验服务由 cmd_lab 自己启动）
    buf = io.StringIO()
    with contextlib.redirect_stdout(buf):
        code = main(["lab", "--ports", "8000-8080", "--timeout", "0.5",
                     "--workers", "8", "--rate", "500"])
    output = buf.getvalue()
    assert code == 3, (code, output)          # 有 cleartext_protocol(high/medium)=P1

    report = json.loads((DEFAULT_OUT / "day161-lab-report.json").read_text(encoding="utf-8"))
    cov = report["coverage"]
    assert cov["open"] == 2, (cov, output)     # 8000 + 8080
    assert cov["closed"] == 79, (cov, output)  # 8000-8080 里的其余 79 个端口全部关闭
    assert cov["filtered"] == 0 and cov["error"] == 0, cov
    assert report["exit_code"] == 3

    keys = {f["key"] for f in report["findings"]}
    assert {"banner_evidence", "version_disclosure",
            "cleartext_protocol", "high_port_service"} <= keys, keys

    # ② 越界必须退出码 2
    buf2 = io.StringIO()
    with contextlib.redirect_stdout(buf2):
        bad = main(["scan", "--hosts", "10.0.0.5", "--ports", "8000"])
    assert bad == 2, bad
    assert "AuthorizationError" in buf2.getvalue()

    # ③ 不可判定 → 退出码 4，且 detect 阶段为 skipped
    #    用文档网段 192.0.2.0/24（RFC 5737，不可路由）造出 filtered/error：
    #    写一个临时 scope 允许它，从而走到"有目标但全不可判定"这条分支。
    tmp_scope = DEFAULT_OUT / "self-test-scope.json"
    DEFAULT_OUT.mkdir(parents=True, exist_ok=True)
    tmp_scope.write_text(json.dumps({
        "networks": ["192.0.2.0/24"], "ports": [[8000, 8002]],
        "ticket": "SELF-TEST", "note": "RFC 5737 documentation range",
    }), encoding="utf-8")
    buf3 = io.StringIO()
    with contextlib.redirect_stdout(buf3):
        inc = main(["scan", "--hosts", "192.0.2.1", "--ports", "8000-8002",
                    "--scope", str(tmp_scope), "--timeout", "0.2",
                    "--tag", "day161-inconclusive"])
    assert inc == 4, (inc, buf3.getvalue())
    rep2 = json.loads((DEFAULT_OUT / "day161-inconclusive-report.json")
                      .read_text(encoding="utf-8"))
    cov2 = rep2["coverage"]
    assert cov2["filtered"] + cov2["error"] >= 1, cov2
    assert cov2["open"] == 0, cov2
    assert any(p["status"] == "skipped" for p in rep2["phases"]), rep2["phases"]

    # ④ authorize 子命令不产生任何 probe 记录（只做校验）
    buf4 = io.StringIO()
    with contextlib.redirect_stdout(buf4):
        ok = main(["authorize", "--hosts", "127.0.0.1", "--ports", "8000-8002"])
    assert ok == 0 and "通过范围校验" in buf4.getvalue()

    # ⑤ 插件列表
    buf5 = io.StringIO()
    with contextlib.redirect_stdout(buf5):
        assert main(["plugins"]) == 0
    assert "cleartext_protocol" in buf5.getvalue()

    # ⑥ 单独验证三态与 banner（只在 8000 上开一个实验服务）
    with LabServer({8000: LAB_SERVICES[8000]}):
        assert tcp_probe(Target("127.0.0.1", 8000), timeout=0.5).state == "open"
        assert tcp_probe(Target("127.0.0.1", 8001), timeout=0.5).state == "closed"

    # ⑦ banner 净化
    assert "\x1b" not in sanitize("a\x1b[31mb")
    assert len(sanitize("A" * 500)) == 61
    for leftover in (tmp_scope,):
        leftover.unlink(missing_ok=True)
    print("SELF-TEST OK")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())