浅念
eebff9f786
Replace the Phase-0 placeholder password encryption with real RSA
(PKCS#1 v1.5). encrypt_password now builds an RSA public key from the
getSecureKey modulus/exponent (RSAPublicNumbers, e=int(hex,16),
n=int(hex,16)), encrypts the UTF-8 password, and returns hex (documented
assumption -- base64 is the one-line fallback noted in the docstring).
get_secure_key now returns the inner data dict to match the new contract;
pc_login keeps {account, password} with a comment on the username/password
and account/pwd fallbacks if pcLogin rejects.
Extend normalize_wechat_type with the Phase-0 confirmed codes:
2 -> image, 3 -> image_legacy, 4 -> video, 10 -> system (1 -> text and
-2 -> divider unchanged; 6/9/13/15/21/22 stay unknown until confirmed).
Tests:
- encrypt_password round-trip: generate an RSA keypair, build secure_key
from its public numbers, encrypt, then decrypt with the private key and
assert the original password comes back. Pins 512-hex output for a
2048-bit key. No network.
- Parser parametrize table extended for the new codes (and a downstream
message-fixture assertion updated from unknown -> image_legacy).
Suite: 87 -> 97 passing. ruff check clean.
Co-Authored-By: Claude <noreply@anthropic.com>
362 lines
14 KiB
Python
362 lines
14 KiB
Python
"""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()))
|