account-manager/scripts/service/chrome_profile_prefs.py

"""Chrome/Edge 独立 Profile：启用「继续上次会话」，让 Session Cookie 可跨次打开保留。"""

from __future__ import annotations

import json
import os
from typing import Any, Dict

from util.logging_config import get_skill_logger

RESTORE_CONTINUE_WHERE_LEFT_OFF = 1

# Chromium ≤ roughly mid-era: session files lived under Default/
_LEGACY_SESSION_FILES = (
    "Current Session",
    "Last Session",
    "Current Tabs",
    "Last Tabs",
)


def profile_has_restorable_session(profile_dir: str) -> bool:
    """True when the profile has on-disk session data Chrome can restore.

    Modern Chrome/Edge write ``Default/Sessions/Session_*`` (and ``Tabs_*``).
    Older builds used ``Default/Current Session`` etc. Either form counts.
    """
    default_dir = os.path.join(os.path.abspath(profile_dir), "Default")
    if any(
        os.path.isfile(os.path.join(default_dir, name))
        for name in _LEGACY_SESSION_FILES
    ):
        return True

    sessions_dir = os.path.join(default_dir, "Sessions")
    if not os.path.isdir(sessions_dir):
        return False
    try:
        names = os.listdir(sessions_dir)
    except OSError:
        return False
    return any(
        name.startswith("Session_") or name.startswith("Tabs_")
        for name in names
        if os.path.isfile(os.path.join(sessions_dir, name))
    )


def preferences_path(profile_dir: str) -> str:
    return os.path.join(os.path.abspath(profile_dir), "Default", "Preferences")


def local_state_path(profile_dir: str) -> str:
    return os.path.join(os.path.abspath(profile_dir), "Local State")


def _patch_preferences_dict(data: Dict[str, Any]) -> Dict[str, Any]:
    before = int((data.get("session") or {}).get("restore_on_startup") or 0)
    exit_before = str((data.get("profile") or {}).get("exit_type") or "")

    session = data.setdefault("session", {})
    session["restore_on_startup"] = RESTORE_CONTINUE_WHERE_LEFT_OFF
    session.pop("startup_urls", None)

    profile = data.setdefault("profile", {})
    if str(profile.get("exit_type") or "").lower() == "crashed":
        profile["exit_type"] = "Normal"
    profile["exited_cleanly"] = True

    browser = data.setdefault("browser", {})
    browser["has_seen_welcome_page"] = True
    browser["check_default_browser"] = False

    return {"restore_before": before, "exit_type_before": exit_before}


def _patch_local_state(profile_dir: str, log) -> Dict[str, Any]:
    path = local_state_path(profile_dir)
    if not os.path.isfile(path):
        return {"ok": True, "skipped": True, "path": path}
    try:
        with open(path, "r", encoding="utf-8") as fh:
            data = json.load(fh)
    except Exception as exc:
        log.warning("chrome_local_state_read_failed path=%s err=%s", path, exc)
        return {"ok": False, "path": path, "error": str(exc)}

    profile = data.setdefault("profile", {})
    info = profile.setdefault("info_cache", {}).setdefault("Default", {})
    info["active_time"] = float(info.get("active_time") or 0.0) or 1.0
    exited = profile.get("exited_cleanly", True)
    profile["exited_cleanly"] = True
    try:
        with open(path, "w", encoding="utf-8") as fh:
            json.dump(data, fh, ensure_ascii=False, separators=(",", ":"))
    except Exception as exc:
        log.warning("chrome_local_state_write_failed path=%s err=%s", path, exc)
        return {"ok": False, "path": path, "error": str(exc), "exited_cleanly_before": exited}

    log.info("chrome_local_state_patched path=%s exited_cleanly %s->True", path, exited)
    return {"ok": True, "path": path, "exited_cleanly_before": exited}


def ensure_session_restore(profile_dir: str) -> Dict[str, Any]:
    """启动浏览器前写入 session.restore_on_startup=1；清 startup_urls；Local State 干净退出。"""
    log = get_skill_logger()
    profile_abs = os.path.abspath(profile_dir)
    prefs = preferences_path(profile_abs)
    os.makedirs(os.path.dirname(prefs), exist_ok=True)
    data: Dict[str, Any] = {}
    existed = os.path.isfile(prefs)
    if existed:
        try:
            with open(prefs, "r", encoding="utf-8") as fh:
                data = json.load(fh)
        except Exception:
            log.exception("chrome_prefs_read_failed path=%s", prefs)
            data = {}

    meta = _patch_preferences_dict(data)
    before = int(meta["restore_before"])
    exit_before = str(meta["exit_type_before"] or "")

    try:
        with open(prefs, "w", encoding="utf-8") as fh:
            json.dump(data, fh, ensure_ascii=False, separators=(",", ":"))
    except Exception as exc:
        log.warning("chrome_prefs_write_failed path=%s err=%s", prefs, exc)
        return {
            "ok": False,
            "path": prefs,
            "existed": existed,
            "restore_before": before,
            "restore_after": before,
            "exit_type_before": exit_before,
            "error": str(exc),
        }

    after = int((data.get("session") or {}).get("restore_on_startup") or 0)
    exit_after = str((data.get("profile") or {}).get("exit_type") or "")
    local = _patch_local_state(profile_abs, log)
    log.info(
        "chrome_session_restore_ensured profile_dir=%s prefs=%s existed=%s "
        "restore_on_startup %s->%s exit_type %s->%s local_state=%s",
        profile_abs,
        prefs,
        existed,
        before,
        after,
        exit_before or "-",
        exit_after or "-",
        local,
    )
    return {
        "ok": True,
        "path": prefs,
        "existed": existed,
        "restore_before": before,
        "restore_after": after,
        "exit_type_before": exit_before,
        "exit_type_after": exit_after,
        "local_state": local,
    }