chrome.py

import os
import sys
import json
import socket
import subprocess
import time
import platform
import logging
import requests
import websocket
from typing import Optional, Tuple

logger = logging.getLogger(__name__)

DEFAULT_CDP_URL = "http://localhost:9222"


class BrowserInfo:
    def __init__(self, web_socket_debugger_url: str):
        self.webSocketDebuggerUrl = web_socket_debugger_url


def start_chrome(debug_port: int = 0) -> Tuple[Optional[subprocess.Popen], str]:
    """启动 Chrome 浏览器

    Args:
        debug_port: 远程调试端口，默认为0表示自动获取

    Returns:
        (chrome_process, user_data_dir) 元组
    """
    home_dir = os.path.expanduser("~")
    user_data_dir = os.path.join(home_dir, "ChromeProfile")

    if debug_port == 0:
        debug_port_env = os.getenv("DEBUG_PORT", "")
        if debug_port_env:
            debug_port = int(debug_port_env)
        if debug_port == 0:
            debug_port = 9222

    if is_port_open("localhost", debug_port):
        return None, user_data_dir

    chrome_path = find_chrome_path()
    if not chrome_path:
        raise RuntimeError("Chrome browser not found")

    cmd = [
        chrome_path,
        f"--remote-debugging-port={debug_port}",
        "--remote-allow-origins=*",
        "--remote-debugging-address=0.0.0.0",  # 可选：允许外部访问
        "--no-first-run",
        "--no-default-browser-check",
        "--disable-gpu",
        "--disable-extensions",
        "--disable-plugins",
        "--disable-sync",
        f"--user-data-dir={user_data_dir}",
        # 静默日志输出的关键参数
        "--log-level=3",  # 只显示致命错误
        "--silent-startup",  # 静默启动
        "--disable-dev-shm-usage",  # 减少崩溃
        "--disable-logging",  # 禁用日志记录
        "--disable-ipc-flooding-protection",  # 减少日志
    ]

    # 将 stdout 和 stderr 重定向到空设备
    if platform.system() == "Windows":
        null_device = "NUL"
    else:
        null_device = "/dev/null"

    try:
        process = subprocess.Popen(
            cmd,
            stdout=open(null_device, "w"),
            stderr=open(null_device, "w")
        )
        logger.info(f"Started Chrome with PID: {process.pid}")
        return process, user_data_dir
    except Exception as e:
        raise RuntimeError(f"Failed to start Chrome: {e}")


def shutdown_chrome(process: Optional[subprocess.Popen], base_url: str = ""):
    """关闭 Chrome 浏览器

    Args:
        process: Chrome 进程对象
        base_url: CDP 基础 URL
    """
    if not base_url:
        base_url = DEFAULT_CDP_URL

    # WebSocket Browser.close
    if close_chrome_via_cdp(base_url):
        logger.info("通过 Browser.close 成功关闭")
        return

    # 最终方案: SIGTERM
    logger.info("CDP 关闭失败，回退到 SIGTERM...")
    graceful_shutdown(process)


def close_chrome_via_cdp(base_url: str) -> bool:
    """使用 WebSocket 发送 Browser.close 命令

    Args:
        base_url: CDP 基础 URL

    Returns:
        是否成功关闭
    """
    cdp_websocket_url = get_browser_websocket_url(base_url)
    if not cdp_websocket_url:
        return False

    try:
        ws = websocket.create_connection(cdp_websocket_url, timeout=5)
        # 发送 Browser.close 命令
        close_cmd = json.dumps({
            "id": 1,
            "method": "Browser.close"
        })
        ws.send(close_cmd)
        ws.close()
        logger.info("已发送 Browser.close 命令")
        return True
    except Exception as e:
        logger.error(f"WebSocket error: {e}")
        return False


def get_browser_websocket_url(base_url: str) -> Optional[str]:
    """获取浏览器 WebSocket 调试 URL

    Args:
        base_url: CDP 基础 URL

    Returns:
        WebSocket URL 或 None
    """
    try:
        resp = requests.get(f"{base_url}/json/version", timeout=5)
        if resp.status_code == 200:
            data = resp.json()
            return data.get("webSocketDebuggerUrl")
    except Exception as e:
        logger.error(f"Failed to get browser websocket URL: {e}")
    return None


def graceful_shutdown(process: Optional[subprocess.Popen]):
    """优雅关闭 Chrome 进程

    Args:
        process: Chrome 进程对象
    """
    if not process or process.poll() is not None:
        logger.info("Chrome 进程不存在")
        return

    logger.info("发送 SIGTERM 以优雅关闭 Chrome...")
    if platform.system() == "Windows":
        # Windows 不支持 SIGTERM，尝试其他方式或直接 Kill
        process.kill()
    else:
        process.terminate()

    # 等待退出
    try:
        process.wait(timeout=10)
        logger.info("Chrome 已优雅退出")
    except subprocess.TimeoutExpired:
        logger.info("超时，强制终止")
        process.kill()


def find_chrome_path() -> Optional[str]:
    """查找 Chrome 浏览器可执行文件路径

    Returns:
        Chrome 可执行文件路径或 None
    """
    system = platform.system()
    if system == "Windows":
        return find_chrome_on_windows()
    elif system == "Darwin":
        return find_chrome_on_macos()
    else:
        return find_chrome_on_unix_like()


def find_chrome_on_windows() -> Optional[str]:
    """在 Windows 上查找 Chrome

    Returns:
        Chrome 可执行文件路径或 None
    """
    # 检查常见安装路径
    common_paths = [
        r"C:\Program Files\Google\Chrome\Application\chrome.exe",
        r"C:\Program Files (x86)\Google\Chrome\Application\chrome.exe",
    ]

    for path in common_paths:
        if os.path.exists(path):
            return path

    # 尝试从注册表查找
    try:
        import winreg
        key = winreg.OpenKey(
            winreg.HKEY_LOCAL_MACHINE,
            r"SOFTWARE\Microsoft\Windows\CurrentVersion\App Paths\chrome.exe"
        )
        chrome_path, _ = winreg.QueryValueEx(key, None)
        winreg.CloseKey(key)

        if os.path.exists(chrome_path):
            return os.path.abspath(chrome_path)
    except Exception as e:
        logger.error(f"Failed to find Chrome in registry: {e}")

    return None


def find_chrome_on_macos() -> Optional[str]:
    """在 macOS 上查找 Chrome

    Returns:
        Chrome 可执行文件路径或 None
    """
    path = "/Applications/Google Chrome.app/Contents/MacOS/Google Chrome"
    if os.path.exists(path):
        return path
    return None


def find_chrome_on_unix_like() -> Optional[str]:
    """在 Unix-like 系统上查找 Chrome

    Returns:
        Chrome 可执行文件路径或 None
    """
    paths = []

    # 添加 PATH 环境变量中的目录
    path_env = os.getenv("PATH", "")
    if path_env:
        paths.extend(path_env.split(":"))

    # 添加常见安装目录
    paths.extend([
        "/usr/bin",
        "/usr/local/bin",
        "/opt/google/chrome/bin",
    ])

    chrome_exe = "google-chrome"  # or "chrome" depending on your system

    for dir_path in paths:
        full_path = os.path.join(dir_path, chrome_exe)
        if os.path.exists(full_path):
            return os.path.abspath(full_path)

    return None


def is_port_open(host: str, port: int) -> bool:
    """检查本地指定端口是否已经开启（监听）

    Args:
        host: 主机地址
        port: 端口号

    Returns:
        端口是否开启
    """
    try:
        sock = socket.create_connection((host, port), timeout=2)
        sock.close()
        return True
    except (socket.timeout, ConnectionRefusedError, OSError):
        return False