"""End-to-end smoke check for the yunke WeChat data mirror (Task 1.9). Wires the existing pieces together against the REAL yunke CRM (read-only) and a local SQLite DB, then prints a summary. It exists to prove the plumbing (session -> scraper -> upsert -> DB) holds end-to-end before the full sync pipeline is built. Auth modes (mutually exclusive; session mode preferred for now) -------------------------------------------------------------- 1. **Existing-session mode (most reliable).** Set ``YK_USERID``, ``YK_COMPANY``, and optionally ``YK_COOKIE``. The script builds a :class:`YunkeClient` directly from that captured session (``userid=`` / ``company=`` / ``cookies=``) and BYPASSES login. This avoids the remaining ``pcLogin`` body-field-name inference in mode 2. 2. **Login mode.** Set ``YUNKE_USERNAME`` and ``YUNKE_PASSWORD`` (and leave the session vars unset). The script calls :meth:`YunkeClient.login`, which runs the ``getSecureKey`` -> ``pcLogin`` flow. Password encryption is real RSA (PKCS#1 v1.5) using the ``getSecureKey`` modulus/exponent — see :func:`app.session.login.encrypt_password`. The remaining inference is the ``pcLogin`` request-body field names; if the real CRM rejects the body, the script catches that, prints a clear message, and exits with code 3. If NEITHER set of env vars is present, the script prints usage and exits with code 2. What it does (kept tiny — read-only) ------------------------------------ * Creates all ORM tables on the smoke DB (``Base.metadata.create_all``). * Scrapes up to ~10 WeChat friends (:func:`scrape_wechat_friends`, first page only), upserts them via :func:`upsert_customers`, prints the count and the first friend (friend_id / friend_nickname / owner_account_id). * For the first scraped friend, scrapes up to ~50 of their messages from the last 7 days (:func:`scrape_wechat_messages`), upserts via :func:`upsert_messages`, prints the count and a 3-row sample. * Verifies rows landed by querying the smoke DB back, and prints a final summary. The script NEVER writes to the upstream CRM — every upstream call is a list / search endpoint. It writes only to the local smoke DB. How to run ---------- From ``yunke-server/`` with the venv active (``.venv/Scripts/activate`` on Windows, ``source .venv/bin/activate`` on *nix):: # Existing-session mode (works today — controller supplies the values): export YK_USERID=aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa export YK_COMPANY=8m0atd export YK_COOKIE="JSESSIONID=deadbeef; other=foo" python -m scripts.smoke_wechat # Login mode (real RSA encryption; pcLogin body field names inferred): export YUNKE_USERNAME=user export YUNKE_PASSWORD=pass python -m scripts.smoke_wechat The smoke DB defaults to ``./smoke.db`` (SQLite); override the URL with ``SMOKE_DB_URL`` (e.g. ``sqlite+aiosqlite:///./smoke.db`` or any other SQLAlchemy async URL). On a 401 / expired session, the script prints a helpful message and exits non-zero (code 4) instead of dumping a traceback. Exit codes: ``0`` success · ``2`` no auth env · ``3`` login failed (e.g. an inferred ``pcLogin`` body field name the real CRM rejects) · ``4`` upstream auth/session error mid-scrape. """ from __future__ import annotations import asyncio import json import os import sys import time from typing import Any import httpx import sqlalchemy as sa from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine from app.config import get_settings from app.db import models # noqa: F401 (register tables on Base.metadata) from app.db.base import Base from app.db.models import Customer, Message from app.scraper.wechat import scrape_wechat_friends, scrape_wechat_messages from app.session.yunke_client import YunkeClient from app.storage.repos import upsert_customers, upsert_messages # Caps that keep the smoke run tiny (≈10 customers, ≈50 messages). MAX_FRIENDS = 10 MAX_MESSAGES = 50 # Last 7 days in ms — backward window for the message scraper. MESSAGE_WINDOW_MS = 7 * 86_400_000 DEFAULT_SMOKE_DB_URL = "sqlite+aiosqlite:///./smoke.db" def _parse_cookie_env(raw: str) -> dict[str, str]: """Parse the ``YK_COOKIE`` env var into a ``{name: value}`` dict. Accepts either: * a raw ``Cookie:`` header (``"JSESSIONID=abc; foo=bar"``), or * a JSON object (``'{"JSESSIONID": "abc"}'``). Empty / unparsable input yields ``{}`` (the caller is fine with no cookies — the ``userid`` + ``company`` headers may be enough on their own). """ raw = (raw or "").strip() if not raw: return {} # JSON first — handy for structured config. if raw.startswith("{"): try: obj = json.loads(raw) except ValueError: return {} if not isinstance(obj, dict): return {} return {str(k): str(v) for k, v in obj.items()} # Otherwise treat as a raw header (semicolon-separated name=value pairs). out: dict[str, str] = {} for part in raw.split(";"): part = part.strip() if not part or "=" not in part: continue name, _, value = part.partition("=") out[name.strip()] = value.strip() return out def build_client_from_env() -> YunkeClient | None: """Construct a :class:`YunkeClient` from env vars (pure / network-free). Selection rules (in priority order): 1. **Existing-session mode** — if EITHER of ``YK_USERID`` / ``YK_COMPANY`` is set, BOTH must be set (``YK_COOKIE`` optional). Constructs the client with that session and bypasses login. 2. **Login mode** — if ``YUNKE_USERNAME`` and ``YUNKE_PASSWORD`` are both set, constructs a credentialed client. The caller must call ``await client.login()`` to actually authenticate. 3. Otherwise returns ``None`` so the caller can print usage and exit. Raises ``ValueError`` if only one of ``YK_USERID`` / ``YK_COMPANY`` is set (a clear caller bug, not something to silently paper over). Pure / network-free — safe to unit-test with ``monkeypatch.setenv``. """ base_url = os.environ.get("YUNKE_BASE_URL") or get_settings().yunke_base_url userid = os.environ.get("YK_USERID", "").strip() company = os.environ.get("YK_COMPANY", "").strip() cookie_raw = os.environ.get("YK_COOKIE", "") if userid or company: # Session mode — BOTH required. if not (userid and company): raise ValueError( "YK_USERID and YK_COMPANY must be set together " "(got only one)." ) return YunkeClient( base_url=base_url, userid=userid, company=company, cookies=_parse_cookie_env(cookie_raw), ) username = os.environ.get("YUNKE_USERNAME", "").strip() password = os.environ.get("YUNKE_PASSWORD", "") if username and password: return YunkeClient( base_url=base_url, username=username, password=password, ) return None def _print_usage(stream=sys.stderr) -> None: stream.write( "smoke_wechat: missing auth env vars.\n" "Set EITHER existing-session env (YK_USERID + YK_COMPANY [+ YK_COOKIE])\n" " OR login env (YUNKE_USERNAME + YUNKE_PASSWORD).\n" "See `python -m scripts.smoke_wechat --help` or the module docstring.\n" ) async def main() -> int: """Run the smoke check. Returns a process exit code (0 = success).""" client = build_client_from_env() if client is None: _print_usage() return 2 # Login mode: try to log in. The pcLogin body field names are still an # inference, so a real-CRM rejection is possible — catch and explain # instead of tracebacking. if not client.logged_in: try: await client.login() except Exception as exc: # noqa: BLE001 — any failure is fatal-but-explained sys.stderr.write( f"\n[smoke] login failed (real RSA encryption is in place; " f"the likely cause is an inferred pcLogin body field name): " f"{exc!r}\n" "Use existing-session mode (YK_USERID / YK_COMPANY / YK_COOKIE) " "if login keeps failing.\n" ) await client.aclose() return 3 smoke_db_url = os.environ.get("SMOKE_DB_URL", DEFAULT_SMOKE_DB_URL) engine = create_async_engine(smoke_db_url, future=True) session_maker = async_sessionmaker(engine, expire_on_commit=False) try: # 1) Create all ORM tables on the smoke DB. async with engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) customers_upserted = 0 messages_upserted = 0 first_friend: dict[str, Any] | None = None sample_msgs: list[dict[str, Any]] = [] # 2) Customers — consume only the first page worth (~10 records). print( f"[smoke] scraping up to {MAX_FRIENDS} friends from " f"{client.base_url} ..." ) friends: list[dict[str, Any]] = [] try: agen = scrape_wechat_friends(client, page_size=MAX_FRIENDS) try: async for rec in agen: friends.append(rec) if len(friends) >= MAX_FRIENDS: break finally: # Explicit cleanup — `async for` does not synchronously close # the generator on `break`. await agen.aclose() except httpx.HTTPStatusError as exc: sys.stderr.write( f"\n[smoke] upstream auth/session error scraping friends: " f"HTTP {exc.response.status_code}.\n" "Session likely expired — refresh YK_COOKIE / YK_USERID / " "YK_COMPANY and retry.\n" ) return 4 if friends: async with session_maker() as session: customers_upserted = await upsert_customers(session, friends) await session.commit() first_friend = friends[0] if first_friend is None: print("[smoke] no friends returned; nothing more to do.") _print_summary(session_maker, customers_upserted, messages_upserted) return 0 print(f"[smoke] customers upserted: {customers_upserted}") print( "[smoke] first friend: " f"friend_id={first_friend.get('friend_id')!r} " f"friend_nickname={first_friend.get('friend_nickname')!r} " f"owner_account_id={first_friend.get('owner_account_id')!r}" ) # 3) Messages for the first friend — last 7 days, capped at 50. owner = first_friend.get("owner_account_id") or "" friend_id = first_friend.get("friend_id") or "" if not (owner and friend_id): print( "[smoke] first friend missing owner_account_id/friend_id; " "skipping messages." ) else: end_ms = int(time.time() * 1000) start_ms = end_ms - MESSAGE_WINDOW_MS print( f"[smoke] scraping up to {MAX_MESSAGES} messages for " f"owner={owner!r} friend={friend_id!r} over the last 7 days ..." ) msgs: list[dict[str, Any]] = [] try: mgen = scrape_wechat_messages( client, wechat_id=owner, friend_id=friend_id, start=start_ms, end=end_ms, ) try: async for rec in mgen: msgs.append(rec) if len(msgs) >= MAX_MESSAGES: break finally: await mgen.aclose() except httpx.HTTPStatusError as exc: sys.stderr.write( f"\n[smoke] upstream auth/session error scraping messages: " f"HTTP {exc.response.status_code}.\n" "Session likely expired — refresh and retry.\n" ) return 4 if msgs: async with session_maker() as session: messages_upserted = await upsert_messages(session, msgs) await session.commit() sample_msgs = msgs[:3] print(f"[smoke] messages upserted: {messages_upserted}") for rec in sample_msgs: text = (rec.get("text") or "") print( f" msg_id={rec.get('msg_id')!r} type={rec.get('type')!r} " f"type_raw={rec.get('type_raw')!r} ts={rec.get('timestamp')!r} " f"text={text[:40]!r}" ) # 4) Verify rows landed by querying the smoke DB back. _print_summary(session_maker, customers_upserted, messages_upserted) return 0 finally: await client.aclose() await engine.dispose() async def _print_summary( session_maker: async_sessionmaker, customers_upserted: int, messages_upserted: int, ) -> None: """Print the final summary, including DB row counts read back from SQLite.""" async with session_maker() as session: db_customers = await session.scalar( sa.select(sa.func.count()).select_from(Customer) ) db_messages = await session.scalar( sa.select(sa.func.count()).select_from(Message) ) print("-" * 60) print("[smoke] SUMMARY") print(f" customers upserted : {customers_upserted}") print(f" messages upserted : {messages_upserted}") print(f" customer rows in DB : {db_customers}") print(f" message rows in DB : {db_messages}") print("-" * 60) if __name__ == "__main__": sys.exit(asyncio.run(main()))