#!/usr/bin/env python3
"""
将 SQL 导出的 API 调用记录（CSV 或 JSON）解密 request_params，输出 JSON 文件。

依赖: pip install pycryptodome

用法:
  python decrypt_api_calls.py -i combmy01_calls.csv -o combmy01_decrypted.json
  python decrypt_api_calls.py -i query_1-2026-06-27_32042.json -o combmy01_decrypted.json

输入必须包含字段: request_params, secret_key（或通过 --secret-key 传入）
可选字段会原样写入输出: id, transaction_id, user_id, product_code, status, cost, start_at, end_at, created_at
输出默认不包含 secret_key。
"""

from __future__ import annotations

import argparse
import base64
import csv
import json
import sys
from pathlib import Path
from typing import Any

try:
    from Crypto.Cipher import AES
except ImportError:
    print("请先安装依赖: pip install pycryptodome", file=sys.stderr)
    sys.exit(1)


OPTIONAL_FIELDS = (
    "id",
    "transaction_id",
    "user_id",
    "product_code",
    "status",
    "cost",
    "start_at",
    "end_at",
    "created_at",
)


def aes_decrypt(cipher_b64: str, key_hex: str) -> Any:
    """与 tyapi-server internal/shared/crypto/crypto.go AesDecrypt 逻辑一致。"""
    key = bytes.fromhex(key_hex.strip())
    raw = base64.b64decode(cipher_b64.strip())
    if len(raw) < 16:
        raise ValueError("ciphertext too short")

    iv, ciphertext = raw[:16], raw[16:]
    plain = AES.new(key, AES.MODE_CBC, iv).decrypt(ciphertext)

    pad = plain[-1]
    if pad < 1 or pad > 16 or pad > len(plain):
        raise ValueError("invalid padding")
    for i in range(pad):
        if plain[-1 - i] != pad:
            raise ValueError("invalid padding")
    plain = plain[:-pad]

    text = plain.decode("utf-8")
    return json.loads(text)


def normalize_row(row: dict) -> dict[str, str]:
    """兼容带 BOM 或大小写不一致的表头；值统一转为字符串。"""
    out: dict[str, str] = {}
    for k, v in row.items():
        key = (k or "").strip().lower().lstrip("\ufeff")
        if v is None:
            out[key] = ""
        elif isinstance(v, (dict, list)):
            out[key] = json.dumps(v, ensure_ascii=False)
        else:
            out[key] = str(v).strip()
    return out


def process_row(row: dict, default_secret_key: str | None) -> dict[str, Any]:
    row = normalize_row(row)

    encrypted = row.get("request_params", "")
    secret_key = row.get("secret_key") or default_secret_key or ""

    if not encrypted:
        raise ValueError("request_params 为空")
    if not secret_key:
        raise ValueError("缺少 secret_key（输入字段或 --secret-key 参数）")

    item: dict[str, Any] = {}
    for field in OPTIONAL_FIELDS:
        if field in row and row[field] != "":
            item[field] = row[field]

    try:
        item["request_params"] = aes_decrypt(encrypted, secret_key)
    except Exception as exc:
        item["request_params"] = None
        item["decrypt_error"] = str(exc)

    return item


def read_text(path: Path) -> str:
    for encoding in ("utf-8-sig", "utf-8", "gbk"):
        try:
            return path.read_text(encoding=encoding)
        except UnicodeDecodeError:
            continue
    raise ValueError(f"无法读取文件编码: {path}")


def read_rows(path: Path) -> list[dict]:
    suffix = path.suffix.lower()
    if suffix == ".json":
        data = json.loads(read_text(path))
        if isinstance(data, list):
            return data
        if isinstance(data, dict):
            for key in ("items", "data", "rows", "records"):
                if key in data and isinstance(data[key], list):
                    return data[key]
            raise ValueError("JSON 根节点应为数组，或包含 items/data/rows/records 数组字段")
        raise ValueError("不支持的 JSON 结构")

    if suffix == ".csv":
        for encoding in ("utf-8-sig", "utf-8", "gbk"):
            try:
                with path.open("r", encoding=encoding, newline="") as f:
                    return list(csv.DictReader(f))
            except UnicodeDecodeError:
                continue
        raise ValueError(f"无法读取 CSV 编码: {path}")

    # 无扩展名时自动探测
    text = read_text(path).lstrip()
    if text.startswith("[") or text.startswith("{"):
        data = json.loads(text)
        if isinstance(data, list):
            return data
        if isinstance(data, dict):
            for key in ("items", "data", "rows", "records"):
                if key in data and isinstance(data[key], list):
                    return data[key]
        raise ValueError("不支持的 JSON 结构")

    for encoding in ("utf-8-sig", "utf-8", "gbk"):
        try:
            with path.open("r", encoding=encoding, newline="") as f:
                return list(csv.DictReader(f))
        except UnicodeDecodeError:
            continue
    raise ValueError(f"无法识别输入格式: {path}")


def main() -> None:
    parser = argparse.ArgumentParser(description="解密 API 调用记录中的 request_params（支持 CSV / JSON）")
    parser.add_argument("-i", "--input", required=True, help="SQL 导出的 CSV 或 JSON 文件路径")
    parser.add_argument("-o", "--output", default="api_calls_decrypted.json", help="输出 JSON 文件路径")
    parser.add_argument(
        "--secret-key",
        help="全局 secret_key（若 CSV 无 secret_key 列，或每行相同可省略 CSV 中该列）",
    )
    args = parser.parse_args()

    input_path = Path(args.input)
    if not input_path.is_file():
        print(f"输入文件不存在: {input_path}", file=sys.stderr)
        sys.exit(1)

    rows = read_rows(input_path)
    if not rows:
        print("输入文件无数据行", file=sys.stderr)
        sys.exit(1)

    result = []
    failed = 0
    for idx, row in enumerate(rows, start=1):
        try:
            result.append(process_row(row, args.secret_key))
            if result[-1].get("decrypt_error"):
                failed += 1
        except Exception as exc:
            failed += 1
            result.append({"row": idx, "decrypt_error": str(exc), "request_params": None})

    output_path = Path(args.output)
    output_path.write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")

    ok = len(result) - failed
    print(f"完成: 共 {len(result)} 条, 成功解密 {ok} 条, 失败 {failed} 条")
    print(f"输出: {output_path.resolve()}")


if __name__ == "__main__":
    main()