cmshoppe/app/chrome.py

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

from __future__ import annotations

import json
import os
import re
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 _safe_shortcut_name(value) -> str:
    text = str(value or "").strip()
    if not text:
        raise ChromeLaunchError("快捷方式名称不能为空")
    text = re.sub(r'[<>:"/\\|?*\x00-\x1f]+', "_", text)
    text = re.sub(r"\s+", " ", text).strip(" .")
    if not text:
        raise ChromeLaunchError("快捷方式名称不能为空")
    return text


def _default_shortcut_name(account) -> str:
    alias = _account_value(account, "alias")
    account_name = _account_value(account, "account_name")
    name = alias or account_name or _required_account_value(account, "slug")
    return _safe_shortcut_name(f"cmshopee-{name}")


def _desktop_dir() -> str:
    return os.path.join(os.path.expanduser("~"), "Desktop")


def _ps_single_quote(value) -> str:
    return "'" + str(value).replace("'", "''") + "'"


def create_shortcut(account, shortcut_path=None, desktop_dir=None, name=None, config=None) -> str:
    """Create a Windows .lnk shortcut for one account's Chrome session."""

    args = build_launch_args(account, config=config)
    chrome_path = os.path.abspath(args[0])
    arguments = subprocess.list2cmdline(args[1:])
    if shortcut_path is None:
        root = os.path.abspath(desktop_dir or _desktop_dir())
        shortcut_name = _safe_shortcut_name(name) if name else _default_shortcut_name(account)
        shortcut_path = os.path.join(root, shortcut_name + ".lnk")
    shortcut_path = os.path.abspath(str(shortcut_path))
    if not shortcut_path.lower().endswith(".lnk"):
        shortcut_path += ".lnk"
    os.makedirs(os.path.dirname(shortcut_path), exist_ok=True)

    working_directory = os.path.dirname(chrome_path) or os.getcwd()
    script = "; ".join(
        [
            "$shell = New-Object -ComObject WScript.Shell",
            f"$shortcut = $shell.CreateShortcut({_ps_single_quote(shortcut_path)})",
            f"$shortcut.TargetPath = {_ps_single_quote(chrome_path)}",
            f"$shortcut.Arguments = {_ps_single_quote(arguments)}",
            f"$shortcut.WorkingDirectory = {_ps_single_quote(working_directory)}",
            f"$shortcut.IconLocation = {_ps_single_quote(chrome_path)}",
            "$shortcut.Save()",
        ]
    )
    try:
        subprocess.run(
            [
                "powershell",
                "-NoProfile",
                "-ExecutionPolicy",
                "Bypass",
                "-Command",
                script,
            ],
            check=True,
            capture_output=True,
            text=True,
        )
    except OSError as exc:
        raise ChromeLaunchError(f"生成快捷方式失败: {exc}") from exc
    except subprocess.CalledProcessError as exc:
        error = (exc.stderr or exc.stdout or str(exc)).strip()
        raise ChromeLaunchError(f"生成快捷方式失败: {error}") from exc
    return shortcut_path


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)