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.
* **Anti-risk-control / stealth** (see §6.1 风控规避): every request also
  carries browser-identical headers (UA / sec-ch-ua / accept / accept-language
  / referer) so automated traffic looks like the normal web app; a shared
  jittered :class:`app.scraper.base.RateLimiter` paces each request; a
  per-run volume cap guards against runaway loops; and risk-control signals
  (HTTP 429 or a 200 body carrying a 风控/限流 marker) trigger exponential
  backoff + retry.
* 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. The 401 path is deliberately distinct from the 429/limit
  path: a 401 is a session death (-> relogin), a 429 is throttling
  (-> backoff); they are never confused.

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

import asyncio
from collections.abc import Mapping
from typing import Any

import httpx

from app.config import get_settings
from app.scraper.base import RateLimiter
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",
}

# Browser-identical headers captured from the real CRM web app. Sent on EVERY
# request (merged ahead of the auth headers; per-request overrides win) so the
# automated client presents the same fingerprint as a normal Chrome-on-Windows
# user. ``referer`` defaults to the app root; callers can override it per call
# when they know the precise cms path. See §6.1 风控规避.
_BROWSER_HEADERS: dict[str, str] = {
    "user-agent": (
        "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
        "(KHTML, like Gecko) Chrome/150.0.0.0 Safari/537.36"
    ),
    "accept": "application/json, text/plain, */*",
    "accept-language": "zh-CN",
    "sec-ch-ua": '"Not;A=Brand";v="8", "Chromium";v="150", "Google Chrome";v="150"',
    "sec-ch-ua-mobile": "?0",
    "sec-ch-ua-platform": '"Windows"',
    "referer": "https://crm2.yunkecn.com/",
}

# Substrings that, when found in a HTTP 200 response body, indicate the CRM is
# throttling / risk-controlling us even though it returned 200 (a common
# Chinese-CRM idiom). Kept deliberately conservative — the explicit markers
# observed in 风控 messages. See §6.1 风控规避.
_RISK_BODY_MARKERS: tuple[str, ...] = ("频繁", "风控", "限制", "稍后")


class RequestCapReached(Exception):
    """Raised when a run exceeds ``max_requests_per_run`` upstream calls.

    Carries the cap and the count so callers can surface a useful message.
    The exception is a control-flow signal for orchestrators (a run that hits
    the cap should stop scheduling more work, not crash the process).
    """

    def __init__(self, cap: int, count: int) -> None:
        self.cap = cap
        self.count = count
        super().__init__(
            f"request cap reached: {count} requests issued, cap is {cap}"
        )


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,
        # --- stealth / anti-risk-control knobs (override Settings defaults) ---
        rate_limiter: RateLimiter | None = None,
        max_requests_per_run: int | None = None,
        backoff_base: float | None = None,
        backoff_max: float | None = None,
        backoff_attempts: int | 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`).

        Stealth knobs default to the conservative values declared on
        :class:`app.config.Settings`; pass any of them explicitly to override
        (tests inject a no-op ``rate_limiter`` and small backoff to stay fast).
        """
        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)

        # --- stealth configuration ----------------------------------------
        # Defaults come from Settings; explicit args win. Tests pass explicit
        # no-op values so the suite doesn't sleep on the conservative prod
        # defaults.
        s = get_settings()
        self._rate_limiter = (
            rate_limiter
            if rate_limiter is not None
            else RateLimiter(
                min_interval=s.request_min_interval, jitter=s.request_jitter
            )
        )
        self._max_requests = (
            s.max_requests_per_run if max_requests_per_run is None else max_requests_per_run
        )
        self._backoff_base = s.risk_backoff_base if backoff_base is None else backoff_base
        self._backoff_max = s.risk_backoff_max if backoff_max is None else backoff_max
        self._backoff_attempts = (
            s.risk_backoff_attempts if backoff_attempts is None else backoff_attempts
        )
        self._request_count = 0

    # -- 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

    @property
    def request_count(self) -> int:
        """Number of upstream requests issued through this client so far."""
        return self._request_count

    # -- 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,
        **kwargs: Any,
    ) -> httpx.Response:
        """Send an authed request with the full stealth layer applied.

        Per request:

        1. **Volume cap** — if ``max_requests_per_run > 0`` and the counter
           has reached it, raise :class:`RequestCapReached` (fail fast, before
           pacing or sending anything).
        2. **Throttle** — await the shared jittered :class:`RateLimiter` once,
           so successive calls are spaced at the conservative cadence.
        3. **Send + 401 handling** — :meth:`_send_with_relogin` performs the
           actual request (with browser + auth headers) and re-logs-in once on
           a 401/session-expired signal, exactly as before.
        4. **Backoff on risk-control signals** — if the response is a 429 or a
           200 carrying a 风控/限流 marker, sleep
           ``backoff_base * 2**attempt`` (capped at ``backoff_max``) and retry
           up to ``backoff_attempts`` times; raise on persistent throttling.
           This is intentionally separate from the 401 path.

        Per-request ``headers`` override both the browser and the auth
        defaults for this call only (per-request values win).
        """
        # 1. Volume cap.
        if self._max_requests > 0 and self._request_count >= self._max_requests:
            raise RequestCapReached(cap=self._max_requests, count=self._request_count)
        self._request_count += 1

        # 2. Throttle ahead of every request.
        await self._rate_limiter.wait()

        user_headers = dict(headers or {})

        # 3 + 4. Send, then back off on risk-control signals. The 401-relogin
        # behaviour lives inside ``_send_with_relogin`` and is unaffected by
        # this loop — a 401 is never treated as a rate-limit signal.
        last_resp: httpx.Response | None = None
        for attempt in range(self._backoff_attempts + 1):
            resp = await self._send_with_relogin(method, path, user_headers, **kwargs)
            if not _is_rate_limited(resp):
                return resp
            last_resp = resp
            if attempt >= self._backoff_attempts:
                break  # retries exhausted -> raise below
            delay = min(
                self._backoff_base * (2 ** attempt), self._backoff_max
            )
            await asyncio.sleep(delay)

        # Persistent risk-control signal after all retries — surface it.
        assert last_resp is not None  # loop ran at least once
        raise httpx.HTTPStatusError(
            "yunke risk-control signal persisted after"
            f" {self._backoff_attempts} backoff retries",
            request=last_resp.request,
            response=last_resp,
        )

    async def _send_with_relogin(
        self,
        method: str,
        path: str,
        user_headers: Mapping[str, str],
        *,
        _retried: bool = False,
        **kwargs: Any,
    ) -> httpx.Response:
        """Perform one actual send with browser + auth headers and 401-relogin.

        Auto-logs-in on the first call when no session is present. On a
        session-expired signal (401 or body code 401) re-logs-in *once* and
        retries; a persistent 401 raises so a dead session cannot loop. This
        preserves the pre-stealth behaviour verbatim — the only addition is
        the browser headers merged into the outbound request.
        """
        if self.login_result is None:
            await self.login()

        merged = {
            **_BROWSER_HEADERS,
            **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._send_with_relogin(
                method, path, 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.

    Note: this is deliberately disjoint from :func:`_is_rate_limited`. A 401
    is relogin territory; a 429/limit is backoff territory.
    """
    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")


def _is_rate_limited(resp: httpx.Response) -> bool:
    """Heuristic: does this response indicate risk-control / throttling?

    Triggers backoff when:

    * HTTP 429 (the canonical "too many requests" status), OR
    * HTTP 200 with a body whose text contains any of the conservative
      risk-control markers (频繁/风控/限制/稍后). Many Chinese CRMs signal
      throttling inline in a 200 envelope rather than via a proper status.

    Deliberately conservative: we do NOT flag generic non-zero business codes
    (those are ordinary errors, not throttling) and we never flag 401 (that
    is session-expiry, handled by :func:`_is_session_expired` + relogin).
    """
    if resp.status_code == 429:
        return True
    if resp.status_code != 200:
        return False
    text = resp.text
    return any(marker in text for marker in _RISK_BODY_MARKERS)