yunke/yunke-server/tests/test_yunke_client_stealth.py

"""Tests for the anti-risk-control / stealth layer in :class:`YunkeClient`.

These exercise the four stealth concerns that ``request()`` now centralizes:

* **Browser-identical headers** on every request (UA / sec-ch-ua / accept
  / accept-language / referer) merged with the existing auth headers, with
  per-request overrides winning.
* A shared jittered :class:`RateLimiter` awaited ahead of each upstream
  request.
* **Exponential backoff + retry** on risk-control signals (HTTP 429, or a 200
  body carrying a 风控/限流 marker such as 频繁/风控/限制/稍后) — distinct from
  401 session-expiry, which keeps its existing relogin-once behaviour.
* A hard per-run request cap that raises :class:`RequestCapReached`.

Every test injects explicit stealth params (no-op throttle, small backoff) so
the suite stays fast and deterministic regardless of the process-wide
:class:`Settings` defaults (which are deliberately slow in production).
"""

from __future__ import annotations

from unittest.mock import AsyncMock

import httpx
import pytest
import respx
from cryptography.hazmat.primitives.asymmetric import rsa

from app.scraper.base import RateLimiter
from app.session.yunke_client import RequestCapReached, YunkeClient

BASE = "https://crm2.yunkecn.com"
USERID = "a" * 32
COMPANY = "8m0atd"

DATA_PATH = "/wechatweb/wxFriends/friend_list/es"
DATA_URL = f"{BASE}{DATA_PATH}"

# The exact browser-captured header values the client must reproduce.
UA = (
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
    "(KHTML, like Gecko) Chrome/150.0.0.0 Safari/537.36"
)
SEC_CH_UA = '"Not;A=Brand";v="8", "Chromium";v="150", "Google Chrome";v="150"'

# Backoff schedule uses base=2.0 so expected delays are 2.0, 4.0, 8.0, ...
BACKOFF_BASE = 2.0

# A real RSA keypair backing the mocked ``getSecureKey`` endpoint so the real
# RSA step in ``pc_login`` has a valid key to encrypt against (mirrors the
# pattern in tests/test_yunke_client.py). Generated once at import.
_TEST_PRIVATE_KEY = rsa.generate_private_key(public_exponent=65537, key_size=2048)


def _secure_key_from_test_public_key() -> dict[str, str]:
    nums = _TEST_PRIVATE_KEY.public_key().public_numbers()
    return {
        "public_exponent": format(nums.e, "x"),
        "modulus": format(nums.n, "x"),
    }


_TEST_SECURE_KEY = _secure_key_from_test_public_key()


def _patch_backoff_sleep(monkeypatch, sink: list[float]) -> None:
    """Replace ``asyncio.sleep`` in yunke_client with an instant recorder."""

    async def fake_sleep(delay: float) -> None:
        sink.append(delay)

    monkeypatch.setattr("app.session.yunke_client.asyncio.sleep", fake_sleep)


@pytest.fixture()
async def client():
    """A logged-in (session-mode) client with a no-op throttle.

    Session mode (userid + company) skips the login flow entirely. Tests that
    need relogin construct their own credentials-mode client inline.
    """
    c = YunkeClient(
        base_url=BASE,
        userid=USERID,
        company=COMPANY,
        rate_limiter=RateLimiter(min_interval=0.0, jitter=0.0),
    )
    try:
        yield c
    finally:
        await c.aclose()


# ---------------------------------------------------------------------------
# 1. Browser-identical headers on every request
# ---------------------------------------------------------------------------


@respx.mock
async def test_browser_headers_present_on_every_request(client: YunkeClient):
    route = respx.post(DATA_URL).mock(
        return_value=httpx.Response(200, json={"ok": True})
    )

    await client.post(DATA_PATH, json={"pageNum": 1})

    sent = route.calls.last.request.headers
    assert sent["user-agent"] == UA
    assert sent["accept"] == "application/json, text/plain, */*"
    assert sent["accept-language"] == "zh-CN"
    assert sent["sec-ch-ua"] == SEC_CH_UA
    assert sent["sec-ch-ua-mobile"] == "?0"
    assert sent["sec-ch-ua-platform"] == '"Windows"'
    assert sent["referer"] == "https://crm2.yunkecn.com/"
    # The existing auth headers still ride along unchanged.
    assert sent["source"] == "yunkepc"
    assert sent["channel"] == "4"
    assert sent["userid"] == USERID
    assert sent["company"] == COMPANY
    assert sent["content-type"] == "application/json"


@respx.mock
async def test_per_request_header_overrides_browser_defaults(client: YunkeClient):
    """A caller-supplied header wins over the browser + auth defaults."""
    route = respx.post(DATA_URL).mock(
        return_value=httpx.Response(200, json={"ok": True})
    )

    await client.post(
        DATA_PATH, json={}, headers={"referer": "https://crm2.yunkecn.com/cms/x"}
    )

    assert route.calls.last.request.headers["referer"] == (
        "https://crm2.yunkecn.com/cms/x"
    )


# ---------------------------------------------------------------------------
# 2. Rate limiter awaited once per request
# ---------------------------------------------------------------------------


@respx.mock
async def test_rate_limiter_wait_awaited_once_per_request(client: YunkeClient):
    respx.post(DATA_URL).mock(return_value=httpx.Response(200, json={"ok": True}))

    spy = AsyncMock()
    client._rate_limiter.wait = spy  # type: ignore[method-assign]

    await client.post(DATA_PATH, json={})

    assert spy.await_count == 1


# ---------------------------------------------------------------------------
# 3. Risk-control signals -> exponential backoff + retry
# ---------------------------------------------------------------------------


@respx.mock
async def test_429_triggers_backoff_retry_then_success(
    monkeypatch, client: YunkeClient
):
    delays: list[float] = []
    _patch_backoff_sleep(monkeypatch, delays)

    route = respx.post(DATA_URL).mock(
        side_effect=[
            httpx.Response(429, json={"msg": "请求过于频繁"}),
            httpx.Response(429, json={"msg": "请求过于频繁"}),
            httpx.Response(200, json={"data": [{"id": 1}]}),
        ]
    )

    resp = await client.post(DATA_PATH, json={})

    assert resp.status_code == 200
    assert resp.json() == {"data": [{"id": 1}]}
    # Initial try + two retries after backoff.
    assert route.call_count == 3
    # base * 2**0 = 2.0, base * 2**1 = 4.0
    assert delays == [2.0, 4.0]


@respx.mock
async def test_200_with_risk_marker_triggers_backoff(monkeypatch, client: YunkeClient):
    """A 200 body carrying a 风控/限流 marker must also trigger backoff — many
    Chinese CRMs signal throttling inline rather than via a real 429."""
    _patch_backoff_sleep(monkeypatch, [])

    route = respx.post(DATA_URL).mock(
        side_effect=[
            httpx.Response(200, json={"code": 500, "msg": "操作过于频繁"}),
            httpx.Response(200, json={"data": []}),
        ]
    )

    resp = await client.post(DATA_PATH, json={})

    assert resp.status_code == 200
    assert resp.json() == {"data": []}
    assert route.call_count == 2


@respx.mock
async def test_persistent_429_raises_after_backoff_attempts(monkeypatch):
    """When the upstream keeps signaling risk past ``risk_backoff_attempts``,
    the client gives up and raises an :class:`httpx.HTTPStatusError`."""
    delays: list[float] = []
    _patch_backoff_sleep(monkeypatch, delays)

    c = YunkeClient(
        base_url=BASE,
        userid=USERID,
        company=COMPANY,
        rate_limiter=RateLimiter(min_interval=0.0, jitter=0.0),
        backoff_attempts=2,
        backoff_base=BACKOFF_BASE,
    )
    try:
        respx.post(DATA_URL).mock(
            return_value=httpx.Response(429, json={"msg": "请求过于频繁"})
        )
        with pytest.raises(httpx.HTTPStatusError):
            await c.post(DATA_PATH, json={})
        # 1 initial attempt + 2 retries = 3 upstream calls, 2 backoff sleeps.
        assert delays == [2.0, 4.0]
    finally:
        await c.aclose()


# ---------------------------------------------------------------------------
# 4. Volume cap -> RequestCapReached
# ---------------------------------------------------------------------------


@respx.mock
async def test_volume_cap_raises_on_nth_request():
    respx.post(DATA_URL).mock(return_value=httpx.Response(200, json={"ok": True}))
    capped = YunkeClient(
        base_url=BASE,
        userid=USERID,
        company=COMPANY,
        rate_limiter=RateLimiter(min_interval=0.0, jitter=0.0),
        max_requests_per_run=2,
    )
    try:
        await capped.post(DATA_PATH, json={})  # 1 -> ok
        await capped.post(DATA_PATH, json={})  # 2 -> ok
        with pytest.raises(RequestCapReached):
            await capped.post(DATA_PATH, json={})  # 3 -> cap hit
    finally:
        await capped.aclose()


@respx.mock
async def test_volume_cap_zero_means_unlimited(client: YunkeClient):
    """``max_requests_per_run=0`` (the default) disables the cap entirely."""
    respx.post(DATA_URL).mock(return_value=httpx.Response(200, json={"ok": True}))
    for _ in range(5):
        await client.post(DATA_PATH, json={})


# ---------------------------------------------------------------------------
# 5. 401 session-expiry still relogs exactly once (existing behaviour preserved)
# ---------------------------------------------------------------------------


@respx.mock
async def test_401_still_relogs_once_with_stealth_active():
    """Stealth must not break the existing 401 -> relogin-once -> retry path.

    A 401 is a session-death (handled by relogin), NOT a risk-control signal,
    so it must not enter the backoff loop.
    """
    secure = respx.get(f"{BASE}/newuc/login/getSecureKey").mock(
        return_value=httpx.Response(200, json={"data": _TEST_SECURE_KEY})
    )
    login = respx.post(f"{BASE}/usercenter/login/pcLogin").mock(
        return_value=httpx.Response(
            200,
            json={"code": 0, "data": {"userid": USERID, "company": COMPANY}},
            headers={"Set-Cookie": "JSESSIONID=deadbeef; Path=/"},
        )
    )
    data = respx.post(DATA_URL).mock(
        side_effect=[
            httpx.Response(401, json={"msg": "session expired"}),
            httpx.Response(200, json={"data": [{"id": 1}]}),
        ]
    )

    c = YunkeClient(
        base_url=BASE,
        username="tester",
        password="hunter2",
        rate_limiter=RateLimiter(min_interval=0.0, jitter=0.0),
    )
    # Prime logged-in so only the 401-driven relogin is counted (mirrors the
    # pattern in test_yunke_client.py).
    from app.session.login import LoginResult

    c.login_result = LoginResult(userid=USERID, company=COMPANY)
    try:
        resp = await c.post(DATA_PATH, json={})
        assert resp.status_code == 200
        assert resp.json() == {"data": [{"id": 1}]}
        assert secure.called
        assert login.call_count == 1
        assert data.call_count == 2
    finally:
        await c.aclose()