#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
模块名称：http_client.py
模块描述：安全测试 HTTP 客户端，封装 requests 库，自动管理 Token、请求签名、日志记录

签名算法（从前端 JS 逆向）:
1. x_random = 随机字符串(8-16位字母数字)
2. x_timestamp = 当前毫秒时间戳
3. sign_str = timestamp + JSON.stringify(body) + random  (有body时)
   sign_str = timestamp + random                        (无body时)
4. sign_hash = SHA256(sign_str)                         → 中间哈希
5. aes_key_source = SHA256(bearer_token 或 随机60位)    → 64字符hex
6. aes_key = aes_key_source[16:32]                      → 32字符hex=16字节
7. aes_iv  = aes_key_source[0:8] + aes_key_source[-8:]  → 16字符hex=8字节
8. x_sign = AES.encrypt(sign_hash, aes_key, aes_iv)     → Base64编码
9. 无 Token 时额外生成 RandomCode 头部（60位随机字符串）

作者：czj
创建日期：2026-07-21
最后修改：2026-07-21
"""

import hashlib
import random
import string
import time
import json
import logging
from typing import Optional, Dict, Any

import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry

# 尝试导入 pycryptodome，不存在则降级
try:
    from Crypto.Cipher import AES
    from Crypto.Util.Padding import pad
    import base64 as b64
    HAS_CRYPTO = True
except ImportError:
    HAS_CRYPTO = False

logger = logging.getLogger(__name__)


class HttpClient:
    """
    HTTP 请求客户端，支持自动 Token 管理和请求签名

    用于安全测试的 HTTP 请求发送，自动处理：
    - 请求签名（AES-CBC）
    - Token 管理
    - 重试机制
    - SSL 验证
    """

    def __init__(self, config: Dict[str, Any]):
        """
        初始化 HTTP 客户端

        Args:
            config: 配置字典，包含:
                - target.base_url: 目标服务器地址
                - target.verify_ssl: 是否验证SSL
                - target.timeout: 请求超时时间
                - accounts: 测试账号配置
                - auth: 认证配置
        """
        self.config = config

        # 目标配置
        target = config.get("target", {})
        self.base_url = target.get("base_url", "")
        self.verify_ssl = target.get("verify_ssl", False)
        self.timeout = target.get("timeout", 30)

        # Token 存储：{账号名: token字符串}
        self._tokens: Dict[str, str] = {}

        # 验证码 uuid 存储
        self._captcha_uuid: Optional[str] = None

        # 创建带重试机制的 session
        self.session = self._create_session()

        # 检查加密库
        if not HAS_CRYPTO:
            logger.warning("pycryptodome 未安装，AES 签名将不可用。请执行: pip install pycryptodome")

    def _create_session(self) -> requests.Session:
        """
        创建带重试机制的 requests Session

        Returns:
            requests.Session: 配置好的会话对象
        """
        session = requests.Session()

        # 重试策略：最多重试 2 次，仅对连接错误重试
        retry_strategy = Retry(
            total=2,
            backoff_factor=0.5,
            status_forcelist=[500, 502, 503, 504],
            allowed_methods=["GET", "POST", "PUT", "DELETE"]
        )
        adapter = HTTPAdapter(max_retries=retry_strategy)
        session.mount("https://", adapter)
        session.mount("http://", adapter)

        # 禁用 SSL 警告（测试环境）
        if not self.verify_ssl:
            import urllib3
            urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

        return session

    @staticmethod
    def encrypt_password(password: str) -> str:
        """
        使用 SHA256 加密密码（与前端加密方式一致）

        Args:
            password: 明文密码

        Returns:
            str: SHA256 十六进制摘要
        """
        return hashlib.sha256(password.encode('utf-8')).hexdigest()

    def _generate_random_string(self, min_len: int = 8, max_len: int = 16, digits: bool = True) -> str:
        """
        生成指定长度的随机字符串

        Args:
            min_len: 最小长度
            max_len: 最大长度
            digits: 是否包含数字

        Returns:
            str: 随机字符串
        """
        length = random.randint(min_len, max_len)
        chars = string.ascii_letters + (string.digits if digits else '')
        return ''.join(random.choice(chars) for _ in range(length))

    @staticmethod
    def _sha256(text: str) -> str:
        """计算 SHA256 哈希值"""
        return hashlib.sha256(text.encode('utf-8')).hexdigest()

    @staticmethod
    def _aes_encrypt(plaintext: str, key_str: str, iv_str: str) -> str:
        """
        AES CBC 模式加密（与前端 CryptoJS 行为一致）
        CryptoJS 传入字符串时按 UTF-8 编码为密钥/IV，不是 hex 解析

        Args:
            plaintext: 明文字符串（SHA256 后的 hex 字符串）
            key_str: 16字符 AES 密钥字符串
            iv_str: 16字符 AES IV 字符串

        Returns:
            str: Base64 编码的密文
        """
        if not HAS_CRYPTO:
            return plaintext  # 降级：无加密库时返回原文

        # CryptoJS 传入字符串时直接按 UTF-8 编码为 bytes
        key_bytes = key_str.encode('utf-8')   # 16 bytes (AES-128)
        iv_bytes = iv_str.encode('utf-8')     # 16 bytes
        cipher = AES.new(key_bytes, AES.MODE_CBC, iv_bytes)
        # 对明文进行 PKCS7 填充
        padded = pad(plaintext.encode('utf-8'), AES.block_size)
        encrypted = cipher.encrypt(padded)
        return b64.b64encode(encrypted).decode('utf-8')

    def _generate_sign(self, body_data: Optional[dict] = None, bearer_token: str = "") -> Dict[str, str]:
        """
        根据前端 JS 逆向的签名算法生成请求签名

        Args:
            body_data: 请求体数据（字典或None）
            bearer_token: Bearer Token 字符串（含 "Bearer " 前缀）

        Returns:
            dict: {x-random, x-timestamp, x-sign, [RandomCode]}
        """
        # 步骤1: 生成随机字符串 x-random（8-16位）
        x_random = self._generate_random_string(8, 16)
        # 步骤2: 毫秒时间戳
        x_timestamp = str(int(time.time() * 1000))

        # 步骤3: 构造签名原文
        if body_data is not None:
            sign_str = x_timestamp + json.dumps(body_data, separators=(',', ':'), ensure_ascii=False) + x_random
        else:
            sign_str = x_timestamp + x_random

        # 步骤4: SHA256 中间哈希
        sign_hash = self._sha256(sign_str)

        # 步骤5-8: AES 加密生成 x-sign
        random_code = None
        if bearer_token:
            # 有 Token 时用 Token 做 AES 密钥源
            aes_key_source = self._sha256(bearer_token)
        else:
            # 无 Token 时用随机60位字符串，并将该字符串作为 RandomCode 头部
            random_code = self._generate_random_string(60, 60)
            aes_key_source = self._sha256(random_code)

        # 步骤6: AES 密钥 = sha256结果的[16:32]（16字节）
        aes_key = aes_key_source[16:32]
        # 步骤7: AES IV = sha256结果的[0:8] + 末尾8字符
        aes_iv = aes_key_source[0:8] + aes_key_source[-8:]

        # 步骤8: AES 加密
        x_sign = self._aes_encrypt(sign_hash, aes_key, aes_iv)

        headers = {
            'X-RANDOM': x_random,
            'X-TIMESTAMP': x_timestamp,
            'X-SIGN': x_sign,
        }
        if random_code:
            headers['RandomCode'] = random_code

        return headers

    def get_captcha_uuid(self) -> Optional[str]:
        """
        获取验证码 UUID（从验证码接口获取）

        Returns:
            str: 验证码 UUID
        """
        auth_config = self.config.get("auth", {})
        captcha_path = auth_config.get("captcha_path", "/platform/api/code")
        url = f"{self.base_url}{captcha_path}"

        try:
            resp = self.session.get(url, verify=self.verify_ssl, timeout=self.timeout)
            if resp.status_code == 200:
                data = resp.json()
                # 从响应中提取 uuid
                uuid_val = data.get('uuid') or data.get('data', {}).get('uuid')
                if uuid_val:
                    self._captcha_uuid = uuid_val
                    logger.debug(f"获取验证码 UUID: {uuid_val}")
                    return uuid_val
        except Exception as e:
            logger.warning(f"获取验证码 UUID 失败: {e}")

        return None

    def login(self, account_key: str = "superadmin") -> Optional[str]:
        """
        使用指定账号登录系统并获取 Token

        Args:
            account_key: 账号配置键名（superadmin/admin/user）

        Returns:
            str: 访问 Token，失败返回 None
        """
        # 如果已有有效 Token，直接返回
        if account_key in self._tokens:
            return self._tokens[account_key]

        accounts = self.config.get("accounts", {})
        account = accounts.get(account_key)
        if not account:
            logger.error(f"未找到账号配置: {account_key}")
            return None

        auth_config = self.config.get("auth", {})
        login_path = auth_config.get("login_path", "/platform/api/auth/login")
        url = f"{self.base_url}{login_path}"
        captcha = accounts.get("captcha", "csba")

        # 获取验证码 UUID
        uuid_val = self.get_captcha_uuid() or ''

        # 构造登录请求体
        encrypted_pwd = self.encrypt_password(account['password'])
        payload = {
            "username": account['username'],
            "password": encrypted_pwd,
            "code": captcha,
            "uuid": uuid_val,
        }

        # 生成签名头部（登录时无 Token，使用随机密钥）
        sign_headers = self._generate_sign(body_data=payload, bearer_token="")

        # 通用请求头
        headers = {
            'Content-Type': 'application/json',
            'Accept': 'application/json, text/plain, */*',
            'referer': '',
        }
        headers.update(sign_headers)

        try:
            logger.info(f"正在登录 [{account_key}]: {account['username']}")
            resp = self.session.post(
                url, json=payload, headers=headers,
                verify=self.verify_ssl, timeout=self.timeout
            )

            if resp.status_code == 200:
                data = resp.json()
                # 尝试从响应中提取 Token
                token = None
                if isinstance(data, dict):
                    token = (
                        data.get('accessToken') or
                        data.get('token') or
                        data.get('access_token') or
                        data.get('data', {}).get('accessToken') or
                        data.get('data', {}).get('token') or
                        data.get('data', {}).get('access_token')
                    )

                if token:
                    self._tokens[account_key] = token
                    logger.info(f"登录成功 [{account_key}], Token: {token[:30]}...")
                    return token
                else:
                    logger.warning(f"登录响应中未找到 Token: {json.dumps(data, ensure_ascii=False)[:300]}")
            else:
                logger.warning(f"登录失败 [{account_key}], HTTP {resp.status_code}: {resp.text[:200]}")

        except Exception as e:
            logger.error(f"登录异常 [{account_key}]: {e}")

        return None

    def get_token(self, account_key: str = "superadmin") -> Optional[str]:
        """
        获取指定账号的 Token（如未登录则自动登录）

        Args:
            account_key: 账号配置键名

        Returns:
            str: Token 字符串
        """
        if account_key in self._tokens:
            return self._tokens[account_key]
        return self.login(account_key)

    def clear_token(self, account_key: Optional[str] = None):
        """
        清除 Token 缓存

        Args:
            account_key: 账号键名，为None则清除所有
        """
        if account_key:
            self._tokens.pop(account_key, None)
        else:
            self._tokens.clear()

    def _get_bearer_token(self, account_key: str) -> str:
        """获取 Bearer Token 字符串（含前缀）"""
        token = self.get_token(account_key)
        if token:
            return f"Bearer {token}"
        return ""

    def _build_headers(self, token: Optional[str] = None, extra_headers: Optional[dict] = None) -> dict:
        """
        构造请求头部

        Args:
            token: 访问 Token
            extra_headers: 额外头部

        Returns:
            dict: 请求头部字典
        """
        headers = {
            'Content-Type': 'application/json',
            'Accept': 'application/json, text/plain, */*',
        }
        if token:
            headers['Authorization'] = f"Bearer {token}"
        if extra_headers:
            headers.update(extra_headers)
        return headers

    def _request_with_sign(
        self,
        method: str,
        path: str,
        json_data: Optional[dict] = None,
        params: Optional[dict] = None,
        token: Optional[str] = None,
        account: str = "superadmin",
        **kwargs
    ) -> Optional[requests.Response]:
        """
        带签名的请求发送（内部方法）

        Args:
            method: HTTP 方法
            path: 接口路径
            json_data: 请求体
            params: URL 查询参数
            token: Token
            account: 账号键名

        Returns:
            requests.Response 或 None
        """
        url = f"{self.base_url}{path}"

        # 获取 Token
        if token is None and account:
            token = self.get_token(account)
        bearer_str = f"Bearer {token}" if token else ""

        # 构造签名（根据方法决定 body_data）
        body_for_sign = None
        if method.lower() in ('post', 'put', 'patch') and json_data is not None:
            body_for_sign = json_data
        elif method.lower() in ('get', 'delete') and params:
            body_for_sign = params

        sign_headers = self._generate_sign(body_data=body_for_sign, bearer_token=bearer_str)

        # 合并头部
        headers = self._build_headers(token=token)
        headers.update(sign_headers)

        logger.debug(f"{method.upper()} {path}")
        try:
            resp = self.session.request(
                method, url,
                json=json_data if method.lower() not in ('get', 'delete') else None,
                params=params if method.lower() in ('get', 'delete') else None,
                headers=headers,
                verify=self.verify_ssl,
                timeout=self.timeout,
                **kwargs
            )
            return resp
        except requests.RequestException as e:
            logger.error(f"{method.upper()} 请求异常 [{path}]: {e}")
            return None

    def get(
        self,
        path: str,
        params: Optional[dict] = None,
        token: Optional[str] = None,
        account: str = "superadmin",
        **kwargs
    ) -> Optional[requests.Response]:
        """
        发送 GET 请求

        Args:
            path: 接口路径
            params: 查询参数
            token: 指定 Token
            account: 账号键名

        Returns:
            requests.Response
        """
        return self._request_with_sign('GET', path, params=params, token=token, account=account, **kwargs)

    def post(
        self,
        path: str,
        json_data: Optional[dict] = None,
        token: Optional[str] = None,
        account: str = "superadmin",
        **kwargs
    ) -> Optional[requests.Response]:
        """
        发送 POST 请求

        Args:
            path: 接口路径
            json_data: 请求体
            token: 指定 Token
            account: 账号键名

        Returns:
            requests.Response
        """
        return self._request_with_sign('POST', path, json_data=json_data, token=token, account=account, **kwargs)

    def put(
        self,
        path: str,
        json_data: Optional[dict] = None,
        token: Optional[str] = None,
        account: str = "superadmin",
        **kwargs
    ) -> Optional[requests.Response]:
        """
        发送 PUT 请求
        """
        return self._request_with_sign('PUT', path, json_data=json_data, token=token, account=account, **kwargs)

    def delete(
        self,
        path: str,
        token: Optional[str] = None,
        account: str = "superadmin",
        **kwargs
    ) -> Optional[requests.Response]:
        """
        发送 DELETE 请求
        """
        return self._request_with_sign('DELETE', path, token=token, account=account, **kwargs)

    def request(
        self,
        method: str,
        path: str,
        token: Optional[str] = None,
        account: str = "superadmin",
        json_data: Optional[dict] = None,
        params: Optional[dict] = None,
        **kwargs
    ) -> Optional[requests.Response]:
        """
        发送自定义 HTTP 方法的请求

        Args:
            method: HTTP 方法（OPTIONS、PATCH、TRACE 等）
            path: 接口路径
            token: 指定 Token
            account: 账号键名
            json_data: 请求体
            params: 查询参数

        Returns:
            requests.Response
        """
        return self._request_with_sign(
            method, path, json_data=json_data,
            params=params, token=token,
            account=account, **kwargs
        )