"""Authenticated HTTP client for the yunke CRM API. :class:`YunkeClient` is the single chokepoint every scraper uses to call the CRM (see design doc §6.1 会话层). It wraps an :class:`httpx.AsyncClient` and: * Auto-logs-in (``getSecureKey`` -> ``pcLogin``) on the first request when no session is present, or when constructed with credentials. * Injects the 5 known auth headers (``userid`` / ``company`` / ``source`` / ``channel`` / ``content-type``) on every request. * Detects HTTP 401 (or a body-level session-expired marker), re-logs-in *once*, and retries. A persistent 401 raises after the single retry — no infinite loop. Phase-0 note: the session is held in-memory only. Future work will persist the encrypted session blob to ``YkAccount.session_blob_enc`` (see ``app.db.models.YkAccount``) so scraper processes can resume an existing session instead of always logging in fresh. """ from __future__ import annotations from collections.abc import Mapping from typing import Any import httpx from app.session.login import ( LoginResult, get_secure_key, pc_login, ) # The 4 constant auth headers. ``userid`` + ``company`` come from the # LoginResult and are merged in separately. _STATIC_AUTH_HEADERS: dict[str, str] = { "source": "yunkepc", "channel": "4", "content-type": "application/json", } class YunkeClient: """Authenticated async HTTP client for the yunke CRM API.""" def __init__( self, base_url: str, *, username: str | None = None, password: str | None = None, userid: str | None = None, company: str | None = None, cookies: Mapping[str, str] | None = None, client: httpx.AsyncClient | None = None, ) -> None: """Construct a client either from credentials or an existing session. Credentials mode (auto-login capable): Pass ``username`` + ``password``; ``login()`` works and 401s will trigger automatic re-login. Existing-session mode (no re-login possible): Pass ``userid`` + ``company`` (optionally ``cookies``). ``login()`` raises if called — the caller has elected to skip credentials, so nothing can re-login on a 401. ``client`` may be supplied to inject a pre-configured :class:`httpx.AsyncClient` (e.g. for custom transports). When omitted a new client bound to ``base_url`` is created and owned by this instance (closed by :meth:`aclose`). """ has_creds = username is not None and password is not None has_session = userid is not None and company is not None if not has_creds and not has_session: raise ValueError( "YunkeClient requires either (username + password) for" " auto-login or (userid + company) for an existing session." ) self.base_url = base_url self.username = username self.password = password self.login_result: LoginResult | None = None self._owns_client = client is None self._client = client or httpx.AsyncClient(base_url=base_url) if has_session: assert userid is not None and company is not None # for mypy self.login_result = LoginResult( userid=userid, company=company, cookies=dict(cookies or {}) ) for name, value in (cookies or {}).items(): self._client.cookies.set(name, value) # -- properties --------------------------------------------------------- @property def client(self) -> httpx.AsyncClient: """The underlying httpx client (exposed for cookie inspection/tests).""" return self._client @property def logged_in(self) -> bool: return self.login_result is not None # -- login -------------------------------------------------------------- async def login(self) -> LoginResult: """Run ``getSecureKey`` -> ``pcLogin`` and store the LoginResult. Raises ``RuntimeError`` if this client was constructed without credentials (existing-session mode) — there is nothing to log in with. """ if self.username is None or self.password is None: raise RuntimeError( "YunkeClient.login() requires username + password; this" " client was constructed from an existing session." ) secure_key = await get_secure_key(self._client) result = await pc_login( self._client, self.username, self.password, secure_key ) self.login_result = result return result # -- authed request ----------------------------------------------------- async def request( self, method: str, path: str, *, headers: Mapping[str, str] | None = None, _retried: bool = False, **kwargs: Any, ) -> httpx.Response: """Send an authed request, auto-logging-in and re-logging-in on 401. Per-request ``headers`` override the auth defaults for this call only. ``_retried`` is an internal guard against infinite re-login loops and must not be set by callers. """ if self.login_result is None: await self.login() user_headers = dict(headers or {}) merged = {**self._auth_headers(), **user_headers} resp = await self._client.request(method, path, headers=merged, **kwargs) if _is_session_expired(resp): if _retried: # Already retried once after a re-login — give up so a # persistently dead session cannot loop forever. raise httpx.HTTPStatusError( "yunke session expired after re-login retry", request=resp.request, response=resp, ) await self.login() return await self.request( method, path, headers=user_headers, _retried=True, **kwargs ) return resp async def get(self, path: str, **kwargs: Any) -> httpx.Response: return await self.request("GET", path, **kwargs) async def post(self, path: str, **kwargs: Any) -> httpx.Response: return await self.request("POST", path, **kwargs) # -- teardown ----------------------------------------------------------- async def aclose(self) -> None: """Close the underlying httpx client if this instance owns it.""" if self._owns_client: await self._client.aclose() # -- internals ---------------------------------------------------------- def _auth_headers(self) -> dict[str, str]: if self.login_result is None: raise RuntimeError("Not logged in") # pragma: no cover return { "userid": self.login_result.userid, "company": self.login_result.company, **_STATIC_AUTH_HEADERS, } def _is_session_expired(resp: httpx.Response) -> bool: """Heuristic: does this response indicate the session is dead? * HTTP 401 -> yes. * HTTP 200 with a body-level ``code`` of ``401`` (a common Chinese-CRM idiom) -> yes. We deliberately stay narrow here — only the explicit auth-failed code, not generic non-zero codes, to avoid re-login storms on ordinary business errors. """ if resp.status_code == 401: return True try: body = resp.json() except (ValueError, httpx.DecodingError): return False return isinstance(body, dict) and body.get("code") in (401, "401")