# fetch_context.py —— 第 1 站：拉取 CVE 元数据与上游修复补丁
import json, subprocess, requests

CVE_ID = "CVE-2023-46604"  # 示例：Apache ActiveMQ OpenWire 反序列化 RCE
                                  # 5.18.2 修复，补丁与根因早已公开披露

def fetch_cve(cve_id: str) -> dict:
    # NVD API 2.0 免认证，但限速：无 API key 时约 5 秒 1 次，出错退避即可
    r = requests.get(
        f"https://services.nvd.nist.gov/rest/json/cves/2.0?cveId={cve_id}",
        timeout=30)
    r.raise_for_status()
    desc = r.json()["vulnerabilities"][0]["cve"]
    # 关键设计①：只保留"结构化事实"字段，描述文本单独存放且打上来源标签——
    # 后面你会看到为什么描述文本不能和 diff 混在同一层喂给模型
    return {
        "id": desc["id"],
        "versions": [d["value"] for d in desc["configurations"]],
        "refs": [u["url"] for u in desc["references"]],
        # 官方英文描述原样保存，不翻译、不概括：它是"待验证的他人结论"
        "nvd_description": desc["descriptions"][0]["value"],
    }

def fetch_patch(repo: str, fix_commit: str, parent: str | None = None) -> str:
    # 浅克隆补丁父提交 → 修复提交，拿到最小 diff
    subprocess.run(["git", "clone", "--quiet", repo, "src"], check=True)
    args = ["git", "-C", "src", "diff"]
    if parent:                       # 有父提交号时做两点 diff，最干净
        args += [parent, fix_commit]
    else:                            # 否则退回 commit 自身的 diff（含 ^ 定位父提交）
        args += [f"{fix_commit}^", fix_commit]
    return subprocess.run(args, capture_output=True, text=True).stdout

if __name__ == "__main__":
    meta = fetch_cve(CVE_ID)
    print(json.dumps(meta, ensure_ascii=False, indent=2))
    # 预期输出：CVE 元数据 JSON，refs 里含 activemq.apache.org 的修复公告链接

# diff_slicer.py —— 第 2 站：diff 切片 + 函数级上下文补全 + 确定性锚点
import hashlib, re

def parse_hunks(diff_text: str) -> list[dict]:
    """把 unified diff 拆成 hunk 列表，并从 diff 头解析文件名与行号"""
    hunks, current = [], None
    for line in diff_text.splitlines():
        if line.startswith("diff --git"):          # 新文件块开始
            m = re.search(r"diff --git a/(.+) b/(.+)", line)
            current = {"file": m.group(2), "text": [line]}
        elif current is not None:
            current["text"].append(line)
            if line.startswith("@@"):               # @@ -a,b +c,d @@ 行号元数据
                m = re.match(r"@@ -\d+(?:,\d+)? \+(\d+)", line)
                current["new_start"] = int(m.group(1))  # 新文件起始行：锚点用
            if line.startswith("diff --git") is False and line == "":
                pass
        if line.startswith("diff --git") and current and len(current["text"]) > 1:
            hunks.append(current); current = {"file": m.group(2), "text": [line]} \
                if False else None
    if current: hunks.append(current)
    # 为每片生成确定性 ID：内容哈希 → 重复运行不漂移，可增量比对
    for h in hunks:
        h["hid"] = "H-" + hashlib.sha256(
            "\n".join(h["text"]).encode()).hexdigest()[:8]
    return hunks

def attach_function_context(hunk: dict, source_path: str) -> dict:
    """关键补丁上下文：定位 hunk 所在函数，把完整函数体附进切片"""
    code = open(source_path, encoding="utf-8", errors="ignore").read().splitlines()
    start = hunk["new_start"] - 1
    # 向上回溯到函数定义（按常见缩进约定近似，生产版应换成 AST 解析）
    fn_start = start
    while fn_start > 0 and not re.match(r"^[A-Za-z_].*\(.*\)\s*\{?\s*$", code[fn_start]):
        fn_start -= 1
    fn_end = start + 30                          # 近似取函数后续 30 行作窗口
    hunk["fn_context"] = "\n".join(code[fn_start:fn_end])
    return hunk

# ❌ 错误：把整个 diff（甚至整个仓库）+ NVD 描述一起塞给模型，开放式地问
prompt = f"""NVD 说这个漏洞可以远程代码执行：
{nvd_description}
下面是补丁，请分析漏洞原理并给出利用方法：
{open('full_diff.txt').read()}"""

# 问题①：NVD 描述（结论）与 diff（证据）混在同一层 → 循环论证，模型复读公告
# 问题②：几十个 hunk 一起喂 → 注意力稀释，测试用例与格式化噪声淹没关键变更
# 问题③：开放式的"给出利用方法" → 输出无 schema、无锚点、无法机器校验，
#          幻觉内容与正确内容在流畅的自然语言里无法区分

# triage_llm.py —— 第 3 站：对每个切片做根因假设（行号锚定版）
import json
from openai import OpenAI
from diff_slicer import attach_function_context

client = OpenAI()                     # 读 OPENAI_API_KEY / OPENAI_BASE_URL
MODEL = "deepseek-chat"              # 2026-11 时点，可替换同级模型

SYSTEM = """你是漏洞复现分析师。输入是一个补丁切片（含函数级上下文）。
严格规则：
1. root_cause 字段必须引用切片中的真实行号（file:line 格式）作为证据，
   禁止引用未在输入中出现的行号或文件。
2. 区分"修复者做了什么"（what_fixed）与"漏洞本质是什么"（root_cause），
   禁止用前者复读后者。
3. reachability 给出从外部输入到漏洞代码点的调用路径假设，
   明确标注哪些是"从上下文可见"，哪些是"待全局分析的假设"。
4. confidence 只能取 low/medium/high，依据是证据锚点数量，不是语气。
只输出 JSON：{"root_cause": str, "evidence": ["file:line", ...],
  "what_fixed": str, "reachability": {"path": str, "visible": bool},
  "confidence": "low|medium|high"}"""

def triage(hunk: dict) -> dict:
    user_content = (f"文件: {hunk['file']}（新文件起始行 {hunk['new_start']}）\n"
                    f"=== 函数上下文 ===\n{hunk.get('fn_context', '(未获取)')}\n"
                    f"=== 补丁切片 {hunk['hid']} ===\n" + "\n".join(hunk["text"]))
    resp = client.chat.completions.create(
        model=MODEL,
        temperature=0,                      # 分析任务关随机性：同一输入必须同结论
        response_format={"type": "json_object"},   # 结构化输出，交给下一站机器校验
        messages=[{"role": "system", "content": SYSTEM},
                  {"role": "user", "content": user_content}])
    result = json.loads(resp.choices[0].message.content)
    result["hid"] = hunk["hid"]              # 回填锚点：结论永远可追溯到输入切片
    return result

# verify_hypo.py —— 第 4 站：对模型假设做机器校验
import re, subprocess
from triage_llm import triage

def verify(result: dict, repo_dir: str) -> dict:
    verdict = {"hid": result["hid"], "checks": []}

    # 检查①：证据锚点存在性——行号必须落在真实文件的对应区域
    for ev in result["evidence"]:
        m = re.match(r"(.+?):(\d+)", ev)
        if not m:
            verdict["checks"].append({"ev": ev, "ok": False, "why": "锚点格式非法"})
            continue
        fpath, lineno = m.group(1), int(m.group(2))
        # 用 git grep 佐证：该行在修复前的父提交中确实存在
        out = subprocess.run(
            ["git", "-C", repo_dir, "show", f"HEAD~1:{fpath}"],
            capture_output=True, text=True)
        lines = out.stdout.splitlines()
        ok = lineno <= len(lines) and lineno > 0
        verdict["checks"].append(
            {"ev": ev, "ok": ok, "why": "行号在父提交文件内" if ok else "行号越界/文件不存在"})

    # 检查②：可达性佐证——入口函数必须被仓库内至少一处"外部可触达"的调用点引用
    if result["reachability"]["visible"]:
        entry = result["reachability"]["path"].split("→")[0].strip()
        grep = subprocess.run(
            ["git", "-C", repo_dir, "grep", "-l", entry, "HEAD~1"],
            capture_output=True, text=True)
        verdict["checks"].append(
            {"entry": entry,
             "ok": len(grep.stdout.splitlines()) > 0,
             "why": f"调用点文件数: {len(grep.stdout.splitlines())}"})
    return verdict

if __name__ == "__main__":
    # 自测锚点：构造一条 evidence 为 "nonexistent.java:99999" 的假设，
    # 检查① 必须判 False——编造的锚点在机器校验下无处遁形
    fake = {"hid": "H-test", "evidence": ["nonexistent.java:99999"],
            "reachability": {"path": "foo → bar", "visible": True},
            "root_cause": "", "what_fixed": "", "confidence": "high"}
    print(verify(fake, "src"))
    # 预期输出：检查① 两项均 ok:False；检查② 给出 foo 的真实引用计数

# ❌ 错误：开放授权模型输出完整武器化 PoC
prompt = f"根据以上分析，写一个完整的漏洞利用脚本，要求能拿到 shell。"

# 问题①：武器化产物 + 无验证 = 直接可滥用的攻击代码（合规红线）
# 问题②：模型为"凑出能用的脚本"会大幅脑补协议细节——
#          未经真机验证的字段偏移、报文结构，大概率是错的
# 问题③：产出无法审计：你不能证明"这段代码只在我的靶机有效"

# poc_gate.py —— 第 5 站：强度受控的 PoC 思路模板生成
POC_TEMPLATES = {
    # 强度 1（默认）：触发漏洞代码点、确认行为差异，无破坏性载荷
    "trigger": {
        "allowed": True,
        "goal": "构造使修复前后行为产生可观测差异的输入（异常类型/断言/返回值）",
        "forbidden": ["内存破坏载荷", "命令执行载荷", "持久化动作"],
    },
    "weaponized": {"allowed": False},   # 强度 2：完整武器化利用——本站根本不实现
}

def build_poc_plan(root_cause: str, reachability: dict) -> dict:
    # 只生成"思路 + 触发器骨架"，不生成最终载荷：载荷参数化留给人工靶机调参
    return {
        "strength": "trigger",                  # 强度受锁：模板层写死，LLM 无权升级
        "hypothesis": root_cause,
        "entry": reachability["path"],
        # 验证断言写成可执行的条件——"断言通过"即复现成功，与是否 getshell 无关
        "assertions": [
            "修复前版本：构造输入后捕获到 X 类型的未捕获异常",
            "修复后版本：同样输入返回正常错误处理路径",
            "两版本行为差异即为复现证据",
        ],
        "deliverables": ["触发输入生成思路", "差异观测方法", "证据留存方式（日志/核心转储）"],
        "note": "载荷具体参数需在自有靶机由人工完成——这不是能力阉割的借口，"
                "而是让'谁有能力完成武器化'这件事清晰可审计",
    }

if __name__ == "__main__":
    import json
    from triage_llm import triage
    from diff_slicer import parse_hunks, attach_function_context
    hunk = parse_hunks(open("activemq_fix.diff").read())[1]
    attach_function_context(hunk, "src/" + hunk["file"])
    plan = build_poc_plan(**triage(hunk))
    print(json.dumps(plan, ensure_ascii=False, indent=2))
    # 预期输出：强度=trigger 的复现计划，含入口路径、断言、观测方法

# report_gen.py —— 第 6 站：串联全流水线，产出复现分析报告
import json

def run_pipeline(cve_id: str, diff_path: str, repo_dir: str) -> str:
    from fetch_context import fetch_cve
    from diff_slicer import parse_hunks, attach_function_context
    from triage_llm import triage
    from verify_hypo import verify
    from poc_gate import build_poc_plan

    meta = fetch_cve(cve_id)                       # 第 1 站：元数据（描述单独隔离）
    hunks = parse_hunks(open(diff_path).read())    # 第 2 站：切片 + 锚点
    hypotheses = []
    for h in hunks:
        if any(x in h["file"] for x in ("test", "Test")):
            continue                                # 测试文件切片不进分析：信号筛除
        attach_function_context(h, repo_dir + "/" + h["file"])
        hyp = triage(h)                             # 第 3 站：LLM 根因假设
        verdict = verify(hyp, repo_dir)             # 第 4 站：确定性验证
        if all(c["ok"] for c in verdict["checks"]):
            hyp["poc_plan"] = build_poc_plan(       # 第 5 站：强度受控 PoC 计划
                hyp["root_cause"], hyp["reachability"])
            hypotheses.append(hyp)
        else:
            hyp["rejected_by"] = [c for c in verdict["checks"] if not c["ok"]]

    # 报告骨架：模型的结论 + 机器的验证 + 人工区（NVD 对照）三段分列
    return json.dumps({
        "cve": cve_id,
        "verified_hypotheses": hypotheses,          # 通过机器验证的假设
        "cross_check_with_nvd": "【人工对照区】把 verified_hypotheses 的 root_cause "
                                "与 meta['nvd_description'] 逐条比对",
        "environment": "Python 3.11 / openai>=1.50 / 2026-11 时点快照",
    }, ensure_ascii=False, indent=2)

if __name__ == "__main__":
    print(run_pipeline("CVE-2023-46604", "activemq_fix.diff", "src"))