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

"""yunke CRM login flow: ``getSecureKey`` -> ``pcLogin``.

Real-API facts (from inspection — see design doc §6.1)
------------------------------------------------------
* ``GET /newuc/login/getSecureKey`` kicks off the login chain.
* ``POST /usercenter/login/pcLogin`` exchanges credentials for a session.
* After login every data request to ``crm2.yunkecn.com`` carries the headers
  ``userid`` / ``company`` / ``source=yunkepc`` / ``channel=4`` /
  ``content-type=application/json`` plus the login session cookie (handled
  by the httpx client's cookie jar, not by this module).

Phase-0 unknowns (clearly marked)
---------------------------------
The exact ``pcLogin`` request body, the password encryption scheme (almost
certainly RSA using the modulus/exponent from ``getSecureKey``), and the
precise response JSON shape were NOT captured. To keep the session layer
end-to-end runnable today:

* :func:`encrypt_password` is a deterministic, documented placeholder. The
  real RSA implementation drops in here later without touching call-sites.
* :func:`pc_login` parses ``userid`` / ``company`` out of a ``data`` object
  with several reasonable key-name variants. When the real shape is captured
  only this parser needs to change.

Tests mock both endpoints, so the placeholders are exercised end-to-end.
"""

from __future__ import annotations

import base64
from dataclasses import dataclass, field
from typing import Any

import httpx

# ---------------------------------------------------------------------------
# Data structures
# ---------------------------------------------------------------------------


@dataclass
class LoginResult:
    """Outcome of a successful yunke login.

    Fields:
        userid:   32-hex user id; sent as the ``userid`` header on every
                  subsequent request.
        company:  company code (e.g. ``"8m0atd"``); sent as the ``company``
                  header.
        cookies:  Snapshot of the session cookies set by ``pcLogin``. The
                  live cookies live on the httpx client's cookie jar; this
                  snapshot is kept so the session can later be persisted
                  (encrypted) into ``YkAccount.session_blob_enc``.
        raw:      The full ``pcLogin`` JSON body, kept for diagnostics and
                  for future field additions without re-rolling the parser.

    Phase-0 assumption: ``userid`` and ``company`` come from the response
    body's ``data`` object. Key-name variants tried in :func:`pc_login`:
    ``userid`` / ``userId`` / ``user_id`` and ``company`` / ``companyCode`` /
    ``company_code``.
    """

    userid: str
    company: str
    cookies: dict[str, str] = field(default_factory=dict)
    raw: dict[str, Any] = field(default_factory=dict)


# ---------------------------------------------------------------------------
# Password encryption — Phase-0 placeholder
# ---------------------------------------------------------------------------


def encrypt_password(plain: str, secure_key: dict[str, Any]) -> str:
    """Encrypt ``plain`` using the key material in ``secure_key``.

    Phase-0 TODO: real RSA encryption per ``getSecureKey``.

    The browser-side crypto was not captured. Based on the endpoint name
    (``getSecureKey``) and the typical Chinese-CRM login idiom, the real
    implementation almost certainly RSA-encrypts the password using a
    modulus/exponent (or PEM public key) returned in ``secure_key`` and
    then base64-encodes the ciphertext. Until that is confirmed by a real
    capture, this function returns base64(UTF-8(plain)) so the login flow
    is end-to-end runnable and deterministic for tests.

    Drop-in contract for the real implementation:
        def encrypt_password(plain: str, secure_key: dict[str, Any]) -> str:
            # 1. Pull modulus/exponent (or PEM) out of secure_key.
            # 2. RSA-OAEP (or PKCS#1 v1.5) encrypt UTF-8(plain).
            # 3. Return base64(ciphertext)."
    """
    # Phase 0 TODO: real RSA encryption per getSecureKey.
    return base64.b64encode(plain.encode("utf-8")).decode("ascii")


# ---------------------------------------------------------------------------
# Endpoint wrappers
# ---------------------------------------------------------------------------


async def get_secure_key(client: httpx.AsyncClient) -> dict[str, Any]:
    """Call ``GET {base}/newuc/login/getSecureKey`` and return the JSON body.

    The exact shape of the secure-key payload is a Phase-0 unknown; this
    function returns the whole decoded JSON so :func:`encrypt_password` (and
    eventually the real RSA step) can pick whatever fields it needs.
    """
    resp = await client.get("/newuc/login/getSecureKey")
    resp.raise_for_status()
    return resp.json()


async def pc_login(
    client: httpx.AsyncClient,
    username: str,
    password: str,
    secure_key: dict[str, Any],
) -> LoginResult:
    """Call ``POST {base}/usercenter/login/pcLogin`` and parse a LoginResult.

    Phase-0 assumptions about the request body (not captured):
        ``{"account": <username>, "password": <encrypt_password(...)>}``.
    Real key names (``account`` vs ``loginName`` vs ``username``) and any
    additional required fields (captcha, deviceId, ...) are to be confirmed;
    only this function needs to change when they are.

    Phase-0 assumptions about the response (not captured): JSON body with a
    ``data`` object containing ``userid`` + ``company``. Several reasonable
    key-name variants are tried (see :func:`_pick`). The session cookie is
    read back from the httpx client's cookie jar — ``pcLogin``'s
    ``Set-Cookie`` lands there automatically.
    """
    body = {
        "account": username,
        "password": encrypt_password(password, secure_key),
    }
    resp = await client.post("/usercenter/login/pcLogin", json=body)
    resp.raise_for_status()
    payload = resp.json()

    data = payload.get("data") if isinstance(payload, dict) else None
    data = data if isinstance(data, dict) else {}

    userid = _pick(data, "userid", "userId", "user_id") or _pick(
        payload, "userid", "userId", "user_id"
    )
    company = _pick(data, "company", "companyCode", "company_code") or _pick(
        payload, "company", "companyCode", "company_code"
    )

    cookies = {c.name: c.value for c in client.cookies.jar}
    return LoginResult(
        userid=userid or "",
        company=company or "",
        cookies=cookies,
        raw=payload if isinstance(payload, dict) else {"data": payload},
    )


def _pick(d: dict[str, Any], *keys: str) -> str | None:
    """Return the first present, truthy value among ``keys`` in ``d``."""
    for k in keys:
        v = d.get(k)
        if v:
            return str(v)
    return None