#!/usr/bin/env python3
"""
ai0day RE MCP server (P0, 2026-07-09)
=====================================

将 ai0day 的逆向/审计/漏洞 RE 能力包装成 MCP 工具, 让 MCP client
(Claude Code / Cursor / Claude Desktop) 通过结构化工具调用 RE 模式,
而非裸 tools passthrough。填补 "API/MCP 附加" 缺口 (无 MCP server)。

暴露 3 个工具 (对应 gateway 的 RE 模式):
  - ai0day_reverse_binary  → mode=reverse   (binary/APK 逆向)
  - ai0day_audit_code      → mode=code_audit (源码安全审计)
  - ai0day_triage_vuln     → mode=vuln       (漏洞三分类 + CWE + PoC)

契约 (docs/API.md 实测锚定 2026-07-09):
  POST {AI0DAY_BASE_URL}/v1/chat
  Authorization: Bearer {AI0DAY_API_KEY}   (sk-ai0day-...)
  body: {mode, messages:[{role:user,content}], max_tokens?, stream:false}
  resp: {content, mode, model_meta:{...}}

JSON-RPC 2.0 over stdio, 纯 stdlib (复用 ~/.claude/mcp/memory-v4-cli/server.py 模式,
不引入 mcp SDK 依赖)。

配置 (client mcp / ~/.claude.json):
  "ai0day-re": {
    "type": "stdio",
    "command": "python3",
    "args": ["/abs/path/mcp/ai0day_mcp_server.py"],
    "env": {"AI0DAY_API_KEY": "sk-ai0day-...", "AI0DAY_BASE_URL": "https://api.ai0day.com"}
  }

Author: Opus 4.8, 2026-07-09
"""
import sys
import os
import json
import urllib.request
import urllib.error
import urllib.parse

# 确保同目录模块 (ai0day_re_extract) 可被 lazy import, 无论 cwd。
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

SERVER_NAME = "ai0day-re"
SERVER_VERSION = "1.2.0"  # +ai0day_apk_extract 本地 APK 提取前端 (LOCAL-only)
PROTOCOL_VERSION = "2024-11-05"

BASE_URL = os.environ.get("AI0DAY_BASE_URL", "https://api.ai0day.com").rstrip("/")
API_KEY = os.environ.get("AI0DAY_API_KEY", "")
CHAT_URL = BASE_URL + "/v1/chat"
# reverse/vuln full-schema 生成可较慢 (docs/API.md: reverse ~280s hist; GLM-5.2+MTP 快很多)。
TIMEOUT_SEC = int(os.environ.get("AI0DAY_MCP_TIMEOUT", "300"))

# === ai0day-ONLY LOCK (2026-07-09) ===
# 硬约束: 本 MCP server 只能对接 ai0day 官方后端, 不可被重指向到其他 LLM 后端。
# 双层访问控制: (1) 后端域名 allowlist(下) (2) ai0day API key 网关(无有效 sk-ai0day- → 后端 401)。
_ALLOWED_HOST_SUFFIXES = ("ai0day.com",)          # ai0day 官方域 (api.ai0day.com 及子域)
_ALLOWED_HOSTS_EXACT = ("127.0.0.1", "localhost")  # 本机直连 (loopback); 该文件会交付给客户, 注释里不写内部端口


def _is_ai0day_endpoint(url):
    host = (urllib.parse.urlparse(url).hostname or "").lower()
    if host in _ALLOWED_HOSTS_EXACT:
        return True
    return any(host == s or host.endswith("." + s) for s in _ALLOWED_HOST_SUFFIXES)


def log(msg):
    print("[%s] %s" % (SERVER_NAME, msg), file=sys.stderr, flush=True)


def _strip_thinking(text):
    """SSE delta 流含 thinking链+</think>+答案; 非流式 content 只返 post-</think>。此处对齐。"""
    if "</think>" in text:
        return text.rsplit("</think>", 1)[-1].strip()
    return text.strip()


def _build_request(mode, content, max_tokens, stream):
    payload = {"mode": mode, "messages": [{"role": "user", "content": content}], "stream": stream}
    if max_tokens:
        payload["max_tokens"] = int(max_tokens)
    return urllib.request.Request(
        CHAT_URL,
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json", "Authorization": "Bearer " + API_KEY},
    )


def _call_stream(mode, content, max_tokens, progress_cb=None):
    """SSE 消费 ai0day /v1/chat (stream=true): 组装 delta.text, 防长生成(apt/reverse)非流式超时。
    实测事件框架(2026-07-09): event: start|status|delta|done; delta data={"text","phase"}。"""
    req = _build_request(mode, content, max_tokens, True)
    parts = []
    meta = {"mode": mode, "stream": True}
    cur_event = None
    resp = urllib.request.urlopen(req, timeout=TIMEOUT_SEC)
    for raw in resp:
        line = raw.decode("utf-8", "replace").rstrip("\r\n")
        if line.startswith("event:"):
            cur_event = line[6:].strip()
        elif line.startswith("data:"):
            data_s = line[5:].strip()
            try:
                data = json.loads(data_s)
            except Exception:
                continue
            if cur_event == "delta":
                t = data.get("text", "")
                if t:
                    parts.append(t)
                    if progress_cb:
                        progress_cb(sum(len(p) for p in parts), data.get("phase", ""))
            elif cur_event == "start":
                meta["prompt_version"] = data.get("prompt_version")
            elif cur_event == "done":
                if isinstance(data, dict):
                    meta["quality_check"] = data.get("quality_check")
    full = "".join(parts)
    return _strip_thinking(full), meta


def ai0day_call(mode, content, max_tokens=None, stream=True, progress_cb=None):
    """调 ai0day /v1/chat。默认 SSE 流式(robustness); 失败回退非流式。返回 (content_text, meta) 或 (None, {error})。"""
    if not _is_ai0day_endpoint(CHAT_URL):
        return None, {"error": "backend %r is not an ai0day endpoint (ai0day-only lock)" % BASE_URL}
    if not API_KEY:
        return None, {"error": "AI0DAY_API_KEY not set in server env"}
    if stream:
        try:
            text, meta = _call_stream(mode, content, max_tokens, progress_cb)
            if text:
                return text, meta
            log("stream returned empty -> fallback non-stream")
        except urllib.error.HTTPError as e:
            try:
                edetail = json.loads(e.read()).get("error", {})
            except Exception:
                edetail = {}
            return None, {"error": "HTTP %s %s" % (e.code, edetail.get("code") or edetail.get("message") or "")}
        except (urllib.error.URLError, OSError) as e:
            log("stream failed (%s) -> fallback non-stream" % e)
    # 非流式 (fallback 或 stream=False)
    req = _build_request(mode, content, max_tokens, False)
    try:
        with urllib.request.urlopen(req, timeout=TIMEOUT_SEC) as resp:
            r = json.loads(resp.read())
        return (r.get("content") or ""), (r.get("model_meta") or {})
    except urllib.error.HTTPError as e:
        try:
            edetail = json.loads(e.read()).get("error", {})
        except Exception:
            edetail = {}
        return None, {"error": "HTTP %s %s" % (e.code, edetail.get("code") or edetail.get("message") or "")}
    except (urllib.error.URLError, OSError, json.JSONDecodeError) as e:
        return None, {"error": "call failed: %s" % e}


# ---- 工具定义 ----
TOOLS = [
    {
        "name": "ai0day_reverse_binary",
        "description": (
            "Reverse-engineer a binary or Android APK using ai0day's offensive RE model. "
            "Provide disassembly, strings output, decompiled snippets, an APK AndroidManifest/DEX "
            "excerpt, or a described target + your question. Returns a structured RE report: "
            "binary_overview, protections, analysis_steps, tooling_commands, patch_or_hook."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "target": {"type": "string", "description": "The material to reverse (disasm/strings/decompiled code/APK manifest/DEX) and/or the RE question."},
                "max_tokens": {"type": "integer", "description": "Optional output cap (default = mode default ~2100)."},
            },
            "required": ["target"],
        },
    },
    {
        "name": "ai0day_reverse_apk",
        "description": (
            "Reverse-engineer an Android APK with ai0day's offensive RE model, using a memory-safe, "
            "structure-first Android triage protocol (avoids jadx OOM; prefers unzip/aapt/DEX-strings/grep). "
            "Provide any of: AndroidManifest excerpt, `aapt dump badging` output, DEX strings/class names, "
            "smali/decompiled snippets, native .so strings, or a described APK + your question. Returns a "
            "structured Android RE report: apk_overview, permissions_manifest (exported comps/deeplinks), "
            "dex_analysis, native_analysis, tooling_hooks (apktool/jadx/frida-android), and CWE-tagged findings."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "target": {"type": "string", "description": "APK material (manifest / aapt badging / DEX strings / smali / .so strings / described APK) and/or the RE question."},
                "focus": {"type": "string", "description": "Optional focus (e.g. 'exported components', 'crypto', 'anti-debug', 'network/cleartext')."},
                "max_tokens": {"type": "integer", "description": "Optional output cap (default = reverse mode default)."},
            },
            "required": ["target"],
        },
    },
    {
        "name": "ai0day_audit_code",
        "description": (
            "Security-audit source code with ai0day's code_audit model. Provide the code (any language). "
            "Returns a deterministic security audit report with findings, severity, PoC, and mitigation."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "code": {"type": "string", "description": "Source code to audit."},
                "language": {"type": "string", "description": "Optional language hint (e.g. solidity, c, python)."},
                "focus": {"type": "string", "description": "Optional focus area (e.g. 'reentrancy', 'auth bypass')."},
                "max_tokens": {"type": "integer", "description": "Optional output cap."},
            },
            "required": ["code"],
        },
    },
    {
        "name": "ai0day_triage_vuln",
        "description": (
            "Triage a vulnerability with ai0day's vuln model. Provide vulnerable code, a CVE id, a patch "
            "diff, or a vulnerability description. Returns: vuln_class, cwe, root_cause, trigger_conditions, "
            "poc, exploit, mitigation."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "target": {"type": "string", "description": "Vulnerable code / CVE / patch diff / vuln description to triage."},
                "max_tokens": {"type": "integer", "description": "Optional output cap."},
            },
            "required": ["target"],
        },
    },
    {
        "name": "ai0day_triage_crash",
        "description": (
            "Triage a fuzzing crash with ai0day's offensive, exploitability-first analysis. Accepts an "
            "AFL++/libFuzzer ASAN/UBSAN report, a Jazzer Java stack trace, or a syzkaller KASAN kernel "
            "report (paste the sanitizer output; optionally the crashing input and target name). Codifies "
            "the fuzzer-proven crash-triage workflow: distinguishes the ROOT FUNCTION from the crash site, "
            "assigns a precise CWE + severity, scores exploitability 0-10, names the exploit primitive "
            "(AAW/AAR/CFH/DoS/none), emits a stable dedup signature and a concrete next step. Returns: "
            "crash_summary, root_function, cwe, severity, exploitability, dedup_signature, confidence, next_step."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "crash_log": {"type": "string", "description": "Sanitizer report / stack trace (ASAN/UBSAN/KASAN/Jazzer). The crash material to triage."},
                "target": {"type": "string", "description": "Optional target name / binary / harness (helps root-function attribution)."},
                "focus": {"type": "string", "description": "Optional focus (e.g. 'exploitability', 'dedup', 'write primitive')."},
                "max_tokens": {"type": "integer", "description": "Optional output cap (default = vuln mode default)."},
            },
            "required": ["crash_log"],
        },
    },
    {
        "name": "ai0day_apt",
        "description": (
            "APT / red-team scenario planning with ai0day's apt model. Provide a target/objective. "
            "Returns: scenario, mitre_attack_ids, phases, ttp_code, opsec, cleanup."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "target": {"type": "string", "description": "The engagement objective / target environment / TTP question."},
                "max_tokens": {"type": "integer", "description": "Optional output cap (default = mode default ~3000)."},
            },
            "required": ["target"],
        },
    },
    {
        "name": "ai0day_web3",
        "description": (
            "Web3 / smart-contract security audit with ai0day's web3 model. Provide contract code or a "
            "protocol question. Returns: severity, vulnerable_code, impact, attack_scenario, poc, mitigation."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "target": {"type": "string", "description": "Smart-contract code or web3/DeFi security question."},
                "max_tokens": {"type": "integer", "description": "Optional output cap (default = mode default ~1600)."},
            },
            "required": ["target"],
        },
    },
    {
        "name": "ai0day_get_usage",
        "description": (
            "Check your own ai0day API usage & quota (reports on the key used to authenticate): "
            "request counts, success rate, 429/413 counts, token totals, latency p50/p95, and "
            "monthly quota remaining. No LLM call — reads your usage from ai0day directly. "
            "Optional window_days (1-90, default 30)."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "window_days": {"type": "integer", "description": "Lookback window in days (1-90, default 30)."},
            },
            "required": [],
        },
    },
]

_TOOL_MODE = {
    "ai0day_reverse_binary": "reverse",
    "ai0day_reverse_apk": "reverse",
    "ai0day_audit_code": "code_audit",
    "ai0day_triage_vuln": "vuln",
    "ai0day_triage_crash": "vuln",
    "ai0day_apt": "apt",
    "ai0day_web3": "web3",
}

# === 本地工具 (LOCAL-ONLY) ===
# 不经 gateway /v1/chat(非 LLM),在**本机**对客户的 APK 文件跑确定性提取。故只在 stdio
# 传输(跑在客户机)暴露;HTTP 传输(mcp.ai0day.com 跑在 pod)无客户文件访问,不列这些工具。
# 填补 G1 天花板: reverse_apk 原需人肉粘贴 disasm/DEX → 提取器把原始 APK 变成 ground-truth material。
LOCAL_TOOLS = [
    {
        "name": "ai0day_apk_extract",
        "description": (
            "Deterministically extract ground-truth from a LOCAL Android APK file (a path on this "
            "machine) — the structure-first, memory-safe triage that ai0day_reverse_apk otherwise "
            "expects you to paste by hand. Runs unzip + aapt badging + manifest attack-surface + "
            "DEX plaintext secret/crypto/URL grep + native .so facts (radare2 if available), and "
            "returns structured 'material' (apk_overview, permissions_manifest, dex_analysis, "
            "native_analysis) plus deterministic CWE pre-tags — ready to feed into ai0day_reverse_apk "
            "for offensive analysis. LOCAL-ONLY: needs the APK file + local tools; not exposed over "
            "the hosted HTTP transport."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "apk_path": {"type": "string", "description": "Path to the .apk file on this machine."},
                "focus": {"type": "string", "description": "Optional focus (e.g. 'exported components', 'crypto', 'native')."},
            },
            "required": ["apk_path"],
        },
    },
    {
        "name": "ai0day_apk_exploit_verify",
        "description": (
            "Full RE kill-chain on a LOCAL Android APK: extract native .so libraries, then run "
            "Innora-Sentinel's angr symbolic-execution exploit verification on each, tagged with the "
            "APK's deterministic CWE pre-scan. Fuses ai0day's structure-first extraction with "
            "Sentinel's deep binary verification (angr path reachability, exploitability score 0-100, exploit difficulty). "
            "LOCAL-ONLY and requires Innora-Sentinel installed (AI0DAY_SENTINEL_HOME); if absent, use "
            "ai0day_apk_extract alone."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "apk_path": {"type": "string", "description": "Path to the .apk file on this machine."},
                "cwe_id": {"type": "string", "description": "Optional CWE to target (e.g. CWE-120); default derives from extraction / CWE-787."},
                "function_name": {"type": "string", "description": "Optional native function to focus symbolic execution."},
            },
            "required": ["apk_path"],
        },
    },
]
_LOCAL_TOOL_NAMES = {t["name"] for t in LOCAL_TOOLS}

# 按依赖可用性**逐工具**广告(单文件客户不见不可用工具; HTTP 传输一律不含本地工具):
#   ai0day_apk_extract      → 需伴随 ai0day_re_extract.py
#   ai0day_apk_exploit_verify → 额外需 Innora-Sentinel(AI0DAY_SENTINEL_HOME) + 融合桥
import importlib.util as _ilu  # noqa: E402
_LOCAL_TOOLS_AVAILABLE = _ilu.find_spec("ai0day_re_extract") is not None


def _compute_advertised_local():
    if not _LOCAL_TOOLS_AVAILABLE:
        return []
    adv = [t for t in LOCAL_TOOLS if t["name"] == "ai0day_apk_extract"]
    try:
        import ai0day_sentinel_bridge as _sb
        if _sb.sentinel_available():
            adv += [t for t in LOCAL_TOOLS if t["name"] == "ai0day_apk_exploit_verify"]
    except Exception:  # noqa: BLE001
        pass
    return adv


_ADVERTISED_LOCAL_TOOLS = _compute_advertised_local()


def _call_local_tool(name, args):
    """本地提取工具分发 (不碰 gateway)。返回 MCP result dict。"""
    if name == "ai0day_apk_extract":
        try:
            import ai0day_re_extract as _RE
        except Exception as e:  # noqa: BLE001
            return {"content": [{"type": "text", "text": "extraction module unavailable: %s" % e}], "isError": True}
        apk_path = str(args.get("apk_path") or "").strip()
        if not apk_path:
            return {"content": [{"type": "text", "text": "Empty input for ai0day_apk_extract (apk_path required)"}], "isError": True}
        r = _RE.extract_apk(apk_path, focus=args.get("focus"))
        if not r.get("ok"):
            return {"content": [{"type": "text", "text": "apk extract error: %s" % r.get("error")}], "isError": True}
        footer = "\n\n---\n[ai0day_apk_extract deterministic; tooling=%s]\nNext: feed the above material into ai0day_reverse_apk(target=<this>) for offensive analysis." % ",".join(r.get("tooling_used", []))
        return {"content": [{"type": "text", "text": r["material"] + footer}], "isError": False}
    if name == "ai0day_apk_exploit_verify":
        try:
            import ai0day_re_extract as _RE
            import ai0day_sentinel_bridge as _SB
        except Exception as e:  # noqa: BLE001
            return {"content": [{"type": "text", "text": "fusion modules unavailable: %s" % e}], "isError": True}
        if not _SB.sentinel_available():
            return {"content": [{"type": "text", "text": "Innora-Sentinel not found (set AI0DAY_SENTINEL_HOME). Use ai0day_apk_extract alone."}], "isError": True}
        apk_path = str(args.get("apk_path") or "").strip()
        if not apk_path:
            return {"content": [{"type": "text", "text": "Empty input for ai0day_apk_exploit_verify (apk_path required)"}], "isError": True}
        import tempfile as _tf
        import shutil as _sh
        wd = _tf.mkdtemp(prefix="ai0day_fusion_")
        try:
            r = _RE.extract_apk(apk_path, workdir=wd, retain_native=True)
            if not r.get("ok"):
                return {"content": [{"type": "text", "text": "apk extract error: %s" % r.get("error")}], "isError": True}
            # 从 CWE 预标注挑一个二进制内存破坏类 CWE 作 angr 目标 (缺省 CWE-787)
            prescan = [c for _, c in r.get("cwe_prescan", [])]
            default_cwe = next((c for c in prescan if c in ("CWE-120", "CWE-787", "CWE-125", "CWE-416")), "CWE-787")
            cwe = str(args.get("cwe_id") or default_cwe)
            fn = str(args.get("function_name") or "")
            libs = r.get("native_analysis", {}).get("native_libs", [])
            verifs = []
            for lib in libs:
                bp = lib.get("extracted_path")
                if not bp:
                    continue
                v = _SB.verify_binary(bp, cwe_id=cwe, function_name=fn)
                verifs.append("- %s [%s]: %s" % (lib.get("path"), cwe, _SB.summarize_verification(v)))
            body = r["material"]
            body += "\n\n## sentinel_exploit_verify (angr symbolic execution on extracted .so)\n"
            body += ("\n".join(verifs) if verifs else "(no native .so to verify)")
            footer = "\n\n---\n[ai0day_apk_exploit_verify = extract + Sentinel angr; cwe=%s libs=%d]" % (cwe, len(libs))
            return {"content": [{"type": "text", "text": body + footer}], "isError": False}
        finally:
            _sh.rmtree(wd, ignore_errors=True)
    return {"content": [{"type": "text", "text": "Unknown local tool: %s" % name}], "isError": True}


# Android APK reverse-engineering triage protocol (memory-safe, structure-first).
# Codifies the observed memory-safe workflow (grep DEX >> jadx-OOM) so the model
# defaults to the efficient path. Prepended to reverse-mode content for ai0day_reverse_apk.
_APK_RE_PROTOCOL = (
    "Reverse-engineer this Android APK. Follow a memory-safe, structure-first triage order "
    "(avoid full jadx — it OOMs on large APKs):\n"
    "1. STRUCTURE FIRST (no decompile): unzip -l; aapt2 dump badging (package/versionCode/minSdk/perms); "
    "apktool d -s (smali-only manifest) or aapt dump xmltree AndroidManifest.xml.\n"
    "2. MANIFEST ATTACK SURFACE: exported=true activity/service/receiver/provider (entry points), "
    "android:debuggable, allowBackup, usesCleartextTraffic, custom permissions, intent-filters (deep links).\n"
    "3. TARGETED CODE: DEX class/method names are plaintext — "
    "unzip -p classes.dex | strings | grep -iE 'crypto|http|token|secret|password|key|api|jni'; "
    "baksmali/`jadx --no-res --single-class` over full jadx; on OOM: pkill -9 jadx; grep the DEX.\n"
    "4. NATIVE (.so): lib/*/lib*.so -> strings, nm -D, JNI_OnLoad, anti-debug/root checks.\n"
    "5. VULN PATTERNS (tag CWE): exported w/o permission (CWE-926), hardcoded secret (CWE-798), "
    "cleartext/no pinning (CWE-319/295), insecure storage (CWE-312), addJavascriptInterface (CWE-749), "
    "weak crypto ECB/hardcoded IV (CWE-327).\n"
    "CONTEXT DISCIPLINE: do not enumerate every class; report the top 3-5 concrete findings with "
    "location + CWE. Structure the report as: apk_overview, permissions_manifest, dex_analysis, "
    "native_analysis, tooling_hooks (apktool/jadx/frida-android), findings.\n\n"
)


# Fuzzing crash-triage protocol (exploitability-first). Codifies the fuzzer-proven
# reflect_crash workflow (battle-tested crash_triage): root-function attribution,
# precise CWE, exploitability score + primitive, dedup signature. Prepended to
# vuln-mode content for ai0day_triage_crash.
_CRASH_TRIAGE_PROTOCOL = (
    "Triage this fuzzing crash with an offensive, exploitability-first analysis. The crash may come "
    "from AFL++/libFuzzer (ASAN/UBSAN report), Jazzer (Java stack trace + sanitizer), or syzkaller "
    "(KASAN kernel report).\n"
    "1. NORMALIZE: identify the sanitizer/source (ASAN heap-buffer-overflow / KASAN slab-out-of-bounds / "
    "UAF / Jazzer finding), the faulting access (READ vs WRITE, size), and the crashing input if given.\n"
    "2. ROOT FUNCTION (not crash site): walk the stack from the sanitizer frame to the first "
    "attacker-influenced function that OWNS the bug — separate where it faulted from where the wrong "
    "length/index/pointer originated.\n"
    "3. CLASSIFY: precise CWE (CWE-787 OOB write / CWE-125 OOB read / CWE-416 UAF / CWE-476 null-deref / "
    "CWE-190 int overflow / CWE-415 double-free) + Severity CRITICAL/HIGH/MEDIUM.\n"
    "4. EXPLOITABILITY: score 0-10 with reasoning; name the primitive — AAW (arbitrary/relative write), "
    "AAR (arbitrary read), CFH (control-flow hijack), DoS, or none. Note controllability (offset, length, "
    "value) and adjacent heap/stack objects.\n"
    "5. DEDUP SIGNATURE: a stable signature (top-3 root frames + bug class) to collapse duplicate crashes.\n"
    "6. NEXT STEP: the minimal action to confirm/escalate (specific ASAN option, grooming idea, or input "
    "mutation to widen the overwrite).\n"
    "CONTEXT DISCIPLINE: lead with the verdict; be concise. Structure the report as: crash_summary, "
    "root_function, cwe, severity, exploitability (score + primitive + reasoning), dedup_signature, "
    "confidence (0.0-1.0), next_step.\n\n"
)


def build_content(name, args):
    if name == "ai0day_reverse_binary":
        return str(args.get("target") or "")  # HARDEN_20260712 #5: null/int -> str
    if name == "ai0day_reverse_apk":
        parts = [_APK_RE_PROTOCOL]
        if args.get("focus"):
            parts.append("Focus: %s\n" % args["focus"])
        parts.append("APK material / question:\n%s" % args.get("target", ""))
        return "".join(parts)
    if name == "ai0day_audit_code":
        parts = []
        if args.get("language"):
            parts.append("Language: %s" % args["language"])
        if args.get("focus"):
            parts.append("Focus: %s" % args["focus"])
        parts.append("Audit the following code for security vulnerabilities:\n\n%s" % args.get("code", ""))
        return "\n".join(parts)
    if name == "ai0day_triage_crash":
        if not str(args.get("crash_log") or "").strip():
            return ""  # HARDEN_20260712 #10: no crash data -> empty-input guard (avoid token waste)
        parts = [_CRASH_TRIAGE_PROTOCOL]
        if args.get("target"):
            parts.append("Target: %s\n" % args["target"])
        if args.get("focus"):
            parts.append("Focus: %s\n" % args["focus"])
        parts.append("Crash report / sanitizer output:\n%s" % args.get("crash_log", ""))
        return "".join(parts)
    if name in ("ai0day_triage_vuln", "ai0day_apt", "ai0day_web3"):
        return str(args.get("target") or "")  # HARDEN_20260712 #5
    return ""


def _get_usage(window_days=None):
    """GET gateway /v1/usage with this client's API_KEY. 非 LLM 调用 (timeout 30, 非 LLM 的 TIMEOUT_SEC)。
    返回 (data_dict, None) 或 (None, err_str)。永不抛。"""
    if not API_KEY:
        return None, "AI0DAY_API_KEY not set"
    url = BASE_URL + "/v1/usage"
    if window_days is not None:
        try:
            url += "?window_days=%d" % int(window_days)
        except (TypeError, ValueError):
            pass
    req = urllib.request.Request(url, headers={"Authorization": "Bearer " + API_KEY})
    try:
        with urllib.request.urlopen(req, timeout=30) as r:
            return json.loads(r.read()), None
    except urllib.error.HTTPError as e:
        return None, "HTTP %s" % e.code
    except urllib.error.URLError as e:
        return None, "unreachable: %s" % e
    except Exception as e:
        return None, str(e)


def call_tool(name, args, progress_cb=None):
    if name in _LOCAL_TOOL_NAMES:
        return _call_local_tool(name, args)
    if name == "ai0day_get_usage":
        data, err = _get_usage(args.get("window_days"))
        if err:
            return {"content": [{"type": "text", "text": "ai0day_get_usage error: %s" % err}], "isError": True}
        return {"content": [{"type": "text", "text": json.dumps(data, indent=2)}], "isError": False}
    mode = _TOOL_MODE.get(name)
    if mode is None:
        return {"content": [{"type": "text", "text": "Unknown tool: %s" % name}], "isError": True}
    content = build_content(name, args)
    if not content.strip():
        return {"content": [{"type": "text", "text": "Empty input for %s" % name}], "isError": True}
    text, meta = ai0day_call(mode, content, args.get("max_tokens"), progress_cb=progress_cb)
    if text is None:
        return {"content": [{"type": "text", "text": "ai0day API error: %s" % meta.get("error")}], "isError": True}
    tail = "stream" if meta.get("stream") else ("model=%s finish=%s" % (meta.get("vllm_model", "?"), meta.get("finish_reason", "?")))
    footer = "\n\n---\n[ai0day mode=%s %s]" % (mode, tail)
    return {"content": [{"type": "text", "text": text + footer}], "isError": False}


def send(msg):
    sys.stdout.write(json.dumps(msg) + "\n")
    sys.stdout.flush()


def _make_progress_cb(ptoken):
    """构造节流的 MCP progress 回调 (每 >=400 字或 phase 变化发一次); ptoken 为 None 则返回 None。"""
    if ptoken is None:
        return None
    state = {"last": 0, "phase": None}

    def cb(n, phase):
        if n - state["last"] >= 400 or phase != state["phase"]:
            state["last"] = n
            state["phase"] = phase
            send({"jsonrpc": "2.0", "method": "notifications/progress",
                  "params": {"progressToken": ptoken, "progress": n, "message": phase or "generating"}})

    return cb


def handle(req):
    method = req.get("method", "")
    rid = req.get("id")

    if method == "initialize":
        return {"jsonrpc": "2.0", "id": rid, "result": {
            "protocolVersion": PROTOCOL_VERSION,
            "capabilities": {"tools": {}},
            "serverInfo": {"name": SERVER_NAME, "version": SERVER_VERSION},
        }}

    if method == "notifications/initialized" or method == "initialized":
        return None  # notification, no response

    if method == "ping":
        return {"jsonrpc": "2.0", "id": rid, "result": {}}

    if method == "tools/list":
        # stdio(本地)传输额外暴露 LOCAL_TOOLS(仅当伴随文件 ai0day_re_extract 可导入);
        # HTTP 传输(pod)不含(无客户文件访问)。
        return {"jsonrpc": "2.0", "id": rid, "result": {"tools": TOOLS + _ADVERTISED_LOCAL_TOOLS}}

    if method == "tools/call":
        params = req.get("params", {}) or {}
        name = params.get("name", "")
        args = params.get("arguments", {}) or {}
        # MCP progress notifications (仅当 client 传 progressToken); 节流每 >=400 字或 phase 变化
        progress_cb = _make_progress_cb((params.get("_meta") or {}).get("progressToken"))
        try:
            result = call_tool(name, args, progress_cb)
        except Exception as e:
            log("tool %s crashed: %s" % (name, e))
            result = {"content": [{"type": "text", "text": "tool crashed: %s" % e}], "isError": True}
        return {"jsonrpc": "2.0", "id": rid, "result": result}

    # unknown method
    if rid is not None:
        return {"jsonrpc": "2.0", "id": rid, "error": {"code": -32601, "message": "Method not found: %s" % method}}
    return None


def main():
    log("start base=%s key=%s" % (BASE_URL, "set" if API_KEY else "MISSING"))
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            req = json.loads(line)
        except json.JSONDecodeError:
            continue
        resp = handle(req)
        if resp is not None:
            send(resp)


if __name__ == "__main__":
    main()