cmshoppe/app/chrome.py

"""Chrome launcher helpers for per-account CDP sessions."""

from __future__ import annotations

import json
import os
import subprocess
import time
import urllib.error
import urllib.request

from . import appconfig
from . import config as account_config


CDP_HOST = "127.0.0.1"


class ChromeLaunchError(RuntimeError):
    """Raised when Chrome launch configuration is invalid."""


def _account_value(account, field, default=None):
    if isinstance(account, dict):
        return account.get(field, default)
    return getattr(account, field, default)


def _required_account_value(account, field):
    value = _account_value(account, field)
    if value is None or str(value).strip() == "":
        raise ChromeLaunchError(f"账号缺少字段: {field}")
    return value


def _debug_port(account) -> int:
    value = int(_required_account_value(account, "debug_port"))
    if value <= 0 or value > 65535:
        raise ChromeLaunchError("debug_port 必须在 1-65535 范围内")
    return value


def _user_data_dir(account, config=None) -> str:
    existing = _account_value(account, "user_data_dir")
    if existing:
        path = os.path.abspath(str(existing))
        os.makedirs(path, exist_ok=True)
        return path

    slug = _account_value(account, "slug")
    if not slug:
        slug = account_config.make_slug(_required_account_value(account, "alias"))
    return account_config.ensure_user_data_dir(slug, config=config)


def build_launch_args(account, config=None) -> list:
    """Build Chrome command arguments for one account."""

    chrome_path = appconfig.chrome_path(config)
    if not chrome_path:
        raise ChromeLaunchError("chrome_path 不能为空")
    port = _debug_port(account)
    user_data_dir = _user_data_dir(account, config=config)
    return [
        chrome_path,
        f"--remote-debugging-port={port}",
        "--remote-allow-origins=*",
        f"--user-data-dir={user_data_dir}",
    ]


def launch_chrome(account, config=None) -> subprocess.Popen:
    """Launch Chrome for one account and return the process handle."""

    args = build_launch_args(account, config=config)
    try:
        return subprocess.Popen(args)
    except OSError as exc:
        raise ChromeLaunchError(f"启动 Chrome 失败: {exc}") from exc


def _json_version_url(port, host=CDP_HOST):
    return f"http://{host}:{int(port)}/json/version"


def _fetch_json_version(port, host=CDP_HOST, timeout=1.0):
    request = urllib.request.Request(_json_version_url(port, host=host), method="GET")
    opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))
    with opener.open(request, timeout=timeout) as response:
        if response.status != 200:
            raise ChromeLaunchError(f"CDP 返回状态码: {response.status}")
        body = response.read()
    return json.loads(body.decode("utf-8"))


def is_running(port, host=CDP_HOST) -> bool:
    """Return whether a Chrome CDP endpoint responds on /json/version."""

    try:
        _fetch_json_version(port, host=host, timeout=1.0)
        return True
    except (OSError, urllib.error.URLError, json.JSONDecodeError, ChromeLaunchError):
        return False


def wait_debug_ready(port, timeout=60, host=CDP_HOST) -> bool:
    """Poll until the Chrome CDP endpoint is ready or timeout expires."""

    deadline = time.monotonic() + float(timeout)
    while True:
        if is_running(port, host=host):
            return True
        if time.monotonic() >= deadline:
            return False
        time.sleep(0.25)