yunke/yunke-server/app/session/yunke_client.py

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