hy3d/hunyuan3dweb/browser/sniffer.py

#!/usr/bin/env python3
"""
腾讯混元3D API 拦截分析工具 (CloakBrowser)
利用持久化登录状态，自动捕获所有 API 请求和响应
"""

import json
import os
import shutil
import sys
import tempfile
import time
from datetime import datetime
from cloakbrowser import launch_persistent_context

from ..config import get_profile_dir, prepare_profile_dir

PROFILE_DIR = str(get_profile_dir())
API_LOG_FILE = "./api_requests.log.json"


def main():
    # 使用临时文件缓冲日志，避免长时间嗅探导致内存无限增长
    log_buffer = tempfile.NamedTemporaryFile(mode="w+", suffix=".jsonl", delete=False, encoding="utf-8")

    def log_request(request):
        url = request.url
        # 只关注同域 API 和关键第三方接口
        if "/api/" in url or "hunyuan" in url:
            entry = {
                "time": datetime.now().isoformat(),
                "type": "request",
                "method": request.method,
                "url": url,
                "headers": dict(request.headers) if hasattr(request, "headers") else {},
            }
            # 尝试获取 POST body
            if request.method in ("POST", "PUT", "PATCH") and hasattr(request, "post_data"):
                try:
                    entry["body"] = request.post_data
                except Exception:
                    pass
            log_buffer.write(json.dumps(entry, ensure_ascii=False) + "\n")
            log_buffer.flush()
            print(f"[REQ] {request.method} {url}")

    def log_response(response):
        url = response.url
        if "/api/" in url or "hunyuan" in url:
            entry = {
                "time": datetime.now().isoformat(),
                "type": "response",
                "status": response.status,
                "url": url,
            }
            # 尝试读取响应体
            try:
                # 只读取 JSON 响应
                content_type = response.headers.get("content-type", "")
                if "json" in content_type:
                    body = response.json()
                    entry["body"] = body
                else:
                    text = response.text()
                    # 限制文本长度，避免过大
                    entry["body_preview"] = text[:500]
            except Exception as e:
                entry["body_error"] = str(e)
            log_buffer.write(json.dumps(entry, ensure_ascii=False) + "\n")
            log_buffer.flush()
            print(f"[RES] {response.status} {url}")

    if not os.path.exists(PROFILE_DIR):
        print(f"错误: 未找到持久化目录 {PROFILE_DIR}")
        print("请先运行 hunyuan3dweb-login 完成登录")
        sys.exit(1)

    profile_dir, temp_dir = prepare_profile_dir()

    context = None
    try:
        context = launch_persistent_context(profile_dir, headless=False)
        page = context.new_page()

        page.on("request", log_request)
        page.on("response", log_response)

        print("正在打开腾讯混元3D并捕获 API...")
        page.goto("https://3d.hunyuan.tencent.com/")
        time.sleep(3)

        # 检查是否已登录
        login_btn = page.locator("button").filter(has_text="登录").first
        if login_btn.count() > 0:
            print("\n警告: 当前未检测到登录状态，API 可能返回未授权")
        else:
            print("\n已检测到登录状态，开始捕获 API...")

        print("\n你可以手动在浏览器中操作（切换页面、生成3D等）")
        print("所有 API 请求会实时打印在终端中")
        print("按 Enter 停止捕获并保存结果...\n")
        input()
    finally:
        if context is not None:
            try:
                context.close()
            except Exception:
                pass
        if temp_dir is not None:
            try:
                shutil.rmtree(temp_dir)
            except Exception:
                pass

    # 从临时文件读取并保存为 JSON 数组
    try:
        log_buffer.flush()
        log_buffer.seek(0)
        logs = [json.loads(line) for line in log_buffer if line.strip()]
    finally:
        log_buffer.close()
        os.unlink(log_buffer.name)

    # 保存日志
    with open(API_LOG_FILE, "w", encoding="utf-8") as f:
        json.dump(logs, f, ensure_ascii=False, indent=2)

    print(f"\n共捕获 {len([l for l in logs if l['type'] == 'request'])} 个请求")
    print(f"结果已保存到: {os.path.abspath(API_LOG_FILE)}")


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n已取消")
        sys.exit(0)