api.py

from fastapi import FastAPI, Request, Response, WebSocket, WebSocketDisconnect, Query
import httpx
import uvicorn
import asyncio
import websockets
import json
import subprocess
import threading
import time
import os
import random
from typing import Dict
from loguru import logger
from starlette.responses import StreamingResponse
import socket

app = FastAPI()

# 浏览器实例管理
BROWSERS: Dict[str, dict] = {}  # {browser_id: {"process": Popen, "ws": str, "port": int, "status": str, "info": dict}}
# 浏览器实例状态: open, closed
PORT_RANGE = (9300, 9700)
CHROME_PATH = "google-chrome"  # Linux下可用命令，或自定义绝对路径
PROFILE_BASE = "/tmp/profiles"


def get_system_usage():
    cpu = subprocess.check_output("top -bn1 | grep 'Cpu(s)'", shell=True).decode().strip()
    mem = subprocess.check_output("free -m | grep Mem", shell=True).decode().strip()
    # 字段映射
    cpu_key_map = {
        "us": "user",
        "sy": "system",
        "ni": "nice",
        "id": "idle",
        "wa": "iowait",
        "hi": "hardware_irq",
        "si": "software_irq",
        "st": "steal"
    }
    # 解析CPU信息
    cpu_parts = cpu.split(":")[1].split(",")
    cpu_info = {}
    for part in cpu_parts:
        value, key = part.strip().split(" ", 1)
        key = key.strip().replace("%", "")
        mapped_key = cpu_key_map.get(key, key)
        cpu_info[mapped_key] = value.strip() + "%"
    # 解析内存信息
    mem_values = mem.split()[1:4]
    mem_total, mem_used, mem_free = mem_values
    mem_available = int(mem_total) - int(mem_used)
    mem_info = {
        "Total": f"{mem_total}MB",
        "Used": f"{mem_used}MB",
        "Free": f"{mem_free}MB",
        "Available": f"{mem_available}MB"
    }
    logger.info(f"CPU: {cpu_info}, Memory: {mem_info}")
    info = {"CPU": cpu_info, "Memory": mem_info}
    return info


def get_free_port():
    used_ports = {b["port"] for b in BROWSERS.values()}
    for _ in range(100):
        port = random.randint(*PORT_RANGE)
        if port not in used_ports:
            return port
    raise RuntimeError("No free port available")


def gen_browser_id():
    return str(int(time.time() * 1000)) + str(random.randint(1000, 9999))


def get_default_launch_args():
    return [
        "--disable-gpu",
        "--disable-dev-shm-usage",
        "--disable-software-rasterizer",
        "--disable-extensions",
        "--disable-background-networking",
        "--disable-default-apps",
        "--disable-sync",
        "--disable-translate",
        "--disable-features=TranslateUI",
        "--no-first-run",
        "--no-default-browser-check",
        "--remote-allow-origins=*",
        "--headless"
    ]


def wait_port(host, port, timeout=5.0):
    """等待端口可用"""
    start = time.time()
    while time.time() - start < timeout:
        try:
            with socket.create_connection((host, port), timeout=0.5):
                return True
        except Exception:
            time.sleep(0.1)
    return False


@app.get("/system/usage")
async def system_usage():
    """
    获取系统资源使用情况
    """
    try:
        usage_info = get_system_usage()
        return {"code": 0, "msg": "success", "data": usage_info}
    except Exception as e:
        logger.error(f"获取系统资源使用情况失败: {e}")
        return {"code": 1, "msg": "failed", "error": str(e)}


# BitBrowser风格API
@app.post("/browser/open")
async def open_browser(request: Request):
    data = await request.json()
    browser_id = data.get("id", '')
    if browser_id in BROWSERS:
        logger.info(f"Browser ID {browser_id}: {BROWSERS[browser_id]}")
        if BROWSERS[browser_id]["status"] == "open":
            return {"code": 0, "msg": "already opened",
                    "data": {"id": browser_id, "port": BROWSERS[browser_id]["port"]}}
        else:
            port = get_free_port()
            BROWSERS[browser_id]["port"] = port
    else:
        # 新建浏览器实例
        browser_id = gen_browser_id()
        port = get_free_port()
    profile_dir = f"{PROFILE_BASE}/{browser_id}"
    os.makedirs(profile_dir, exist_ok=True)
    # 支持自定义启动参数并去重
    launch_args = data.get("launchArgs", "")
    args = [CHROME_PATH, f"--remote-debugging-port={port}", f"--user-data-dir={profile_dir}"]
    # 加入默认参数
    logger.info(f"使用默认参数: {get_default_launch_args()}")
    args.extend(get_default_launch_args())
    if launch_args:
        # 拆分参数并去重，保留顺序
        base_args = args[1:]  # 除去chrome可执行文件
        extra_args = [a for a in launch_args.split() if a]
        all_args = base_args + extra_args
        seen = set()
        deduped_args = []
        for arg in all_args:
            key = arg.split('=')[0].strip() if '=' in arg else arg.strip()
            if key not in seen:
                seen.add(key)
                deduped_args.append(arg)
        args = [CHROME_PATH] + deduped_args
    proc = subprocess.Popen(args, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    if not wait_port("127.0.0.1", port, timeout=60):
        proc.terminate()
        return {"code": 2, "msg": f"chrome端口{port}未就绪"}
    async with httpx.AsyncClient() as client:
        resp = await client.get(f"http://127.0.0.1:{port}/json/version")
        if resp.status_code == 200:
            browser_info = resp.json()
            ws_url = browser_info.get("webSocketDebuggerUrl").replace("ws://127.0.0.1", "wss://0.0.0.0")
            logger.info(f"Browser opened with ID: {browser_id}, Info: {browser_info}")
    BROWSERS[browser_id] = {"process": proc, "ws": ws_url, "port": port, "status": "open", "info": browser_info}
    return {"code": 0, "msg": "success", "data": {"id": browser_id, "ws": ws_url}, "info": browser_info}


@app.post("/browser/close")
async def close_browser(request: Request):
    data = await request.json()
    browser_id = data.get("id")
    b = BROWSERS.get(browser_id)
    if not b:
        return {"code": 1, "msg": "not found"}
    b["process"].terminate()
    b["status"] = "closed"
    return {"code": 0, "msg": "closed", "data": {"id": browser_id}}


@app.post("/browser/delete")
async def delete_browser(request: Request):
    data = await request.json()
    browser_id = data.get("id")
    b = BROWSERS.pop(browser_id, None)
    if not b:
        return {"code": 1, "msg": "not found"}
    try:
        b["process"].terminate()
    except Exception:
        pass
    # 删除用户数据目录
    profile_dir = f"{PROFILE_BASE}/{browser_id}"
    if os.path.exists(profile_dir):
        import shutil
        shutil.rmtree(profile_dir, ignore_errors=True)
    return {"code": 0, "msg": "deleted", "data": {"id": browser_id}}


@app.post("/browser/update")
async def update_browser(request: Request):
    data = await request.json()
    # 兼容 BitBrowser 结构
    launch_args = ""
    if isinstance(data.get("browserFingerPrint"), dict):
        launch_args = data["browserFingerPrint"].get("launchArgs", "")
    elif "launchArgs" in data:
        launch_args = data["launchArgs"]
    # 合并 launchArgs 到 data
    if launch_args:
        data["launchArgs"] = launch_args
    return await open_browser(Request({**request.scope, "body": json.dumps(data).encode()}, receive=request._receive))


@app.post("/browser/ports")
async def browser_ports():
    # 返回 {id: port} 格式
    data = {k: str(v["ws"]) for k, v in BROWSERS.items() if v["status"] == "open"}
    return {"success": True, "data": data}


@app.post("/health")
async def health():
    return {"code": 0, "msg": "ok"}


CDP_PATHS = [
    "/json",
    "/json/list",
    "/json/version",
    "/json/protocol",
]


def get_browser_by_id(browser_id: str):
    b = BROWSERS.get(browser_id)
    if not b or b["status"] != "open":
        return None
    return b


async def find_browser_by_target_id(target_id: str):
    """
    根据 targetId 智能查找对应的浏览器实例

    1. 检查 target_id 是否是一个 browser_id
    2. 如果不是，查询所有浏览器实例的页面列表，找到匹配 target_id 的页面所在浏览器
    3. 如果找不到，返回第一个可用的浏览器实例
    """
    # 检查是否直接是 browser_id
    if target_id in BROWSERS:
        return BROWSERS[target_id]

    # 查询所有浏览器实例，找到匹配的页面
    for browser_id, browser in BROWSERS.items():
        if browser["status"] != "open":
            continue
        try:
            # 尝试获取浏览器的页面列表
            port = browser["port"]
            async with httpx.AsyncClient() as client:
                resp = await client.get(f"http://127.0.0.1:{port}/json/list")
                if resp.status_code == 200:
                    pages = resp.json()
                    # 检查页面是否包含 target_id
                    for page in pages:
                        if page.get("id") == target_id:
                            return browser
        except Exception as e:
            logger.error(f"查询浏览器 {browser_id} 页面列表失败: {e}")

    # 如果找不到，返回第一个可用的浏览器实例
    for browser_id, browser in BROWSERS.items():
        if browser["status"] == "open":
            return browser

    return None


@app.api_route("/json", methods=["GET"])
@app.api_route("/json/list", methods=["GET"])
@app.api_route("/json/version", methods=["GET"])
@app.api_route("/json/protocol", methods=["GET"])
async def cdp_native_proxy(request: Request):
    # 这个接口只返回最后一个浏览器实例的内容
    # 1. 获取最后一个浏览器实例
    if not BROWSERS:
        return Response(content="No browser instance", status_code=404)
    browser_id = next(reversed(BROWSERS.keys()))
    b = get_browser_by_id(browser_id)
    if not b:
        return Response(content="browser not found", status_code=404)
    port = b["port"]
    # 2. 构造目标URL
    path = request.url.path
    url = f"http://127.0.0.1:{port}{path}"
    # 3. 代理请求
    headers = dict(request.headers)
    headers["host"] = "127.0.0.1"
    body = await request.body()
    async with httpx.AsyncClient(follow_redirects=True) as client:
        resp = await client.request(
            request.method,
            url,
            headers=headers,
            content=body,
            params=request.query_params
        )
    return Response(
        content=resp.content,
        status_code=resp.status_code,
        headers=dict(resp.headers)
    )


@app.websocket("/devtools/{tab_type}/{target_id}")
async def cdp_native_ws_proxy(websocket: WebSocket, tab_type: str, target_id: str):
    if not target_id:
        await websocket.close()
        return
    # 智能查找 target_id 对应的浏览器实例
    browser = await find_browser_by_target_id(target_id)
    if not browser:
        await websocket.close(code=1008, reason="无法找到有效的浏览器实例")
        return

    await websocket.accept()
    tab_type = "page" if tab_type == "page" else "browser"
    port = browser["port"]
    target_url = f"ws://127.0.0.1:{port}/devtools/{tab_type}/{target_id}"
    logger.info(f"Forwarding to: {target_url}")

    try:
        async with websockets.connect(target_url) as target_ws:
            async def forward_client_to_server():
                try:
                    while True:
                        data = await websocket.receive_text()
                        await target_ws.send(data)
                except WebSocketDisconnect:
                    pass

            async def forward_server_to_client():
                try:
                    while True:
                        response = await target_ws.recv()
                        await websocket.send_text(response)
                except websockets.exceptions.ConnectionClosed:
                    pass

            await asyncio.gather(
                forward_client_to_server(),
                forward_server_to_client()
            )
    except Exception as e:
        logger.error(f"WebSocket代理错误: {e}")
    finally:
        await websocket.close()


if __name__ == "__main__":
    uvicorn.run("api:app", host="0.0.0.0", port=8000, reload=True)