cmshoppe/app/chrome.py

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

from __future__ import annotations

import json
import os
import re
import shutil
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"
CHROME_APP_PATHS_KEY = r"SOFTWARE\Microsoft\Windows\CurrentVersion\App Paths\chrome.exe"


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


def _normalize_chrome_path(value) -> str:
    """Return a filesystem-ready Chrome path, or an empty string."""

    text = str(value or "").strip()
    if len(text) >= 2 and text[0] == text[-1] and text[0] in {"\"", "'"}:
        text = text[1:-1].strip()
    if not text:
        return ""
    text = os.path.expandvars(os.path.expanduser(text))
    return os.path.abspath(text)


def _load_winreg():
    """Return winreg only on Windows so non-Windows callers stay harmless."""

    if os.name != "nt":
        return None
    try:
        import winreg
    except ImportError:
        return None
    return winreg


def _registry_view_flags(winreg):
    flags = [0]
    for name in ("KEY_WOW64_64KEY", "KEY_WOW64_32KEY"):
        value = getattr(winreg, name, 0)
        if value and value not in flags:
            flags.append(value)
    return flags


def _close_registry_key(key):
    close = getattr(key, "Close", None) or getattr(key, "close", None)
    if callable(close):
        try:
            close()
        except OSError:
            pass


def _registry_chrome_candidates():
    """Yield Chrome App Paths values without exposing registry failures."""

    winreg = _load_winreg()
    if winreg is None:
        return

    key_read = getattr(winreg, "KEY_READ", 0)
    hives = (
        getattr(winreg, "HKEY_CURRENT_USER", None),
        getattr(winreg, "HKEY_LOCAL_MACHINE", None),
    )
    seen = set()
    for hive in hives:
        if hive is None:
            continue
        for view_flag in _registry_view_flags(winreg):
            marker = (hive, view_flag)
            if marker in seen:
                continue
            seen.add(marker)
            key = None
            try:
                key = winreg.OpenKey(
                    hive,
                    CHROME_APP_PATHS_KEY,
                    0,
                    key_read | view_flag,
                )
                value, _value_type = winreg.QueryValueEx(key, "")
            except (AttributeError, OSError):
                continue
            finally:
                if key is not None:
                    _close_registry_key(key)
            if value:
                yield value


def _standard_chrome_candidates():
    """Yield the supported Windows Chrome install locations."""

    suffix = ("Google", "Chrome", "Application", "chrome.exe")
    for variable in ("ProgramFiles", "ProgramFiles(x86)", "LOCALAPPDATA"):
        base = os.environ.get(variable, "").strip()
        if base:
            yield os.path.join(base, *suffix)


def is_configured_chrome_path_valid(value) -> bool:
    """Return whether a configured Chrome executable can be launched safely."""

    normalized = _normalize_chrome_path(value)
    if normalized and os.path.isfile(normalized):
        return True
    text = str(value or "").strip().strip("\"")
    return text.lower() == "chrome.exe" and bool(shutil.which(text))


def detect_chrome_path() -> str:
    """Find an installed Google Chrome executable without probing other browsers."""

    for candidate in _registry_chrome_candidates():
        normalized = _normalize_chrome_path(candidate)
        if normalized and os.path.isfile(normalized):
            return normalized
    for candidate in _standard_chrome_candidates():
        normalized = _normalize_chrome_path(candidate)
        if normalized and os.path.isfile(normalized):
            return normalized
    path_candidate = shutil.which("chrome.exe")
    normalized = _normalize_chrome_path(path_candidate)
    if normalized and os.path.isfile(normalized):
        return normalized
    return ""


def ensure_configured_chrome_path(config=None, config_path=None) -> dict:
    """Persist an auto-detected path only when the configured path is unusable."""

    path = config_path or (config or {}).get("config_path") or appconfig.CONFIG_PATH
    current = appconfig.load_config(path) if config is None else config
    configured_path = appconfig.chrome_path(current)
    if is_configured_chrome_path_valid(configured_path):
        return {
            "changed": False,
            "config": current,
            "path": configured_path,
            "message": "",
        }

    detected_path = detect_chrome_path()
    if not detected_path:
        return {
            "changed": False,
            "config": current,
            "path": configured_path,
            "message": "",
        }

    updated = dict(current)
    updated["chrome_path"] = detected_path
    saved = appconfig.save_config(updated, path=path)
    return {
        "changed": True,
        "config": saved,
        "path": detected_path,
        "message": f"已自动定位 Chrome：{detected_path}",
    }


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)