ksjsb/core/h5_kws.py

"""Pure-Python KWS/WebWeapon fallback ticket generator.

The Nebula H5 bundle initializes WebWeapon roughly as:

- product cookie: ``kwpsecproductname=kuaishou-growth``;
- fallback fingerprint: AES-CBC(``...|nonce8``, key=iv ``K8wm...``)
  wrapped as ``K + b64[:4] + W + b64[4:-2] + F + b64[-2:]``;
- fallback signature token: random 64-char ``kwssectoken`` plus
  AES-CBC(``...|kwssectoken[:8]``, key=iv ``H4t...``) wrapped with ``S``.

When the online Web DFP config succeeds, APP may replace these with
server-provided ``secToken`` + sign-script output.  This module implements the
local fallback path so a fresh ``ksck`` does not need stale copied KWS cookies.
"""

from __future__ import annotations

import base64
import json
import os
import re
import secrets
import subprocess
import time
from dataclasses import dataclass
from pathlib import Path
from random import Random
from typing import Any, Mapping, Protocol
from urllib.parse import quote

from core.h5_kww_alg import kwf_aes_cbc_decrypt, kwf_aes_cbc_encrypt


KWS_PRODUCT_NAME = "kuaishou-growth"
KWS_FINGERPRINT_KEY = "K8wm5PvY9nX7qJc2"
KWS_TOKEN_KEY = "H4tL6rNd3vB9xM5k"
KWS_CONFIG_KEY = "webweaponconfigs"
KWS_CONFIG_URL = "https://gdfp.gifshow.com/s/w/c"
KWS_RANDOM_ALPHABET = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"
KWS_URL_SLICE_LEN = 80
KWS_SCRIPT_CODE_RE = re.compile(r"^[0-9a-z]{64}$")
DEFAULT_KWS_SIGN_RUNNER = Path(__file__).with_name("h5_kws_sign.mjs")
DEFAULT_KWS_SIGN_SCRIPT = Path(__file__).with_name("kws-11-0.0.1-obfuscated.5e0a90af726d8a7e.js")


class _RandLike(Protocol):
    def randrange(self, stop: int) -> int: ...


@dataclass(frozen=True)
class H5KwsTicket:
    kwpsecproductname: str
    kwssectoken: str
    kwscode: str
    kwfv1: str
    kww: str
    fingerprint_plain: str
    token_plain: str

    def cookie_fields(self, *, include_fingerprint: bool = True) -> dict[str, str]:
        fields = {
            "kwpsecproductname": self.kwpsecproductname,
            "kwssectoken": self.kwssectoken,
            "kwscode": self.kwscode,
        }
        if include_fingerprint:
            fields["kwfv1"] = self.kwfv1
        return fields


@dataclass(frozen=True)
class H5KwsConfigRequest:
    product_name: str
    ts: int
    did: str
    plain: str
    data: str
    url: str = KWS_CONFIG_URL

    def body(self) -> dict[str, str]:
        return {"data": self.data}


def _now_ms() -> int:
    return int(time.time() * 1000)


def _random_text(length: int, rng: _RandLike | None = None) -> str:
    alphabet = KWS_RANDOM_ALPHABET
    if rng is None:
        return "".join(secrets.choice(alphabet) for _ in range(length))
    return "".join(alphabet[rng.randrange(len(alphabet))] for _ in range(length))


def _encode_uri(value: str) -> str:
    """Match JavaScript ``encodeURI`` for the URL prefix used by WebWeapon."""

    return quote(str(value), safe=";/?:@&=+$,#-_.!~*'()")


def _encrypt_webweapon_b64(plain: str, key_text: str) -> str:
    key = key_text.encode("utf-8")
    encrypted = kwf_aes_cbc_encrypt(plain.encode("utf-8"), key=key, iv=key)
    return base64.b64encode(encrypted).decode("ascii")


def webweapon_aes_encrypt_b64(
    plain: str,
    key_text: str = KWS_CONFIG_KEY,
) -> str:
    """AES-CBC/PKCS7 + Base64 used by WebWeapon config payloads."""

    return _encrypt_webweapon_b64(str(plain), key_text)


def webweapon_aes_decrypt_b64(
    encrypted_b64: str,
    key_text: str = KWS_CONFIG_KEY,
) -> str:
    """Decrypt WebWeapon AES-CBC/PKCS7 Base64 text."""

    key = key_text.encode("utf-8")
    try:
        encrypted = base64.b64decode(encrypted_b64, validate=True)
        plain = kwf_aes_cbc_decrypt(encrypted, key=key, iv=key)
    except Exception as exc:
        raise ValueError("invalid WebWeapon AES payload") from exc
    return plain.decode("utf-8")


def _wrap_webweapon_value(encrypted_b64: str, marker: str) -> str:
    return f"K{encrypted_b64[:4]}W{encrypted_b64[4:-2]}{marker}{encrypted_b64[-2:]}"


def is_h5_kws_script_code(value: str) -> bool:
    """Validate the 64-char code emitted by the KWS ``signUrl`` script."""

    return bool(KWS_SCRIPT_CODE_RE.fullmatch(str(value or "")))


def _short_runner_error(stdout: str, stderr: str) -> str:
    detail = (stderr or stdout or "").strip()
    if len(detail) <= 500:
        return detail
    return detail[:240] + "...(truncated)..." + detail[-240:]


def run_h5_kws_sign_script(
    *,
    script_path: str | Path | None = None,
    runner: str | Path = DEFAULT_KWS_SIGN_RUNNER,
    node_bin: str | None = None,
    timeout: int = 10,
) -> str:
    """Run the local KWS ``signUrl`` VM script and return ``kwscode``.

    This is the script-equivalent path for the online WebWeapon config mode:
    ``/s/w/c`` gives ``secToken`` + ``signUrl``; loading that sign script calls
    ``window.kwscb(code)``.  The Node runner supplies a deterministic browser
    fixture and captures the callback result.
    """

    runner_path = Path(runner)
    script = Path(script_path) if script_path is not None else DEFAULT_KWS_SIGN_SCRIPT
    if not runner_path.exists():
        raise FileNotFoundError(f"KWS sign runner not found: {runner_path}")
    if not script.exists():
        raise FileNotFoundError(f"KWS sign script not found: {script}")

    node = node_bin or os.environ.get("NODE_BIN") or "node"
    proc = subprocess.run(
        [node, str(runner_path), str(script)],
        text=True,
        encoding="utf-8",
        capture_output=True,
        timeout=timeout,
        check=False,
    )
    if proc.returncode != 0:
        raise RuntimeError(f"KWS sign runner failed: {_short_runner_error(proc.stdout, proc.stderr)}")
    try:
        data = json.loads(proc.stdout)
    except json.JSONDecodeError as exc:
        raise RuntimeError(f"KWS sign runner returned invalid JSON: {_short_runner_error(proc.stdout, proc.stderr)}") from exc
    if not data.get("ok"):
        raise RuntimeError(f"KWS sign runner failed: {data.get('error')}")
    kwscode = str(data.get("kwscode") or "")
    if not is_h5_kws_script_code(kwscode):
        raise RuntimeError(f"KWS sign runner returned invalid kwscode length/shape: {kwscode!r}")
    return kwscode


def build_h5_kws_script_ticket(
    *,
    sec_token: str,
    kwscode: str | None = None,
    product_name: str = KWS_PRODUCT_NAME,
    script_path: str | Path | None = None,
    runner: str | Path = DEFAULT_KWS_SIGN_RUNNER,
    node_bin: str | None = None,
    timeout: int = 10,
) -> dict[str, str]:
    """Build KWS cookie fields from server ``secToken`` and sign script code."""

    token = str(sec_token or "")
    if not token:
        raise ValueError("sec_token must be non-empty")
    code = kwscode
    if not code:
        script = Path(script_path) if script_path is not None else DEFAULT_KWS_SIGN_SCRIPT
        try:
            from core.h5_kws_vm import kwscode_from_known_h5_kws_script

            code = kwscode_from_known_h5_kws_script(script)
        except Exception:
            code = None
    if not code:
        code = run_h5_kws_sign_script(
            script_path=script_path,
            runner=runner,
            node_bin=node_bin,
            timeout=timeout,
        )
    if not is_h5_kws_script_code(code):
        raise ValueError("kwscode must be a 64-char lowercase alnum KWS script code")
    return {
        "kwpsecproductname": product_name,
        "kwssectoken": token,
        "kwscode": code,
    }


def build_h5_kws_default_ticket(
    *,
    url: str,
    did: str,
    product_name: str = KWS_PRODUCT_NAME,
    now_ms: int | None = None,
    fingerprint_nonce: str | None = None,
    sec_token: str | None = None,
    rng: Random | None = None,
) -> H5KwsTicket:
    """Build WebWeapon ``getDefaultData(true)`` ticket fields.

    ``fingerprint_nonce`` corresponds to the random 8-char suffix used for the
    fallback ``kwfv1``.  ``sec_token`` corresponds to the 64-char
    ``kwssectoken``; its first 8 chars are folded into ``kwscode``.
    """

    ts = _now_ms() if now_ms is None else int(now_ms)
    fp_nonce = fingerprint_nonce if fingerprint_nonce is not None else _random_text(8, rng)
    if len(fp_nonce) != 8:
        raise ValueError("fingerprint_nonce must be 8 characters")

    token = sec_token if sec_token is not None else _random_text(64, rng)
    if len(token) != 64:
        raise ValueError("sec_token must be 64 characters")

    url_prefix = _encode_uri(str(url)[:KWS_URL_SLICE_LEN])
    common = f"{url_prefix}|{did}|{product_name}|{ts}|"
    fingerprint_plain = common + fp_nonce
    token_plain = common + token[:8]

    kwfv1 = _wrap_webweapon_value(
        _encrypt_webweapon_b64(fingerprint_plain, KWS_FINGERPRINT_KEY),
        "F",
    )
    kwscode = _wrap_webweapon_value(
        _encrypt_webweapon_b64(token_plain, KWS_TOKEN_KEY),
        "S",
    )
    return H5KwsTicket(
        kwpsecproductname=product_name,
        kwssectoken=token,
        kwscode=kwscode,
        kwfv1=kwfv1,
        kww=kwfv1,
        fingerprint_plain=fingerprint_plain,
        token_plain=token_plain,
    )


def build_h5_kws_config_request_data(
    *,
    did: str,
    product_name: str = KWS_PRODUCT_NAME,
    ts_ms: int | None = None,
) -> H5KwsConfigRequest:
    """Build the encrypted ``/s/w/c`` config POST body.

    WebWeapon sends compact JSON in this exact key order:
    ``productName`` → ``ts`` → ``did``.  The JSON string is encrypted with
    AES-CBC where key and IV are both ``webweaponconfigs``.
    """

    ts = _now_ms() if ts_ms is None else int(ts_ms)
    payload = {
        "productName": product_name,
        "ts": ts,
        "did": str(did),
    }
    plain = json.dumps(payload, ensure_ascii=False, separators=(",", ":"))
    return H5KwsConfigRequest(
        product_name=product_name,
        ts=ts,
        did=str(did),
        plain=plain,
        data=webweapon_aes_encrypt_b64(plain, KWS_CONFIG_KEY),
    )


def decrypt_h5_kws_config_response_data_rsp(data_rsp: str) -> dict[str, Any]:
    """Decrypt the ``dataRsp`` field returned by ``/s/w/c``."""

    plain = webweapon_aes_decrypt_b64(data_rsp, KWS_CONFIG_KEY)
    try:
        decoded = json.loads(plain)
    except json.JSONDecodeError as exc:
        raise ValueError("WebWeapon config response is not JSON") from exc
    if not isinstance(decoded, dict):
        raise ValueError("WebWeapon config response must be a JSON object")
    return decoded


def decrypt_h5_kws_config_response(response: Mapping[str, Any]) -> dict[str, Any]:
    """Decrypt a full ``/s/w/c`` JSON response object containing ``dataRsp``."""

    data_rsp = response.get("dataRsp")
    if not isinstance(data_rsp, str) or not data_rsp:
        raise ValueError("WebWeapon config response is missing dataRsp")
    return decrypt_h5_kws_config_response_data_rsp(data_rsp)


__all__ = [
    "DEFAULT_KWS_SIGN_RUNNER",
    "DEFAULT_KWS_SIGN_SCRIPT",
    "H5KwsTicket",
    "H5KwsConfigRequest",
    "KWS_CONFIG_KEY",
    "KWS_CONFIG_URL",
    "KWS_FINGERPRINT_KEY",
    "KWS_PRODUCT_NAME",
    "KWS_RANDOM_ALPHABET",
    "KWS_SCRIPT_CODE_RE",
    "KWS_TOKEN_KEY",
    "KWS_URL_SLICE_LEN",
    "build_h5_kws_script_ticket",
    "build_h5_kws_config_request_data",
    "build_h5_kws_default_ticket",
    "decrypt_h5_kws_config_response",
    "decrypt_h5_kws_config_response_data_rsp",
    "is_h5_kws_script_code",
    "run_h5_kws_sign_script",
    "webweapon_aes_decrypt_b64",
    "webweapon_aes_encrypt_b64",
]