diff --git a/yunke-server/scripts/scrape_via_page.py b/yunke-server/scripts/scrape_via_page.py index 0037cc5..c0cf86b 100644 --- a/yunke-server/scripts/scrape_via_page.py +++ b/yunke-server/scripts/scrape_via_page.py @@ -1185,6 +1185,92 @@ def _check_auth_envelope(resp): return resp +# ── cookie 自动续期(2026-09-03)────────────────────────────────────────── # +# 缓存文件优先于 compose 注入的 YK_COOKIE(自动续期成功时写入,永远比 env 新); +# 实体撞 CookieExpiredError → 浏览器表单重登 → 落缓存 → 整实体重跑一次 +# (upsert 幂等 + 水位增量,重跑安全);二次过期直接抛出,绝不循环。 + + +def _cookie_cache_path(): + from pathlib import Path + + override = os.environ.get("YK_COOKIE_CACHE") + if override: + return Path(override) + from app.config import get_settings + + return Path(get_settings().media_dir) / ".yk_cookie_cache" + + +def _load_cookie() -> str: + """缓存文件 > YK_COOKIE env > settings(.env)。空串=三者皆无。""" + try: + cached = _cookie_cache_path().read_text(encoding="utf-8").strip() + if cached: + return cached + except OSError: + pass + cookie = os.environ.get("YK_COOKIE", "") + if not cookie: + from app.config import get_settings + + cookie = get_settings().yk_cookie + return cookie or "" + + +def _save_cookie_cache(cookie: str) -> None: + """尽力而为持久化:路径不可写(本地无挂载)只告警不炸。""" + try: + p = _cookie_cache_path() + p.parent.mkdir(parents=True, exist_ok=True) + p.write_text(cookie, encoding="utf-8") + except OSError as exc: + logger.warning(f"cookie 缓存写入失败({exc}),新 cookie 仅本次进程生效") + + +async def _run_with_auto_relogin(entity, cookie: str, *, relogin=None): + """跑一个 entity(cookie 参数化);撞过期自动重登换 cookie 重跑一次。""" + if relogin is None: + relogin = _relogin + try: + return await entity(cookie) + except CookieExpiredError: + logger.warning("cookie 过期 → 浏览器表单自动重登…") + new_cookie = await relogin() + _save_cookie_cache(new_cookie) + logger.info("自动续期成功,新 cookie 已缓存;重跑本实体") + return await entity(new_cookie) + + +async def _relogin() -> str: + """浏览器表单登录云客换取新 cookie(2026-09-03 实测一次成功)。 + + 选择器来自真实页面:tab 文本「密码登录」、placeholder 请输入手机号/ + 请输入密码、按钮「登 录」;成功判据 = URL 跳到 /cms/home/* 且 + document.cookie 出现 userToken。 + """ + from app.config import get_settings + + s = get_settings() + if not s.yunke_username or not s.yunke_password: + raise CookieExpiredError("自动重登失败:未配置 YUNKE_USERNAME/YUNKE_PASSWORD") + try: + async with PageScraper("") as scraper: + page = await scraper.context.new_page() + await page.goto(f"{YUNKE_BASE}/cms/auth/login", wait_until="networkidle", timeout=30000) + await page.click("text=密码登录") + await page.fill('input[placeholder="请输入手机号"]', s.yunke_username) + await page.fill('input[placeholder="请输入密码"]', s.yunke_password) + await page.click('button:has-text("登 录")') + await page.wait_for_url("**/cms/home/**", timeout=20000) + cookie = await page.evaluate("() => document.cookie") + except Exception as exc: + raise CookieExpiredError(f"自动重登失败:{exc}") from exc + if "userToken=" not in cookie: + raise CookieExpiredError("自动重登失败:登录后未见 userToken cookie") + return cookie + + def _rows_flat(body: dict) -> list[dict]: """Extract rows from the flat envelope: ``data.data[]`` (wxUser/list, queryWeChatVoiceList, call/list).""" @@ -2207,17 +2293,21 @@ async def main(): parser.add_argument("--full-history", action="store_true", help="消息类实体强制全量翻历史(默认增量水位)") args = parser.parse_args() - cookie = os.environ.get("YK_COOKIE", "") - if not cookie: - # Try reading from .env - from app.config import get_settings - cookie = get_settings().yk_cookie + cookie = _load_cookie() if not cookie: logger.error("YK_COOKIE not set in .env or environment") return logger.info(f"Starting page-based scrape: entity={args.entity}") + async def _run(c: str) -> None: + await _dispatch_entity(args, c) + + # cookie 过期自动续期:重登→落缓存→整实体重跑(2026-09-03) + await _run_with_auto_relogin(_run, cookie) + + +async def _dispatch_entity(args, cookie: str) -> None: if args.entity == "friends": await scrape_friends(cookie, args.max_employees) elif args.entity == "groups": diff --git a/yunke-server/tests/test_auto_relogin.py b/yunke-server/tests/test_auto_relogin.py new file mode 100644 index 0000000..179007a --- /dev/null +++ b/yunke-server/tests/test_auto_relogin.py @@ -0,0 +1,106 @@ +"""cookie 自动续期(2026-09-03):缓存优先级 + 撞过期自动重登重跑。 + +机制::func:`_load_cookie` 优先读缓存文件(自动续期成功时写入, +MEDIA_DIR/.yk_cookie_cache)——比 compose 注入的 YK_COOKIE 新;实体运行撞 +:class:`CookieExpiredError` 时 :func:`_run_with_auto_relogin` 调浏览器表单 +重登拿新 cookie、落缓存、整实体重跑一次(幂等,水位增量天然安全)。 +二次过期不再重试(防死循环)。 +""" + +from __future__ import annotations + +import os + +import pytest + +from scripts import scrape_via_page as sv + + +@pytest.fixture +def cache_file(tmp_path, monkeypatch): + """把 cookie 缓存指到临时目录。""" + path = tmp_path / ".yk_cookie_cache" + monkeypatch.setenv("YK_COOKIE_CACHE", str(path)) + return path + + +# ── 缓存:读/写/优先级 ────────────────────────────────────────────────────── # + + +def test_load_cookie_prefers_cache_over_env(cache_file, monkeypatch): + monkeypatch.setenv("YK_COOKIE", "env-cookie") + cache_file.write_text("cache-cookie", encoding="utf-8") + assert sv._load_cookie() == "cache-cookie" + + +def test_load_cookie_falls_back_to_env_when_no_cache(cache_file, monkeypatch): + monkeypatch.setenv("YK_COOKIE", "env-cookie") + assert sv._load_cookie() == "env-cookie" + + +def test_save_then_load_roundtrip(cache_file): + sv._save_cookie_cache("fresh-cookie") + assert cache_file.read_text(encoding="utf-8").strip() == "fresh-cookie" + assert sv._load_cookie() == "fresh-cookie" + + +def test_save_cache_unwritable_path_does_not_raise(monkeypatch): + # 持久化是尽力而为:路径不可写(本地无挂载等)不能炸,回退 env 语义仍在 + monkeypatch.setenv("YK_COOKIE_CACHE", "/nonexistent-root/x/y/.cookie") + sv._save_cookie_cache("whatever") # 不应抛 + + +# ── 撞过期自动重登重跑 ───────────────────────────────────────────────────── # + + +async def test_retry_after_relogin(cache_file): + calls: list[str] = [] + + async def entity(cookie: str) -> str: + calls.append(cookie) + if cookie == "old": + raise sv.CookieExpiredError("云客 cookie 已过期") + return f"ok-with-{cookie}" + + async def fake_relogin() -> str: + return "new" + + result = await sv._run_with_auto_relogin(entity, "old", relogin=fake_relogin) + assert result == "ok-with-new" + assert calls == ["old", "new"] # 旧 cookie 失败一次,新 cookie 重跑成功 + assert sv._load_cookie() == "new" # 新 cookie 已落缓存 + + +async def test_no_relogin_when_other_error(cache_file): + async def entity(cookie: str): + raise RuntimeError("业务炸了") + + async def fake_relogin() -> str: + raise AssertionError("不应触发重登") + + with pytest.raises(RuntimeError, match="业务炸了"): + await sv._run_with_auto_relogin(entity, "old", relogin=fake_relogin) + + +async def test_second_expiry_propagates_no_loop(cache_file): + """重登后仍过期 → 直接抛出,绝不无限循环。""" + + async def always_expired(cookie: str): + raise sv.CookieExpiredError("还是过期") + + async def fake_relogin() -> str: + return "renewed" + + with pytest.raises(sv.CookieExpiredError): + await sv._run_with_auto_relogin(always_expired, "old", relogin=fake_relogin) + + +async def test_relogin_failure_propagates(cache_file): + async def entity(cookie: str): + raise sv.CookieExpiredError("过期") + + async def failing_relogin() -> str: + raise sv.CookieExpiredError("自动重登失败(无账密/验证码)") + + with pytest.raises(sv.CookieExpiredError, match="自动重登失败"): + await sv._run_with_auto_relogin(entity, "old", relogin=failing_relogin)