AutoRpa/rpa_selection/search_yaoex.py

from __future__ import annotations

import base64
import time
from binascii import Error as BinasciiError
from dataclasses import dataclass
from typing import Any

import requests
from Crypto.Cipher import AES
from Crypto.Util.Padding import unpad

from http_util import post_json_with_retry
from mall_config import (
    YAOEX_MALL_NAME,
    YAOEX_PRICE_DECRYPT_BASE_KEY,
    YAOEX_SEARCH_API_URL,
)
from rules import Candidate, normalize_vendor, parse_expiry_months
from util import normalize_text, parse_float_or_none


@dataclass
class YaoexAuth:
    token: str
    user_id: str
    buyer_code: str


def build_yaoex_search_request(
    auth: YaoexAuth, query_key: str
) -> tuple[dict[str, Any], dict[str, Any]]:
    timestamp = str(int(time.time() * 1000))
    common_part: dict[str, Any] = {
        "traderName": "yaoex_pc",
        "trader": "pc",
        "closesignature": "yes",
        "signature_method": "md5",
        "signature": "",
        "timestamp": timestamp,
        "token": auth.token,
        "userToken": auth.token,
    }
    form_payload = {
        **common_part,
        "userId": auth.user_id,
        "roleId": "",
        "userType": "",
        "buyerCode": auth.buyer_code,
        "nowPage": 1,
        "per": 20,
        "keyword": query_key,
        "catSearchId": "",
        "specs": "",
        "factoryIds": "",
        "sellerCodes": "",
        "sellerFilterMode": 1,
        "sortColumn": "default",
        "sortMode": "default",
        "ver": 1,
        "stock_mode": 1,
        "showExtendCard": "true",
        "needDinnerPrice": "true",
        "limitStart": "",
        "limitEnd": "",
        "deadLineStart": 12,
        "deadLineEnd": "",
        "filterDtos": "",
        "showWholePurchase": "true",
    }
    return common_part, form_payload


def build_yaoex_price_decrypt_key(user_id: Any) -> bytes:
    user_text = normalize_text(user_id)
    if not user_text:
        return YAOEX_PRICE_DECRYPT_BASE_KEY.encode("utf-8")
    dynamic_key = YAOEX_PRICE_DECRYPT_BASE_KEY[:10] + user_text[:6].rjust(6, "0")
    return dynamic_key.encode("utf-8")


def decrypt_yaoex_price(raw_price: Any, user_id: Any) -> Any:
    if raw_price is None:
        return ""
    if isinstance(raw_price, (int, float)):
        return raw_price

    text = normalize_text(raw_price)
    if not text:
        return ""

    try:
        return float(text)
    except ValueError:
        pass

    try:
        cipher_bytes = base64.b64decode(text)
        key_bytes = build_yaoex_price_decrypt_key(user_id)
        plain_bytes = AES.new(key_bytes, AES.MODE_ECB).decrypt(cipher_bytes)
        plain_text = unpad(plain_bytes, AES.block_size).decode("utf-8").strip()
    except (BinasciiError, ValueError, UnicodeDecodeError):
        return text

    if plain_text == "":
        return 0
    try:
        return float(plain_text)
    except ValueError:
        return plain_text


def query_yaoex_products(
    session: requests.Session,
    auth: YaoexAuth,
    query_key: str,
) -> tuple[list[dict[str, Any]], str]:
    params, payload = build_yaoex_search_request(auth, query_key)
    headers = {
        "Accept": "application/json, text/plain, */*",
        "Content-Type": "application/x-www-form-urlencoded",
        "Origin": "https://mall.yaoex.com",
        "Referer": "https://mall.yaoex.com/",
        "X-Request-Agent": "Axios",
        "X-Requested-With": "XMLHttpRequest",
        "User-Agent": (
            "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
            "(KHTML, like Gecko) Chrome/152.0.0.0 Safari/537.36"
        ),
    }
    payload_json = post_json_with_retry(
        session,
        YAOEX_SEARCH_API_URL,
        platform_name=YAOEX_MALL_NAME,
        query_key=query_key,
        params=params,
        data=payload,
        headers=headers,
    )
    data_section = payload_json.get("data") or {}
    products = data_section.get("shopProducts") or []
    if not isinstance(products, list):
        raise RuntimeError(f"1号药城接口返回 shopProducts 结构异常：{type(products)}")
    response_user_id = normalize_text(data_section.get("userId")) or auth.user_id
    return products, response_user_id


def map_yaoex_products(products: list[dict[str, Any]], user_id: Any) -> list[Candidate]:
    candidates: list[Candidate] = []
    if products:
        sample = products[0]
        expiry_keys = {
            key: sample.get(key)
            for key in sample
            if any(token in str(key).lower() for token in ("dead", "valid", "expir", "period", "effect"))
        }
        print(f"1号药城效期相关字段示例：{expiry_keys}")
    for product in products:
        vendor_raw = normalize_text(product.get("vendorName"))
        expiry_raw = ""
        for key in ("deadLine", "deadLineDesc", "validDate", "expireDate", "period"):
            expiry_raw = normalize_text(product.get(key))
            if expiry_raw:
                break
        raw_price = decrypt_yaoex_price(product.get("price", ""), user_id)
        candidates.append(
            Candidate(
                mall_name=YAOEX_MALL_NAME,
                sku_id=normalize_text(product.get("spuCode")),
                name=normalize_text(product.get("shortName")),
                spec=normalize_text(product.get("spec")),
                price=parse_float_or_none(raw_price),
                expiry_raw=expiry_raw,
                expiry_months=parse_expiry_months(expiry_raw),
                vendor_raw=vendor_raw,
                vendor_norm=normalize_vendor(vendor_raw),
                stock=parse_float_or_none(product.get("stockCount")),
                extra={"stockCountDesc": product.get("stockCountDesc", "")},
            )
        )
    return candidates


def search_candidates(
    auth: YaoexAuth,
    query_key: str,
    session: requests.Session | None = None,
    cache: dict[str, list[Candidate]] | None = None,
) -> list[Candidate]:
    store = cache if cache is not None else {}
    if query_key in store:
        return store[query_key]
    own_session = session or requests.Session()
    try:
        products, user_id = query_yaoex_products(own_session, auth, query_key)
        mapped = map_yaoex_products(products, user_id)
    finally:
        if session is None:
            own_session.close()
    store[query_key] = mapped
    return mapped