xhs_crawler_claude/app/auth_manager.py

"""auth_manager.py

登录态管理（单账号）
- 使用 Playwright storage_state 保存到 data/auth_state.json
- 额外保存元数据到 data/auth_meta.json（登录时间/最后验证时间等）
- 提供 verify_login / get_login_status 等方法

设计目标：
1) 不依赖“文件存在”判断登录，而是用 Playwright 打开页面做轻量验证
2) 单账号长期复用：只要不掉线，就一直用
"""

from __future__ import annotations

import json
import os
from dataclasses import dataclass
from datetime import datetime, timedelta
from typing import Optional, Tuple

from playwright.async_api import Browser, BrowserContext, TimeoutError as PlaywrightTimeoutError

from app import config


AUTH_META_FILE = os.path.join(config.DATA_DIR, "auth_meta.json")


@dataclass
class LoginStatus:
    logged_in: bool
    login_time: Optional[str] = None
    last_verified_time: Optional[str] = None
    message: Optional[str] = None


def _utc_iso() -> str:
    # 使用本地时间 ISO 字符串，方便 UI 展示
    return datetime.now().isoformat(timespec="seconds")


def load_auth_meta() -> dict:
    if not os.path.exists(AUTH_META_FILE):
        return {}
    try:
        with open(AUTH_META_FILE, "r", encoding="utf-8") as f:
            return json.load(f) or {}
    except Exception:
        return {}


def save_auth_meta(meta: dict) -> None:
    tmp = AUTH_META_FILE + ".tmp"
    with open(tmp, "w", encoding="utf-8") as f:
        json.dump(meta, f, ensure_ascii=False, indent=2)
    os.replace(tmp, AUTH_META_FILE)


async def _new_context(browser: Browser) -> BrowserContext:
    """基于已保存的 storage_state 创建 context。"""
    if os.path.exists(config.AUTH_STATE_FILE):
        return await browser.new_context(storage_state=config.AUTH_STATE_FILE)
    return await browser.new_context()


async def open_login_page(browser: Browser) -> None:
    """打开登录页面（通常需要用户在 VNC/有界面环境里扫码）。"""
    context = await _new_context(browser)
    page = await context.new_page()
    await page.goto(config.LOGIN_URL, wait_until="domcontentloaded")
    # 不关闭页面：让用户扫码
    # 注意：context/page 生命周期由浏览器进程维护；任务结束时统一关闭 browser 即可


async def verify_login(browser: Browser, timeout_sec: int = 20) -> Tuple[bool, str]:
    """验证当前 storage_state 是否仍然有效。

    策略（尽量稳，避免依赖某个易变 selector）：
    1) 打开首页；
    2) 若跳转到 /login 或出现明显登录按钮/文案，则判定未登录；
    3) 若未跳登录页，并且存在一定数量 cookie，则判定登录有效（启发式）。
    """
    if not os.path.exists(config.AUTH_STATE_FILE):
        return False, "未找到登录态文件，请先登录"

    context = await _new_context(browser)
    page = await context.new_page()

    try:
        await page.goto(config.XIAOHONGSHU_BASE_URL + "/", wait_until="domcontentloaded", timeout=timeout_sec * 1000)
        await page.wait_for_timeout(800)

        url = page.url or ""
        if "/login" in url:
            return False, "已跳转到登录页（登录态可能已失效）"

        # 启发式：页面上是否出现明显“登录”入口
        login_text_hits = await page.locator("text=登录").count()
        if login_text_hits > 0:
            # 有些已登录页面也会出现“登录”字样（比如弹窗），再做一次 cookie 辅助判断
            cookies = await context.cookies()
            if len(cookies) < 3:
                return False, "检测到登录入口且 cookie 数量过少，疑似未登录"

        cookies = await context.cookies()
        if len(cookies) < 3:
            return False, "cookie 数量过少，疑似未登录"

        # 通过：写入 meta
        meta = load_auth_meta()
        if not meta.get("login_time"):
            meta["login_time"] = _utc_iso()
        meta["last_verified_time"] = _utc_iso()
        save_auth_meta(meta)

        # 刷新保存一次 state（让 cookie 续期信息尽可能落盘）
        try:
            await context.storage_state(path=config.AUTH_STATE_FILE)
        except Exception:
            # 不要因为保存失败影响验证结果
            pass

        return True, "ok"
    except PlaywrightTimeoutError:
        return False, "验证登录超时（网络或风控），请稍后重试"
    except Exception as e:
        return False, f"验证登录异常: {e}"
    finally:
        try:
            await page.close()
        except Exception:
            pass
        try:
            await context.close()
        except Exception:
            pass


async def get_login_status(browser: Optional[Browser]) -> LoginStatus:
    """获取登录状态（如果 browser 可用则做轻量验证，否则仅返回本地 meta 状态）。"""
    meta = load_auth_meta()
    if not os.path.exists(config.AUTH_STATE_FILE):
        return LoginStatus(logged_in=False, message="未登录")

    if browser is None:
        return LoginStatus(
            logged_in=True if meta.get("login_time") else True,
            login_time=meta.get("login_time"),
            last_verified_time=meta.get("last_verified_time"),
            message="仅本地状态（未进行在线验证）"
        )

    ok, msg = await verify_login(browser)
    meta = load_auth_meta()
    return LoginStatus(
        logged_in=ok,
        login_time=meta.get("login_time"),
        last_verified_time=meta.get("last_verified_time"),
        message=None if ok else msg
    )


def clear_login_state() -> None:
    """清理登录态。"""
    for p in [config.AUTH_STATE_FILE, AUTH_META_FILE]:
        try:
            if os.path.exists(p):
                os.remove(p)
        except Exception:
            pass