"""L2b YARA 规则与二进制分析引擎。

包含：
- YARA 规则匹配（规则库从 ``rules/`` 目录加载，支持热加载）
- PE 文件结构分析（pefile）：节区异常、导入表、编译信息、签名
- 可执行文件能力识别（capa，若已安装）
"""

from __future__ import annotations

import os
import threading
from typing import Any, Dict, List, Optional

from app.engines.base import Engine, EngineResult, Finding, ScanTarget, Severity

BINARY_EXTS = {".exe", ".dll", ".sys", ".scr", ".com", ".cpl", ".ocx", ".msi", ".apk", ".jar", ".so", ".bin", ".elf"}


# ---------- pefile 失败时的手工 PE 分析 fallback ----------

# PE Section Characteristics 常量（与 pefile.SECTION_CHARACTERISTICS 一致）
_PE_SCN_MEM_EXECUTE = 0x20000000
_PE_SCN_MEM_WRITE = 0x80000000

# 声明扩展名 → 该格式自己的魔数。这些容器本就不是 PE，用 pefile 解析必然抛错；
# 把"格式完全正常的安装包/共享库"记成 PeParseFailed 会给每个 .msi/.jar/.so 白送 10 分。
_NATIVE_MAGIC: Dict[str, Any] = {
    ".msi": (b"\xd0\xcf\x11\xe0",),
    ".msp": (b"\xd0\xcf\x11\xe0",),
    ".mst": (b"\xd0\xcf\x11\xe0",),
    ".jar": (b"PK\x03\x04", b"PK\x05\x06"),
    ".war": (b"PK\x03\x04", b"PK\x05\x06"),
    ".apk": (b"PK\x03\x04", b"PK\x05\x06"),
    ".so": (b"\x7fELF",),
    ".elf": (b"\x7fELF",),
}


def _is_native_container(path: str, ext: str) -> bool:
    """文件内容是否就是其声明格式（非 PE）。"""
    expected = _NATIVE_MAGIC.get(ext)
    if not expected:
        return False
    try:
        with open(path, "rb") as fh:
            head = fh.read(4)
    except OSError:
        return False
    return any(head.startswith(m) for m in expected)


def _manual_pe_analysis(path: str) -> List[Finding]:
    """pefile 解析失败时，手工搜索 PE\\0\\0 签名并解析节表。

    覆盖场景：假 PE / 伪装 PE / e_lfanew 偏移错误。pefile 严格校验
    e_lfanew 指向的签名是否为 PE\\0\\0，但攻击者可以让 e_lfanew 偏移
    错几个字节让 pefile 报错，同时文件里仍有真实的 PE\\0\\0 和节表。
    """
    import struct

    findings: List[Finding] = []
    try:
        with open(path, "rb") as fh:
            data = fh.read()
    except OSError:
        return findings

    # 1) MZ 头存在？（PE 文件必需）
    if len(data) < 2 or data[:2] != b"MZ":
        return findings  # 根本不是 PE

    # 2) 尝试从 e_lfanew 取 PE 偏移；不对就全局搜索 PE\0\0
    pe_off = -1
    if len(data) >= 0x40:
        try:
            e_lfanew = struct.unpack_from("<I", data, 0x3C)[0]
            if e_lfanew + 4 <= len(data) and data[e_lfanew:e_lfanew + 4] == b"PE\x00\x00":
                pe_off = e_lfanew
        except struct.error:
            pass

    if pe_off < 0:
        pe_off = data.find(b"PE\x00\x00")
    if pe_off < 0:
        return findings  # 没找到 PE 签名

    # 3) COFF header 在 PE\0\0 后 4 字节
    coff_off = pe_off + 4
    if coff_off + 24 > len(data):
        return findings
    num_sections = struct.unpack_from("<H", data, coff_off + 2)[0]
    size_opt_header = struct.unpack_from("<H", data, coff_off + 16)[0]
    section_table_off = coff_off + 20 + size_opt_header

    if num_sections == 0 or num_sections > 96:  # 合理上限
        return findings
    if section_table_off + num_sections * 40 > len(data):
        return findings

    # 4) 逐节扫描 RWX 属性
    rwx_sections = []
    for i in range(num_sections):
        sec_off = section_table_off + i * 40
        name = data[sec_off:sec_off + 8].rstrip(b"\x00")
        characteristics = struct.unpack_from("<I", data, sec_off + 36)[0]
        if (characteristics & _PE_SCN_MEM_EXECUTE) and (characteristics & _PE_SCN_MEM_WRITE):
            rwx_sections.append(name.decode("utf-8", errors="replace"))

    if rwx_sections:
        findings.append(
            Finding(
                engine="binary", layer="L2B",
                severity=Severity.HIGH.value, rule="RWXSectionDetected",
                detail=f"手工 PE 分析检测到可写且可执行的危险节：{', '.join(rwx_sections)}",
                weight=50.0,
            )
        )

    return findings


# ---------- 模块级全局 YARA 规则单例 ----------
# 所有 YaraEngine 实例共享同一个编译结果，避免重复 compile
# rules_dir 不同时（极少见）才分叉加载
_YARA_GLOBAL: Dict[str, Dict[str, Any]] = {}
_YARA_LOCK = threading.Lock()


def _get_global_yara(rules_dir: str) -> Optional[Any]:
    """获取全局共享的 YARA Rules 对象（自动按目录 mtime 重载）。"""
    with _YARA_LOCK:
        cached = _YARA_GLOBAL.get(rules_dir)
        import yara as _yr
        newest = 0.0
        files: List[str] = []
        if os.path.isdir(rules_dir):
            for name in sorted(os.listdir(rules_dir)):
                if name.endswith((".yar", ".yara")):
                    path = os.path.join(rules_dir, name)
                    files.append(path)
                    newest = max(newest, os.path.getmtime(path))

        if cached and cached.get("mtime") == newest:
            return cached.get("rules")

        if not files:
            _YARA_GLOBAL[rules_dir] = {"rules": None, "mtime": newest}
            return None

        try:
            compiled = _yr.compile(filepaths={os.path.basename(p): p for p in files})
        except Exception:
            _YARA_GLOBAL[rules_dir] = {"rules": None, "mtime": newest}
            return None

        _YARA_GLOBAL[rules_dir] = {"rules": compiled, "mtime": newest}
        return compiled


def _discover_rules_dir() -> str:
    """自动探测 YARA 规则目录（优先环境变量 > /app/rules > 项目根 rules/）。"""
    env_dir = os.environ.get("MESHSAFE_RULES_DIR", "").strip()
    if env_dir and os.path.isdir(env_dir):
        return env_dir

    for candidate in ("/app/rules", "rules", "./rules", "../rules"):
        if os.path.isdir(candidate):
            return os.path.abspath(candidate)

    return env_dir or "/app/rules"


class YaraEngine(Engine):
    """YARA 规则匹配引擎（使用全局共享的编译结果）。"""

    name = "yara"
    layer = "L2B"

    def __init__(self, rules_dir: str | None = None) -> None:
        # 显式传入时才用传入值，否则走自动探测
        self.rules_dir = rules_dir if rules_dir else _discover_rules_dir()

    def _load_rules(self) -> Any:
        """委托给全局单例。"""
        try:
            import yara  # noqa: F401 — 只是探测可用性
        except Exception:  # noqa: BLE001
            return None
        return _get_global_yara(self.rules_dir)

    def scan(self, target: ScanTarget, ctx: Any) -> EngineResult:
        result = EngineResult(engine=self.name)
        rules = self._load_rules()
        if rules is None:
            result.ok = False
            result.error = "YARA 规则库不可用或编译失败"
            return result

        try:
            # 本环境的 yara-python 4.5.4 在 Match 对象上默认就带 meta，
            # 且 rules.match() 不接受 with_meta 关键字（传了会整引擎报错）。
            matches = rules.match(target.path)
        except Exception as exc:  # noqa: BLE001
            result.ok = False
            result.error = f"YARA 匹配失败：{exc}"
            return result

        for match in matches:
            meta = getattr(match, "meta", None) or {}
            severity, weight = _judge_yara_rule(match.rule, match.tags, meta)
            # 覆盖：如果 ctx 里的策略为该 YARA 规则定义了精细权重，用它
            try:
                policy = ctx.policy_engine.policy
                yara_weights = getattr(policy, "yara_weights", {}) or {}
                if match.rule in yara_weights:
                    weight = float(yara_weights[match.rule])
            except Exception:
                pass  # ctx 没有 policy 字段时保持默认 weight
            desc = str(meta.get("description", "")).strip()
            tags_text = ", ".join(match.tags)
            result.findings.append(
                Finding(
                    engine=self.name,
                    layer=self.layer,
                    severity=severity,
                    rule=match.rule,
                    detail=(
                        f"YARA 规则 {match.rule} 命中"
                        + (f"（{desc}）" if desc else f"（标签：{tags_text or '无'}）")
                    ),
                    weight=weight,
                    meta={"tags": list(match.tags)},
                )
            )
        # rule_set_version 用全局 compiled 对象的 id（每次 reload 变）
        try:
            _r = _get_global_yara(self.rules_dir)
            result.meta["rule_set_version"] = str(id(_r)) if _r else "unavailable"
        except Exception:
            result.meta["rule_set_version"] = "unknown"
        return result

    def health(self) -> Dict[str, Any]:
        rules = self._load_rules()
        ok = rules is not None
        return {
            "engine": self.name,
            "version": "rules" if ok else "none",
            "ok": ok,
            "status": "ok" if ok else "degraded",
            "error": None if ok else "YARA 规则库未加载",
            "info": f"rules_dir={self.rules_dir}" if not ok else "rules loaded",
        }


class BinaryEngine(Engine):
    """PE/二进制文件结构分析引擎。"""

    name = "binary"
    layer = "L2B"

    def scan(self, target: ScanTarget, ctx: Any) -> EngineResult:
        result = EngineResult(engine=self.name)
        ext = (target.declared_ext or "").lower()
        if ext not in BINARY_EXTS:
            return result

        # 内容确实是声明的那个非 PE 容器（MSI=CFB、JAR/APK=ZIP、SO/ELF=ELF）→ 不做 PE 分析
        if _is_native_container(target.path, ext):
            return result

        try:
            import pefile
        except Exception:  # noqa: BLE001
            result.ok = False
            result.error = "pefile 未安装"
            return result

        try:
            pe = pefile.PE(target.path, fast_load=False)
        except Exception as exc:  # noqa: BLE001
            result.findings.append(
                Finding(
                    engine=self.name,
                    layer=self.layer,
                    severity=Severity.MEDIUM.value,
                    rule="PeParseFailed",
                    detail=f"PE 结构解析失败：{exc}",
                    weight=10.0,
                )
            )
            # --- pefile 失败时尝试手工 PE 分析（覆盖假 PE / 伪装 PE）---
            manual = _manual_pe_analysis(target.path)
            if manual:
                result.findings.extend(manual)
            return result

        try:
            result.findings.extend(self._analyze(pe, target))
        finally:
            try:
                pe.close()
            except Exception:  # noqa: BLE001
                pass

        return result

    def _analyze(self, pe: Any, target: ScanTarget) -> List[Finding]:
        findings: List[Finding] = []

        # 1) 节区熵值：加密/加壳的节区熵接近 8
        for section in getattr(pe, "sections", []) or []:
            try:
                name = section.Name.rstrip(b"\x00").decode("utf-8", errors="replace")
                entropy = section.get_entropy()
            except Exception:  # noqa: BLE001
                continue
            if entropy >= 7.2:
                findings.append(
                    Finding(
                        engine=self.name,
                        layer=self.layer,
                        severity=Severity.MEDIUM.value,
                        rule="HighEntropySection",
                        detail=f"节区 {name} 熵值 {entropy:.2f}，可能经过加密或加壳",
                        weight=20.0,
                    )
                )

        # 2) 导入表能力信号：只做佐证，不做独立定罪
        #    Windows 自身的内核/外壳组件（kernel32、shell32、wer）必然导入
        #    VirtualAllocEx/WriteProcessMemory 这类 API——它们就是这些 API 的实现方。
        #    真正的注入判定由 YARA C1_PE_ProcessInjection（导入名邻接窗口）负责，
        #    这里按「命中个数」而不是「命中即高危」计分，避免正常软件被判恶意。
        tier1, tier2, tier3 = self._suspicious_imports(pe)
        if len(tier1) >= 2:
            findings.append(
                Finding(
                    engine=self.name,
                    layer=self.layer,
                    severity=Severity.HIGH.value,
                    rule="SuspiciousImports",
                    detail=f"导入表同时包含多个注入/后门级 API：{', '.join(tier1[:8])}",
                    weight=20.0,
                )
            )
        elif len(tier1) == 1:
            findings.append(
                Finding(
                    engine=self.name,
                    layer=self.layer,
                    severity=Severity.LOW.value,
                    rule="SuspiciousImports",
                    detail=f"导入表含单个注入/后门级 API（系统组件常态，需其他信号佐证）：{tier1[0]}",
                    weight=5.0,
                )
            )
        if len(tier2) >= 3:
            findings.append(
                Finding(
                    engine=self.name,
                    layer=self.layer,
                    severity=Severity.LOW.value,
                    rule="CapabilityImports",
                    detail=f"导入表同时具备多项危险能力：{', '.join(tier2[:8])}",
                    weight=10.0,
                )
            )
        elif tier3:
            findings.append(
                Finding(
                    engine=self.name,
                    layer=self.layer,
                    severity=Severity.INFO.value,
                    rule="AntiDebugImports",
                    detail=f"导入表含反调试/环境探测 API（正常软件普遍使用）：{', '.join(tier3[:8])}",
                    weight=0.0,
                )
            )

        # 3) 无签名可执行文件
        if not self._has_signature(pe):
            findings.append(
                Finding(
                    engine=self.name,
                    layer=self.layer,
                    severity=Severity.MEDIUM.value,
                    rule="UnsignedExecutable",
                    detail="可执行文件未携带数字签名（内网自研工具常态，仅辅助人工判断）",
                    weight=10.0,
                )
            )

        # 4) 编译时间戳异常（明显晚于其他节区时间，或为 1970/未来时间）
        timestamp = getattr(pe, "FILE_HEADER", None)
        if timestamp is not None:
            try:
                ts = int(timestamp.TimeDateStamp)
                if ts == 0 or ts > 4102444800:  # 2100 年之后
                    findings.append(
                        Finding(
                            engine=self.name,
                            layer=self.layer,
                            severity=Severity.LOW.value,
                            rule="SuspiciousTimestamp",
                            detail="PE 编译时间戳异常，可能被伪造",
                            weight=0.0,
                        )
                    )
            except Exception:  # noqa: BLE001
                pass

        # 5) UPX 壳显式检测 —— 检查节名 UPX0 / UPX1 / UPX!
        upx_sections = []
        for section in getattr(pe, "sections", []) or []:
            try:
                name = section.Name.rstrip(b"\x00").decode("utf-8", errors="replace")
                if name.upper().startswith("UPX"):
                    upx_sections.append(name)
            except Exception:
                continue
        if upx_sections:
            findings.append(
                Finding(
                    engine=self.name,
                    layer=self.layer,
                    severity=Severity.MEDIUM.value,
                    rule="UPXPackerDetected",
                    detail=f"检测到 UPX 加壳节：{', '.join(upx_sections)}",
                    weight=30.0,
                )
            )

        # 6) RWX 可写可执行节 —— 高危（绕过 DEP 的典型手法）
        from pefile import SECTION_CHARACTERISTICS  # noqa: PLC0415
        MEM_EXECUTE = SECTION_CHARACTERISTICS["IMAGE_SCN_MEM_EXECUTE"]
        MEM_WRITE = SECTION_CHARACTERISTICS["IMAGE_SCN_MEM_WRITE"]
        rwx_sections = []
        for section in getattr(pe, "sections", []) or []:
            try:
                flags = section.Characteristics
                if (flags & MEM_EXECUTE) and (flags & MEM_WRITE):
                    name = section.Name.rstrip(b"\x00").decode("utf-8", errors="replace")
                    rwx_sections.append(name)
            except Exception:
                continue
        if rwx_sections:
            findings.append(
                Finding(
                    engine=self.name,
                    layer=self.layer,
                    severity=Severity.HIGH.value,
                    rule="RWXSectionDetected",
                    detail=f"检测到可写且可执行的危险节：{', '.join(rwx_sections)}",
                    weight=40.0,
                )
            )

        # 7) 32/64 位识别（info 级别，辅助判断）
        machine = None
        pe32 = None
        try:
            machine = int(pe.FILE_HEADER.Machine)
            pe32 = (pe.OPTIONAL_HEADER.Magic == 0x10B)  # PE32
        except Exception:
            pass
        if machine is not None:
            try:
                from pefile import MACHINE_TYPES  # noqa: PLC0415
                machine_name = MACHINE_TYPES.get(machine, f"UNKNOWN(0x{machine:04x})")
            except Exception:
                machine_name = f"0x{machine:04x}"
            bits = "32-bit" if pe32 else "64-bit"
            findings.append(
                Finding(
                    engine=self.name,
                    layer=self.layer,
                    severity=Severity.LOW.value,
                    rule="PEArchInfo",
                    detail=f"PE 架构：{machine_name} / {bits}",
                    weight=0.0,
                )
            )

        # 8) 调试信息（pdb 路径或 debug 节）
        has_debug_dir = False
        has_debug_section = False
        try:
            debug_dir = pe.OPTIONAL_HEADER.DATA_DIRECTORY[6]  # IMAGE_DIRECTORY_ENTRY_DEBUG
            if debug_dir.VirtualAddress and debug_dir.Size:
                has_debug_dir = True
        except Exception:
            pass
        for section in getattr(pe, "sections", []) or []:
            try:
                name = section.Name.rstrip(b"\x00").decode("utf-8", errors="replace").lower()
                if ".debug" in name or name.startswith("debug") or name.startswith(".pdb"):
                    has_debug_section = True
                    break
            except Exception:
                continue
        if has_debug_dir or has_debug_section:
            findings.append(
                Finding(
                    engine=self.name,
                    layer=self.layer,
                    severity=Severity.LOW.value,
                    rule="DebugInformationPresent",
                    detail="PE 包含调试信息（PDB 路径或 .debug 节），可用于逆向定位",
                    weight=0.0,
                )
            )

        return findings

    def _suspicious_imports(self, pe: Any) -> "tuple[List[str], List[str], List[str]]":
        """导入 API 按恶意指向分档。

        返回 (tier1, tier2, tier3)：
          tier1 注入/后门级 API —— 单独出现即可疑
          tier2 危险能力 API    —— 正常软件也常用，需多项共存才计分
          tier3 反调试/探测 API —— 信息项，不计分（大量合法商业软件带加壳自保会用）
        """
        tier1_set = {
            "VirtualAllocEx", "WriteProcessMemory", "CreateRemoteThread",
            "NtUnmapViewOfSection", "SetWindowsHookEx", "SetWindowsHookExA",
            "SetWindowsHookExW", "URLDownloadToFileA", "URLDownloadToFileW",
            "QueueUserAPC", "NtQueueApcThread", "RtlCreateUserThread",
        }
        tier2_set = {
            "OpenProcess", "WinHttpOpen", "WinHttpOpenRequest", "WinHttpSendRequest",
            "CryptEncrypt", "CryptDeriveKey", "AdjustTokenPrivileges",
            "LookupPrivilegeValueA", "LookupPrivilegeValueW", "OpenProcessToken",
            "VirtualProtectEx", "ReadProcessMemory", "SuspendThread", "GetAsyncKeyState",
            "SetWinEventHook", "HttpSendRequestA", "InternetWriteFile",
            "RegSetValueExA", "RegSetValueExW",
            # 以下 API 恶意软件常用，但任何"要装服务/要启动进程/要联网"的正常
            # 软件都必然导入它们，不足以单独定罪
            "WinExec", "ShellExecuteExA", "ShellExecuteExW", "CreateServiceA",
            "CreateServiceW", "InternetOpenUrlA", "InternetOpenUrlW",
            "CreateNamedPipeA", "DefineDosDeviceW",
        }
        tier3_set = {
            "IsDebuggerPresent", "CheckRemoteDebuggerPresent", "OutputDebugStringA",
            "GetTickCount", "GetTickCount64", "QueryPerformanceCounter",
            "GetProcessHeap", "FindWindowA", "FindWindowW",
        }
        found1: List[str] = []
        found2: List[str] = []
        found3: List[str] = []
        try:
            for entry in getattr(pe, "DIRECTORY_ENTRY_IMPORT", []) or []:
                for imp in entry.imports:
                    if not imp.name:
                        continue
                    name = imp.name.decode("utf-8", errors="replace")
                    if name in tier1_set:
                        found1.append(name)
                    elif name in tier2_set:
                        found2.append(name)
                    elif name in tier3_set:
                        found3.append(name)
        except Exception:  # noqa: BLE001
            pass
        return sorted(set(found1)), sorted(set(found2)), sorted(set(found3))

    def _has_signature(self, pe: Any) -> bool:
        """检查 PE 是否包含 Authenticode 签名目录。"""
        try:
            directory = pe.OPTIONAL_HEADER.DATA_DIRECTORY[4]  # IMAGE_DIRECTORY_ENTRY_SECURITY
            return bool(directory.VirtualAddress and directory.Size)
        except Exception:  # noqa: BLE001
            return False


_YARA_RISK_SEVERITY = {
    "critical": Severity.CRITICAL.value,
    "high": Severity.HIGH.value,
    "medium": Severity.MEDIUM.value,
    "low": Severity.LOW.value,
}
_YARA_RISK_WEIGHT = {"critical": 85.0, "high": 60.0, "medium": 30.0, "low": 10.0}


def _judge_yara_rule(rule: str, tags: List[str], meta: Dict[str, Any] | None = None) -> tuple[str, float]:
    """依据规则 meta.risk_level 判定严重度（其次才是标签/规则名启发式）。

    v0.2.3 只看标签与规则名，规则库里的 risk_level 被完全忽略，
    结果 65 分的 C1_PE_RAT 显示成"低危"，权重与严重度两套口径。
    """
    risk = str((meta or {}).get("risk_level", "")).strip().lower()
    if risk in _YARA_RISK_SEVERITY:
        return _YARA_RISK_SEVERITY[risk], _YARA_RISK_WEIGHT[risk]

    tag_set = {t.lower() for t in tags}
    text = f"{rule} {' '.join(tags)}".lower()

    if "critical" in tag_set or any(k in text for k in ("malware", "ransomware", "backdoor", "c2", "exploit")):
        return Severity.CRITICAL.value, 85.0
    if "high" in tag_set or any(k in text for k in ("trojan", "dropper", "obfusc", "shellcode")):
        return Severity.HIGH.value, 60.0
    if "medium" in tag_set or any(k in text for k in ("suspicious", "pua", "packer")):
        return Severity.MEDIUM.value, 30.0
    if "suspicious" in tag_set:
        return Severity.MEDIUM.value, 30.0
    return Severity.LOW.value, 10.0