competition.py

#!/usr/bin/env python3
"""
competition.py — TCH 比赛编排器 v2
基于官方 API，支持并行 worker 调度、实例生命周期管理、Hint 策略。

三层架构:
  Layer 0: watchdog.sh (保活)
  Layer 1: competition.py (确定性编排器) ← 本文件
  Layer 2: tchkiller (渗透智能体)

Usage:
    # 正式比赛模式
    python3 competition.py --server <HOST> --token <TOKEN>
    python3 competition.py --server <HOST> --token <TOKEN> --workers 3

    # Mock 模式 (本地测试)
    python3 competition.py --mock
    python3 competition.py --mock --targets targets.txt
    python3 competition.py --mock --challenges challenges.json

    # 恢复上次中断的比赛
    python3 competition.py --server <HOST> --token <TOKEN> --resume

    # 紧急: 关闭所有运行中的赛题实例
    python3 competition.py --server <HOST> --token <TOKEN> --stop-all

    # 查看当前状态
    python3 competition.py --status

    # 试运行 (不实际执行)
    python3 competition.py --mock --dry-run
"""
from __future__ import annotations

import argparse
import asyncio
import json
import sys
import time
from datetime import datetime
from pathlib import Path

ROOT = Path(__file__).resolve().parent
COMP_TIMELINE = ROOT / "comp_timeline.jsonl"
sys.path.insert(0, str(ROOT))

from comp.platform import MockPlatform, TCHPlatform, Platform, Challenge
from comp.state import StateManager
from comp.strategy import Strategy
from comp.worker_pool import WorkerPool


def comp_timeline_append(entry: dict):
    """追加一条竞赛级事件到 comp_timeline.jsonl"""
    try:
        entry.setdefault("ts", time.time())
        entry.setdefault("time", datetime.now().strftime("%H:%M:%S"))
        with open(COMP_TIMELINE, "a", encoding="utf-8") as f:
            f.write(json.dumps(entry, ensure_ascii=False) + "\n")
    except Exception:
        pass


# ---------------------------------------------------------------------------
# 主编排
# ---------------------------------------------------------------------------

async def run_competition(platform: Platform, state_mgr: StateManager,
                          max_workers: int = 3,
                          extra_args: list[str] | None = None,
                          dry_run: bool = False):
    """比赛主循环"""
    state = state_mgr.state
    strategy = Strategy(state)
    pool = WorkerPool(
        platform=platform,
        state_mgr=state_mgr,
        strategy=strategy,
        max_workers=max_workers,
        extra_args=extra_args,
        timeline_fn=comp_timeline_append,
    )

    print("\n" + "=" * 60)
    print(f"🏆 TCH Competition Runner v2")
    print(f"   Workers: {pool.max_workers} | Dry-run: {dry_run}")
    print("=" * 60)

    # 初始赛题获取
    try:
        status = platform.get_challenges()
    except Exception as e:
        print(f"❌ 无法获取赛题: {e}")
        return

    challenges = status.challenges
    if not challenges:
        print("❌ 没有可用赛题")
        return

    state.current_level = status.current_level
    state_mgr.sync_from_platform(challenges)

    comp_timeline_append({
        "event": "session_start",
        "total_challenges": status.total_challenges,
        "solved": status.solved_challenges,
        "current_level": status.current_level,
        "workers": max_workers,
    })

    print(f"\n📋 赛题: {status.total_challenges} 道 (已解: {status.solved_challenges})")
    print(f"🔑 当前关卡: Level {status.current_level}")

    # 按关卡/难度显示
    levels = sorted(set(c.level for c in challenges))
    for lv in levels:
        lv_chs = [c for c in challenges if c.level == lv]
        easy = sum(1 for c in lv_chs if c.difficulty == "easy")
        medium = sum(1 for c in lv_chs if c.difficulty == "medium")
        hard = sum(1 for c in lv_chs if c.difficulty == "hard")
        solved = sum(1 for c in lv_chs if c.solved)
        print(f"   Level {lv}: {len(lv_chs)} 题 "
              f"(E:{easy} M:{medium} H:{hard}) "
              f"[已解: {solved}]")
        for c in sorted(lv_chs, key=lambda x: x.difficulty):
            flag_info = f"🚩{c.flag_got_count}/{c.flag_count}"
            score_info = f"💰{c.total_got_score}/{c.total_score}"
            status_icon = "✅" if c.solved else "⬚"
            print(f"     {status_icon} {c.title} ({c.difficulty}) "
                  f"| {flag_info} | {score_info}")
            if c.description:
                print(f"        📝 {c.description}")

    print(f"\n{strategy.summary(challenges)}")

    # 运行 worker pool
    stats = await pool.run(dry_run=dry_run)

    # 最终报告
    try:
        final_status = platform.get_challenges()
        challenges = final_status.challenges
    except Exception:
        pass

    print("\n" + "=" * 60)
    print("📋 最终报告")
    print("=" * 60)
    print(strategy.summary(challenges))
    print(f"\n  执行统计:")
    print(f"    总执行: {stats.total_executed}")
    print(f"    成功: {stats.total_success}")
    print(f"    失败: {stats.total_failed}")
    print(f"    超时: {stats.total_timeout}")
    print(f"    总耗时: {stats.total_duration_s / 60:.1f} min")
    print(f"    总花费: ${state.total_cost_usd:.2f}")

    # 保存最终报告
    report_file = ROOT / "comp_report.json"
    total_score = sum(c.total_got_score for c in challenges)
    report = {
        "timestamp": datetime.now().isoformat(),
        "total_score": total_score,
        "total_challenges": len(challenges),
        "solved_challenges": sum(1 for c in challenges if c.solved),
        "current_level": state.current_level,
        "flags_submitted": state.total_flags,
        "total_cost_usd": state.total_cost_usd,
        "stats": {
            "total_executed": stats.total_executed,
            "total_success": stats.total_success,
            "total_failed": stats.total_failed,
            "total_timeout": stats.total_timeout,
            "total_duration_s": stats.total_duration_s,
        },
        "challenges": {},
    }
    for c in challenges:
        cs = state.challenges.get(c.code)
        report["challenges"][c.code] = {
            "title": c.title,
            "level": c.level,
            "difficulty": c.difficulty,
            "solved": c.solved,
            "score": f"{c.total_got_score}/{c.total_score}",
            "flags": f"{c.flag_got_count}/{c.flag_count}",
            "attempts": cs.attempt_count if cs else 0,
            "hint_viewed": c.hint_viewed,
        }
    report_file.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"\n  📄 报告已保存: {report_file}")


# ---------------------------------------------------------------------------
# stop-all 紧急操作
# ---------------------------------------------------------------------------

def stop_all_instances(platform: Platform):
    """紧急停止所有运行中的赛题实例"""
    print("\n🚨 紧急操作: 停止所有运行中的实例...")
    try:
        count = platform.stop_all()
        print(f"✅ 已停止 {count} 个实例")
    except Exception as e:
        print(f"❌ 停止失败: {e}")


# ---------------------------------------------------------------------------
# 状态查看
# ---------------------------------------------------------------------------

def show_status(state_file: Path):
    """查看当前比赛状态"""
    if not state_file.exists():
        print("❌ 没有比赛状态文件")
        return

    sm = StateManager(state_file)
    s = sm.state

    print("\n" + "=" * 60)
    print("📊 Competition Status")
    print("=" * 60)
    print(f"  Session:  {s.session_id}")
    print(f"  Start:    {s.start_time}")
    print(f"  Level:    {s.current_level}")
    print(f"  Score:    {s.total_score}")
    print(f"  Flags:    {s.total_flags}")
    print(f"  Solved:   {s.total_challenges_solved}")
    print(f"  Cost:     ${s.total_cost_usd:.2f}")
    print()

    # 按 level 分组
    by_level: dict[int, list] = {}
    for cid, cs in s.challenges.items():
        by_level.setdefault(cs.level, []).append(cs)

    for lv in sorted(by_level.keys()):
        print(f"  Level {lv}:")
        for cs in sorted(by_level[lv], key=lambda x: x.difficulty):
            if cs.solved:
                icon = "✅"
            elif cs.skip:
                icon = "⏭️"
            elif cs.attempt_count > 0:
                icon = "🔄"
            else:
                icon = "⬜"
            diff_tag = {"easy": "E", "medium": "M", "hard": "H"}.get(cs.difficulty, "?")
            print(f"    {icon} [{diff_tag}] {cs.name or cs.code}")
            if cs.attempt_count > 0:
                print(f"          Attempts: {cs.attempt_count} | "
                      f"Flags: {len(cs.flags_submitted)}/{cs.flag_count} | "
                      f"Cost: ${cs.total_cost_usd:.2f}"
                      f"{' | Hint: ✅' if cs.hint_viewed else ''}")
        print()

    if s.log:
        print("  📝 最近日志:")
        for entry in s.log[-10:]:
            print(f"     {entry}")


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def parse_args():
    p = argparse.ArgumentParser(
        description="TCH Competition Runner v2 — 比赛编排器",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # 正式比赛
  python3 competition.py --server 10.0.0.1:8080 --token <TOKEN> --workers 3
  python3 competition.py --server 10.0.0.1:8080 --token <TOKEN> --resume

  # 紧急停止所有实例
  python3 competition.py --server 10.0.0.1:8080 --token <TOKEN> --stop-all

  # 本地测试
  python3 competition.py --mock
  python3 competition.py --mock --workers 1 --dry-run
  python3 competition.py --mock --challenges tests/challenges.json

  # 查看状态
  python3 competition.py --status
        """,
    )

    # 平台选择
    group = p.add_argument_group("平台配置")
    group.add_argument("--mock", action="store_true", help="使用 MockPlatform 进行本地测试")
    group.add_argument("--server", type=str, help="比赛平台地址 (HOST:PORT)")
    group.add_argument("--token", type=str, help="Agent Token")

    # Mock 选项
    mock_group = p.add_argument_group("Mock 选项")
    mock_group.add_argument("--targets", type=str, help="目标列表文件 (每行一个)")
    mock_group.add_argument("--challenges", type=str, help="赛题 JSON 文件")

    # 调度选项
    sched_group = p.add_argument_group("调度选项")
    sched_group.add_argument("--workers", type=int, default=3, choices=[1, 2, 3],
                             help="并发 worker 数量 (1-3，默认: 3)")

    # 操作模式
    ops_group = p.add_argument_group("操作模式")
    ops_group.add_argument("--resume", action="store_true", help="从上次中断处恢复")
    ops_group.add_argument("--status", action="store_true", help="查看当前比赛状态")
    ops_group.add_argument("--stop-all", action="store_true",
                           help="紧急: 关闭所有运行中的赛题实例")
    ops_group.add_argument("--dry-run", action="store_true",
                           help="试运行，不实际执行 tchkiller")
    ops_group.add_argument("--reset", action="store_true",
                            help="清除本地状态文件，重新开始")
    ops_group.add_argument("--reset-range", action="store_true",
                             help="远程重置靶场状态（清空提交记录，仅 test_range 支持）")
    ops_group.add_argument("--state-file", type=str, default="comp_state.json",
                           help="状态文件路径 (默认: comp_state.json)")

    # 透传给 tchkiller 的参数
    pass_group = p.add_argument_group("tchkiller 透传参数")
    pass_group.add_argument("--model", default=None, help="主模型")
    pass_group.add_argument("--provider", default=None, help="API provider")
    pass_group.add_argument("--no-team", action="store_true", help="禁用 Agent Teams")
    pass_group.add_argument("--no-orchestrator", action="store_true", help="禁用决策 agent")
    pass_group.add_argument("--browser", action="store_true", help="启用浏览器")

    return p.parse_args()


def build_platform(args) -> Platform:
    """根据参数构建平台实例"""
    if args.mock:
        if args.challenges:
            return MockPlatform(challenges_file=args.challenges)
        elif args.targets:
            targets = Path(args.targets).read_text(encoding="utf-8").strip().splitlines()
            targets = [t.strip() for t in targets if t.strip() and not t.startswith("#")]
            return MockPlatform(targets=targets)
        else:
            return MockPlatform()
    elif args.server and args.token:
        return TCHPlatform(server_host=args.server, agent_token=args.token)
    else:
        print("❌ 请指定平台: --mock (测试) 或 --server + --token (正式比赛)")
        sys.exit(1)


def main():
    args = parse_args()
    state_file = ROOT / args.state_file

    # 状态查看
    if args.status:
        show_status(state_file)
        return

    # 紧急停止
    if args.stop_all:
        platform = build_platform(args)
        stop_all_instances(platform)
        return

    # 构建平台
    platform = build_platform(args)

    # 构建透传参数
    extra_args = []
    if args.model:
        extra_args.extend(["--model", args.model])
    if args.provider:
        extra_args.extend(["--provider", args.provider])
    if args.browser:
        extra_args.append("--browser")
    if args.no_team:
        extra_args.append("--no-team")
    if args.no_orchestrator:
        extra_args.append("--no-orchestrator")

    # Worker 数 vs Provider 配置文件检查
    # Worker 0 → providers.json, Worker 1 → providers2.json, Worker 2 → providers3.json
    max_workers = args.workers
    if max_workers > 1:
        provider_files = ["providers.json"]
        for i in range(2, max_workers + 1):
            cfg = ROOT / f"providers{i}.json"
            if cfg.exists():
                provider_files.append(f"providers{i}.json")
        n_distinct = len(provider_files)
        if n_distinct < max_workers:
            shared = max_workers - n_distinct
            print(f"  ⚠️ {shared} 个 worker 将共享 providers.json"
                  f"（缺少 {', '.join(f'providers{i}.json' for i in range(n_distinct + 1, max_workers + 1))}）")
            print(f"  💡 建议: 创建独立的 provider 配置文件以避免 API 并发冲突")

    # 初始化状态
    if args.reset and state_file.exists():
        state_file.unlink()
        print("🗑️  已清除本地状态文件")
    if args.reset_range:
        platform.reset_range()
    state_mgr = StateManager(state_file)
    if not args.resume:
        state_mgr.state.session_id = f"comp-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
        state_mgr.state.start_time = datetime.now().isoformat()
    else:
        if not state_file.exists():
            print("⚠️ 没有找到上次的状态文件，将启动新会话")
            state_mgr.state.session_id = f"comp-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
            state_mgr.state.start_time = datetime.now().isoformat()
        else:
            print(f"🔄 恢复会话: {state_mgr.state.session_id}")
    state_mgr.save()

    # 运行
    asyncio.run(run_competition(
        platform=platform,
        state_mgr=state_mgr,
        max_workers=max_workers,
        extra_args=extra_args,
        dry_run=args.dry_run,
    ))


if __name__ == "__main__":
    main()