yunke/yunke-server/scripts/smoke_wechat.py

"""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"
# Base URL used when neither ``YUNKE_BASE_URL`` env nor Settings supply one.
# Mirrors the Settings default so build_client_from_env never hands httpx a
# ``None`` base_url (which would raise TypeError at client construction).
DEFAULT_YUNKE_BASE_URL = "https://crm2.yunkecn.com"


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).

    Reads ``os.environ`` ONLY — never :func:`get_settings`. The smoke script's
    ``main()`` calls :func:`_push_settings_into_environ` first so a ``.env``
    file is honored (Settings -> os.environ, once at startup); tests call this
    function directly without the push, so they stay a pure function of
    ``os.environ`` and can be exercised with ``monkeypatch.setenv``.

    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", "").strip() or DEFAULT_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 _push_settings_into_environ() -> None:
    """Honor ``.env`` by pushing non-empty Settings values into ``os.environ``.

    The smoke script is run by a controller that supplies creds via ``.env``,
    but :func:`build_client_from_env` reads ``os.environ`` only (so it stays
    unit-testable with ``monkeypatch.setenv``). This helper bridges the two:
    it reads the process-wide :func:`get_settings` singleton (which loads
    ``.env``) and copies each non-empty session/login value into
    ``os.environ`` **only if that var is not already set there** — real env
    wins, exactly as the original ``os.environ.get(...) or settings.*`` chain
    did, but evaluated ONCE at smoke startup instead of inside the pure
    builder.

    Synchronous and side-effectful by design; called as the first line of
    :func:`main`. Not called by the unit tests, so they remain hermetic.
    """
    settings = get_settings()
    _ENV_MAP = {
        "YUNKE_USERNAME": settings.yunke_username,
        "YUNKE_PASSWORD": settings.yunke_password,
        "YK_USERID": settings.yk_userid,
        "YK_COMPANY": settings.yk_company,
        "YK_COOKIE": settings.yk_cookie,
    }
    for key, value in _ENV_MAP.items():
        if value and not os.environ.get(key):
            os.environ[key] = value
    # YUNKE_BASE_URL: only push when the setting diverges from the default and
    # the env var is not already set (avoids overwriting an explicit real env).
    if settings.yunke_base_url and not os.environ.get("YUNKE_BASE_URL"):
        os.environ["YUNKE_BASE_URL"] = settings.yunke_base_url


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)."""
    # Honor .env (via Settings) by lifting non-empty values into os.environ
    # BEFORE the pure, os.environ-only builder runs. Real env vars win; this
    # only fills gaps. Done once at startup so build_client_from_env stays a
    # pure function of os.environ (and stays unit-testable with setenv/delenv).
    _push_settings_into_environ()
    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.")
            await _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.
        await _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()))