#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
模块名称：security_executor.py
模块描述：安全测试执行引擎，负责执行 API 安全测试用例

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

import json
import base64
import logging
import time
from dataclasses import dataclass, field
from datetime import datetime
from typing import Optional, List, Dict, Any, Callable

from app.executors.http_client import HttpClient
from app.executors.auth_helper import AuthHelper
from app.config import settings

logger = logging.getLogger(__name__)


# ==================== 数据类 ====================

@dataclass
class SecurityStepResult:
    """
    安全测试步骤结果

    Attributes:
        order (int): 步骤顺序
        name (str): 步骤名称
        action (str): 动作类型
        status (str): 状态 (pending/running/passed/failed/skipped)
        duration (float): 耗时（秒）
        log (str): 执行日志
        error (Optional[str]): 错误信息
        is_vulnerable (bool): 是否存在漏洞
        level (str): 风险等级
        request_info (str): 请求信息
        response_info (str): 响应信息
        start_time (datetime): 开始时间
        end_time (datetime): 结束时间
    """

    order: int
    name: str = ""
    action: str = ""
    status: str = "pending"
    duration: float = 0.0
    log: str = ""
    error: Optional[str] = None
    is_vulnerable: bool = False
    level: str = "info"
    request_info: str = ""
    response_info: str = ""
    start_time: Optional[datetime] = None
    end_time: Optional[datetime] = None


@dataclass
class SecurityCaseResult:
    """
    安全测试用例结果

    Attributes:
        case_id (str): 用例ID
        test_id (str): OWASP测试编号
        name (str): 用例名称
        status (str): 状态 (passed/failed/skipped/error)
        is_vulnerable (bool): 是否存在漏洞
        level (str): 风险等级
        description (str): 漏洞描述
        request_info (str): 请求信息
        response_info (str): 响应信息
        fix_suggestion (str): 修复建议
        duration (float): 执行耗时（秒）
        steps (list): 步骤结果列表
        error (Optional[str]): 错误信息
        metadata (dict): 额外元数据
        start_time (datetime): 开始时间
        end_time (datetime): 结束时间
    """

    case_id: str = ""
    test_id: str = ""
    name: str = ""
    status: str = "pending"
    is_vulnerable: bool = False
    level: str = "info"
    description: str = ""
    request_info: str = ""
    response_info: str = ""
    fix_suggestion: str = ""
    duration: float = 0.0
    steps: List[SecurityStepResult] = field(default_factory=list)
    error: Optional[str] = None
    metadata: Dict[str, Any] = field(default_factory=dict)
    start_time: Optional[datetime] = None
    end_time: Optional[datetime] = None


# ==================== 断言引擎 ====================

class AssertionEngine:
    """
    断言引擎，评估测试响应是否存在漏洞
    """

    # 权限拒绝相关的错误码
    AUTH_DENIED_CODES = {
        "401", "403", "40301", "40101",  # 通用鉴权错误
        "B0027",   # 权限不足。无法请求接口
        "B0017",   # 接口请求方式错误
        "B0002",   # 验证码失效
        "A0076",   # 无效token
        "A0301",   # token已过期
    }

    # 限流相关错误码/关键词
    RATE_LIMIT_CODES = {
        "429", "B0029", "A0302",
    }
    RATE_LIMIT_KEYWORDS = [
        "请求过于频繁", "请稍后再试", "限流", "rate limit",
        "too many requests", "throttl", "频率限制",
    ]

    # 权限拒绝相关的消息关键词
    DENY_KEYWORDS = [
        "权限不足", "不允许访问", "拒绝访问",
        "未授权", "请登录", "重新登录", "已退出",
        "无效token", "token无效", "token已过期",
        "AccessDenied", "用户不存在",
    ]

    @classmethod
    def is_success(cls, response) -> bool:
        """
        判断响应是否表示业务成功（即越权成功）

        Args:
            response: requests.Response 对象

        Returns:
            bool: True 表示接口返回了正常数据（存在越权风险）
        """
        if response is None:
            return False

        # 401/403 直接表示鉴权拦截
        if response.status_code in (401, 403):
            return False

        # 非 200 状态码不算成功
        if response.status_code != 200:
            return False

        try:
            data = response.json()

            # 1. 检查 success 字段
            if "success" in data and not data["success"]:
                return False

            # 2. 检查业务错误码
            code = data.get("code", data.get("status", data.get("errcode")))
            if code is not None:
                if str(code) in cls.AUTH_DENIED_CODES:
                    return False

            # 3. 检查 message/msg 中的权限拒绝关键词
            message = str(data.get("message", data.get("msg", "")))
            if message:
                if any(kw in message for kw in cls.DENY_KEYWORDS):
                    return False

            # 4. 检查是否有有效数据返回
            result_data = data.get("data")
            if result_data is None or result_data == {} or result_data == []:
                # 允许 success=True 但 data 为空的情况（如删除/修改操作）
                if data.get("success") is True:
                    return True
                return False

            return True

        except Exception:
            pass

        return False

    @classmethod
    def is_rate_limited(cls, response) -> bool:
        """
        判断响应是否表示触发了限流

        Args:
            response: requests.Response 对象

        Returns:
            bool: True 表示触发了限流
        """
        if response is None:
            return False

        # 429 状态码直接判定为限流
        if response.status_code == 429:
            return True

        try:
            data = response.json()
            code = str(data.get("code", data.get("status", "")))
            if code in cls.RATE_LIMIT_CODES:
                return True
            message = str(data.get("message", data.get("msg", "")))
            if any(kw in message for kw in cls.RATE_LIMIT_KEYWORDS):
                return True
        except Exception:
            pass

        # 检查响应头中的限流标志
        resp_headers = dict(response.headers)
        if any(h in resp_headers for h in ["X-RateLimit-Remaining", "X-RateLimit-Limit", "Retry-After"]):
            return True

        return False

    @classmethod
    def evaluate_assertions(cls, response, assertions: List[Dict], duration: float = 0.0) -> tuple:
        """
        评估断言列表

        Args:
            response: requests.Response 对象
            assertions: 断言配置列表
            duration: 请求耗时（秒），用于 response_time_limit

        Returns:
            tuple: (is_vulnerable, description)
        """
        if response is None:
            return False, "无响应"

        is_vulnerable = False
        descriptions = []

        for assertion in assertions:
            atype = assertion.get("type", "")
            expect = assertion.get("expect")
            desc = assertion.get("description", "")

            if atype == "status_code_not_equal":
                # 状态码不应该等于某个值（如越权访问应返回403，不应返回200）
                if response.status_code == expect:
                    is_vulnerable = True
                    descriptions.append(f"状态码={response.status_code}，{desc}")

            elif atype == "status_code_equal":
                # 状态码应该等于某个值
                if response.status_code != expect:
                    is_vulnerable = True
                    descriptions.append(f"预期状态码{expect}，实际{response.status_code}")

            elif atype == "response_code_not_contain":
                # 响应业务码不应包含某些值
                try:
                    data = response.json()
                    code = str(data.get("code", data.get("status", "")))
                    if code not in [str(e) for e in expect]:
                        is_vulnerable = True
                        descriptions.append(f"业务码={code}，{desc}")
                except Exception:
                    pass

            elif atype == "response_code_contain":
                # 响应业务码应包含某个值
                try:
                    data = response.json()
                    code = str(data.get("code", data.get("status", "")))
                    if code not in [str(e) for e in expect]:
                        is_vulnerable = True
                        descriptions.append(f"预期业务码{expect}，实际{code}")
                except Exception:
                    pass

            elif atype == "body_contains":
                # 响应体应包含某字符串
                if expect not in response.text:
                    is_vulnerable = True
                    descriptions.append(f"响应体不包含'{expect}'，{desc}")

            elif atype == "body_not_contains":
                # 响应体不应包含某字符串
                if expect in response.text:
                    is_vulnerable = True
                    descriptions.append(f"响应体包含'{expect}'，{desc}")

            elif atype == "is_success":
                # 响应应表示业务成功（用于验证越权）
                if cls.is_success(response) == expect:
                    is_vulnerable = True
                    descriptions.append(desc)

            elif atype == "response_size_limit":
                # 响应体大小不应超过限制（KB）
                max_kb = assertion.get("expect_max_kb", 1024)
                if response.text:
                    actual_kb = len(response.text.encode('utf-8')) / 1024
                    if actual_kb > max_kb:
                        is_vulnerable = True
                        descriptions.append(f"响应大小={actual_kb:.1f}KB，超过限制{max_kb}KB，{desc}")

            elif atype == "response_time_limit":
                # 响应时间不应超过限制（秒）
                max_seconds = assertion.get("expect_max_seconds", 10)
                if duration > max_seconds:
                    is_vulnerable = True
                    descriptions.append(f"响应耗时={duration:.2f}s，超过限制{max_seconds}s，{desc}")

            elif atype == "rate_limit_exists":
                # 检测是否存在速率限制（期望有限流则expect=True）
                expect_limited = expect if expect is not None else True
                is_limited = False
                # 429 状态码直接判定为限流
                if response.status_code == 429:
                    is_limited = True
                else:
                    try:
                        data = response.json()
                        code = str(data.get("code", data.get("status", "")))
                        if code in cls.RATE_LIMIT_CODES:
                            is_limited = True
                        message = str(data.get("message", data.get("msg", "")))
                        if any(kw in message for kw in cls.RATE_LIMIT_KEYWORDS):
                            is_limited = True
                    except Exception:
                        pass
                    # 检查响应头中的限流标志
                    resp_headers = dict(response.headers)
                    if any(h in resp_headers for h in ["X-RateLimit-Remaining", "X-RateLimit-Limit", "Retry-After"]):
                        is_limited = True
                if is_limited != expect_limited:
                    is_vulnerable = True
                    descriptions.append(f"限流状态={is_limited}，期望={expect_limited}，{desc}")

            elif atype == "header_not_contains":
                # 响应头不应包含某些字段
                header_names = expect if isinstance(expect, list) else [expect]
                resp_headers = {k.lower(): v for k, v in response.headers.items()}
                for hname in header_names:
                    if hname.lower() in resp_headers:
                        is_vulnerable = True
                        descriptions.append(f"响应头包含'{hname}'，{desc}")

            elif atype == "header_contains":
                # 响应头应存在某些字段（如安全响应头），缺失则视为漏洞
                header_names = expect if isinstance(expect, list) else [expect]
                resp_headers = {k.lower(): v for k, v in response.headers.items()}
                for hname in header_names:
                    if hname.lower() not in resp_headers:
                        is_vulnerable = True
                        descriptions.append(f"响应头缺失'{hname}'，{desc}")

            elif atype == "cors_allow_wildcard":
                # CORS 配置错误：Access-Control-Allow-Origin 值为 * 视为漏洞
                resp_headers = {k.lower(): v for k, v in response.headers.items()}
                acao = resp_headers.get("access-control-allow-origin")
                acac = resp_headers.get("access-control-allow-credentials")
                if acao == "*":
                    is_vulnerable = True
                    descriptions.append(f"CORS 配置错误：Access-Control-Allow-Origin={acao}，{desc}")
                elif (acao and acac and acac.lower() == "true"):
                    # 允许任意来源且携带凭据是最危险的组合
                    is_vulnerable = True
                    descriptions.append(f"CORS 配置错误：ACAO={acao} 且允许凭据，{desc}")

        return is_vulnerable, "; ".join(descriptions) if descriptions else "断言通过"


# ==================== 执行器 ====================

class SecurityExecutor:
    """
    安全测试执行器

    负责执行 API 安全测试用例，不依赖浏览器，使用 requests 库发送 HTTP 请求。
    与 PlaywrightExecutor 并列，通过 case_type=security 区分调用。
    """

    def __init__(self, config: Dict[str, Any]):
        """
        初始化安全测试执行器

        Args:
            config: 安全测试配置字典（来自 SecurityConfig.to_client_config()）
        """
        self.config = config
        self.client = HttpClient(config)
        self.auth = AuthHelper(self.client)
        self.results: List[SecurityCaseResult] = []
        self._is_running = False
        self._stop_requested = False
        self._progress_callback: Optional[Callable] = None

    def start(self) -> bool:
        """
        初始化执行环境，登录所有账号

        Returns:
            bool: 是否成功初始化
        """
        logger.info("初始化安全测试执行环境...")

        # 登录所有测试账号
        tokens = self.auth.login_all_accounts()
        success_count = sum(1 for v in tokens.values() if v)
        logger.info(f"登录完成: {success_count}/{len(tokens)} 个账号成功")

        if success_count == 0:
            logger.error("所有账号登录失败，无法继续测试")
            return False

        self._is_running = True
        return True

    def stop(self):
        """停止执行"""
        self._stop_requested = True
        self._is_running = False
        logger.info("安全测试执行器已停止")

    def set_progress_callback(self, callback: Callable):
        """
        设置进度回调函数

        Args:
            callback: 回调函数，签名为 callback(case_id, status, result)
        """
        self._progress_callback = callback

    @classmethod
    def _generate_fake_token(cls) -> str:
        """
        生成伪造的 Token（用于测试伪造 Token 认证绕过）

        Returns:
            str: 伪造的 JWT Token（alg=none 无签名）
        """
        # 使用 alg=none 的 JWT 格式：base64头部.base64载荷.空签名
        header = base64.b64encode(json.dumps({"alg": "none", "typ": "JWT"}).encode()).decode()
        payload = base64.b64encode(json.dumps({
            "sub": "fake_user", "exp": 0, "iat": 0
        }).encode()).decode()
        return f"{header}.{payload}."

    @classmethod
    def _extract_from_response(cls, response, json_path: str) -> Any:
        """
        从响应 JSON 中按路径提取值

        Args:
            response: requests.Response 对象
            json_path: 如 "data.records[0].id"

        Returns:
            Any: 提取的值，失败返回 None
        """
        if response is None:
            return None
        try:
            data = response.json()
            # 按 "." 分割路径段
            parts = json_path.split(".")
            current = data
            for part in parts:
                # 处理数组索引：如 "records[0]"
                if "[" in part and part.endswith("]"):
                    name, idx_str = part.split("[")
                    idx = int(idx_str[:-1])
                    current = current.get(name, [])[idx] if isinstance(current, dict) else None
                else:
                    current = current.get(part) if isinstance(current, dict) else None
                if current is None:
                    return None
            return current
        except Exception:
            return None

    @classmethod
    def _substitute_vars(cls, value: Any, var_dict: Dict[str, Any]) -> Any:
        """
        替换字符串中的 ${var_name} 占位符

        Args:
            value: 原始值（str/dict/list）
            var_dict: 变量字典

        Returns:
            Any: 替换后的值
        """
        if isinstance(value, str):
            for key, val in var_dict.items():
                placeholder = f"${{{key}}}"
                if placeholder in value:
                    value = value.replace(placeholder, str(val) if val is not None else "")
            return value
        elif isinstance(value, dict):
            return {k: cls._substitute_vars(v, var_dict) for k, v in value.items()}
        elif isinstance(value, list):
            return [cls._substitute_vars(v, var_dict) for v in value]
        return value

    def execute_case(self, case: Dict[str, Any]) -> SecurityCaseResult:
        """
        执行单个安全测试用例

        Args:
            case: 用例配置字典，包含:
                - id: 用例ID
                - name: 用例名称
                - steps: 测试步骤定义（JSON）

        Returns:
            SecurityCaseResult: 执行结果
        """
        start_time = datetime.now()
        case_id = case.get("id", "")
        case_name = case.get("name", "")

        result = SecurityCaseResult(
            case_id=case_id,
            name=case_name,
            start_time=start_time,
        )

        steps_config = case.get("steps", {})

        # 提取测试配置
        test_type = steps_config.get("test_type", "api_security")
        if test_type != "api_security":
            result.status = "skipped"
            result.error = f"不支持的测试类型: {test_type}"
            result.end_time = datetime.now()
            return result

        target = steps_config.get("target", {})
        auth_config = steps_config.get("auth", {})
        request_config = steps_config.get("request", {})
        assertions = steps_config.get("assertions", [])
        vulnerability_info = steps_config.get("vulnerability", {})

        # 填充漏洞信息
        result.test_id = vulnerability_info.get("id", "")
        result.level = vulnerability_info.get("level", "info")
        result.description = vulnerability_info.get("description", "")
        result.fix_suggestion = vulnerability_info.get("fix_suggestion", "")
        result.metadata = {
            "vuln_name": vulnerability_info.get("name", ""),
        }

        try:
            # ==================== 前置步骤：提取资源 ID ====================
            pre_steps = steps_config.get("pre_steps", [])
            extracted_vars: Dict[str, Any] = {}

            for pre_step in pre_steps:
                pre_action = pre_step.get("action", "")
                if pre_action == "get_resource_id":
                    pre_account = pre_step.get("account", "user")
                    pre_path = pre_step.get("path", "/")
                    pre_params = pre_step.get("params", {})
                    pre_extract = pre_step.get("extract", {})

                    pre_token = self.auth.get_token(pre_account)
                    logger.info(f"前置步骤: GET {pre_path} (账号: {pre_account})")

                    pre_response = self.client.get(
                        path=pre_path,
                        params=pre_params,
                        token=pre_token,
                        account=pre_account,
                    )

                    if pre_response:
                        for var_name, json_path in pre_extract.items():
                            value = self._extract_from_response(pre_response, json_path)
                            if value is not None:
                                extracted_vars[var_name] = value
                                logger.info(f"  提取变量 {var_name} = {value}")
                            else:
                                logger.warning(f"  提取变量 {var_name} 失败: {json_path}")
                    else:
                        logger.warning(f"前置步骤请求失败: {pre_path}")

            # ==================== Token 处理 ====================
            account = auth_config.get("account", "user")
            use_fake_token = auth_config.get("use_fake_token", False)
            after_logout = auth_config.get("after_logout", False)

            token = self.auth.get_token(account)

            # 处理 fake token
            if use_fake_token:
                token = self._generate_fake_token()
                logger.info(f"使用伪造 Token 测试认证绕过")
                result.metadata["token_type"] = "fake"

            # 处理 logout 后使用旧 Token
            elif after_logout and token:
                # 从配置或知识库获取注销路径
                auth_cfg = self.config.get("auth", {})
                logout_path = auth_cfg.get("logout_path", "/platform/api/auth/logout")
                logger.info(f"注销后测试: POST {logout_path}")
                self.client.post(path=logout_path, token=token, account=account)
                # Token 已被注销，但保留变量中的值用于请求
                logger.info("已注销，使用旧 Token 发送请求")
                result.metadata["token_type"] = "post_logout"

            if auth_config.get("required", False) and not token:
                result.status = "error"
                result.error = f"账号 {account} 登录失败，无法获取 Token"
                result.level = "info"
                return result

            # ==================== 替换变量 ====================
            method = target.get("method", "GET").upper()
            path = target.get("path", "/")
            params = request_config.get("params", {})
            body = request_config.get("body")

            # 替换 ${var_name} 占位符
            if extracted_vars:
                path = self._substitute_vars(path, extracted_vars)
                params = self._substitute_vars(params, extracted_vars)
                body = self._substitute_vars(body, extracted_vars)

            request_info = f"请求方法: {method}\n请求路径: {path}"
            if params:
                request_info += f"\n查询参数: {params}"
            if body:
                request_info += f"\n请求体: {body}"
            request_info += f"\n使用账号: {account}"
            if use_fake_token:
                request_info += "\nToken类型: 伪造Token"
            if after_logout:
                request_info += "\nToken类型: 注销后Token"

            result.request_info = request_info

            # ==================== 限流测试（批量快速请求） ====================
            rate_limit_test = steps_config.get("rate_limit_test", {})
            if rate_limit_test:
                request_count = rate_limit_test.get("request_count", 10)
                interval_ms = rate_limit_test.get("interval_ms", 100)
                expect_limited = rate_limit_test.get("expect_limited", True)

                logger.info(f"限流测试: {request_count} 次请求, 间隔 {interval_ms}ms")
                rate_limited_responses = 0
                total_duration = 0.0
                last_response = None

                for i in range(request_count):
                    if self._stop_requested:
                        break
                    step_start = time.time()
                    resp = self.client.request(
                        method=method,
                        path=path,
                        token=token,
                        account=account,
                        json_data=body,
                        params=params,
                    )
                    step_duration = time.time() - step_start
                    total_duration += step_duration
                    last_response = resp

                    if resp and AssertionEngine.is_rate_limited(resp):
                        rate_limited_responses += 1

                    if i < request_count - 1:
                        time.sleep(interval_ms / 1000.0)

                step_duration = total_duration

                if last_response:
                    result.response_info = (
                        f"限流测试: {request_count}次请求, "
                        f"{rate_limited_responses}次被限流\n"
                        f"最后一次响应状态码: {last_response.status_code}\n"
                        f"响应体: {last_response.text[:500]}"
                    )
                else:
                    result.response_info = "无响应"

                # 限流测试的断言：使用最后一次响应评估
                is_vulnerable, desc = AssertionEngine.evaluate_assertions(
                    last_response, assertions, duration=step_duration
                )

                # 补充限流测试的描述
                if not is_vulnerable and expect_limited and rate_limited_responses == 0:
                    is_vulnerable = True
                    desc = f"限流测试: {request_count}次请求均未被限流"
                elif is_vulnerable and expect_limited and rate_limited_responses > 0:
                    desc = f"限流测试: {rate_limited_responses}/{request_count}次被限流，存在速率限制"

                result.is_vulnerable = is_vulnerable
                result.duration = step_duration
                result.status = "passed" if not is_vulnerable else "failed"

                if is_vulnerable:
                    result.description = vulnerability_info.get("description", desc)
                    logger.warning(f"[{result.test_id}] 发现漏洞: {result.description}")
                else:
                    logger.info(f"[{result.test_id}] 安全: {case_name}")

            # ==================== 普通请求 ====================
            else:
                step_start = time.time()
                response = self.client.request(
                    method=method,
                    path=path,
                    token=token,
                    account=account,
                    json_data=body,
                    params=params,
                )
                step_duration = time.time() - step_start

                # 构建响应信息
                if response:
                    result.response_info = (
                        f"状态码: {response.status_code}\n"
                        f"响应体: {response.text[:500]}"
                    )
                else:
                    result.response_info = "无响应"

                # 评估断言
                is_vulnerable, desc = AssertionEngine.evaluate_assertions(
                    response, assertions, duration=step_duration
                )

                result.is_vulnerable = is_vulnerable
                result.duration = step_duration
                result.status = "passed" if not is_vulnerable else "failed"

                if is_vulnerable:
                    result.description = vulnerability_info.get("description", desc)
                    logger.warning(f"[{result.test_id}] 发现漏洞: {result.description}")
                else:
                    logger.info(f"[{result.test_id}] 安全: {case_name}")

            # 通知进度
            if self._progress_callback:
                self._progress_callback(case_id, result.status, result)

        except Exception as e:
            result.status = "error"
            result.error = str(e)
            logger.error(f"执行用例 {case_id} 异常: {e}")

        result.end_time = datetime.now()
        result.duration = (result.end_time - result.start_time).total_seconds()

        return result

    def execute_cases(
        self,
        cases: List[Dict[str, Any]],
        progress_callback: Optional[Callable] = None
    ) -> List[SecurityCaseResult]:
        """
        批量执行安全测试用例

        Args:
            cases: 用例配置列表
            progress_callback: 进度回调函数

        Returns:
            list: SecurityCaseResult 列表
        """
        if progress_callback:
            self.set_progress_callback(progress_callback)

        self.results = []
        total = len(cases)

        for idx, case in enumerate(cases, 1):
            if self._stop_requested:
                logger.info("收到停止请求，中断执行")
                break

            logger.info(f"执行用例 [{idx}/{total}]: {case.get('name', '')}")
            result = self.execute_case(case)
            self.results.append(result)

            # 请求间隔（防止触发限流）
            limits = self.config.get("limits", {})
            interval = limits.get("request_interval", 0.5)
            time.sleep(interval)

        return self.results

    def get_summary(self) -> Dict[str, int]:
        """
        获取执行摘要统计

        Returns:
            dict: 统计数据
        """
        summary = {
            "total": len(self.results),
            "critical": 0,
            "high": 0,
            "medium": 0,
            "low": 0,
            "info": 0,
            "passed": 0,
            "failed": 0,
            "error": 0,
        }

        for r in self.results:
            if r.is_vulnerable:
                level = r.level
                if level in summary:
                    summary[level] += 1
                summary["failed"] += 1
            else:
                summary["passed"] += 1

            if r.status == "error":
                summary["error"] += 1

        return summary

    @property
    def is_running(self) -> bool:
        """是否正在执行"""
        return self._is_running