MeshSate/app/engines/keywords.py

"""内网特征词表引擎（P1）。

数据驱动：词表从 policy.yaml 的 ``keyword_rules`` 段加载，不在代码中硬编码。
支持三类匹配模式：

* ``contains`` — 子串匹配（区分大小写）
* ``regex`` — 正则表达式匹配
* ``any_of`` — 多关键词中任一命中

用途：识别内网特有模板、密级标识、命名规律等——辅助宏例外判断和内容审计方向。
"""

from __future__ import annotations

import logging
import os
import re
from typing import Any, Dict, List

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

logger = logging.getLogger(__name__)


class KeywordEngine(Engine):
    """内网特征词表匹配引擎。"""

    # ---- 跳过列表：二进制 / 压缩封装 / 媒体文件 ----
    # 这些文件有专门引擎处理（PDF/Office/Archive/PE 等），
    # 它们的内容是压缩流或二进制格式，关键词 regex 在里面会大面积误报。
    # 短关键词 (nc/wce/curl/...) 在压缩流里等于伪随机字节必然碰撞。
    BINARY_SKIP_EXTS = frozenset({
        # PDF / Office
        ".pdf", ".docx", ".docm", ".doc", ".xlsx", ".xlsm", ".xls",
        ".pptx", ".pptm", ".ppt", ".ppsx", ".pps", ".potx",
        # 压缩包
        ".rar", ".zip", ".7z", ".tar", ".gz", ".bz2", ".xz", ".tgz",
        ".iso", ".img", ".dmg",
        # 图片
        ".png", ".jpg", ".jpeg", ".gif", ".bmp", ".tiff", ".tif", ".webp",
        ".ico", ".svg",
        # 视频
        ".mp4", ".avi", ".mkv", ".mov", ".flv", ".wmv", ".webm",
        # 音频
        ".mp3", ".wav", ".flac", ".ogg", ".aac", ".wma", ".m4a",
    })

    name = "keywords"
    layer = "L2B"
    version = "1.0"

    def __init__(self, keyword_rules: List[Dict[str, Any]] | None = None) -> None:
        """初始化。keyword_rules 来自 PolicyStore.current().keyword_rules。"""
        self._rules = keyword_rules or []
        # 预编译正则
        self._compiled: List[Dict[str, Any]] = []
        for rule in self._rules:
            pattern = rule.get("pattern", "")
            if rule.get("match") == "regex" and pattern:
                try:
                    compiled = re.compile(pattern, re.IGNORECASE if rule.get("ignore_case") else 0)
                except re.error:
                    logger.warning("词表规则正则编译失败: %s", rule.get("name"))
                    continue
                self._compiled.append({**rule, "_re": compiled})
            else:
                self._compiled.append(rule)

    # ------------------------------------------------------------------- 接口实现

    def is_enabled(self) -> bool:
        return bool(self._rules)

    def scan(self, target: ScanTarget, ctx: Any) -> EngineResult:
        result = EngineResult(engine=self.name)
        if not self._rules:
            return result

        findings: List[Finding] = []

        # ---- 跳过二进制/压缩封装文件 ----
        # 这些文件内容是压缩流或纯二进制，短关键词 regex 在里面必然误报
        ext = (target.declared_ext or "").lower()
        if not ext and target.path:
            ext = os.path.splitext(target.path)[1].lower()
        if ext in self.BINARY_SKIP_EXTS:
            return result

        try:
            text = self._extract_text(target.path, target.declared_ext)
        except Exception:  # noqa: BLE001
            return result

        if not text:
            return result

        # ---- 预处理：合并换行符以支持跨行关键词 ----
        # 攻击者会把 Invoke-WebRequest 这种关键字用 CRLF 拆成
        # "Invoke-Web\r\nRequest" 来绕过基于单行的匹配。
        # 这里把所有 \r\n / \r / \n 替换为空，让跨行关键词重新拼接。
        # 注意：不影响纯文本/多行语义，我们匹配的都是单行关键词或跨行命令。
        text = text.replace("\r\n", "").replace("\r", "").replace("\n", "")

        for rule in self._compiled:
            match_mode = rule.get("match", "contains")
            rule_name = rule.get("name", "keyword")
            weight = rule.get("weight", 30.0)
            severity_str = rule.get("severity", "medium")
            severity = getattr(Severity, severity_str.upper(), Severity.MEDIUM).value

            hits: List[str] = []

            if match_mode == "regex":
                compiled = rule.get("_re")
                if compiled:
                    for m in compiled.finditer(text):
                        snippet = text[max(0, m.start() - 10): m.end() + 10]
                        hits.append(snippet.strip())
                        if len(hits) >= 3:
                            break

            elif match_mode == "any_of":
                keywords = rule.get("keywords", [])
                for kw in keywords:
                    if kw.lower() in text.lower():
                        hits.append(kw)
                        if len(hits) >= 3:
                            break

            elif match_mode == "contains":
                pattern = rule.get("pattern", "")
                if pattern and (pattern.lower() in text.lower()):
                    hits.append(pattern)

            if hits:
                findings.append(
                    Finding(
                        engine=self.name,
                        layer=self.layer,
                        severity=severity,
                        rule=rule_name,
                        detail=f"词表命中：{rule_name}（模式={match_mode}），命中 {len(hits)} 处：{hits[:2]}",
                        weight=float(weight),
                    )
                )

        result.findings = findings
        result.meta["scanned_chars"] = len(text)
        return result

    # ------------------------------------------------------------------- 内部辅助

    @staticmethod
    def _extract_text(path: str, ext: str) -> str:
        """从文件中提取可匹配的文本（简化版，够用就行）。

        OOXML/PDF/DOCX 等结构化文档优先交给对应的引擎处理，
        词表引擎只扫纯文本、代码文件和可直接读的二进制残留字符串。
        """
        text_exts = {".txt", ".csv", ".json", ".xml", ".html", ".htm",
                     ".md", ".log", ".ini", ".cfg", ".yaml", ".yml",
                     ".sh", ".bash", ".zsh", ".py", ".pl", ".rb", ".js", ".ts",
                     ".bat", ".cmd", ".ps1", ".vbs", ".sql", ".go", ".java",
                     ".c", ".cpp", ".h", ".cs"}

        ext_lower = (ext or "").lower()
        if ext_lower in text_exts:
            try:
                with open(path, "r", encoding="utf-8", errors="replace") as fh:
                    return fh.read(256 * 1024)
            except OSError:
                return ""

        # 二进制文件：提取可打印字符串（≥4 字符）
        try:
            with open(path, "rb") as fh:
                data = fh.read(512 * 1024)
        except OSError:
            return ""

        strings: List[str] = []
        current = bytearray()
        for byte in data:
            if 32 <= byte <= 126 or byte in (9,):  # 可打印或 tab
                current.append(byte)
            else:
                if len(current) >= 4:
                    strings.append(current.decode("ascii", errors="replace"))
                current = bytearray()
        if len(current) >= 4:
            strings.append(current.decode("ascii", errors="replace"))

        return "\n".join(strings)