#!/usr/bin/env python3
"""文件验证模块 - 魔数验证、恶意内容扫描、文件名安全检查，防止上传攻击。"""

import logging
import re
from collections.abc import Callable
from typing import TypedDict, cast

logger = logging.getLogger(__name__)


class ImageTypeInfo(TypedDict, total=False):
    """图片类型信息（用于魔数识别）"""

    extension: str
    mime_type: str
    description: str
    additional_check: Callable[[bytes], bool]


class FileValidationResult(TypedDict):
    """文件验证结果结构（用于类型检查与 IDE 提示）"""

    valid: bool
    file_type: str | None
    mime_type: str | None
    extension: str | None
    size: int
    warnings: list[str]
    errors: list[str]


# 图片格式魔数字典：{魔数字节: {extension, mime_type, description, additional_check?}}
IMAGE_MAGIC_NUMBERS: dict[bytes, ImageTypeInfo] = {
    # PNG格式
    b"\x89\x50\x4e\x47\x0d\x0a\x1a\x0a": {
        "extension": ".png",
        "mime_type": "image/png",
        "description": "PNG图片",
    },
    # JPEG格式 (多种变体)
    b"\xff\xd8\xff\xe0": {
        "extension": ".jpg",
        "mime_type": "image/jpeg",
        "description": "JPEG图片 (JFIF)",
    },
    b"\xff\xd8\xff\xe1": {
        "extension": ".jpg",
        "mime_type": "image/jpeg",
        "description": "JPEG图片 (EXIF)",
    },
    b"\xff\xd8\xff\xe2": {
        "extension": ".jpg",
        "mime_type": "image/jpeg",
        "description": "JPEG图片 (Canon)",
    },
    b"\xff\xd8\xff\xe3": {
        "extension": ".jpg",
        "mime_type": "image/jpeg",
        "description": "JPEG图片 (Samsung)",
    },
    b"\xff\xd8\xff\xdb": {
        "extension": ".jpg",
        "mime_type": "image/jpeg",
        "description": "JPEG图片 (标准)",
    },
    # GIF格式
    b"\x47\x49\x46\x38\x37\x61": {
        "extension": ".gif",
        "mime_type": "image/gif",
        "description": "GIF图片 (87a)",
    },
    b"\x47\x49\x46\x38\x39\x61": {
        "extension": ".gif",
        "mime_type": "image/gif",
        "description": "GIF图片 (89a)",
    },
    # WebP格式
    b"\x52\x49\x46\x46": {
        "extension": ".webp",
        "mime_type": "image/webp",
        "description": "WebP图片",
        "additional_check": lambda data: data[8:12] == b"WEBP",
    },
    # BMP格式
    b"\x42\x4d": {
        "extension": ".bmp",
        "mime_type": "image/bmp",
        "description": "BMP图片",
    },
    # TIFF格式
    b"\x49\x49\x2a\x00": {
        "extension": ".tiff",
        "mime_type": "image/tiff",
        "description": "TIFF图片 (Little Endian)",
    },
    b"\x4d\x4d\x00\x2a": {
        "extension": ".tiff",
        "mime_type": "image/tiff",
        "description": "TIFF图片 (Big Endian)",
    },
    # ICO格式
    b"\x00\x00\x01\x00": {
        "extension": ".ico",
        "mime_type": "image/x-icon",
        "description": "ICO图标",
    },
    # SVG格式 (XML开头)
    b"\x3c\x3f\x78\x6d\x6c": {
        "extension": ".svg",
        "mime_type": "image/svg+xml",
        "description": "SVG矢量图",
        "additional_check": lambda data: b"<svg" in data[:1024].lower(),
    },
}

# 危险文件扩展名黑名单：可执行文件、脚本、打包文件
DANGEROUS_EXTENSIONS = {
    ".exe",
    ".bat",
    ".cmd",
    ".com",
    ".scr",
    ".pif",
    ".vbs",
    ".js",
    ".jar",
    ".msi",
    ".dll",
    ".sys",
    ".drv",
    ".ocx",
    ".cpl",
    ".inf",
    ".reg",
    ".ps1",
    ".sh",
    ".bash",
    ".zsh",
    ".fish",
    ".py",
    ".pl",
    ".rb",
    ".php",
    ".asp",
    ".jsp",
    ".war",
    ".ear",
    ".deb",
    ".rpm",
    ".dmg",
    ".pkg",
    ".app",
}

# 恶意内容正则模式：JavaScript/PHP/Shell/SQL 注入特征
MALICIOUS_PATTERNS = [
    # JavaScript代码模式
    rb"<script[^>]*>",
    rb"javascript:",
    rb"eval\s*\(",
    rb"document\.write",
    rb"window\.location",
    # PHP代码模式
    rb"<\?php",
    rb"<\?=",
    rb"eval\s*\(",
    rb"system\s*\(",
    rb"exec\s*\(",
    # Shell命令模式
    rb"#!/bin/",
    rb"rm\s+-rf",
    rb"wget\s+",
    rb"curl\s+",
    # SQL注入模式
    rb"union\s+select",
    rb"drop\s+table",
    rb"insert\s+into",
    rb"delete\s+from",
]


class FileValidationError(Exception):
    """文件验证异常"""

    pass


class FileValidator:
    """文件验证器 - 魔数验证、恶意内容扫描、文件名安全检查。"""

    # 【优化】类级别常量：危险字符集合（所有实例共享）
    _DANGEROUS_CHARS = frozenset(["<", ">", ":", '"', "|", "?", "*", "\0"])

    def __init__(self, max_file_size: int = 10 * 1024 * 1024):  # 10MB
        """初始化并预编译恶意内容正则"""
        # 验证max_file_size参数
        if max_file_size <= 0:
            raise ValueError(f"max_file_size 必须为正数，当前值: {max_file_size}")

        self.max_file_size = max_file_size
        # 【优化】预编译正则并缓存 decoded pattern 字符串
        self.compiled_patterns = []
        for pattern in MALICIOUS_PATTERNS:
            compiled = re.compile(pattern, re.IGNORECASE)
            pattern_str = pattern.decode("utf-8", errors="ignore")
            self.compiled_patterns.append((compiled, pattern_str))

    def validate_file(
        self,
        file_data: bytes | None,
        filename: str,
        declared_mime_type: str | None = None,
    ) -> FileValidationResult:
        """验证文件安全性，返回 {valid, file_type, mime_type, extension, size, warnings, errors}"""
        # 验证输入参数
        if not filename or not filename.strip():
            return {
                "valid": False,
                "file_type": None,
                "mime_type": None,
                "extension": None,
                "size": 0,
                "warnings": [],
                "errors": ["文件名为空"],
            }

        if file_data is None:
            return {
                "valid": False,
                "file_type": None,
                "mime_type": None,
                "extension": None,
                "size": 0,
                "warnings": [],
                "errors": ["文件数据为空（None）"],
            }

        result: FileValidationResult = {
            "valid": False,
            "file_type": None,
            "mime_type": None,
            "extension": None,
            "size": len(file_data),
            "warnings": [],
            "errors": [],
        }

        try:
            # 1. 基础检查
            self._validate_basic_properties(file_data, filename, result)

            # 2. 魔数验证
            detected_type = self._validate_magic_number(file_data, result)

            # 3. 文件名验证
            self._validate_filename(filename, result)

            # 4. MIME类型一致性检查
            if declared_mime_type:
                self._validate_mime_consistency(
                    declared_mime_type, detected_type, result
                )

            # 5. 恶意内容扫描
            self._scan_malicious_content(file_data, result)

            # 6. 最终验证结果
            result["valid"] = len(result["errors"]) == 0

            if result["valid"]:
                logger.info(f"文件验证通过: {filename} ({result['file_type']})")
            else:
                logger.warning(f"文件验证失败: {filename}, 错误: {result['errors']}")

        except Exception as e:
            logger.error(f"文件验证过程中出错: {e}", exc_info=True)
            result["errors"].append(f"验证过程异常: {str(e)}")
            result["valid"] = False

        return result

    def _validate_basic_properties(
        self, file_data: bytes, filename: str, result: FileValidationResult
    ) -> None:
        """检查文件大小、文件名长度、危险扩展名"""
        # 检查文件大小
        if len(file_data) == 0:
            result["errors"].append("文件为空")
            # 不再提前return，继续检查文件名安全性

        if len(file_data) > self.max_file_size:
            result["errors"].append(
                f"文件大小超过限制: {len(file_data)} > {self.max_file_size}"
            )

        # 检查文件名长度
        if len(filename) > 255:
            result["errors"].append("文件名过长")

        # 【优化】使用 rsplit 代替 Path，避免创建对象
        # 原逻辑：Path(filename).suffix.lower()
        # 优化后：提取 '.' 后的扩展名，保留 '.' 前缀
        parts = filename.rsplit(".", 1)
        file_ext = ("." + parts[1]).lower() if len(parts) > 1 else ""

        # 检查危险扩展名
        if file_ext and file_ext in DANGEROUS_EXTENSIONS:
            result["errors"].append(f"危险的文件扩展名: {file_ext}")

    def _validate_magic_number(
        self, file_data: bytes, result: FileValidationResult
    ) -> ImageTypeInfo | None:
        """通过魔数识别真实文件类型（PNG/JPEG 快速路径优化）"""
        detected_type: ImageTypeInfo | None = None

        # 【优化】快速路径：优先检查最常见的格式（PNG、JPEG）
        # PNG 魔数检查（约占 40% 图片上传）
        if file_data.startswith(b"\x89\x50\x4e\x47\x0d\x0a\x1a\x0a"):
            detected_type = cast(
                ImageTypeInfo,
                {
                    "extension": ".png",
                    "mime_type": "image/png",
                    "description": "PNG图片",
                },
            )

        # JPEG 魔数检查（约占 50% 图片上传）
        # 所有 JPEG 变体的前 3 字节都是 \xff\xd8\xff
        elif file_data.startswith(b"\xff\xd8\xff"):
            detected_type = cast(
                ImageTypeInfo,
                {
                    "extension": ".jpg",
                    "mime_type": "image/jpeg",
                    "description": "JPEG图片",
                },
            )

        # 快速路径命中，直接返回
        if detected_type:
            result["file_type"] = detected_type["description"]
            result["mime_type"] = detected_type["mime_type"]
            result["extension"] = detected_type["extension"]
            return detected_type

        # 【优化】慢速路径：跳过已在快速路径检查的 PNG 和 JPEG 格式
        # PNG 魔数：\x89\x50\x4e\x47\x0d\x0a\x1a\x0a
        # JPEG 魔数（5个变体）：\xff\xd8\xff\xe0, \xff\xd8\xff\xe1, \xff\xd8\xff\xe2,
        #                        \xff\xd8\xff\xe3, \xff\xd8\xff\xdb
        _SKIP_MAGIC_BYTES = {
            b"\x89\x50\x4e\x47\x0d\x0a\x1a\x0a",  # PNG
            b"\xff\xd8\xff\xe0",  # JPEG JFIF
            b"\xff\xd8\xff\xe1",  # JPEG EXIF
            b"\xff\xd8\xff\xe2",  # JPEG Canon
            b"\xff\xd8\xff\xe3",  # JPEG Samsung
            b"\xff\xd8\xff\xdb",  # JPEG 标准
        }

        # 慢速路径：检查其他所有格式（跳过快速路径已检查的）
        for magic_bytes, type_info in IMAGE_MAGIC_NUMBERS.items():
            # 【优化】跳过快速路径已检查的格式
            if magic_bytes in _SKIP_MAGIC_BYTES:
                continue

            if file_data.startswith(magic_bytes):
                # 额外检查添加错误处理
                if "additional_check" in type_info:
                    try:
                        if not type_info["additional_check"](file_data):
                            continue
                    except Exception as e:
                        logger.warning(
                            f"额外检查失败: {type_info.get('description', 'Unknown')} - {e}",
                            exc_info=True,
                        )
                        continue

                detected_type = type_info
                result["file_type"] = type_info["description"]
                result["mime_type"] = type_info["mime_type"]
                result["extension"] = type_info["extension"]
                break

        if not detected_type:
            result["errors"].append("无法识别的文件格式或不支持的文件类型")

        return detected_type

    def _validate_filename(self, filename: str, result: FileValidationResult) -> None:
        """检查路径遍历、特殊字符、隐藏文件"""
        # 检查空文件名或只包含空格/点的文件名
        stripped_name = filename.strip()
        if not stripped_name or stripped_name == "." or stripped_name == "..":
            result["errors"].append("文件名无效（空或只包含点）")

        # 检查路径遍历攻击
        if ".." in filename or "/" in filename or "\\" in filename:
            result["errors"].append("文件名包含非法字符")

        # 【优化】使用类级别 frozenset 和反转循环顺序
        # 原逻辑：any(char in filename for char in dangerous_chars) O(n * m)
        # 优化后：any(char in _DANGEROUS_CHARS for char in filename) O(n)
        if any(char in self._DANGEROUS_CHARS for char in filename):
            result["warnings"].append("文件名包含特殊字符")

        # 检查隐藏文件
        if filename.startswith("."):
            result["warnings"].append("隐藏文件")

    def _validate_mime_consistency(
        self,
        declared_mime: str,
        detected_type: ImageTypeInfo | None,
        result: FileValidationResult,
    ):
        """检查声明的 MIME 类型与检测结果是否一致"""
        if not detected_type:
            return

        # 提取MIME类型的主类型（忽略参数部分）
        # 例如："image/png; charset=utf-8" → "image/png"
        declared_main_type = declared_mime.split(";")[0].strip().lower()
        detected_main_type = detected_type["mime_type"].lower()

        if declared_main_type != detected_main_type:
            result["warnings"].append(
                f"MIME类型不一致: 声明={declared_mime}, 检测={detected_type['mime_type']}"
            )

    def _scan_malicious_content(
        self, file_data: bytes, result: FileValidationResult
    ) -> None:
        """扫描前 64KB 检测恶意代码特征"""
        # 只扫描文件的前64KB，避免性能问题
        scan_data = file_data[: 64 * 1024]

        # 遍历所有模式，报告所有匹配
        for compiled, pattern_str in self.compiled_patterns:
            if compiled.search(scan_data):
                # 【优化】使用预先 decoded 的 pattern_str，避免重复 decode
                result["errors"].append(f"检测到可疑内容模式: {pattern_str}")
                # 不短路，继续检查其他模式


# 【优化】模块级单例：预创建默认 FileValidator 实例，避免重复初始化
# 所有 validate_uploaded_file() 调用共享此实例，避免重复编译正则表达式
_default_validator = FileValidator()


def validate_uploaded_file(
    file_data: bytes | None, filename: str, mime_type: str | None = None
) -> FileValidationResult:
    """便捷函数：使用默认单例验证文件"""
    return _default_validator.validate_file(file_data, filename, mime_type)


def is_safe_image_file(file_data: bytes, filename: str) -> bool:
    """便捷函数：返回文件是否通过验证"""
    result = validate_uploaded_file(file_data, filename)
    return result["valid"] and len(result["errors"]) == 0
