mcp_server.py

#!/usr/bin/env python3
"""tchkiller MCP 工具服务器 — 独立 STDIO 进程模式

CC（包括 Agent Teams 子代理）通过 stdio 与本进程通信。
每个 CC 实例启动自己的 mcp_server.py 副本。
"""

import asyncio
import json
import os
import re
import sys
import fcntl
from datetime import datetime
from pathlib import Path

from mcp.server import Server
from mcp.types import Tool, TextContent

# ── 环境变量 ──────────────────────────────────────────────
OUTPUT_DIR = os.environ.get("TCHKILLER_OUTPUT_DIR", "./output")
TARGET = os.environ.get("TCHKILLER_TARGET", "unknown")
SCOPE = os.environ.get("TCHKILLER_SCOPE", "")
DEBUG_LOG = Path("/tmp/tchkiller_mcp_debug.log")

# ── vulndb 路径 ──────────────────────────────────────────
VULNDB_DIR = Path(__file__).parent / "vulndb"
VULNDB_SQLITE = Path(__file__).parent / "vulndb.sqlite"


def _log(tag: str, msg: str):
    try:
        ts = datetime.now().strftime("%H:%M:%S.%f")[:-3]
        pid = os.getpid()
        with open(DEBUG_LOG, "a", encoding="utf-8") as f:
            f.write(f"[{ts}] [pid={pid}] [{tag}] {msg}\n")
    except Exception:
        pass


def _evidence_dir() -> Path:
    p = Path(OUTPUT_DIR) / "evidence"
    p.mkdir(parents=True, exist_ok=True)
    return p


def _next_evidence_id() -> int:
    """基于已有文件计算下一个 ID（多进程安全）"""
    d = _evidence_dir()
    existing = [f.name for f in d.glob("[0-9][0-9][0-9]-*.md")]
    if not existing:
        return 1
    ids = []
    for name in existing:
        try:
            ids.append(int(name[:3]))
        except ValueError:
            pass
    return max(ids) + 1 if ids else 1


# ── MCP Server ────────────────────────────────────────────
server = Server("tchkiller-tools", version="0.6.0")

# ── 全局别名归一化 ────────────────────────────────────────
# 把模型可能用的各种参数名统一到 canonical 名称
# key=canonical, value=接受的别名列表 (优先级从高到低)
TOOL_ALIASES: dict[str, dict[str, list[str]]] = {
    "evidence_read": {
        "id": ["id", "path", "file_path", "file", "name", "query", "keyword", "filename"],
    },
    "evidence_save": {
        "title": ["title", "name", "vuln_name"],
        "description": ["description", "desc", "detail", "content"],
    },
    "report_vuln": {
        "title": ["title", "vuln_name", "name"],
        "severity": ["severity", "vuln_type", "vuln_severity", "level"],
        "url": ["url", "vuln_url", "target", "endpoint"],
        "detail": ["detail", "description", "evidence", "desc"],
    },
    "read_skill": {
        "id": ["id", "skill", "skill_id", "name"],
    },
    "read_vuln": {
        "id": ["id", "vuln_id", "cve", "name"],
    },
    "locate_tool": {
        "name": ["name", "tool_name", "tool", "query", "keyword"],
    },
}


def _resolve_aliases(tool_name: str, args: dict) -> dict:
    """将模型传入的别名参数归一化到 canonical 名称。
    不修改原始 dict，返回归一化后的新 dict。"""
    aliases = TOOL_ALIASES.get(tool_name)
    if not aliases:
        return args
    resolved = dict(args)
    for canonical, alt_names in aliases.items():
        if canonical in resolved and resolved[canonical]:
            continue  # canonical 已存在且有值，跳过
        for alt in alt_names:
            if alt != canonical and alt in resolved and resolved[alt]:
                resolved[canonical] = resolved.pop(alt)
                break
    return resolved


@server.list_tools()
async def list_tools() -> list[Tool]:
    return [
        Tool(
            name="evidence_save",
            description="保存漏洞证据。必填: title, description。强烈建议提供: payload(攻击载荷/请求), response(服务器响应)。可选: severity, target, endpoint, param",
            inputSchema={
                "type": "object",
                "properties": {
                    "title": {"type": "string", "description": "漏洞名称"},
                    "description": {"type": "string", "description": "漏洞描述"},
                    "severity": {"type": "string", "description": "严重程度: CRITICAL/HIGH/MEDIUM/LOW"},
                    "target": {"type": "string", "description": "目标URL"},
                    "payload": {"type": "string", "description": "攻击载荷或完整请求（强烈建议提供）"},
                    "response": {"type": "string", "description": "服务器响应关键部分（强烈建议提供）"},
                    "endpoint": {"type": "string", "description": "漏洞端点路径"},
                    "param": {"type": "string", "description": "漏洞参数名"},
                    "method": {"type": "string", "description": "HTTP 方法"},
                    "evidence_type": {"type": "string", "description": "证据类型 (xss/sqli/ssrf/idor/etc)"},
                    "vuln_name": {"type": "string", "description": "漏洞名称 (title 的别名)"},
                },
            },
        ),
        Tool(
            name="evidence_list",
            description="列出已保存的所有漏洞证据",
            inputSchema={"type": "object", "properties": {}},
        ),
        Tool(
            name="evidence_read",
            description="读取指定证据的完整内容。传入证据编号(如 '1', '3')或文件名关键词(如 'SSRF', 'IDOR')。",
            inputSchema={
                "type": "object",
                "properties": {
                    "id": {"type": "string", "description": "证据编号(如 '1', '003')或文件名关键词(如 'SSRF', '侦察')"},
                    "path": {"type": "string", "description": "同 id，也接受文件路径"},
                    "file_path": {"type": "string", "description": "同 id，也接受完整路径"},
                    "file": {"type": "string", "description": "同 id，也接受文件名"},
                    "name": {"type": "string", "description": "同 id，也接受证据名称"},
                    "query": {"type": "string", "description": "同 id，也接受搜索关键词"},
                },
            },
        ),
        Tool(
            name="report_vuln",
            description="报告发现的漏洞。必填: title(漏洞名称), severity(CRITICAL/HIGH/MEDIUM/LOW), url(漏洞URL)。可选: detail(详细描述)。也接受别名: vuln_name=title, vuln_type=severity, vuln_url=url",
            inputSchema={
                "type": "object",
                "properties": {
                    "title": {"type": "string", "description": "漏洞名称 (别名: vuln_name)"},
                    "severity": {"type": "string", "description": "严重程度: CRITICAL/HIGH/MEDIUM/LOW (别名: vuln_type)"},
                    "url": {"type": "string", "description": "漏洞URL (别名: vuln_url)"},
                    "detail": {"type": "string", "description": "详细描述 (别名: description)"},
                    "vuln_name": {"type": "string", "description": "漏洞名称 (title的别名)"},
                    "vuln_type": {"type": "string", "description": "漏洞类型/严重程度 (severity的别名)"},
                    "vuln_url": {"type": "string", "description": "漏洞URL (url的别名)"},
                    "vuln_param": {"type": "string", "description": "漏洞参数"},
                    "vuln_severity": {"type": "string", "description": "严重程度 (severity的别名)"},
                    "description": {"type": "string", "description": "详细描述 (detail的别名)"},
                },
            },
        ),
        Tool(
            name="list_vulns",
            description="列出当前已发现的所有漏洞",
            inputSchema={"type": "object", "properties": {}},
        ),
        Tool(
            name="check_scope",
            description="检查目标是否在授权范围内",
            inputSchema={
                "type": "object",
                "properties": {
                    "target": {"type": "string", "description": "要检查的目标（不传则使用主目标）"},
                },
            },
        ),
        Tool(
            name="get_targets",
            description="获取当前任务的目标列表和授权范围",
            inputSchema={"type": "object", "properties": {}},
        ),
        Tool(
            name="list_skills",
            description="搜索可用的渗透技巧。传入关键词（如 sqli, xss, ssti, recon），返回匹配的技巧列表。",
            inputSchema={
                "type": "object",
                "properties": {
                    "keyword": {"type": "string", "description": "搜索关键词"},
                },
                "required": [],
            },
        ),
        Tool(
            name="read_skill",
            description="读取指定渗透技巧的完整内容。传入技巧 ID（如 sql-injection-methodology）。",
            inputSchema={
                "type": "object",
                "properties": {
                    "id": {"type": "string", "description": "技巧 ID"},
                    "skill": {"type": "string", "description": "同 id，也接受技巧名"},
                    "skill_id": {"type": "string", "description": "同 id"},
                    "name": {"type": "string", "description": "同 id，也接受技巧名称"},
                },
            },
        ),
        Tool(
            name="locate_tool",
            description="在本地文件系统搜索渗透工具的安装路径。仅当命令执行失败或 which 找不到工具时使用。"
                        "搜索范围: PATH、/pentest/、/usr/local/bin/、/usr/bin/。",
            inputSchema={
                "type": "object",
                "properties": {
                    "name": {
                        "type": "string",
                        "description": "工具名称（如 noPac, zerologon, pyGPOAbuse, passthecert, k8spider）",
                    },
                    "tool_name": {"type": "string", "description": "同 name"},
                    "tool": {"type": "string", "description": "同 name"},
                },
                "required": [],
            },
        ),
        Tool(
            name="search_vulndb",
            description="搜索漏洞知识库。根据产品名、关键词、CVE 编号搜索已知漏洞的利用方法。当指纹识别发现已知产品（如通达OA、致远OA、用友NC、Shiro、Fastjson、ThinkPHP、Spring、Log4j、Tomcat 等）时，务必搜索知识库获取精确的利用 payload。",
            inputSchema={
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "搜索关键词（产品名、CVE编号、漏洞类型等），如: 'shiro 反序列化'、'CVE-2021-44228'、'通达OA'",
                    },
                    "product": {
                        "type": "string",
                        "description": "精确产品名过滤（可选），如: tongda-oa, shiro, fastjson, thinkphp",
                    },
                    "severity": {
                        "type": "string",
                        "description": "严重性过滤（可选）: CRITICAL, HIGH, MEDIUM, LOW",
                    },
                },
                "required": [],
            },
        ),
        Tool(
            name="read_vuln",
            description="读取漏洞知识库中指定漏洞的完整利用详情（payload、利用步骤、验证方法）。先用 search_vulndb 搜索，再用此工具获取详情。",
            inputSchema={
                "type": "object",
                "properties": {
                    "id": {
                        "type": "string",
                        "description": "漏洞 ID（如 CVE-2020-1957、Shiro-550、CNVD-2019-34895）",
                    },
                    "vuln_id": {"type": "string", "description": "同 id"},
                    "cve": {"type": "string", "description": "同 id，也接受 CVE 编号"},
                    "name": {"type": "string", "description": "同 id，也接受漏洞名称"},
                },
            },
        ),
        Tool(
            name="http_request",
            description="发送 HTTP 请求，自动管理 Cookie/Session（同一 host:port 的 Cookie 自动保持）。"
                        "适合需要登录态保持、CSRF token 传递、复杂 payload（XML/XXE/SOAP）的场景。"
                        "登录流程示例：1) GET /login.php 获取 token 2) POST /login.php 提交凭据 — Cookie 自动携带，无需手动管理。",
            inputSchema={
                "type": "object",
                "properties": {
                    "url": {"type": "string", "description": "目标 URL"},
                    "method": {"type": "string", "description": "HTTP 方法 (GET/POST/PUT/DELETE/PATCH)，默认 GET"},
                    "headers": {
                        "type": "object",
                        "description": "请求头字典，如 {\"Content-Type\": \"application/xml\"}",
                        "additionalProperties": {"type": "string"},
                    },
                    "body": {"type": "string", "description": "请求体（原始字符串，不做转义处理）"},
                    "timeout": {"type": "integer", "description": "超时秒数，默认 15"},
                    "follow_redirects": {"type": "boolean", "description": "是否跟随重定向，默认 true"},
                    "new_session": {"type": "boolean", "description": "设为 true 时清除该 host 的所有 Cookie，开启全新会话。用于测试未认证访问。默认 false"},
                },
                "required": ["url"],
            },
        ),
        Tool(
            name="interactive_session",
            description="管理交互式终端会话（基于 tmux）。用于任何需要持续交互的场景：反弹 shell (nc)、SSH 登录、"
                        "数据库客户端 (mysql/redis-cli)、调试器、FTP 等。每个 session 独立运行，支持发送命令和读取输出。"
                        "action: start(启动新会话运行命令) | send(发送命令+Enter) | send_raw(发送原始按键如Ctrl-C) | read(读取当前屏幕) | close(关闭会话) | list(列出所有会话)",
            inputSchema={
                "type": "object",
                "properties": {
                    "action": {
                        "type": "string",
                        "description": "操作类型: start|send|send_raw|read|close|list",
                        "enum": ["start", "send", "send_raw", "read", "close", "list"],
                    },
                    "session_name": {
                        "type": "string",
                        "description": "会话名称（如 'nc_listener', 'ssh_target', 'mysql'）。默认 'default'",
                    },
                    "command": {
                        "type": "string",
                        "description": "start: 启动命令（如 'nc -lvp 4444'）; send: 要执行的命令; send_raw: 原始按键（如 'C-c' 发送 Ctrl-C）",
                    },
                    "wait": {
                        "type": "number",
                        "description": "发送命令后等待输出的秒数（默认 2，最大 30）。反弹 shell 等慢操作建议设 5-10",
                    },
                    "lines": {
                        "type": "integer",
                        "description": "read 操作读取的历史行数（默认 100）",
                    },
                },
                "required": ["action"],
            },
        ),
    ]


@server.call_tool()
async def call_tool(name: str, arguments: dict) -> list[TextContent]:
    # 全局别名归一化 — 在 handler 之前统一参数名
    arguments = _resolve_aliases(name, arguments)
    _log(name, f"CALLED args={arguments}")
    try:
        if name == "evidence_save":
            return await _handle_evidence_save(arguments)
        elif name == "evidence_list":
            return await _handle_evidence_list(arguments)
        elif name == "evidence_read":
            return await _handle_evidence_read(arguments)
        elif name == "report_vuln":
            return await _handle_report_vuln(arguments)
        elif name == "list_vulns":
            return await _handle_list_vulns(arguments)
        elif name == "check_scope":
            return await _handle_check_scope(arguments)
        elif name == "get_targets":
            return await _handle_get_targets(arguments)
        elif name == "list_skills":
            return await _handle_list_skills(arguments)
        elif name == "read_skill":
            return await _handle_read_skill(arguments)
        elif name == "locate_tool":
            return await _handle_locate_tool(arguments)
        elif name == "search_vulndb":
            return await _handle_search_vulndb(arguments)
        elif name == "read_vuln":
            return await _handle_read_vuln(arguments)
        elif name == "http_request":
            return await _handle_http_request(arguments)
        elif name == "interactive_session":
            return await _handle_interactive_session(arguments)
        else:
            return [TextContent(type="text", text=f"❌ Unknown tool: {name}")]
    except Exception as e:
        _log(name, f"ERROR: {e}")
        return [TextContent(type="text", text=f"❌ 工具错误: {e}")]


# ── evidence_save ─────────────────────────────────────────

def _normalize_evidence_key(title: str, url: str, param: str = "") -> str:
    """归一化 evidence 去重 key: vuln_type:path?param"""
    from urllib.parse import urlparse, parse_qs
    vtype = _classify_vuln_type(title)
    parsed = urlparse(url)
    path = parsed.path.rstrip("/").lower()
    # 合并 URL query 参数和显式 param 字段
    params = sorted(parse_qs(parsed.query).keys())
    if param and param.lower() not in params:
        params.append(param.lower())
        params.sort()
    param_key = ",".join(params) if params else ""
    key = f"{vtype}:{path}"
    if param_key:
        key += f"?{param_key}"
    return key


def _find_existing_evidence(evidence_dir: Path, norm_key: str) -> Path | None:
    """查找已有的同 endpoint+vuln_type 的 evidence 文件（用于去重合并）"""
    index_path = evidence_dir / ".evidence_keys.json"
    if not index_path.exists():
        return None
    try:
        mapping = json.loads(index_path.read_text(encoding="utf-8"))
        filename = mapping.get(norm_key)
        if filename:
            fp = evidence_dir / filename
            if fp.exists():
                return fp
    except Exception:
        pass
    return None


def _save_evidence_key(evidence_dir: Path, norm_key: str, filename: str):
    """记录 evidence 去重 key → filename 映射"""
    index_path = evidence_dir / ".evidence_keys.json"
    mapping = {}
    if index_path.exists():
        try:
            mapping = json.loads(index_path.read_text(encoding="utf-8"))
        except Exception:
            pass
    mapping[norm_key] = filename
    index_path.write_text(json.dumps(mapping, ensure_ascii=False, indent=2), encoding="utf-8")


async def _handle_evidence_save(args: dict) -> list[TextContent]:
    title = (args.get("title") or args.get("vuln_name") or args.get("name") or "").strip()
    desc = (args.get("description") or args.get("desc") or args.get("detail") or args.get("content") or "").strip()
    if not title and not desc:
        return [TextContent(type="text",
            text="❌ 缺少必要参数。用法: evidence_save(title='漏洞名', description='描述')\n"
                 "别名: title=name/vuln_name, description=desc/detail/content")]
    if not title:
        title = desc.split('\n')[0][:80] or "Unknown"
    if not desc:
        desc = title
    severity = (args.get("severity") or "MEDIUM").strip().upper()
    target = (args.get("target") or TARGET).strip()
    payload = (args.get("payload") or "").strip() or "(未提供)"
    response = (args.get("response") or "").strip() or "(未提供)"
    endpoint = (args.get("endpoint") or "").strip()
    param = (args.get("param") or "").strip()
    method = (args.get("method") or "").strip()
    evidence_type = (args.get("evidence_type") or "").strip()

    evidence_dir = _evidence_dir()

    # ── 去重检查: 同一 endpoint + vuln_type 合并到已有文件 ──
    url_for_key = endpoint or target
    norm_key = _normalize_evidence_key(title, url_for_key, param)
    existing = _find_existing_evidence(evidence_dir, norm_key)
    if existing:
        ts = datetime.now().strftime("%Y%m%d-%H%M%S")
        append_section = f"\n\n---\n## [合并] {ts}\n\n{desc}\n"
        if payload and payload != "(未提供)":
            append_section += f"\n### Payload\n```\n{payload}\n```\n"
        if response and response != "(未提供)":
            append_section += f"\n### 响应\n```\n{response[:1000]}\n```\n"
        with open(existing, "a", encoding="utf-8") as f:
            f.write(append_section)
        _log("evidence_save", f"MERGED into {existing.name} (key={norm_key})")
        return [TextContent(type="text", text=f"📎 证据已合并到: {existing.name}")]

    # ── 新建 evidence 文件 ──
    eid = _next_evidence_id()
    ts = datetime.now().strftime("%Y%m%d-%H%M%S")
    safe_name = title.replace(" ", "_").replace("/", "_")[:50]
    filename = f"{eid:03d}-{safe_name}.md"
    filepath = evidence_dir / filename

    emoji = {"CRITICAL": "🔴", "HIGH": "🟠", "MEDIUM": "🟡", "LOW": "🟢"}.get(severity, "⚪")

    # 构建额外信息行
    extra_lines = []
    if endpoint:
        extra_lines.append(f"- **端点**: {endpoint}")
    if param:
        extra_lines.append(f"- **参数**: {param}")
    if method:
        extra_lines.append(f"- **方法**: {method}")
    if evidence_type:
        extra_lines.append(f"- **类型**: {evidence_type}")
    extra_section = "\n".join(extra_lines)
    if extra_section:
        extra_section = "\n" + extra_section

    content = f"""# {emoji} {title} [{severity}]

- **目标**: {target}
- **严重性**: {severity}
- **时间**: {ts}{extra_section}

## 描述
{desc}

## Payload
```
{payload}
```

## 服务器响应
```
{response[:2000]}
```
"""
    filepath.write_text(content, encoding="utf-8")
    _save_evidence_key(evidence_dir, norm_key, filename)
    _update_evidence_index(evidence_dir)
    _log("evidence_save", f"SAVED {filepath}")
    return [TextContent(type="text", text=f"✅ 证据已保存: {filepath} ({emoji} {severity})")]


# ── evidence_list ─────────────────────────────────────────
async def _handle_evidence_list(args: dict) -> list[TextContent]:
    evidence_dir = _evidence_dir()
    files = sorted(f for f in evidence_dir.glob("*.md") if f.name != "index.md")
    if not files:
        return [TextContent(type="text", text="📭 暂无漏洞证据")]
    lines = [f"📁 已保存 {len(files)} 条证据:\n"]
    for f in files:
        lines.append(f"  - {f}")
    return [TextContent(type="text", text="\n".join(lines))]


# ── evidence_read ─────────────────────────────────────────
async def _handle_evidence_read(args: dict) -> list[TextContent]:
    query = (args.get("id") or "").strip()
    if not query:
        return [TextContent(type="text", text="❌ 请提供证据编号(如 '1')或关键词(如 'SSRF')")]

    # 如果传入的是完整路径，提取文件名作为匹配依据
    if "/" in query:
        query = query.rsplit("/", 1)[-1]
    # 去掉 .md 后缀
    if query.endswith(".md"):
        query = query[:-3]

    evidence_dir = _evidence_dir()
    files = sorted(f for f in evidence_dir.glob("*.md") if f.name != "index.md")
    if not files:
        return [TextContent(type="text", text="📭 暂无证据")]

    matched = None

    # 精确文件名匹配 (带或不带 .md)
    for f in files:
        if f.name == query or f.name == query + ".md" or f.stem == query:
            matched = f
            break

    # 编号匹配 (001, 1)
    if not matched:
        try:
            num = int(query)
            prefix = f"{num:03d}-"
            for f in files:
                if f.name.startswith(prefix):
                    matched = f
                    break
        except ValueError:
            pass

    # 关键词模糊匹配
    if not matched:
        query_lower = query.lower()
        for f in files:
            if query_lower in f.name.lower():
                matched = f
                break

    if not matched:
        names = "\n".join(f"  - {f.name}" for f in files[:20])
        return [TextContent(type="text",
            text=f"❌ 未找到匹配 '{query}' 的证据。可用证据:\n{names}")]

    content = matched.read_text(encoding="utf-8", errors="replace")
    if len(content) > 8000:
        content = content[:8000] + f"\n\n... (截断，共 {len(content)} 字符)"
    return [TextContent(type="text", text=f"📄 {matched.name}\n\n{content}")]


# ── report_vuln helpers ───────────────────────────────────

_VULN_TYPE_MAP = {
    "sqli": ["sql injection", "sqli", "sql注入", "authentication bypass"],
    "xss_reflected": ["reflected xss", "反射型xss", "反射型 xss"],
    "xss_stored": ["stored xss", "存储型xss", "存储型 xss"],
    "xss_dom": ["dom xss", "dom-based xss"],
    "idor": ["idor", "越权", r"unauthorized.*access", r"account.*disclosure", "insecure direct"],
    "csrf": ["csrf", "cross-site request"],
    "ssrf": ["ssrf", "server-side request"],
    "path_traversal": ["path traversal", "路径遍历", "lfi", "rfi", "file inclusion"],
    "info_disclosure": ["information disclosure", "信息泄露", "info leak", "swagger", "stack trace", "directory listing"],
    "business_logic": ["business logic", "业务逻辑", r"zero.*transfer", r"negative.*amount", "金额操纵"],
    "auth_bypass": ["authentication bypass", r"cookie.*manipulation", "认证绕过"],
    "priv_escalation": ["privilege escalation", "权限提升"],
    "weak_cred": ["weak credential", "default credential", "弱口令", "弱密码", "brute force"],
    "missing_header": [r"missing.*header", "security header", "clickjacking"],
    "file_upload": ["file upload", "文件上传", "unrestricted upload", "webshell"],
    "rce": ["remote code", "command injection", "命令注入", "代码执行", "rce"],
    "xxe": ["xxe", "xml external", "xml 外部实体"],
    "deserialization": ["deserialization", "反序列化", "unserialize"],
    "open_redirect": ["open redirect", "url redirect", "开放重定向"],
    "cors": ["cors", "cross-origin"],
}


def _classify_vuln_type(title: str) -> str:
    title_lower = title.lower()
    for vtype, keywords in _VULN_TYPE_MAP.items():
        for kw in keywords:
            if re.search(kw, title_lower):
                return vtype
    return title_lower[:30]


def _normalize_vuln_key(title: str, url: str) -> str:
    """归一化去重 key: vuln_type:endpoint_path?param_names"""
    from urllib.parse import urlparse, parse_qs
    parsed = urlparse(url)
    path = parsed.path.rstrip("/").lower()
    params = sorted(parse_qs(parsed.query).keys())
    param_key = ",".join(params) if params else ""
    vtype = _classify_vuln_type(title)
    key = f"{vtype}:{path}"
    if param_key:
        key += f"?{param_key}"
    return key


def _lookup_evidence(vuln_id: str) -> dict | None:
    """从 evidence 文件中交叉引用漏洞信息"""
    evidence_dir = _evidence_dir()
    for f in evidence_dir.glob("*.md"):
        if f.name == "index.md":
            continue
        if vuln_id in f.name:
            text = f.read_text(encoding="utf-8", errors="replace")
            info = {}
            for line in text.split("\n"):
                if line.startswith("# "):
                    info["title"] = re.sub(r"^#\s*[🔴🟠🟡🟢⚪]\s*", "", line).split("[")[0].strip()
                if "**目标**" in line:
                    info["url"] = line.split(":", 1)[-1].strip() if ":" in line else ""
                if "**严重性**" in line:
                    info["severity"] = line.split(":", 1)[-1].strip().upper() if ":" in line else "MEDIUM"
            if info:
                return info
    return None


_URL_RE = re.compile(r'https?://[^\s<>"\']+')


def _extract_url_from_text(text: str) -> str:
    m = _URL_RE.search(text)
    return m.group(0).rstrip(".,;)") if m else ""


# ── report_vuln ───────────────────────────────────────────
async def _handle_report_vuln(args: dict) -> list[TextContent]:
    _log("report_vuln", f"args={args}")

    vuln_type = (args.get("title") or args.get("name") or args.get("vuln_name") or args.get("vuln_type") or "").strip()
    severity = (args.get("severity") or "MEDIUM").strip().upper()
    url = (args.get("url") or args.get("target") or args.get("endpoint") or args.get("vuln_url") or "").strip()
    detail = (args.get("detail") or args.get("description") or args.get("evidence") or "").strip()
    vuln_param = (args.get("vuln_param") or "").strip()

    # severity 归一化 (fuzzy match)
    _SEV_CANONICAL = ("CRITICAL", "HIGH", "MEDIUM", "LOW")
    if severity not in _SEV_CANONICAL:
        _sev_up = severity.upper()
        _matched = False
        for _canon in _SEV_CANONICAL:
            if _canon in _sev_up or _sev_up.startswith(_canon[:3]):
                severity = _canon
                _matched = True
                break
        if not _matched:
            severity = "MEDIUM"

    # ── 容错层 0: JSON blob 解析 ──
    _raw_vuln = args.get("vuln") or args.get("vulnerability") or ""
    if isinstance(_raw_vuln, str) and _raw_vuln.startswith("{"):
        try:
            blob = json.loads(_raw_vuln)
            if isinstance(blob, dict):
                if not vuln_type:
                    vuln_type = (blob.get("title") or blob.get("name") or blob.get("type") or "").strip()[:120]
                if not url:
                    url = (blob.get("url") or blob.get("endpoint") or blob.get("target") or "").strip()
                if not detail:
                    detail = _raw_vuln
                if severity == "MEDIUM":
                    s = (blob.get("severity") or "").strip().upper()
                    if s in _SEV_CANONICAL:
                        severity = s
                _log("report_vuln", f"parsed JSON blob: title={vuln_type}, url={url}")
        except (ValueError, TypeError):
            if not vuln_type and len(_raw_vuln) > 5:
                vuln_type = _raw_vuln[:80]
    elif isinstance(_raw_vuln, str) and _raw_vuln and not vuln_type:
        vuln_type = _raw_vuln[:80]

    # ── 容错层 1: vuln_id 交叉引用 evidence ──
    if (not vuln_type or vuln_type == "Unknown") and args.get("vuln_id"):
        ev_info = _lookup_evidence(str(args["vuln_id"]))
        if ev_info:
            vuln_type = vuln_type if (vuln_type and vuln_type != "Unknown") else ev_info.get("title", "")
            if not url:
                url = ev_info.get("url", "")
            if not detail:
                detail = ev_info.get("description", "")
            if severity == "MEDIUM" and ev_info.get("severity"):
                severity = ev_info["severity"]

    # ── 容错层 2: 从 description 降级提取 title ──
    if not vuln_type and detail:
        first_line = detail.split('\n')[0].split('。')[0].split('. ')[0][:80]
        if len(first_line) > 5:
            vuln_type = first_line

    # ── 容错层 3: 从所有参数值中智能提取 URL ──
    if not url:
        for key, val in args.items():
            if key in ("title", "severity", "detail", "vuln_id"):
                continue
            if isinstance(val, str):
                found = _extract_url_from_text(val)
                if found:
                    url = found
                    break
    if not url and detail:
        found = _extract_url_from_text(detail)
        if found:
            url = found

    # ── 容错层 4: 最终 fallback ──
    vuln_type = vuln_type or "Unknown"
    url = url or "unknown"

    if vuln_type == "Unknown" and url == "unknown":
        _log("report_vuln", f"rejected: no usable fields in {list(args.keys())}")
        return [TextContent(type="text",
            text="❌ 缺少必要参数。正确用法: report_vuln(title='漏洞名', severity='HIGH', url='http://target/path')\n"
                 "可选: detail='详细描述'")]

    full_detail = detail
    if vuln_param:
        full_detail = f"参数: {vuln_param}\n{detail}" if detail else f"参数: {vuln_param}"

    # 去重（归一化 key）
    norm_key = _normalize_vuln_key(vuln_type, url)
    vulns_path = Path(OUTPUT_DIR) / "vulns.json"
    vulns = _read_vulns_json(vulns_path)

    for v in vulns:
        existing_key = _normalize_vuln_key(v.get("title", ""), v.get("url", ""))
        if existing_key == norm_key:
            # 合并 detail
            if full_detail and full_detail not in (v.get("detail") or ""):
                v["detail"] = (v.get("detail") or "") + f"\n\n[合并] {vuln_type}:\n{full_detail}"
                _write_vulns_json(vulns_path, vulns)
            _log("report_vuln", f"dedup merge: {norm_key}")
            return [TextContent(type="text", text=f"⚠️ 漏洞已记录 (合并): {vuln_type} @ {url}")]

    vuln = {
        "id": len(vulns) + 1,
        "title": vuln_type,
        "type": vuln_type,
        "severity": severity,
        "url": url,
        "detail": full_detail,
        "time": datetime.now().isoformat(),
    }
    vulns.append(vuln)
    _write_vulns_json(vulns_path, vulns)

    # 同时保存 evidence
    evidence_dir = _evidence_dir()
    eid = _next_evidence_id()
    ts = datetime.now().strftime("%Y%m%d-%H%M%S")
    safe_name = vuln_type.replace(" ", "_").replace("/", "_")[:50]
    filename = f"{eid:03d}-{safe_name}.md"
    filepath = evidence_dir / filename

    emoji = {"CRITICAL": "🔴", "HIGH": "🟠", "MEDIUM": "🟡", "LOW": "🟢"}.get(severity, "⚪")

    content = f"""# {emoji} {vuln_type} [{severity}]

- **目标**: {url}
- **严重性**: {severity}
- **时间**: {ts}

## 描述
{full_detail or '(无详细描述)'}
"""
    filepath.write_text(content, encoding="utf-8")
    _update_evidence_index(evidence_dir)
    _log("report_vuln", f"SAVED vuln #{vuln['id']}: {vuln_type}")
    return [TextContent(type="text", text=f"{emoji} 漏洞 #{vuln['id']} 已记录: {vuln_type} [{severity}] @ {url}")]


# ── list_vulns ────────────────────────────────────────────
async def _handle_list_vulns(args: dict) -> list[TextContent]:
    vulns_path = Path(OUTPUT_DIR) / "vulns.json"
    vulns = _read_vulns_json(vulns_path)
    if not vulns:
        return [TextContent(type="text", text="📭 暂未发现漏洞")]

    severity_order = {"CRITICAL": 0, "HIGH": 1, "MEDIUM": 2, "LOW": 3}
    sorted_vulns = sorted(vulns, key=lambda v: severity_order.get(v.get("severity", ""), 9))

    lines = [f"📊 已发现 {len(vulns)} 个漏洞:\n"]
    lines.append("| # | 严重性 | 类型 | URL |")
    lines.append("|---|--------|------|-----|")
    for v in sorted_vulns:
        emoji = {"CRITICAL": "🔴", "HIGH": "🟠", "MEDIUM": "🟡", "LOW": "🟢"}.get(v.get("severity", ""), "⚪")
        lines.append(f"| {v['id']} | {emoji} {v['severity']} | {v['type']} | {v['url']} |")
    return [TextContent(type="text", text="\n".join(lines))]


# ── check_scope ───────────────────────────────────────────
async def _handle_check_scope(args: dict) -> list[TextContent]:
    from urllib.parse import urlparse
    import ipaddress

    target = args.get("target", "").strip() or os.environ.get("TCHKILLER_TARGET", "").strip()
    scope = [s.strip() for s in SCOPE.split(",") if s.strip()]

    if not scope:
        return [TextContent(type="text", text="⚠️ 未配置授权范围")]

    def normalize(t):
        if "://" in t:
            return urlparse(t).hostname or t
        if ":" in t and not t.startswith("["):
            return t.split(":")[0]
        return t

    host = normalize(target)
    for entry in scope:
        entry_host = normalize(entry)
        if host == entry_host or host.endswith(f".{entry_host}"):
            return [TextContent(type="text", text=f"✅ {target} 在授权范围内")]
        try:
            if ipaddress.ip_address(host) in ipaddress.ip_network(entry, strict=False):
                return [TextContent(type="text", text=f"✅ {target} 在授权范围内")]
        except ValueError:
            pass

    return [TextContent(type="text", text=f"❌ {target} 超出授权范围！授权范围: {', '.join(scope)}")]


# ── get_targets ───────────────────────────────────────────
async def _handle_get_targets(args: dict) -> list[TextContent]:
    scope = [s.strip() for s in SCOPE.split(",") if s.strip()]
    mode = os.environ.get("TCHKILLER_MODE", "auto")
    lines = [f"🎯 主目标: {TARGET}", f"📋 模式: {mode}", "🔒 授权范围:"]
    for s in scope:
        lines.append(f"  - {s}")
    return [TextContent(type="text", text="\n".join(lines))]


# ── list_skills ───────────────────────────────────────────
def _get_skills_dir() -> Path:
    candidates = [
        Path(OUTPUT_DIR) / ".claude" / "skills",
        Path(__file__).parent / ".claude" / "skills",
    ]
    for d in candidates:
        if d.exists():
            return d
    return candidates[0]


def _load_frontmatter(skill_md: Path) -> dict:
    text = skill_md.read_text(encoding="utf-8")
    if not text.startswith("---"):
        return {"name": skill_md.parent.name, "content": text}
    parts = text.split("---", 2)
    if len(parts) < 3:
        return {"name": skill_md.parent.name, "content": text}
    meta = {}
    for line in parts[1].strip().split("\n"):
        if ":" in line:
            key, val = line.split(":", 1)
            meta[key.strip()] = val.strip().strip('"').strip("'")
    return {
        "name": meta.get("name", skill_md.parent.name),
        "description": meta.get("description", ""),
        "tags": meta.get("tags", ""),
    }


async def _handle_list_skills(args: dict) -> list[TextContent]:
    keyword = (args.get("keyword") or "").strip().lower()
    skills_dir = _get_skills_dir()

    if not skills_dir.exists():
        return [TextContent(type="text", text=f"❌ Skills 目录不存在: {skills_dir}")]

    results = []
    for skill_path in sorted(skills_dir.iterdir()):
        skill_md = skill_path / "SKILL.md" if skill_path.is_dir() else None
        if not skill_md or not skill_md.exists():
            continue
        meta = _load_frontmatter(skill_md)
        skill_id = skill_path.name
        if keyword:
            searchable = f"{skill_id} {meta.get('name','')} {meta.get('description','')} {meta.get('tags','')}".lower()
            if keyword not in searchable:
                continue
        desc = meta.get("description", "")
        if len(desc) > 80:
            desc = desc[:80] + "..."
        results.append(f"  - **{skill_id}**: {desc}")

    if not results:
        msg = f"📭 未找到匹配 '{keyword}' 的技巧" if keyword else "📭 无可用技巧"
        return [TextContent(type="text", text=msg)]

    header = f"🔍 找到 {len(results)} 个技巧" + (f" (关键词: {keyword})" if keyword else "") + ":\n"
    return [TextContent(type="text", text=header + "\n".join(results))]


# ── read_skill ────────────────────────────────────────────
async def _handle_read_skill(args: dict) -> list[TextContent]:
    skill_id = (args.get("id") or "").strip()
    if not skill_id:
        return [TextContent(type="text", text="❌ 请提供技巧 ID，用 list_skills 查看可用列表")]

    skills_dir = _get_skills_dir()
    skill_md = skills_dir / skill_id / "SKILL.md"
    if not skill_md.exists():
        return [TextContent(type="text", text=f"❌ 技巧 '{skill_id}' 不存在。用 list_skills 搜索可用技巧")]

    content = skill_md.read_text(encoding="utf-8")
    if len(content) > 8000:
        content = content[:8000] + f"\n\n... (内容已截断，共 {len(content)} 字符)"
    return [TextContent(type="text", text=f"📖 技巧: {skill_id}\n\n{content}")]


# ── 辅助函数 ──────────────────────────────────────────────
def _read_vulns_json(path: Path) -> list:
    if not path.exists():
        return []
    try:
        fd = os.open(str(path), os.O_RDONLY)
        try:
            fcntl.flock(fd, fcntl.LOCK_SH)
            data = path.read_text(encoding="utf-8")
            return json.loads(data) if data.strip() else []
        finally:
            fcntl.flock(fd, fcntl.LOCK_UN)
            os.close(fd)
    except Exception:
        return []


def _write_vulns_json(path: Path, vulns: list):
    path.parent.mkdir(parents=True, exist_ok=True)
    fd = os.open(str(path), os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o644)
    try:
        fcntl.flock(fd, fcntl.LOCK_EX)
        os.write(fd, json.dumps(vulns, indent=2, ensure_ascii=False).encode("utf-8"))
    finally:
        fcntl.flock(fd, fcntl.LOCK_UN)
        os.close(fd)


def _update_evidence_index(evidence_dir: Path):
    files = sorted(f for f in evidence_dir.glob("*.md") if f.name != "index.md")
    lines = ["# 漏洞证据索引\n", f"生成时间: {datetime.now().isoformat()}\n"]
    lines.append(f"共发现 **{len(files)}** 个漏洞\n")
    lines.append("| # | 文件 | 类型 |")
    lines.append("|---|------|------|")
    for i, f in enumerate(files, 1):
        name = f.stem.split("-", 1)[-1].replace("_", " ") if "-" in f.stem else f.stem
        lines.append(f"| {i} | [{f.name}](./{f.name}) | {name} |")
    (evidence_dir / "index.md").write_text("\n".join(lines) + "\n", encoding="utf-8")


# ── vulndb 辅助函数 ──────────────────────────────────────
def _get_vulndb():
    """获取 vulndb SQLite 连接，不存在则尝试构建"""
    import sqlite3
    if not VULNDB_SQLITE.exists():
        if not VULNDB_DIR.exists():
            return None
        try:
            _build_vulndb_index()
        except Exception:
            return None
    if not VULNDB_SQLITE.exists():
        return None
    return sqlite3.connect(str(VULNDB_SQLITE))


def _build_vulndb_index():
    """从 Markdown 构建 FTS5 索引"""
    import sqlite3
    import re
    conn = sqlite3.connect(str(VULNDB_SQLITE))
    conn.execute("DROP TABLE IF EXISTS vulns")
    conn.execute("DROP TABLE IF EXISTS vulns_fts")
    conn.execute("""
        CREATE TABLE vulns (
            id TEXT PRIMARY KEY, title TEXT, product TEXT, vendor TEXT,
            version_affected TEXT, severity TEXT, tags TEXT, fingerprint TEXT,
            filepath TEXT, content TEXT
        )
    """)
    conn.execute("""
        CREATE VIRTUAL TABLE vulns_fts USING fts5(
            title, product, tags, fingerprint, content,
            tokenize='unicode61'
        )
    """)
    count = 0
    for md_file in sorted(VULNDB_DIR.rglob("*.md")):
        raw = md_file.read_text(encoding="utf-8", errors="replace")
        if not raw.startswith("---"):
            continue
        parts = raw.split("---", 2)
        if len(parts) < 3:
            continue
        meta = {}
        for line in parts[1].strip().split("\n"):
            if ":" in line:
                key, val = line.split(":", 1)
                meta[key.strip()] = val.strip().strip('"').strip("'")
        vuln_id = meta.get("id", md_file.stem)
        title = meta.get("title", md_file.stem)
        product = meta.get("product", md_file.parent.name)
        tags_raw = meta.get("tags", "")
        tags = re.sub(r"[\[\]]", "", tags_raw)
        fingerprint_raw = meta.get("fingerprint", "")
        fingerprint = re.sub(r"[\[\]\"]", "", fingerprint_raw)
        rel_path = str(md_file.relative_to(VULNDB_DIR))
        body = parts[2].strip()
        conn.execute(
            "INSERT OR REPLACE INTO vulns VALUES (?,?,?,?,?,?,?,?,?,?)",
            (vuln_id, title, product, meta.get("vendor", ""),
             meta.get("version_affected", ""), meta.get("severity", "MEDIUM"),
             tags, fingerprint, rel_path, body),
        )
        conn.execute(
            "INSERT INTO vulns_fts (rowid, title, product, tags, fingerprint, content) "
            "VALUES ((SELECT rowid FROM vulns WHERE id=?), ?, ?, ?, ?, ?)",
            (vuln_id, title, product, tags, fingerprint, body),
        )
        count += 1
    conn.commit()
    conn.close()
    _log("vulndb", f"Index built: {count} entries")
    return count


# ── locate_tool ───────────────────────────────────────────
_TOOL_SEARCH_DIRS = ["/pentest", "/usr/local/bin", "/usr/bin", "/usr/sbin",
                     str(Path.home() / ".local" / "bin")]

async def _handle_locate_tool(args: dict) -> list[TextContent]:
    tool_name = (args.get("name") or args.get("tool_name")
                 or args.get("tool") or "").strip()
    if not tool_name:
        return [TextContent(type="text", text="❌ 请提供工具名称（name 参数）")]

    _log("locate_tool", f"searching: {tool_name}")
    results = []

    # 1. which (PATH 搜索)
    import shutil as _shutil
    which_path = _shutil.which(tool_name)
    if which_path:
        results.append(f"PATH: {which_path}")

    # 2. 白名单目录搜索
    import subprocess
    for search_dir in _TOOL_SEARCH_DIRS:
        if not Path(search_dir).is_dir():
            continue
        try:
            proc = subprocess.run(
                ["find", search_dir, "-maxdepth", "3",
                 "-iname", f"*{tool_name}*", "-type", "f"],
                capture_output=True, text=True, timeout=5
            )
            for line in proc.stdout.strip().split("\n"):
                if line and line not in results:
                    results.append(line)
        except Exception:
            pass

    if results:
        _log("locate_tool", f"found {len(results)} results")
        return [TextContent(type="text",
                text=f"🔍 工具 '{tool_name}' 搜索结果:\n" + "\n".join(f"  - {r}" for r in results[:15]))]
    else:
        _log("locate_tool", "not found")
        return [TextContent(type="text",
                text=f"❌ 未找到工具 '{tool_name}'。搜索范围: PATH, {', '.join(_TOOL_SEARCH_DIRS)}")]


# ── search_vulndb ────────────────────────────────────────
async def _handle_search_vulndb(args: dict) -> list[TextContent]:
    query = (args.get("query") or "").strip()
    product = (args.get("product") or "").strip()
    severity = (args.get("severity") or "").strip().upper()

    _log("search_vulndb", f"query={query} product={product} severity={severity}")

    if not query and not product:
        return [TextContent(type="text", text="❌ 请提供搜索关键词（query）或产品名（product）")]

    conn = _get_vulndb()
    if not conn:
        return _vulndb_fallback_search(query, product)

    try:
        results = []
        if query:
            try:
                c = conn.execute(
                    "SELECT v.id, v.title, v.product, v.severity, v.version_affected, v.filepath "
                    "FROM vulns_fts fts JOIN vulns v ON v.rowid = fts.rowid "
                    "WHERE vulns_fts MATCH ? ORDER BY rank LIMIT 20",
                    (query,),
                )
                results = c.fetchall()
            except Exception:
                pass

        if not results and product:
            c = conn.execute(
                "SELECT id, title, product, severity, version_affected, filepath "
                "FROM vulns WHERE product LIKE ? ORDER BY severity",
                (f"%{product}%",),
            )
            results = c.fetchall()

        if not results and query:
            like = f"%{query}%"
            c = conn.execute(
                "SELECT id, title, product, severity, version_affected, filepath FROM vulns "
                "WHERE id LIKE ? OR title LIKE ? OR product LIKE ? OR tags LIKE ? "
                "OR fingerprint LIKE ? OR version_affected LIKE ? "
                "ORDER BY severity LIMIT 20",
                (like, like, like, like, like, like),
            )
            results = c.fetchall()

        # 如果完整查询无匹配，按空格拆分后选最具特异性的词搜索
        if not results and query and " " in query:
            words = [w.strip() for w in query.split() if len(w.strip()) >= 2]
            # 统计每个词的匹配数，选匹配最少但 >0 的词（最具特异性）
            word_hits = []
            for word in words:
                wlike = f"%{word}%"
                c = conn.execute(
                    "SELECT COUNT(*) FROM vulns "
                    "WHERE id LIKE ? OR title LIKE ? OR product LIKE ? OR tags LIKE ? "
                    "OR fingerprint LIKE ? OR version_affected LIKE ?",
                    (wlike, wlike, wlike, wlike, wlike, wlike),
                )
                cnt = c.fetchone()[0]
                if cnt > 0:
                    word_hits.append((cnt, word))
            # 按匹配数升序（最少匹配 = 最特异的词优先）
            word_hits.sort()
            for _, word in word_hits:
                wlike = f"%{word}%"
                c = conn.execute(
                    "SELECT id, title, product, severity, version_affected, filepath FROM vulns "
                    "WHERE id LIKE ? OR title LIKE ? OR product LIKE ? OR tags LIKE ? "
                    "OR fingerprint LIKE ? OR version_affected LIKE ? "
                    "ORDER BY severity LIMIT 20",
                    (wlike, wlike, wlike, wlike, wlike, wlike),
                )
                results = c.fetchall()
                if results:
                    break

        conn.close()

        if severity:
            results = [r for r in results if r[3] == severity]

        if not results:
            return [TextContent(type="text", text=f"📭 未找到匹配 '{query or product}' 的漏洞。知识库可能未收录。")]

        lines = [f"🔍 找到 {len(results)} 条漏洞:\n"]
        for r in results:
            vuln_id, title, prod, sev, ver, fp = r
            sev_icon = {"CRITICAL": "🔴", "HIGH": "🟠", "MEDIUM": "🟡", "LOW": "🟢"}.get(sev, "⚪")
            lines.append(f"  {sev_icon} **{vuln_id}** — {title}")
            lines.append(f"    产品: {prod} | 影响: {ver} | 严重性: {sev}")
            lines.append(f"    → 调用 `read_vuln(id=\"{vuln_id}\")` 获取完整利用步骤和 payload")
            lines.append("")
        return [TextContent(type="text", text="\n".join(lines))]

    except Exception as e:
        conn.close()
        _log("search_vulndb", f"ERROR: {e}")
        return _vulndb_fallback_search(query, product)


def _vulndb_fallback_search(query: str, product: str) -> list[TextContent]:
    """无 SQLite 时降级为文件遍历搜索"""
    if not VULNDB_DIR.exists():
        return [TextContent(type="text", text="❌ 漏洞知识库目录不存在")]
    query_lower = (query or "").lower()
    product_lower = (product or "").lower()
    results = []
    for md_file in sorted(VULNDB_DIR.rglob("*.md")):
        content = md_file.read_text(encoding="utf-8", errors="replace").lower()
        if product_lower and product_lower not in str(md_file).lower():
            continue
        if query_lower and query_lower not in content and query_lower not in md_file.stem.lower():
            continue
        title = md_file.stem
        for line in md_file.read_text(encoding="utf-8", errors="replace").split("\n"):
            if line.startswith("title:"):
                title = line.split(":", 1)[1].strip().strip('"')
                break
        results.append((md_file.stem, title))
    if not results:
        return [TextContent(type="text", text=f"📭 未找到匹配 '{query or product}' 的漏洞（降级搜索）")]
    lines = [f"🔍 找到 {len(results)} 条漏洞（降级搜索）:\n"]
    for vuln_id, title in results[:20]:
        lines.append(f"  - **{vuln_id}** — {title}")
        lines.append(f"    → 调用 `read_vuln(id=\"{vuln_id}\")` 获取详情")
    return [TextContent(type="text", text="\n".join(lines))]


# ── read_vuln ────────────────────────────────────────────
async def _handle_read_vuln(args: dict) -> list[TextContent]:
    vuln_id = (args.get("id") or "").strip()
    _log("read_vuln", f"id={vuln_id}")
    if not vuln_id:
        return [TextContent(type="text", text="❌ 请提供漏洞 ID，先用 search_vulndb 搜索")]

    filepath = None
    conn = _get_vulndb()
    if conn:
        try:
            c = conn.execute("SELECT filepath FROM vulns WHERE id = ?", (vuln_id,))
            row = c.fetchone()
            if row:
                filepath = VULNDB_DIR / row[0]
            conn.close()
        except Exception:
            conn.close()

    if not filepath or not filepath.exists():
        vuln_id_lower = vuln_id.lower()
        for md_file in VULNDB_DIR.rglob("*.md"):
            if vuln_id_lower in md_file.stem.lower():
                filepath = md_file
                break

    if not filepath or not filepath.exists():
        return [TextContent(type="text", text=f"❌ 漏洞 '{vuln_id}' 不存在。用 search_vulndb 搜索可用漏洞")]

    content = filepath.read_text(encoding="utf-8", errors="replace")
    if content.startswith("---"):
        parts = content.split("---", 2)
        if len(parts) >= 3:
            content = parts[2].strip()

    if len(content) > 10000:
        content = content[:10000] + f"\n\n... (内容已截断，共 {len(content)} 字符)"

    return [TextContent(type="text", text=f"📋 漏洞详情: {vuln_id}\n\n{content}")]


# ── flag 自动捕获 ─────────────────────────────────────────

_FLAG_RE = re.compile(r'flag\{[a-zA-Z0-9_\-]+\}', re.IGNORECASE)
_AUTO_SAVED_FLAGS: set[str] = set()  # 去重：同一 flag 只保存一次

def _auto_save_flag(resp_text: str, url: str):
    """从 HTTP 响应中自动提取 flag 并保存到 evidence（防止 agent 遗漏）"""
    flags = _FLAG_RE.findall(resp_text)
    if not flags:
        return
    for flag in flags:
        if flag in _AUTO_SAVED_FLAGS:
            continue
        _AUTO_SAVED_FLAGS.add(flag)
        try:
            evidence_dir = _evidence_dir()
            ts = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"auto_flag_{ts}.md"
            filepath = evidence_dir / filename
            content = f"""# 🚩 自动捕获 Flag

- **Flag**: `{flag}`
- **来源 URL**: `{url}`
- **捕获时间**: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}

> 此 flag 由 http_request 工具自动检测并保存。
"""
            filepath.write_text(content, encoding="utf-8")
            _log("auto_flag", f"AUTO-SAVED flag from {url}: {flag} → {filepath}")
        except Exception as e:
            _log("auto_flag", f"Failed to auto-save flag: {e}")


# ── http_request ──────────────────────────────────────────

# 全局 Cookie 管理器 — 按 host:port 隔离，跨请求自动保持登录态
_cookie_jars: dict[str, "http.cookiejar.CookieJar"] = {}

def _get_cookie_jar(url: str, new_session: bool = False) -> "http.cookiejar.CookieJar":
    """获取或创建 host:port 粒度的 CookieJar"""
    import http.cookiejar
    from urllib.parse import urlparse
    parsed = urlparse(url)
    port = parsed.port or (443 if parsed.scheme == "https" else 80)
    key = f"{parsed.hostname}:{port}"
    if new_session or key not in _cookie_jars:
        _cookie_jars[key] = http.cookiejar.CookieJar()
    return _cookie_jars[key]


async def _handle_http_request(args: dict) -> list[TextContent]:
    """发送 HTTP 请求 — 自动管理 Cookie/Session"""
    import urllib.request
    from urllib.request import Request
    from urllib.error import HTTPError, URLError
    import ssl

    url = args.get("url", "")
    if not url:
        return [TextContent(type="text", text="❌ 缺少 url 参数")]

    method = args.get("method", "GET").upper()
    headers = args.get("headers", {})
    body = args.get("body")
    timeout = args.get("timeout", 15)
    follow_redirects = args.get("follow_redirects", True)
    new_session = args.get("new_session", False)

    cookie_jar = _get_cookie_jar(url, new_session)

    # 构建请求
    data = body.encode("utf-8") if body else None
    req = Request(url, data=data, method=method)
    for k, v in headers.items():
        req.add_header(k, v)
    if data and "Content-Type" not in headers and "content-type" not in headers:
        req.add_header("Content-Type", "application/x-www-form-urlencoded")

    # 忽略 SSL 验证
    ctx = ssl.create_default_context()
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE

    # 构建 opener（带 cookie 管理 + 可选禁用重定向）
    handlers: list = [
        urllib.request.HTTPCookieProcessor(cookie_jar),
        urllib.request.HTTPSHandler(context=ctx),
    ]
    if not follow_redirects:
        class NoRedirectHandler(urllib.request.HTTPRedirectHandler):
            def redirect_request(self, req, fp, code, msg, headers, newurl):
                return None
        handlers.append(NoRedirectHandler)
    opener = urllib.request.build_opener(*handlers)

    try:
        resp = opener.open(req, timeout=timeout)
        status = resp.status
        resp_headers = dict(resp.headers)
        resp_body = resp.read()
        try:
            resp_text = resp_body.decode("utf-8", errors="replace")
        except Exception:
            resp_text = repr(resp_body[:2000])

    except HTTPError as e:
        status = e.code
        resp_headers = dict(e.headers) if e.headers else {}
        try:
            resp_text = e.read().decode("utf-8", errors="replace")
        except Exception:
            resp_text = str(e)
    except URLError as e:
        return [TextContent(type="text", text=f"❌ 连接失败: {e.reason}")]
    except Exception as e:
        return [TextContent(type="text", text=f"❌ 请求错误: {e}")]

    # 截断过长响应
    body_len = len(resp_text)
    if body_len > 8000:
        resp_text = resp_text[:8000] + f"\n... (截断, 共 {body_len} 字符)"

    # 格式化重要响应头
    header_lines = []
    for k in ["Content-Type", "Location", "Set-Cookie", "Server", "X-Powered-By"]:
        if k in resp_headers:
            header_lines.append(f"  {k}: {resp_headers[k]}")

    # 摘要行
    ctype = resp_headers.get("Content-Type", "")
    has_flag = "flag{" in resp_text.lower() or "flag-{" in resp_text.lower()
    summary = f"HTTP {status} | {body_len} bytes | {ctype.split(';')[0]}"
    if has_flag:
        summary += " | ⚠️ CONTAINS FLAG"
        # 自动提取并保存 flag 到 evidence（防止 agent 遗漏）
        _auto_save_flag(resp_text, url)

    # 显示当前 session cookie 状态（帮助 agent 感知登录态）
    session_cookies = [c for c in cookie_jar if "sess" in c.name.lower() or "token" in c.name.lower()]
    if session_cookies:
        cookie_info = ", ".join(f"{c.name}={c.value[:12]}..." for c in session_cookies[:3])
        summary += f" | 🍪 {cookie_info}"

    result = summary + "\n"
    if header_lines:
        result += "\n".join(header_lines) + "\n"
    result += f"\n{resp_text}"

    return [TextContent(type="text", text=result)]

# ── interactive_session (tmux) ────────────────────────────

import shutil
import subprocess
import time

_TMUX_BIN = shutil.which("tmux")
_TMUX_PREFIX = "tck_"  # session name prefix to avoid conflicts
_ACTIVE_SESSIONS: dict[str, dict] = {}  # name -> {created, last_read, command}

def _tmux(*args: str, timeout: int = 5) -> tuple[int, str]:
    """Run a tmux command synchronously. Returns (returncode, stdout)."""
    if not _TMUX_BIN:
        return (1, "tmux not found in PATH")
    cmd = [_TMUX_BIN] + list(args)
    try:
        r = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
        return (r.returncode, r.stdout + r.stderr)
    except subprocess.TimeoutExpired:
        return (1, "tmux command timed out")
    except Exception as e:
        return (1, str(e))


def _session_name(name: str) -> str:
    """Ensure session name has our prefix."""
    if not name.startswith(_TMUX_PREFIX):
        return _TMUX_PREFIX + name
    return name


def _clean_ansi(text: str) -> str:
    """Strip ANSI escape sequences from captured pane output."""
    return re.sub(r'\x1b\[[0-9;]*[a-zA-Z]', '', text)


async def _handle_interactive_session(args: dict) -> list[TextContent]:
    action = args.get("action", "").lower()
    name = args.get("session_name", "default")
    sess = _session_name(name)

    if action == "start":
        command = args.get("command", "")
        if not command:
            return [TextContent(type="text", text="❌ 'command' is required for start")]

        # Kill existing session with same name if any
        _tmux("kill-session", "-t", sess)
        time.sleep(0.2)

        # Start new detached tmux session running the command
        rc, out = _tmux("new-session", "-d", "-s", sess, "-x", "200", "-y", "50", command)
        if rc != 0:
            return [TextContent(type="text", text=f"❌ Failed to start session: {out}")]

        _ACTIVE_SESSIONS[name] = {
            "created": datetime.now().isoformat(),
            "command": command,
            "tmux_name": sess,
        }
        _log("interactive_session", f"START session={sess} command={command}")

        # Wait a moment for the program to initialize, then capture initial output
        wait = min(args.get("wait", 2), 30)
        time.sleep(wait)

        rc, output = _tmux("capture-pane", "-t", sess, "-p", "-S", "-100")
        output = _clean_ansi(output).strip()

        return [TextContent(type="text", text=f"✅ Session '{name}' started\n$ {command}\n\n{output}")]

    elif action == "send":
        command = args.get("command", "")
        if not command:
            return [TextContent(type="text", text="❌ 'command' is required for send")]

        # Verify session exists
        rc, _ = _tmux("has-session", "-t", sess)
        if rc != 0:
            return [TextContent(type="text", text=f"❌ Session '{name}' not found. Use action='start' first.")]

        # Send keys
        _tmux("send-keys", "-t", sess, command, "Enter")
        _log("interactive_session", f"SEND session={sess} command={command}")

        # Wait for output
        wait = min(args.get("wait", 2), 30)
        time.sleep(wait)

        # Capture pane content
        rc, output = _tmux("capture-pane", "-t", sess, "-p", "-S", "-100")
        output = _clean_ansi(output).strip()

        return [TextContent(type="text", text=output)]

    elif action == "read":
        rc, _ = _tmux("has-session", "-t", sess)
        if rc != 0:
            return [TextContent(type="text", text=f"❌ Session '{name}' not found.")]

        lines = args.get("lines", 100)
        rc, output = _tmux("capture-pane", "-t", sess, "-p", "-S", f"-{lines}")
        output = _clean_ansi(output).strip()

        return [TextContent(type="text", text=output if output else "(empty — no output yet)")]

    elif action == "send_raw":
        # Send raw keys without Enter (for Ctrl-C, special keys, etc.)
        keys = args.get("command", "")
        if not keys:
            return [TextContent(type="text", text="❌ 'command' is required for send_raw")]

        rc, _ = _tmux("has-session", "-t", sess)
        if rc != 0:
            return [TextContent(type="text", text=f"❌ Session '{name}' not found.")]

        _tmux("send-keys", "-t", sess, keys)
        _log("interactive_session", f"SEND_RAW session={sess} keys={keys}")

        wait = min(args.get("wait", 1), 30)
        time.sleep(wait)

        rc, output = _tmux("capture-pane", "-t", sess, "-p", "-S", "-100")
        output = _clean_ansi(output).strip()

        return [TextContent(type="text", text=output)]

    elif action == "close":
        rc, out = _tmux("kill-session", "-t", sess)
        _ACTIVE_SESSIONS.pop(name, None)
        _log("interactive_session", f"CLOSE session={sess}")
        if rc != 0:
            return [TextContent(type="text", text=f"⚠️ Session '{name}' may not exist: {out}")]
        return [TextContent(type="text", text=f"✅ Session '{name}' closed")]

    elif action == "list":
        rc, out = _tmux("list-sessions", "-F", "#{session_name} #{session_created} #{pane_current_command}")
        if rc != 0:
            return [TextContent(type="text", text="No active tmux sessions")]

        # Filter to our sessions only
        lines = []
        for line in out.strip().split("\n"):
            if line.startswith(_TMUX_PREFIX):
                display_name = line.replace(_TMUX_PREFIX, "", 1)
                lines.append(f"  • {display_name}")

        if not lines:
            return [TextContent(type="text", text="No active interactive sessions")]
        return [TextContent(type="text", text="Active sessions:\n" + "\n".join(lines))]

    else:
        return [TextContent(type="text", text=f"❌ Unknown action '{action}'. Use: start|send|send_raw|read|close|list")]


# ── Main ──────────────────────────────────────────────────
async def main():
    from mcp.server.stdio import stdio_server
    from mcp.server import InitializationOptions
    from mcp.types import ServerCapabilities, ToolsCapability

    _log("server", f"Starting STDIO MCP server (OUTPUT_DIR={OUTPUT_DIR}, TARGET={TARGET})")

    init_options = InitializationOptions(
        server_name="tchkiller-tools",
        server_version="0.7.0",
        capabilities=ServerCapabilities(
            tools=ToolsCapability(),
        ),
    )

    async with stdio_server() as (read_stream, write_stream):
        await server.run(read_stream, write_stream, init_options)


if __name__ == "__main__":
    asyncio.run(main())