MeshSate/app/engines/extractor.py

"""递归安全解包引擎。

安全约束（全部强制执行，不依赖调用方自觉）：
- 解包深度上限
- 解压总大小上限与单文件大小上限
- 压缩比上限（防压缩炸弹）
- 路径穿越拦截（Zip Slip）
- 符号链接/硬链接逃逸拦截
- 解包目标固定在独立临时目录，文件名重写为安全形式（不直接使用归档内原始名）

解包出来的文件会加入待检测队列，实现"递归检测"，避免只扫外层被绕过。
"""

from __future__ import annotations

import os
import re
import shutil
import tarfile
import tempfile
import zipfile
from dataclasses import dataclass, field
from typing import Any, List, Optional

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

ARCHIVE_EXTS = {".zip", ".jar", ".war", ".apk", ".epub", ".docx", ".xlsx", ".pptx", ".odt", ".ods", ".odp"}
TAR_EXTS = {".tar", ".gz", ".tgz", ".bz2", ".xz", ".tbz2", ".txz"}

_SAFE_NAME_RE = re.compile(
    r"[^A-Za-z0-9._一-鿿"  # ASCII + CJK 统一汉字
    r"\U0001F300-\U0001FAFF"  # Emoji 表情符号
    r"\U0001F000-\U0001F2FF"  # 符号和象形文字扩展 A
    r"\U00002600-\U000027BF"  # 杂项符号 + 装饰符号
    r"\U00002B00-\U00002BFF"  # 箭头等
    r"-]+"
)


@dataclass
class ExtractionReport:
    """解包结果。"""

    files: List[ScanTarget] = field(default_factory=list)
    findings: List[Finding] = field(default_factory=list)
    total_bytes: int = 0
    truncated: bool = False

    def to_dict(self) -> dict:
        return {
            "file_count": len(self.files),
            "total_bytes": self.total_bytes,
            "truncated": self.truncated,
        }


class SafeExtractor:
    """递归安全解包器。"""

    def __init__(self, ctx: Any) -> None:
        self.ctx = ctx
        self.settings = ctx.settings
        self.engine_name = "extractor"

    def extract(self, target: ScanTarget, workdir: str) -> ExtractionReport:
        report = ExtractionReport()
        ext = (target.declared_ext or "").lower()

        if ext in ARCHIVE_EXTS or zipfile.is_zipfile(target.path):
            self._extract_zip(target, workdir, report)
        elif ext in TAR_EXTS or tarfile.is_tarfile(target.path):
            self._extract_tar(target, workdir, report)

        return report

    # ------------------------------------------------------------------ 工具

    def _sanitize_name(self, raw_name: str) -> str:
        """把归档内原始文件名改写为安全的相对路径。

        去掉目录层级与危险字符，避免原始名称中的 ``..``、绝对路径、
        Windows 保留设备名等造成问题。原始名称仍保留在逻辑路径中以便追溯。
        """
        name = raw_name.replace("\\", "/")
        base = os.path.basename(name)
        base = _SAFE_NAME_RE.sub("_", base).strip("._")
        if not base:
            base = "unnamed"
        # 防止 Windows 保留设备名
        stem, dot, suffix = base.partition(".")
        reserved = {
            "con", "prn", "aux", "nul",
            *(f"com{i}" for i in range(1, 10)),
            *(f"lpt{i}" for i in range(1, 10)),
        }
        if stem.lower() in reserved:
            stem = f"_{stem}"
        return f"{stem}{dot}{suffix}" if dot else stem

    def _target_dir(self, workdir: str, target: ScanTarget) -> str:
        """为每个文件分配独立的解包子目录。"""
        digest = abs(hash(target.logical_path or target.filename)) % 100000
        path = os.path.join(workdir, f"ex_{target.depth}_{digest}")
        os.makedirs(path, exist_ok=True)
        return path

    def _budget_exhausted(self, report: ExtractionReport) -> bool:
        return report.total_bytes >= self.settings.max_extract_total_bytes

    def _record_limit(self, report: ExtractionReport, limit: int = 5000) -> bool:
        if len(report.files) >= limit:
            report.truncated = True
            return True
        return False

    def _add_finding(
        self,
        report: ExtractionReport,
        severity: str,
        rule: str,
        detail: str,
        weight: float = 0.0,
    ) -> None:
        report.findings.append(
            Finding(
                engine=self.engine_name,
                layer="L0",
                severity=severity,
                rule=rule,
                detail=detail,
                weight=weight,
            )
        )

    # ------------------------------------------------------------------ ZIP

    def _extract_zip(self, target: ScanTarget, workdir: str, report: ExtractionReport) -> None:
        out_dir = self._target_dir(workdir, target)
        max_ratio = self.settings.max_extract_ratio
        max_total = self.settings.max_extract_total_bytes
        max_depth = self.settings.max_extract_depth

        try:
            with zipfile.ZipFile(target.path) as zf:
                infos = zf.infolist()
                for info in infos:
                    if target.depth + 1 > max_depth:
                        self._add_finding(
                            report,
                            Severity.MEDIUM.value,
                            "ExtractDepthExceeded",
                            f"嵌套层级超过上限 {max_depth}，深层内容未展开",
                            15.0,
                        )
                        report.truncated = True
                        break

                    if self._budget_exhausted(report) or self._record_limit(report):
                        self._add_finding(
                            report,
                            Severity.MEDIUM.value,
                            "ExtractBudgetExceeded",
                            "解包数量或体积达到上限，后续内容未展开",
                            10.0,
                        )
                        break

                    if info.is_dir():
                        continue

                    name = info.filename
                    if _is_traversal(name):
                        self._add_finding(
                            report,
                            Severity.CRITICAL.value,
                            "PathTraversalBlocked",
                            f"已拦截路径穿越条目：{name}",
                            50.0,
                        )
                        continue

                    # 压缩炸弹判定：按声明大小与实际膨胀比
                    if info.compress_size > 0:
                        ratio = info.file_size / info.compress_size
                        if ratio > max_ratio:
                            self._add_finding(
                                report,
                                Severity.HIGH.value,
                                "ArchiveBombBlocked",
                                f"成员 {name} 压缩比 {ratio:.1f}:1 超过阈值，已跳过",
                                35.0,
                            )
                            continue

                    if info.file_size > self.settings.max_upload_bytes:
                        self._add_finding(
                            report,
                            Severity.MEDIUM.value,
                            "MemberTooLarge",
                            f"成员 {name} 大小超过单文件上限，已跳过",
                            10.0,
                        )
                        continue

                    safe_name = self._sanitize_name(name)
                    dest = os.path.join(out_dir, safe_name)
                    try:
                        written = 0
                        overflow = False
                        with zf.open(info) as src, open(dest, "wb") as out:
                            while True:
                                chunk = src.read(1024 * 1024)
                                if not chunk:
                                    break
                                written += len(chunk)
                                if written > self.settings.max_upload_bytes:
                                    out.close()
                                    os.remove(dest)
                                    self._add_finding(
                                        report,
                                        Severity.MEDIUM.value,
                                        "MemberTooLarge",
                                        f"成员 {name} 实际解压超限，已丢弃",
                                        10.0,
                                    )
                                    overflow = True
                                    break
                                out.write(chunk)
                        # 成员必须回到扫描队列，否则压缩包/OOXML 里的内容永远不参与检测
                        if not overflow:
                            report.total_bytes += written
                            report.files.append(
                                ScanTarget(
                                    path=dest,
                                    filename=safe_name,
                                    declared_ext=os.path.splitext(safe_name)[1].lower(),
                                    logical_path=f"{target.logical_path}!/{name}",
                                    depth=target.depth + 1,
                                    is_extracted=True,
                                )
                            )
                    except Exception as exc:  # noqa: BLE001
                        self._add_finding(
                            report,
                            Severity.LOW.value,
                            "ExtractMemberFailed",
                            f"成员 {name} 解压失败：{exc}",
                            0.0,
                        )
        except zipfile.BadZipFile as exc:
            self._add_finding(
                report,
                Severity.MEDIUM.value,
                "ArchiveBroken",
                f"ZIP 结构损坏：{exc}",
                10.0,
            )
        except Exception as exc:  # noqa: BLE001
            self._add_finding(
                report,
                Severity.LOW.value,
                "ExtractFailed",
                f"解包失败：{exc}",
                0.0,
            )

    # ------------------------------------------------------------------ TAR

    def _extract_tar(self, target: ScanTarget, workdir: str, report: ExtractionReport) -> None:
        out_dir = self._target_dir(workdir, target)
        max_depth = self.settings.max_extract_depth

        try:
            with tarfile.open(target.path) as tf:
                for member in tf.getmembers():
                    if target.depth + 1 > max_depth:
                        self._add_finding(
                            report,
                            Severity.MEDIUM.value,
                            "ExtractDepthExceeded",
                            f"嵌套层级超过上限 {max_depth}",
                            15.0,
                        )
                        report.truncated = True
                        break

                    if self._budget_exhausted(report) or self._record_limit(report):
                        self._add_finding(
                            report,
                            Severity.MEDIUM.value,
                            "ExtractBudgetExceeded",
                            "解包数量或体积达到上限",
                            10.0,
                        )
                        break

                    name = member.name
                    if _is_traversal(name) or member.issym() or member.islnk():
                        self._add_finding(
                            report,
                            Severity.CRITICAL.value,
                            "PathTraversalBlocked",
                            f"已拦截路径穿越或链接条目：{name}",
                            50.0,
                        )
                        continue

                    if not member.isfile():
                        continue

                    if member.size > self.settings.max_upload_bytes:
                        self._add_finding(
                            report,
                            Severity.MEDIUM.value,
                            "MemberTooLarge",
                            f"成员 {name} 大小超过单文件上限，已跳过",
                            10.0,
                        )
                        continue

                    safe_name = self._sanitize_name(name)
                    dest = os.path.join(out_dir, safe_name)
                    try:
                        src = tf.extractfile(member)
                        if src is None:
                            continue
                        with src, open(dest, "wb") as out:
                            written = 0
                            overflow = False
                            while True:
                                chunk = src.read(1024 * 1024)
                                if not chunk:
                                    break
                                written += len(chunk)
                                if written > self.settings.max_upload_bytes:
                                    out.close()
                                    os.remove(dest)
                                    self._add_finding(
                                        report,
                                        Severity.MEDIUM.value,
                                        "MemberTooLarge",
                                        f"成员 {name} 实际解压超限，已丢弃",
                                        10.0,
                                    )
                                    overflow = True
                                    break
                                out.write(chunk)
                        if not overflow:
                            report.total_bytes += written
                            report.files.append(
                                ScanTarget(
                                    path=dest,
                                    filename=safe_name,
                                    declared_ext=os.path.splitext(safe_name)[1].lower(),
                                    logical_path=f"{target.logical_path}!/{name}",
                                    depth=target.depth + 1,
                                    is_extracted=True,
                                )
                            )
                    except Exception as exc:  # noqa: BLE001
                        self._add_finding(
                            report,
                            Severity.LOW.value,
                            "ExtractMemberFailed",
                            f"成员 {name} 解压失败：{exc}",
                            0.0,
                        )
        except tarfile.TarError as exc:
            self._add_finding(
                report,
                Severity.MEDIUM.value,
                "ArchiveBroken",
                f"TAR 结构异常：{exc}",
                10.0,
            )
        except Exception as exc:  # noqa: BLE001
            self._add_finding(
                report,
                Severity.LOW.value,
                "ExtractFailed",
                f"解包失败：{exc}",
                0.0,
            )


def _is_traversal(name: str) -> bool:
    normalized = name.replace("\\", "/")
    if normalized.startswith("/"):
        return True
    if len(normalized) > 1 and normalized[1] == ":":
        return True
    return ".." in normalized.split("/")


def cleanup_dir(path: str) -> None:
    """安全删除临时目录。"""
    try:
        shutil.rmtree(path, ignore_errors=True)
    except Exception:  # noqa: BLE001
        pass


def make_workdir(root: str, prefix: str = "scan_") -> str:
    os.makedirs(root, exist_ok=True)
    return tempfile.mkdtemp(prefix=prefix, dir=root)