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

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

Real-API facts (from Phase-0 reconnaissance — see
``docs/superpowers/specs/phase0-findings.md`` §1)
------------------------------------------------------
* ``GET /newuc/login/getSecureKey`` returns the RSA public key as
  ``{"data": {"public_exponent": "10001", "modulus": "<hex>"}}``.
* The password is **RSA-encrypted (PKCS#1 v1.5)** with that public key
  (exponent ``0x10001`` = 65537, modulus = the returned hex) — see
  :func:`encrypt_password`.
* ``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).

Remaining inferences (clearly marked)
-------------------------------------
The exact ``pcLogin`` request-body field names and the precise response JSON
shape were not fully captured. To keep the session layer end-to-end runnable:

* :func:`pc_login` sends ``{"account", "password"}`` (password = the RSA-hex).
  If ``pcLogin`` rejects, the next variants to try are ``{"username",
  "password"}`` / ``{"account", "pwd"}`` — see phase0-findings §1.
* :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 flow is exercised end-to-end.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

import httpx
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicNumbers

# ---------------------------------------------------------------------------
# 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 — RSA (PKCS#1 v1.5), Phase-0 confirmed
# ---------------------------------------------------------------------------


def encrypt_password(plain: str, secure_key: dict[str, Any]) -> str:
    """RSA-encrypt ``plain`` (PKCS#1 v1.5) and return the ciphertext as hex.

    ``secure_key`` is the parsed ``data`` object from
    ``GET /newuc/login/getSecureKey`` (Phase-0 confirmed shape)::

        {"public_exponent": "10001", "modulus": "<hex>"}

    Both fields are hex strings. ``public_exponent`` is typically ``"10001"``
    (65537); ``modulus`` is the RSA modulus. An :class:`RSAPublicNumbers`
    public key is built from ``int(modulus, 16)`` / ``int(exponent, 16)``,
    used to PKCS#1 v1.5-encrypt UTF-8(plain), and the ciphertext is returned
    as a lowercase hex string.

    Output-format assumption: **hex**. Phase-0 confirmed the public-key shape
    but not the ciphertext encoding; hex is the common Chinese-CRM convention
    and is what this implementation uses. If ``pcLogin`` ever rejects the
    body, base64 is the first alternative to try — both are one-line changes
    here (``ciphertext.hex()`` <-> ``base64.b64encode(ciphertext).decode()``).

    Raises ``ValueError`` if ``secure_key`` is missing either field.
    """
    exp_hex = secure_key.get("public_exponent")
    mod_hex = secure_key.get("modulus")
    if not exp_hex or not mod_hex:
        raise ValueError(
            "secure_key must contain 'public_exponent' and 'modulus' (hex"
            " strings) — got: "
            f"{secure_key!r}"
        )

    e = int(str(exp_hex), 16)
    n = int(str(mod_hex), 16)
    public_key = RSAPublicNumbers(e, n).public_key()
    ciphertext = public_key.encrypt(plain.encode("utf-8"), padding.PKCS1v15())
    return ciphertext.hex()


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


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

    Phase-0 confirmed return shape::

        {"code": 10000, "data": {"public_exponent": "10001", "modulus": "<hex>"}}

    Returns the inner ``data`` dict (the ``public_exponent`` / ``modulus`` key
    material) so callers — in particular :func:`encrypt_password` — receive
    the key directly. If the wrapper is ever absent, the whole payload is
    returned as a defensive fallback.
    """
    resp = await client.get("/newuc/login/getSecureKey")
    resp.raise_for_status()
    payload = resp.json()
    if isinstance(payload, dict) and isinstance(payload.get("data"), dict):
        return payload["data"]
    return payload


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.

    Request body inference (field names NOT yet captured — see
    phase0-findings §1)::

        {"account": <username>, "password": <encrypt_password(...)>}

    where ``password`` is the RSA-hex from :func:`encrypt_password`. If
    ``pcLogin`` rejects this, the next variants to try are
    ``{"username", "password"}`` or ``{"account", "pwd"}``; any additional
    required fields (captcha, deviceId, ...) land here too. Only this
    function changes when the real shape is captured.

    Response inference (NOT yet 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