ksjsb/core/dfp_sign.py

"""DFP / unifiedId atlasSign input builders and 10405 sign wrappers."""

from __future__ import annotations

from dataclasses import dataclass
import hashlib
from typing import Mapping

from .atlas_sign import atlas_sign, atlas_sign_to_digest_hex
from .enc_data import ZT_OUTER_CONFIGS
from .sig3 import kwsg_10418_digest24_unmix


DFP_SDK_ID = "7e46b28a-8c93-4940-8238-4c60e64e3c81"
DFP_SIGN_KEYS = {"sign", "__NS_sig3", "__NS_xfalcon", "__NStokensig", "sig"}
DFP_PAYLOAD_KEYS = ("deviceInfo", "carryInfo", "data")


@dataclass(frozen=True)
class DfpSignMaterial:
    scheme: str
    keys: tuple[str, ...]
    sign_input: str
    excluded_keys: tuple[str, ...] = ()
    note: str = ""

    @property
    def input_len(self) -> int:
        return len(self.sign_input)

    @property
    def input_sha256(self) -> str:
        return hashlib.sha256(self.sign_input.encode("utf-8", errors="ignore")).hexdigest()


def _first_payload_key(params: Mapping[str, str], preferred: tuple[str, ...]) -> str:
    for key in preferred:
        if key in params:
            return key
    for key in DFP_PAYLOAD_KEYS:
        if key in params:
            return key
    raise KeyError("missing DFP encrypted payload field")


def legacy_product_ts_sv_payload(params: Mapping[str, str]) -> DfpSignMaterial:
    """Build `productName + ts + sv + encryptedPayload`.

    Used by `gdfp/report` and `unifiedId/logReport/android`.
    """

    payload_key = _first_payload_key(params, ("deviceInfo", "carryInfo", "data"))
    keys = ("productName", "ts", "sv", payload_key)
    sign_input = (
        params.get("productName", "")
        + params.get("ts", "")
        + params.get("sv", "2")
        + params.get(payload_key, "")
    )
    excluded = tuple(key for key in params if key not in keys and key not in DFP_SIGN_KEYS)
    return DfpSignMaterial(
        scheme="legacy_product_ts_sv_payload",
        keys=keys,
        sign_input=sign_input,
        excluded_keys=excluded,
        note="productName + ts + sv + encrypted payload; append rdid/didtag/ft after sign",
    )


def id_mapping_data_only(params: Mapping[str, str]) -> DfpSignMaterial:
    """Build `data` only for `unifiedId/logReport/idMapping`."""

    payload_key = _first_payload_key(params, ("data",))
    return DfpSignMaterial(
        scheme="id_mapping_data_only",
        keys=(payload_key,),
        sign_input=params.get(payload_key, ""),
        note="idMapping signs only encoded data",
    )


def tree_values_sorted_non_empty_except_sign(params: Mapping[str, str]) -> DfpSignMaterial:
    """Build TreeMap-style non-empty values without separators."""

    keys = tuple(
        key
        for key in sorted(params)
        if key not in DFP_SIGN_KEYS and params.get(key, "") != ""
    )
    sign_input = "".join(params[key] for key in keys)
    excluded = tuple(key for key in params if key not in keys and key not in DFP_SIGN_KEYS)
    return DfpSignMaterial(
        scheme="tree_values_sorted_non_empty_except_sign",
        keys=keys,
        sign_input=sign_input,
        excluded_keys=excluded,
        note="TreeMap key order, non-empty values, no separators",
    )


def build_dfp_sign_material(kind: str, params: Mapping[str, str]) -> DfpSignMaterial:
    if kind in {"gdfp_report", "unified_log_report"}:
        return legacy_product_ts_sv_payload(params)
    if kind == "unified_id_mapping":
        return id_mapping_data_only(params)
    if kind in {"unified_fetch", "unified_repair", "unified_check_repair"}:
        return tree_values_sorted_non_empty_except_sign(params)
    raise ValueError(f"unsupported DFP sign kind: {kind}")


def dfp_atlas_sign(
    sign_input: str | bytes | bytearray,
    counter: int,
    unix_time: int,
    *,
    session_seed: int,
    sdk_id: str = DFP_SDK_ID,
) -> str:
    """Generate the 64hex DFP 10405 atlasSign value."""

    return atlas_sign(
        sign_input,
        sdk_id,
        counter,
        unix_time,
        session_seed=session_seed,
    )


def sign_dfp_form(
    kind: str,
    params: Mapping[str, str],
    counter: int,
    unix_time: int,
    *,
    session_seed: int,
    sdk_id: str = DFP_SDK_ID,
) -> dict[str, str]:
    """Return a signed copy of a DFP/unifiedId form."""

    material = build_dfp_sign_material(kind, params)
    signed = {str(key): str(value) for key, value in params.items()}
    signed["sign"] = dfp_atlas_sign(
        material.sign_input,
        counter,
        unix_time,
        session_seed=session_seed,
        sdk_id=sdk_id,
    )
    return signed


def parse_dfp_atlas_sign(sign_hex64: str, sdk_id: str = DFP_SDK_ID) -> dict:
    digest_hex = atlas_sign_to_digest_hex(sign_hex64, sdk_id)
    parsed = kwsg_10418_digest24_unmix(digest_hex)
    parsed["digest24_hex"] = digest_hex
    parsed["sdk_id"] = sdk_id
    parsed["head8"] = ZT_OUTER_CONFIGS[sdk_id]["head8"].hex()
    return parsed


__all__ = [
    "DFP_PAYLOAD_KEYS",
    "DFP_SDK_ID",
    "DFP_SIGN_KEYS",
    "DfpSignMaterial",
    "build_dfp_sign_material",
    "dfp_atlas_sign",
    "id_mapping_data_only",
    "legacy_product_ts_sv_payload",
    "parse_dfp_atlas_sign",
    "sign_dfp_form",
    "tree_values_sorted_non_empty_except_sign",
]