fws_code/tgebrowser_client.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
TgeBrowser API 客户端：用于通过 REST API 新建/启动浏览器环境，获取调试端口。
文档: https://docs.tgebrowser.com/api.html
"""
from __future__ import annotations

import os
from typing import Any, Optional

import requests

# TgeBrowser 本地 API 默认端口
DEFAULT_API_BASE = "http://127.0.0.1:50326"
# API Key（写死，也可通过环境变量 TGEBROWSER_API_KEY 覆盖）
DEFAULT_API_KEY = "asp_b985921d8464496b2a7273e423d1bd005ec31bf27070ce38"
ENV_API_KEY = "TGEBROWSER_API_KEY"


def get_api_key() -> str:
    """优先从环境变量获取 API Key，未配置时使用默认值。"""
    key = os.environ.get(ENV_API_KEY)
    if key and key.strip():
        return key.strip()
    return DEFAULT_API_KEY


class TgeBrowserClient:
    """TgeBrowser 本地 API 客户端。"""

    def __init__(self, base_url: str = DEFAULT_API_BASE, api_key: Optional[str] = None):
        self.base_url = base_url.rstrip("/")
        self.api_key = api_key or get_api_key()
        self._headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }

    def _post(self, path: str, json: Optional[dict] = None) -> dict:
        resp = requests.post(
            f"{self.base_url}{path}",
            headers=self._headers,
            json=json or {},
            timeout=30,
        )
        resp.raise_for_status()
        data = resp.json()
        if not data.get("success", False):
            msg = data.get("message", "未知错误")
            raise RuntimeError(f"TgeBrowser API 失败: {msg}")
        return data

    def _get(self, path: str) -> dict:
        resp = requests.get(
            f"{self.base_url}{path}",
            headers=self._headers,
            timeout=30,
        )
        resp.raise_for_status()
        data = resp.json()
        if not data.get("success", False):
            msg = data.get("message", "未知错误")
            raise RuntimeError(f"TgeBrowser API 失败: {msg}")
        return data

    def status(self) -> dict:
        """检查 API 服务状态。"""
        return self._get("/api/status")

    def create_browser(
        self,
        browser_name: str = "miguSM_automation",
        start_page_url: Optional[str] = None,
        **kwargs: Any,
    ) -> dict:
        """
        创建新的浏览器环境。
        返回 data.envId 用于后续 start。
        """
        # 使用 TgeBrowser 文档要求的完整结构，避免 400
        start_page_value = [start_page_url] if start_page_url else []
        start_page_mode = "custom" if start_page_value else "none"

        # 指定手机端指纹，随机生成手机 UA
        mobile_fingerprint = {
            "os": "Android",
            "platformVersion": 12,
        }
        payload = {
            "browserName": browser_name,
            "proxy": {"protocol": "direct"},
            "fingerprint": mobile_fingerprint,
            "startInfo": {
                "startPage": {"mode": start_page_mode, "value": start_page_value},
                "otherConfig": {"openConfigPage": False, "checkPage": False},
            },
            **kwargs,
        }
        data = self._post("/api/browser/create", json=payload)
        return data.get("data", {})

    def start_browser(
        self,
        env_id: Optional[int] = None,
        user_index: Optional[int] = None,
        port: Optional[int] = None,
        **kwargs: Any,
    ) -> dict:
        """
        启动浏览器环境。
        返回 data 含 ws、port、http 等，用于 DrissionPage 连接。
        """
        payload: dict[str, Any] = {}
        if env_id is not None:
            payload["envId"] = env_id
        elif user_index is not None:
            payload["userIndex"] = user_index
        else:
            raise ValueError("必须指定 envId 或 userIndex")

        if port is not None:
            payload["port"] = port
        payload.update(kwargs)

        data = self._post("/api/browser/start", json=payload)
        return data.get("data", {})

    def stop_browser(self, env_id: Optional[int] = None, user_index: Optional[int] = None) -> dict:
        """停止浏览器环境。"""
        payload: dict[str, Any] = {}
        if env_id is not None:
            payload["envId"] = env_id
        elif user_index is not None:
            payload["userIndex"] = user_index
        else:
            raise ValueError("必须指定 envId 或 userIndex")
        return self._post("/api/browser/stop", json=payload)

    def create_and_start(
        self,
        browser_name: str = "miguSM_automation",
        start_page_url: Optional[str] = None,
    ) -> dict:
        """
        创建并启动新浏览器，每次调用都会新建一个环境。
        返回含 ws、port、envId 等，供 DrissionPage 连接。
        """
        create_data = self.create_browser(
            browser_name=browser_name,
            start_page_url=start_page_url,
        )
        env_id = create_data.get("envId")
        if env_id is None:
            raise RuntimeError("创建浏览器失败：未返回 envId")

        start_data = self.start_browser(env_id=env_id)
        start_data["envId"] = env_id
        return start_data