#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
2520 PC端主控服务
架构: Browser → PC:7885 (Run_2520.py) → 7884_brain.exe (本地子进程)
                                    → Jetson:7886 (jetson_agent.py via HTTP)
功能:
  1. 直接调用7884_brain.exe进行结构归纳分析和样本生成(本地子进程)
  2. 通过HTTP远程管理Jetson设备上的jetson_agent.py
  3. 通过HTTP上传推送样本文件到Jetson
  4. 提供Web UI界面
"""

import os
import sys
import json
import uuid
import time
import glob
import subprocess
import threading
import webbrowser
from http.server import ThreadingHTTPServer, BaseHTTPRequestHandler
from urllib.parse import urlparse, parse_qs
from urllib.request import Request, urlopen
from urllib.error import URLError, HTTPError
from datetime import datetime

# ==================== 路径与常量 ====================

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
CONFIG_FILE = os.path.join(SCRIPT_DIR, "config.json")
INDEX_HTML = os.path.join(SCRIPT_DIR, "web", "index.html")
RECORDS_DIR = os.path.join(SCRIPT_DIR, "records")

# 确保记录目录存在
os.makedirs(RECORDS_DIR, exist_ok=True)

# 默认配置
DEFAULT_CONFIG = {
    "pc_port": 7885,
    "brain_exe": "../7884_brain.exe",
    "jetson_host": "192.168.1.108",
    "jetson_port": 7886,
    "jetson_user": "nvidia",
    "fuzz_timeout": 30,
    "workspace": "./fuzz_workspace",
    "crash_dir": "./crash_records"
}

# ==================== 全局状态 ====================

config = {}
config_lock = threading.Lock()

# 7884引擎任务队列
task_queue = []
task_lock = threading.Lock()
running_processes = {}
process_lock = threading.Lock()


# ==================== 配置管理 ====================

def load_config():
    """从config.json加载配置，不存在则使用默认值"""
    global config
    if os.path.exists(CONFIG_FILE):
        try:
            with open(CONFIG_FILE, "r", encoding="utf-8") as f:
                saved = json.load(f)
            merged = dict(DEFAULT_CONFIG)
            merged.update(saved)
            config = merged
        except Exception as e:
            print(f"[配置] 加载配置失败，使用默认值: {e}")
            config = dict(DEFAULT_CONFIG)
    else:
        config = dict(DEFAULT_CONFIG)
    # 确保目录存在
    _ensure_dirs()
    # 保存配置（首次会创建文件）
    save_config()


def save_config():
    """保存当前配置到config.json"""
    with config_lock:
        try:
            with open(CONFIG_FILE, "w", encoding="utf-8") as f:
                json.dump(config, f, ensure_ascii=False, indent=2)
        except Exception as e:
            print(f"[配置] 保存配置失败: {e}")


def _ensure_dirs():
    """确保工作目录和崩溃记录目录存在"""
    workspace = config.get("workspace", "./fuzz_workspace")
    crash_dir = config.get("crash_dir", "./crash_records")
    os.makedirs(workspace, exist_ok=True)
    os.makedirs(crash_dir, exist_ok=True)


def get_brain_exe_path():
    """获取7884_brain.exe的绝对路径"""
    brain = config.get("brain_exe", "../7884_brain.exe")
    if not os.path.isabs(brain):
        return os.path.normpath(os.path.join(SCRIPT_DIR, brain))
    return brain


# ==================== 记录管理 ====================

def generate_task_id():
    """生成任务ID"""
    return datetime.now().strftime("%Y%m%d_%H%M%S_") + uuid.uuid4().hex[:6]


def save_record(task_id, record_data):
    """保存任务记录到JSON文件"""
    path = os.path.join(RECORDS_DIR, f"{task_id}.json")
    try:
        with open(path, "w", encoding="utf-8") as f:
            json.dump(record_data, f, ensure_ascii=False, indent=2)
    except Exception as e:
        print(f"[记录] 保存失败: {e}")


def load_records():
    """加载所有任务记录"""
    records = []
    for fp in glob.glob(os.path.join(RECORDS_DIR, "*.json")):
        try:
            with open(fp, "r", encoding="utf-8") as f:
                records.append(json.load(f))
        except Exception:
            pass
    records.sort(key=lambda x: x.get("created_at", ""), reverse=True)
    return records


def delete_record(task_id):
    """删除指定任务记录"""
    path = os.path.join(RECORDS_DIR, f"{task_id}.json")
    if os.path.exists(path):
        os.remove(path)
        return True
    return False


# ==================== 7884引擎辅助函数 ====================

def parse_xml_default_output(xml_file):
    """解析Peach XML中的默认输出路径"""
    try:
        import xml.etree.ElementTree as ET
        tree = ET.parse(xml_file)
        root = tree.getroot()
        for pub in root.iter('Publisher'):
            for p in pub.findall('Param'):
                if p.get('name') == 'FileName':
                    return p.get('value', '')
    except Exception:
        pass
    return ""


def find_peach_xml(output_dir):
    """在输出目录中查找format_schema.xml"""
    for root, dirs, files in os.walk(output_dir):
        for fn in files:
            if fn in ("format_schema_peach.xml", "format_schema.xml"):
                return os.path.join(root, fn).replace("\\", "/")
    return None


# ==================== 7884引擎任务执行 ====================

def worker_loop():
    """工作线程：循环处理待执行的任务"""
    while True:
        task = None
        with task_lock:
            for t in task_queue:
                if t["status"] == "pending":
                    t["status"] = "running"
                    t["started_at"] = datetime.now().isoformat()
                    task = t
                    break
        if not task:
            time.sleep(0.5)
            continue

        try:
            _execute_task(task)
        except Exception as e:
            task["status"] = "failed"
            task["error"] = str(e)
            task["finished_at"] = datetime.now().isoformat()
        save_record(task["task_id"], task)
        with task_lock:
            if task in task_queue:
                task_queue.remove(task)


def _execute_task(task):
    """根据任务模式分发执行"""
    mode = task.get("mode", "analyze")
    if mode == "analyze":
        _run_analyze(task)
    elif mode == "gen":
        _run_gen(task)
    else:
        task["status"] = "failed"
        task["error"] = f"未知模式: {mode}"


def _run_analyze(task):
    """执行结构归纳分析任务"""
    brain_exe = get_brain_exe_path()
    if not os.path.exists(brain_exe):
        task["status"] = "failed"
        task["error"] = f"7884_brain.exe 未找到: {brain_exe}"
        print(f"[分析] 引擎未找到: {brain_exe}")
        return

    input_dir = task.get("input", "")
    output_dir = task.get("output_dir", "")
    if not input_dir or not output_dir:
        task["status"] = "failed"
        task["error"] = "缺少输入或输出路径"
        print(f"[分析] 缺少路径: input={input_dir}, output={output_dir}")
        return

    print(f"[分析] 输入目录: {input_dir} (存在: {os.path.exists(input_dir)})")
    print(f"[分析] 输出目录: {output_dir} (存在: {os.path.exists(output_dir)})")

    if not os.path.exists(input_dir):
        os.makedirs(input_dir, exist_ok=True)
    if not os.path.exists(output_dir):
        os.makedirs(output_dir, exist_ok=True)

    input_files = os.listdir(input_dir) if os.path.isdir(input_dir) else []
    print(f"[分析] 输入文件数: {len(input_files)}")

    # 构建命令行
    cmd = [brain_exe, "--input", input_dir, "--output", output_dir]

    run_mode = task.get("run_mode", "full")
    if run_mode == "cpu-only":
        cmd.extend(["--mode", "cpu-only"])
    elif run_mode == "lightweight":
        cmd.extend(["--mode", "lightweight"])

    if task.get("use_iffa", False):
        cmd.append("--use-iffa")
        iffa_path = task.get("iffa_path", "")
        if iffa_path:
            cmd.extend(["--iffa-path", iffa_path])

    max_iter = task.get("max_iterations", 0)
    if max_iter > 0:
        cmd.extend(["--max-iterations", str(max_iter)])

    patience = task.get("patience", 0)
    if patience > 0:
        cmd.extend(["--patience", str(patience)])

    if task.get("cuda_disable", False):
        cmd.append("--cuda-disable")

    cpu_threads = task.get("cpu_threads", 0)
    if cpu_threads > 0:
        cmd.extend(["--cpu-threads", str(cpu_threads)])

    cmd_str = ' '.join(f'"{c}"' if ' ' in c else c for c in cmd)
    print(f"[分析] 执行命令: {cmd_str}")

    # 启动子进程（Windows下使用CREATE_NO_WINDOW）
    proc = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        creationflags=subprocess.CREATE_NO_WINDOW if sys.platform == "win32" else 0
    )

    with process_lock:
        running_processes[task["task_id"]] = proc

    # 实时读取输出
    stdout_lines = []
    try:
        for line in proc.stdout:
            stdout_lines.append(line)
            print(f"[Brain] {line.rstrip()}")
    except Exception as e:
        print(f"[分析] 读取输出异常: {e}")

    proc.wait(timeout=7200)
    stdout = ''.join(stdout_lines)

    with process_lock:
        running_processes.pop(task["task_id"], None)

    task["returncode"] = proc.returncode
    task["stdout"] = stdout or ""

    print(f"[分析] 进程退出码: {proc.returncode}")
    if proc.returncode != 0:
        print(f"[分析] 最后20行输出:")
        for l in stdout_lines[-20:]:
            print(f"  | {l.rstrip()}")

    if proc.returncode == 0:
        task["status"] = "completed"
        task["peach_xml"] = find_peach_xml(output_dir)
        print(f"[分析] 成功! Peach XML: {task['peach_xml']}")
    else:
        task["status"] = "failed"
        task["error"] = f"进程退出码 {proc.returncode}"

    task["finished_at"] = datetime.now().isoformat()


def _run_gen(task):
    """执行样本生成任务"""
    brain_exe = get_brain_exe_path()
    if not os.path.exists(brain_exe):
        task["status"] = "failed"
        task["error"] = f"7884_brain.exe 未找到: {brain_exe}"
        print(f"[生成] 引擎未找到: {brain_exe}")
        return

    xml_file = task.get("xml_file", "")
    if not xml_file or not os.path.exists(xml_file):
        task["status"] = "failed"
        task["error"] = f"Peach XML 未找到: {xml_file}"
        print(f"[生成] XML未找到: {xml_file}")
        return

    # 构建命令行
    cmd = [brain_exe, "--Gen", xml_file]

    iterations = task.get("iterations", 1000)
    cmd.extend(["-i", str(iterations)])

    output_template = task.get("output_template", "")
    if output_template:
        cmd.extend(["-o", output_template])

    seed = task.get("seed", 0)
    if seed > 0:
        cmd.extend(["-s", str(seed)])

    engine = task.get("engine", "pro")
    if engine == "pro":
        cmd.extend(["--engine", "pro"])

    range_val = task.get("range", "")
    if range_val:
        cmd.extend(["--range", range_val])

    cmd_str = ' '.join(f'"{c}"' if ' ' in c else c for c in cmd)
    print(f"[生成] 执行命令: {cmd_str}")

    # 启动子进程
    proc = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        creationflags=subprocess.CREATE_NO_WINDOW if sys.platform == "win32" else 0
    )

    with process_lock:
        running_processes[task["task_id"]] = proc

    # 实时读取输出
    stdout_lines = []
    try:
        for line in proc.stdout:
            stdout_lines.append(line)
            print(f"[Gen-Brain] {line.rstrip()}")
    except Exception as e:
        print(f"[生成] 读取输出异常: {e}")

    proc.wait(timeout=7200)
    stdout = ''.join(stdout_lines)

    with process_lock:
        running_processes.pop(task["task_id"], None)

    task["returncode"] = proc.returncode
    task["stdout"] = stdout or ""

    print(f"[生成] 进程退出码: {proc.returncode}")
    if proc.returncode != 0:
        print(f"[生成] 最后20行输出:")
        for l in stdout_lines[-20:]:
            print(f"  | {l.rstrip()}")

    task["status"] = "completed" if proc.returncode == 0 else "failed"
    if proc.returncode != 0:
        task["error"] = f"生成进程退出码 {proc.returncode}"
    else:
        print(f"[生成] 成功!")
    task["finished_at"] = datetime.now().isoformat()


# ==================== Jetson通信 (HTTP客户端) ====================

def jetson_request(method, path, body=None, headers=None, timeout=30):
    """
    发送HTTP请求到Jetson设备上的jetson_agent.py
    返回: (status_code, response_dict) 或 (0, error_dict)
    """
    jetson_host = config.get("jetson_host", "192.168.1.108")
    jetson_port = config.get("jetson_port", 7886)
    url = f"http://{jetson_host}:{jetson_port}{path}"

    try:
        req_headers = {}
        if body is not None:
            if isinstance(body, str):
                body = body.encode("utf-8")
            req_headers["Content-Type"] = "application/json"
        if headers:
            req_headers.update(headers)

        req = Request(url, data=body, headers=req_headers, method=method)
        resp = urlopen(req, timeout=timeout)
        resp_data = resp.read().decode("utf-8")
        try:
            return resp.status, json.loads(resp_data)
        except json.JSONDecodeError:
            return resp.status, {"raw": resp_data}

    except HTTPError as e:
        try:
            err_body = e.read().decode("utf-8")
            try:
                err_data = json.loads(err_body)
            except json.JSONDecodeError:
                err_data = {"error": err_body}
            return e.code, err_data
        except Exception:
            return e.code, {"error": str(e)}

    except (URLError, OSError) as e:
        reason = str(e.reason) if hasattr(e, 'reason') else str(e)
        return 0, {"error": f"Jetson连接失败: {reason}", "connected": False}

    except Exception as e:
        return 0, {"error": f"请求异常: {str(e)}", "connected": False}


def upload_file_to_jetson(filepath, remote_dir, filename=None):
    """
    通过HTTP multipart/form-data上传单个文件到Jetson agent
    目标: POST http://{jetson_host}:{jetson_port}/api/upload
    """
    jetson_host = config.get("jetson_host", "192.168.1.108")
    jetson_port = config.get("jetson_port", 7886)
    url = f"http://{jetson_host}:{jetson_port}/api/upload"

    filename = filename or os.path.basename(filepath)
    boundary = uuid.uuid4().hex

    with open(filepath, 'rb') as f:
        file_data = f.read()

    body = (
        f'--{boundary}\r\n'
        f'Content-Disposition: form-data; name="file"; filename="{filename}"\r\n'
        f'Content-Type: application/octet-stream\r\n\r\n'
    ).encode('utf-8') + file_data + (
        f'\r\n--{boundary}\r\n'
        f'Content-Disposition: form-data; name="remote_dir"\r\n\r\n'
        f'{remote_dir}\r\n'
        f'--{boundary}--\r\n'
    ).encode('utf-8')

    req = Request(url, data=body, method='POST')
    req.add_header('Content-Type', f'multipart/form-data; boundary={boundary}')
    resp = urlopen(req, timeout=60)
    return json.loads(resp.read().decode('utf-8'))


def push_samples_to_jetson(local_sample_dir, remote_dir):
    """
    读取本地样本目录中的所有文件，逐个上传到Jetson agent
    返回上传结果列表
    """
    if not os.path.isdir(local_sample_dir):
        return {"success": False, "error": f"本地目录不存在: {local_sample_dir}"}

    files = []
    for fn in os.listdir(local_sample_dir):
        fp = os.path.join(local_sample_dir, fn)
        if os.path.isfile(fp):
            files.append((fn, fp))

    if not files:
        return {"success": False, "error": "本地目录中没有文件"}

    results = []
    success_count = 0
    fail_count = 0

    for fn, fp in files:
        try:
            result = upload_file_to_jetson(fp, remote_dir, fn)
            results.append({"file": fn, "success": True, "result": result})
            success_count += 1
            print(f"[推送] 上传成功: {fn}")
        except Exception as e:
            results.append({"file": fn, "success": False, "error": str(e)})
            fail_count += 1
            print(f"[推送] 上传失败: {fn} - {e}")

    return {
        "success": True,
        "total": len(files),
        "success_count": success_count,
        "fail_count": fail_count,
        "results": results
    }


# ==================== HTTP请求处理器 ====================

class Handler2520(BaseHTTPRequestHandler):
    """2520 PC端主控HTTP请求处理器"""

    def log_message(self, format, *args):
        print(f"[{self.address_string()}] {format % args}")

    def _send_json(self, data, status=200):
        """发送JSON响应，启用CORS"""
        content = json.dumps(data, ensure_ascii=False).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", len(content))
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Methods", "GET, POST, DELETE, OPTIONS")
        self.send_header("Access-Control-Allow-Headers", "Content-Type")
        self.end_headers()
        self.wfile.write(content)

    def _send_file(self, filepath, content_type="text/html"):
        """发送静态文件"""
        if not os.path.exists(filepath):
            self.send_error(404, "File not found")
            return
        with open(filepath, "rb") as f:
            content = f.read()
        self.send_response(200)
        self.send_header("Content-Type", f"{content_type}; charset=utf-8")
        self.send_header("Content-Length", len(content))
        self.send_header("Access-Control-Allow-Origin", "*")
        self.end_headers()
        self.wfile.write(content)

    def _read_body(self):
        """读取请求体"""
        content_length = int(self.headers.get("Content-Length", 0))
        return self.rfile.read(content_length).decode("utf-8")

    def _read_body_json(self):
        """读取请求体并解析为JSON"""
        body = self._read_body()
        try:
            return json.loads(body)
        except (json.JSONDecodeError, UnicodeDecodeError):
            return None

    # ==================== CORS预检 ====================

    def do_OPTIONS(self):
        """处理CORS预检请求"""
        self.send_response(200)
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Methods", "GET, POST, DELETE, OPTIONS")
        self.send_header("Access-Control-Allow-Headers", "Content-Type")
        self.end_headers()

    # ==================== GET路由 ====================

    def do_GET(self):
        parsed = urlparse(self.path)
        path = parsed.path

        # 静态页面
        if path == "/" or path == "/index.html":
            self._send_file(INDEX_HTML)

        # 7884引擎状态
        elif path == "/api/status":
            self._handle_engine_status()

        # 任务队列
        elif path == "/api/queue":
            self._handle_queue()

        # 分析记录列表
        elif path == "/api/records":
            self._handle_records()

        # 单条记录详情
        elif path.startswith("/api/records/"):
            task_id = path.split("/api/records/")[1]
            self._handle_record_detail(task_id)

        # Jetson设备状态
        elif path == "/api/jetson/status":
            self._handle_jetson_status()

        # Jetson目标程序列表
        elif path == "/api/jetson/targets":
            self._proxy_jetson_get("/api/targets")

        # Jetson样本目录列表
        elif path == "/api/jetson/samples":
            self._proxy_jetson_get("/api/samples")

        # Jetson模糊测试状态
        elif path == "/api/jetson/fuzz/status":
            self._proxy_jetson_get("/api/fuzz/status")

        # Jetson崩溃记录列表
        elif path == "/api/jetson/crashes":
            self._proxy_jetson_get("/api/crashes")

        # Jetson单条崩溃记录
        elif path.startswith("/api/jetson/crashes/"):
            crash_id = path[len("/api/jetson/crashes/"):]
            self._proxy_jetson_get(f"/api/crashes/{crash_id}")

        # 配置获取
        elif path == "/api/config":
            self._handle_get_config()

        # ==================== Jetson独有Fuzz模块代理 ====================
        elif path.startswith("/api/module/"):
            self._handle_module_proxy("GET", path, body=None)

        # 导出漏洞报告
        elif path == "/api/export/report":
            generate_and_export_report(self, config)

        else:
            self.send_error(404, "Not found")

    # ==================== POST路由 ====================

    def do_POST(self):
        parsed = urlparse(self.path)
        path = parsed.path

        # 7884引擎 - 结构归纳分析
        if path == "/api/run":
            self._handle_run()

        # 7884引擎 - 样本生成
        elif path == "/api/gen":
            self._handle_gen()

        # 任务取消
        elif path == "/api/queue/cancel":
            self._handle_queue_cancel()

        # Jetson模糊测试启动
        elif path == "/api/jetson/fuzz/start":
            self._proxy_jetson_post("/api/fuzz/start", timeout=300)

        # Jetson模糊测试停止
        elif path == "/api/jetson/fuzz/stop":
            self._proxy_jetson_post("/api/fuzz/stop")

        # 推送样本到Jetson
        elif path == "/api/push-samples":
            self._handle_push_samples()

        # 配置更新
        elif path == "/api/config":
            self._handle_set_config()

        # ==================== Jetson独有Fuzz模块代理 ====================
        elif path.startswith("/api/module/"):
            body = self._read_body()
            self._handle_module_proxy("POST", path, body)

        else:
            self.send_error(404, "Not found")

    # ==================== DELETE路由 ====================

    def do_DELETE(self):
        parsed = urlparse(self.path)
        path = parsed.path

        # 删除分析记录
        if path.startswith("/api/records/"):
            task_id = path.split("/api/records/")[1]
            ok = delete_record(task_id)
            self._send_json({"success": ok})

        # 删除Jetson崩溃记录
        elif path.startswith("/api/jetson/crashes/"):
            crash_id = path[len("/api/jetson/crashes/"):]
            self._proxy_jetson_delete(f"/api/crashes/{crash_id}")

        # ==================== Jetson独有Fuzz模块代理 ====================
        elif path.startswith("/api/module/"):
            self._handle_module_proxy("DELETE", path, body=None)

        else:
            self.send_error(404, "Not found")

    # ==================== 7884引擎接口 ====================

    def _handle_engine_status(self):
        """获取7884引擎状态"""
        with task_lock:
            pending = sum(1 for t in task_queue if t["status"] == "pending")
            running = sum(1 for t in task_queue if t["status"] == "running")

        brain_path = get_brain_exe_path()
        self._send_json({
            "status": "ok",
            "brain_exists": os.path.exists(brain_path),
            "brain_path": brain_path,
            "queue": {"pending": pending, "running": running, "total": len(task_queue)}
        })

    def _handle_queue(self):
        """获取任务队列"""
        with task_lock:
            items = []
            for t in task_queue:
                items.append({
                    "task_id": t["task_id"],
                    "mode": t.get("mode", ""),
                    "status": t["status"],
                    "created_at": t.get("created_at", ""),
                    "started_at": t.get("started_at", ""),
                    "finished_at": t.get("finished_at", ""),
                    "error": t.get("error", ""),
                    "label": t.get("label", "")
                })
        self._send_json({"queue": items})

    def _handle_records(self):
        """获取分析记录列表"""
        records = load_records()
        self._send_json({"records": records})

    def _handle_record_detail(self, task_id):
        """获取单条记录详情"""
        path = os.path.join(RECORDS_DIR, f"{task_id}.json")
        if not os.path.exists(path):
            self._send_json({"error": "记录不存在"}, 404)
            return
        try:
            with open(path, "r", encoding="utf-8") as f:
                data = json.load(f)
            self._send_json(data)
        except Exception as e:
            self._send_json({"error": str(e)}, 500)

    def _handle_run(self):
        """提交结构归纳分析任务"""
        params = self._read_body_json()
        if params is None:
            self._send_json({"error": "无效的JSON请求"}, 400)
            return

        task_id = generate_task_id()
        task = {
            "task_id": task_id,
            "mode": "analyze",
            "status": "pending",
            "created_at": datetime.now().isoformat(),
            "started_at": "",
            "finished_at": "",
            "input": params.get("input", ""),
            "output_dir": params.get("output", ""),
            "run_mode": params.get("run_mode", "full"),
            "use_iffa": params.get("use_iffa", False),
            "iffa_path": params.get("iffa_path", ""),
            "max_iterations": params.get("max_iterations", 0),
            "patience": params.get("patience", 0),
            "cuda_disable": params.get("cuda_disable", False),
            "cpu_threads": params.get("cpu_threads", 0),
            "label": params.get("label", ""),
            "stdout": "",
            "error": "",
            "returncode": None,
            "peach_xml": ""
        }

        with task_lock:
            task_queue.append(task)
        save_record(task_id, task)

        print(f"[API] 新分析任务: {task_id}")
        print(f"[API]   input={params.get('input', '')}")
        print(f"[API]   output={params.get('output', '')}")
        print(f"[API]   run_mode={params.get('run_mode', 'full')}")
        print(f"[API]   label={params.get('label', '')}")

        self._send_json({"success": True, "task_id": task_id})

    def _handle_gen(self):
        """提交样本生成任务"""
        params = self._read_body_json()
        if params is None:
            self._send_json({"error": "无效的JSON请求"}, 400)
            return

        xml_file = params.get("xml_file", "")
        if not xml_file:
            self._send_json({"error": "xml_file 为必填项"}, 400)
            return

        task_id = generate_task_id()
        output_template = params.get("output_template", "")
        gen_output_dir = ""
        if not output_template:
            xml_default = parse_xml_default_output(xml_file)
            if xml_default:
                gen_output_dir = os.path.dirname(xml_default.replace("{0}", "0"))
                os.makedirs(gen_output_dir, exist_ok=True)

        task = {
            "task_id": task_id,
            "mode": "gen",
            "status": "pending",
            "created_at": datetime.now().isoformat(),
            "started_at": "",
            "finished_at": "",
            "xml_file": xml_file,
            "iterations": params.get("iterations", 1000),
            "output_template": output_template,
            "seed": params.get("seed", 0),
            "engine": params.get("engine", "pro"),
            "range": params.get("range", ""),
            "label": params.get("label", ""),
            "output": "",
            "error": "",
            "returncode": None
        }

        with task_lock:
            task_queue.append(task)
        save_record(task_id, task)

        print(f"[API] 新生成任务: {task_id}, xml={xml_file}, iterations={params.get('iterations', 1000)}")

        self._send_json({"success": True, "task_id": task_id, "output_dir": gen_output_dir})

    def _handle_queue_cancel(self):
        """取消任务"""
        params = self._read_body_json()
        if params is None:
            self._send_json({"error": "无效的JSON请求"}, 400)
            return

        task_id = params.get("task_id", "")
        with task_lock:
            for t in task_queue:
                if t["task_id"] == task_id:
                    if t["status"] == "pending":
                        t["status"] = "cancelled"
                        t["finished_at"] = datetime.now().isoformat()
                        save_record(task_id, t)
                        self._send_json({"success": True, "message": "任务已取消"})
                        return
                    elif t["status"] == "running":
                        # 终止运行中的进程
                        with process_lock:
                            proc = running_processes.get(task_id)
                            if proc:
                                try:
                                    proc.kill()
                                except Exception:
                                    pass
                        t["status"] = "cancelled"
                        t["finished_at"] = datetime.now().isoformat()
                        save_record(task_id, t)
                        self._send_json({"success": True, "message": "任务已终止"})
                        return

        self._send_json({"success": False, "error": "任务未找到或无法取消"})

    # ==================== Jetson代理接口 ====================

    def _proxy_jetson_get(self, jetson_path, timeout=30):
        """代理GET请求到Jetson agent"""
        status, data = jetson_request("GET", jetson_path, timeout=timeout)
        if status == 0:
            # Jetson离线，返回502
            self._send_json(data, 502)
        else:
            self._send_json(data, status)

    def _proxy_jetson_post(self, jetson_path, timeout=30):
        """代理POST请求到Jetson agent"""
        body = self._read_body()
        status, data = jetson_request("POST", jetson_path, body=body, timeout=timeout)
        if status == 0:
            self._send_json(data, 502)
        else:
            self._send_json(data, status)

    def _proxy_jetson_delete(self, jetson_path, timeout=30):
        """代理DELETE请求到Jetson agent"""
        status, data = jetson_request("DELETE", jetson_path, timeout=timeout)
        if status == 0:
            self._send_json(data, 502)
        else:
            self._send_json(data, status)

    def _handle_jetson_status(self):
        """获取Jetson设备状态（含连接信息）"""
        status, data = jetson_request("GET", "/api/status", timeout=10)
        if status == 0:
            # Jetson离线
            self._send_json({
                "connected": False,
                "error": data.get("error", "连接失败"),
                "jetson_host": config.get("jetson_host", "192.168.1.108"),
                "jetson_port": config.get("jetson_port", 7886)
            }, 502)
        else:
            # Jetson在线，附加连接信息
            data["connected"] = True
            data["jetson_host"] = config.get("jetson_host", "192.168.1.108")
            data["jetson_port"] = config.get("jetson_port", 7886)
            self._send_json(data, status)

    # ==================== Jetson独有Fuzz模块代理 ====================

    def _handle_module_proxy(self, method, path, body):
        """代理到Jetson的Fuzz模块管理API"""
        try:
            agent_path = path

            headers = {}
            ct = self.headers.get("Content-Type", "")
            if ct:
                headers["Content-Type"] = ct

            if method == "GET":
                status_code, data = jetson_request("GET", agent_path, timeout=30)
            elif method == "POST":
                status_code, data = jetson_request("POST", agent_path, body=body, headers=headers, timeout=300)
            elif method == "DELETE":
                status_code, data = jetson_request("DELETE", agent_path, timeout=30)
            else:
                self._send_json({"error": "不支持的HTTP方法"}, 405)
                return

            if status_code == 0:
                self._send_json({"error": "Jetson连接失败", "connected": False}, 502)
            else:
                self._send_json(data, status_code)
        except Exception as e:
            try:
                self._send_json({"error": f"模块代理错误: {str(e)}"}, 500)
            except:
                pass

    # ==================== 样本推送接口 ====================

    def _handle_push_samples(self):
        """推送本地样本文件到Jetson设备"""
        params = self._read_body_json()
        if params is None:
            self._send_json({"error": "无效的JSON请求"}, 400)
            return

        local_dir = params.get("local_dir", "")
        remote_dir = params.get("remote_dir", "")

        if not local_dir:
            self._send_json({"error": "local_dir 为必填项"}, 400)
            return
        if not remote_dir:
            self._send_json({"error": "remote_dir 为必填项"}, 400)
            return
        if not os.path.isdir(local_dir):
            self._send_json({"error": f"本地目录不存在: {local_dir}"}, 400)
            return

        print(f"[推送] 开始推送样本: {local_dir} → {remote_dir}")

        try:
            result = push_samples_to_jetson(local_dir, remote_dir)
            if not result.get("success", False):
                self._send_json(result, 400)
            else:
                self._send_json(result)
        except Exception as e:
            self._send_json({"success": False, "error": str(e)}, 502)

    # ==================== 配置接口 ====================

    def _handle_get_config(self):
        """获取当前配置"""
        with config_lock:
            safe_config = {
                "pc_port": config.get("pc_port", 7885),
                "brain_exe": config.get("brain_exe", "../7884_brain.exe"),
                "jetson_host": config.get("jetson_host", "192.168.1.108"),
                "jetson_port": config.get("jetson_port", 7886),
                "jetson_user": config.get("jetson_user", "nvidia"),
                "fuzz_timeout": config.get("fuzz_timeout", 30),
                "workspace": config.get("workspace", "./fuzz_workspace"),
                "crash_dir": config.get("crash_dir", "./crash_records")
            }
        self._send_json(safe_config)

    def _handle_set_config(self):
        """更新配置"""
        params = self._read_body_json()
        if params is None:
            self._send_json({"error": "无效的JSON请求"}, 400)
            return

        with config_lock:
            # 允许更新的字段
            updatable = [
                "pc_port", "brain_exe", "jetson_host", "jetson_port",
                "jetson_user", "fuzz_timeout", "workspace", "crash_dir"
            ]
            updated = {}
            for key in updatable:
                if key in params:
                    config[key] = params[key]
                    updated[key] = params[key]

        # 保存到文件
        save_config()
        # 确保目录存在
        _ensure_dirs()

        self._send_json({"success": True, "updated": updated})


# ==================== 进程清理 ====================

def kill_stale_processes(port):
    """启动前清理占用目标端口的旧进程，防止多进程冲突"""
    import subprocess
    try:
        output = subprocess.check_output(
            f'netstat -ano | findstr :{port} | findstr LISTENING',
            shell=True
        ).decode().strip()
        if not output:
            return
        for line in output.split('\n'):
            parts = line.strip().split()
            if len(parts) >= 5:
                try:
                    pid = int(parts[-1])
                    if pid != os.getpid():
                        os.kill(pid, 9)
                        print(f"  已终止旧进程 PID={pid}")
                except (ValueError, OSError):
                    pass
    except subprocess.CalledProcessError:
        pass  # 没有找到旧进程，正常启动

def cleanup_processes():
    """清理所有运行中的子进程和待执行任务"""
    with process_lock:
        for proc in running_processes.values():
            try:
                proc.kill()
            except Exception:
                pass
        running_processes.clear()

    with task_lock:
        for t in task_queue:
            if t["status"] in ("pending", "running"):
                t["status"] = "cancelled"
                t["finished_at"] = datetime.now().isoformat()
                save_record(t["task_id"], t)


# ==================== 漏洞报告导出 ====================

def generate_and_export_report(handler, config):
    """生成并导出漏洞HTML报告"""
    try:
        script_dir = os.path.dirname(os.path.abspath(__file__))
        generator = os.path.join(script_dir, "tools", "bug_report_generator.py")
        output = os.path.join(script_dir, "bug_export.html")

        jetson_host = config.get("jetson_host", "192.168.1.108")
        jetson_user = config.get("jetson_user", "nvidia")
        jetson_pass = config.get("jetson_pass", "nvidia")

        cmd = [
            sys.executable, generator,
            "--jetson", jetson_host,
            "--port", "22",
            "--user", jetson_user,
            "--password", jetson_pass,
            "-o", output,
            "--title", "2520 漏洞报告"
        ]

        result = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
        if result.returncode != 0:
            handler._send_json({"error": f"报告生成失败: {result.stderr}"}, 500)
            return

        with open(output, "rb") as f:
            html_content = f.read()

        handler.send_response(200)
        handler.send_header("Content-Type", "text/html; charset=utf-8")
        handler.send_header("Content-Disposition", "attachment; filename=2520_report.html")
        handler.send_header("Content-Length", len(html_content))
        handler.send_header("Access-Control-Allow-Origin", "*")
        handler.end_headers()
        handler.wfile.write(html_content)

        try:
            os.remove(output)
        except:
            pass

    except subprocess.TimeoutExpired:
        handler._send_json({"error": "报告生成超时"}, 504)
    except Exception as e:
        handler._send_json({"error": str(e)}, 500)


# ==================== 服务器启动 ====================

def run_server():
    """启动2520 PC端主控服务"""
    # 加载配置
    load_config()

    # 启动前清理占用端口的旧进程
    port = config.get("pc_port", 7885)
    kill_stale_processes(port)

    # 启动7884引擎工作线程
    worker_thread = threading.Thread(target=worker_loop, daemon=True)
    worker_thread.start()

    server = ThreadingHTTPServer(("0.0.0.0", port), Handler2520)

    brain_path = get_brain_exe_path()
    jetson_host = config.get("jetson_host", "192.168.1.108")
    jetson_port = config.get("jetson_port", 7886)

    print("=" * 55)
    print("  2520 PC端主控服务")
    print(f"  运行地址: http://0.0.0.0:{port}/")
    print(f"  7884引擎: {brain_path} (存在: {os.path.exists(brain_path)})")
    print(f"  Jetson设备: {jetson_host}:{jetson_port}")
    print("=" * 55)
    print("按 Ctrl+C 停止服务...")

    # 自动打开浏览器
    def open_browser():
        webbrowser.open(f"http://127.0.0.1:{port}/")

    threading.Timer(1.0, open_browser).start()

    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\n正在关闭服务...")
    finally:
        cleanup_processes()
        server.shutdown()
        print("服务已停止")


if __name__ == "__main__":
    run_server()
