MeshSate/app/services/scanner.py

"""扫描编排服务。

职责：
- 组装各检测引擎，执行流水线
- 递归解包并对子文件重复检测
- 汇总信号，交给策略引擎裁决
- 写入数据库与审计日志
- 管理临时文件生命周期
"""

from __future__ import annotations

import os
import shutil
import tempfile
import time
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional

from app.config import get_settings
from app.engines.base import Engine, EngineResult, Finding, ScanTarget, Severity
from app.engines.binary import BinaryEngine, YaraEngine
from app.engines.clamav import ClamAVEngine
from app.engines.documents import DocumentEngine
from app.engines.extractor import SafeExtractor, cleanup_dir
from app.engines.keywords import KeywordEngine
from app.engines.structure import StructureEngine
from app.services import audit
from app.services.policy import PolicyEngine, PolicyStore, Decision

_settings = get_settings()

TEMP_ROOT = os.environ.get("MESHSAFE_TEMP_ROOT", "/tmp/meshsafe")


@dataclass
class ScanContext:
    """一次扫描任务的上下文。"""

    settings: Any
    policy_engine: PolicyEngine
    workdir: str
    task_id: str
    channel: str = "default"
    client_id: str = ""
    client_ip: str = ""


@dataclass
class ScanOutcome:
    """扫描结果汇总。"""

    task_id: str
    verdict: str
    action: str
    score: float
    reason: str
    findings: List[Finding] = field(default_factory=list)
    engine_results: List[EngineResult] = field(default_factory=list)
    extracted: List[Dict[str, Any]] = field(default_factory=list)
    file_info: Dict[str, Any] = field(default_factory=dict)
    meta: Dict[str, Any] = field(default_factory=dict)
    duration_ms: int = 0
    triggered: List[Dict[str, Any]] = field(default_factory=list)
    encrypted: bool = False
    open_risk: Dict[str, Any] = field(default_factory=lambda: {
        "level": "low", "factors": [], "advice": "", "requires_macro_enable": False,
    })

    # ---- 调度/审计元信息（由调用方填充，写入 DB 用） ----
    channel: str = "default"
    client_id: str = ""
    client_ip: str = ""
    mode: str = "sync"  # sync | async

    # ------------------------------------------------------------------ DB 写入

    def persist_to_db(self) -> None:
        """将扫描结果写入 scan_tasks + scan_findings 表。

        DB 写入失败不影响扫描结论返回（fail-open），只打 WARNING 日志。
        """
        import logging
        from datetime import datetime, timezone

        logger = logging.getLogger(__name__)
        try:
            from app.db import SessionLocal
            from app.models import ScanFinding, ScanTask

            db = SessionLocal()
            try:
                task = ScanTask(
                    id=self.task_id,
                    channel=self.channel,
                    client_id=self.client_id,
                    filename=self.file_info.get("filename", ""),
                    declared_ext=self.file_info.get("declared_ext", ""),
                    size=self.file_info.get("size", 0),
                    sha256=self.file_info.get("sha256", ""),
                    md5=self.file_info.get("md5", ""),
                    real_type=self.file_info.get("real_type", ""),
                    status="completed",
                    mode=self.mode,
                    verdict=self.verdict,
                    action=self.action,
                    score=self.score,
                    reason=self.reason,
                    policy_version=self.meta.get("policy_version", ""),
                    detail=self.meta,
                    finished_at=datetime.now(timezone.utc),
                    duration_ms=self.duration_ms,
                )
                db.add(task)

                for f in self.findings:
                    db.add(ScanFinding(
                        task_id=self.task_id,
                        engine=f.engine,
                        layer=f.layer,
                        severity=f.severity,
                        rule=f.rule,
                        detail=f.detail,
                        extracted_path=f.extracted_path,
                        weight=f.weight,
                    ))

                db.commit()
                logger.info(
                    "扫描结果已持久化 task_id=%s findings=%d",
                    self.task_id, len(self.findings),
                )

                # --- 幂等记录写入（upsert，冲突忽略）---
                try:
                    from app.models import IdempotencyRecord
                    sha256 = self.file_info.get("sha256", "")
                    pv = self.meta.get("policy_version", "")
                    if sha256 and pv:
                        existing = db.query(IdempotencyRecord).filter_by(
                            sha256=sha256, policy_version=pv,
                        ).first()
                        if existing is None:
                            db.add(IdempotencyRecord(
                                sha256=sha256,
                                policy_version=pv,
                                task_id=self.task_id,
                            ))
                            db.commit()
                except Exception:  # noqa: BLE001
                    db.rollback()  # 幂等写入失败不影响主记录
            finally:
                db.close()
        except Exception as exc:  # noqa: BLE001
            logger.warning(
                "扫描结果写入 DB 失败（不影响扫描结论）task_id=%s err=%s",
                self.task_id, exc,
            )

    def to_response(self, include_details: bool = True) -> Dict[str, Any]:
        payload: Dict[str, Any] = {
            "task_id": self.task_id,
            "verdict": self.verdict,
            "score": self.score,
            "action": self.action,
            "reason": self.reason,
            "status": "completed",
            "file": self.file_info,
            "meta": self.meta,
            "open_risk": self.open_risk,
        }
        if include_details:
            payload["findings"] = [f.to_dict() for f in self.findings]
            payload["extracted"] = self.extracted
            payload["triggered"] = self.triggered
        return payload


class Scanner:
    """扫描流水线执行器。"""

    def __init__(self, policy_store: PolicyStore) -> None:
        self.policy_store = policy_store
        self.structure_engine = StructureEngine()
        self.document_engine = DocumentEngine()
        self.binary_engine = BinaryEngine()
        self.yara_engine = YaraEngine()  # 自动探测 rules_dir
        # 词表引擎：从当前策略加载 keyword_rules
        self.keyword_engine = KeywordEngine(keyword_rules=policy_store.current().keyword_rules)
        # ClamAVEngine 延迟初始化：首次 scan 时才尝试连接，避免无 ClamAV 环境下启动崩溃
        self._clamav_engine: Optional[ClamAVEngine] = None

    @property
    def clamav_engine(self) -> ClamAVEngine:
        if self._clamav_engine is None:
            self._clamav_engine = ClamAVEngine(
                host=_settings.clamav_host,
                port=_settings.clamav_port,
                timeout=_settings.clamav_timeout,
            )
        return self._clamav_engine

    def engines(self) -> List[Engine]:
        # KeywordEngine 每次都从当前策略重建（支持热重载后立即生效）
        self.keyword_engine = KeywordEngine(
            keyword_rules=self.policy_store.current().keyword_rules
        )
        return [
            self.structure_engine,
            self.clamav_engine,
            self.document_engine,
            self.keyword_engine,
            self.yara_engine,
            self.binary_engine,
        ]

    def _build_idempotent_outcome(
        self,
        idem_row,
        finding_rows,
        *,
        task_id: str,
        channel: str,
        client_id: str,
        filename: str,
    ) -> ScanOutcome:
        """命中幂等记录时，从 DB 组装一个 ScanOutcome 返回。"""
        from app.engines.base import Finding, Severity

        findings = [
            Finding(
                engine=f.engine,
                layer=f.layer,
                severity=Severity(f.severity),
                rule=f.rule,
                detail=f.detail or "",
                weight=f.weight or 50.0,
                extracted_path=f.extracted_path,
                meta={},
            )
            for f in finding_rows
        ]
        return ScanOutcome(
            task_id=task_id,
            channel=channel,
            client_id=client_id,
            filename=filename,
            verdict="known",
            action="pass",  # 幂等返回 pass，调用方按上次结论继续
            message=f"幂等命中（sha256={idem_row.sha256[:12]}... policy={idem_row.policy_version}）",
            findings=findings,
            score=0.0,
            reason="",
            meta={
                "idempotent": True,
                "original_task_id": idem_row.task_id,
                "policy_version": idem_row.policy_version,
            },
        )

    def scan_file(
        self,
        file_path: str,
        original_filename: str,
        *,
        task_id: str,
        channel: str = "default",
        client_id: str = "",
        client_ip: str = "",
        mode: str = "sync",
    ) -> ScanOutcome:
        """执行一次完整扫描。"""
        import logging

        logger = logging.getLogger(__name__)
        started = time.perf_counter()
        policy = self.policy_store.current()
        policy_engine = PolicyEngine(policy)

        logger.info(
            "开始扫描 task_id=%s filename=%s channel=%s client_id=%s",
            task_id, original_filename, channel, client_id,
        )

        # ---------- 幂等检查：同 SHA256 + 策略版本 → 直接返回上次结果 ----------
        try:
            import hashlib
            with open(file_path, "rb") as fh:
                _h = hashlib.sha256()
                while True:
                    _c = fh.read(65536)
                    if not _c:
                        break
                    _h.update(_c)
            root_sha256 = _h.hexdigest()
            pv = policy.version
            from app.db import SessionLocal
            from app.models import IdempotencyRecord
            _db = SessionLocal()
            try:
                _idem = _db.query(IdempotencyRecord).filter_by(
                    sha256=root_sha256, policy_version=pv,
                ).first()
                if _idem is not None:
                    # 命中幂等 → 从 DB 加载上次 ScanTask + ScanFindings 组装结果
                    from app.models import ScanFinding as _SF
                    _findings_rows = _db.query(_SF).filter(_SF.task_id == _idem.task_id).all()
                    logger.info(
                        "幂等命中：sha256=%s policy_version=%s → 复用 task_id=%s",
                        root_sha256[:12] + "...", pv, _idem.task_id,
                    )
                    # 用原始 task_id（不是旧的），但内容复用
                    return self._build_idempotent_outcome(
                        _idem, _findings_rows, task_id=task_id, channel=channel,
                        client_id=client_id, filename=original_filename,
                    )
            finally:
                _db.close()
        except Exception as exc:  # noqa: BLE001
            logger.debug("幂等检查跳过: %s", exc)

        os.makedirs(TEMP_ROOT, exist_ok=True)
        workdir = tempfile.mkdtemp(prefix="scan_", dir=TEMP_ROOT)

        ctx = ScanContext(
            settings=_settings,
            policy_engine=policy_engine,
            workdir=workdir,
            task_id=task_id,
            channel=channel,
            client_id=client_id,
            client_ip=client_ip,
        )

        all_findings: List[Finding] = []
        engine_results: List[EngineResult] = []
        engine_errors: List[str] = []
        extracted_meta: List[Dict[str, Any]] = []
        encrypted = False
        root_info: Dict[str, str] = {"real_type": "", "real_ext": ""}

        root_target = ScanTarget(
            path=file_path,
            filename=original_filename,
            declared_ext=os.path.splitext(original_filename)[1].lower(),
            logical_path=original_filename,
            depth=0,
        )

        try:
            queue: List[ScanTarget] = [root_target]
            processed = 0

            while queue:
                target = queue.pop(0)
                processed += 1
                if processed > 2000:
                    break

                for engine in self.engines():
                    result = engine.safe_scan(target, ctx)
                    engine_results.append(result)
                    if not result.ok and result.error:
                        engine_errors.append(f"{engine.name}: {result.error}")
                    all_findings.extend(result.findings)

                    if engine.name == "structure" and result.meta.get("real_type"):
                        if target.depth == 0:
                            root_info = {
                                "real_type": result.meta.get("real_type", ""),
                                "real_ext": result.meta.get("real_ext", ""),
                            }

                    if engine.name == "document":
                        if any(f.rule in ("PdfEncrypted", "OfficeEncrypted") for f in result.findings):
                            encrypted = True

                # 解包并递归检测子文件
                extractor = SafeExtractor(ctx)
                report = extractor.extract(target, workdir)
                all_findings.extend(report.findings)

                if report.files:
                    for child in report.files:
                        queue.append(child)
                        extracted_meta.append(
                            {
                                "path": child.logical_path,
                                "filename": child.filename,
                                "size": os.path.getsize(child.path) if os.path.exists(child.path) else 0,
                                "depth": child.depth,
                                "sha256": "",
                            }
                        )
        finally:
            cleanup_dir(workdir)

        # --- 哈希计算（根文件 + 所有子文件）---
        def _sha256(path: str) -> str:
            import hashlib
            h = hashlib.sha256()
            with open(path, "rb") as fh:
                while True:
                    chunk = fh.read(65536)
                    if not chunk:
                        break
                    h.update(chunk)
            return h.hexdigest()

        def _md5(path: str) -> str:
            import hashlib
            h = hashlib.md5()
            with open(path, "rb") as fh:
                while True:
                    chunk = fh.read(65536)
                    if not chunk:
                        break
                    h.update(chunk)
            return h.hexdigest()

        root_sha256 = ""
        root_md5 = ""
        if os.path.exists(file_path):
            root_sha256 = _sha256(file_path)
            root_md5 = _md5(file_path)

        # 补齐提取文件的哈希
        for item in extracted_meta:
            child_path = item.get("path", "")
            if os.path.exists(child_path) and not item.get("sha256"):
                try:
                    item["sha256"] = _sha256(child_path)
                except OSError:
                    pass

        # 为每个 finding 附加根文件 SHA256（供宏例外指纹比对用）
        if root_sha256:
            for f in all_findings:
                if f.meta is None:
                    f.meta = {}
                f.meta["file_sha256"] = root_sha256

        size = os.path.getsize(file_path) if os.path.exists(file_path) else 0
        real_type = root_info.get("real_type", "")
        real_ext = root_info.get("real_ext", "")

        decision = policy_engine.decide(
            all_findings,
            filename=original_filename,
            declared_ext=root_target.declared_ext,
            size=size,
            real_type=real_type,
            engine_errors=engine_errors,
            encrypted=encrypted,
            md5=root_info.get("md5", ""),
        )

        # 计算 open_risk（独立于裁决结果，供人阅读）
        open_risk = PolicyEngine.compute_open_risk(
            decision.verdict,
            all_findings,
            declared_ext=root_target.declared_ext,
            real_type=real_type,
        )

        duration_ms = int((time.perf_counter() - started) * 1000)

        logger.info(
            "扫描完成 task_id=%s verdict=%s action=%s score=%.1f open_risk.level=%s "
            "findings=%d triggered=%d duration=%dms",
            task_id, decision.verdict, decision.action, decision.score,
            open_risk["level"], len(all_findings), len(decision.triggered), duration_ms,
        )
        if decision.triggered:
            logger.warning(
                "task_id=%s 触发规则: %s",
                task_id,
                ", ".join(f"{t.get('type', 'unknown')}:{t.get('rule', '?')}" for t in decision.triggered),
            )
        if engine_errors:
            logger.warning("task_id=%s 引擎错误: %s", task_id, "; ".join(engine_errors))

        clamav_meta = next(
            (r.meta for r in engine_results if r.engine == "clamav"), {}
        )

        outcome = ScanOutcome(
            task_id=task_id,
            verdict=decision.verdict,
            action=decision.action,
            score=decision.score,
            reason=decision.reason,
            findings=all_findings,
            engine_results=engine_results,
            extracted=extracted_meta,
            file_info={
                "sha256": root_sha256,
                "md5": root_md5,
                "size": size,
                "real_type": real_type,
                "real_ext": real_ext,
                "declared_ext": root_target.declared_ext,
                "filename": original_filename,
            },
            meta={
                "scan_duration_ms": duration_ms,
                "engines": [e.name for e in self.engines()],
                "engine_errors": engine_errors,
                "clamav": clamav_meta,
                "policy_version": policy.version,
                "policy_sha256": policy.sha256,
                "scanned_objects": processed,
                "scanned_at": datetime.now(timezone.utc).isoformat(),
            },
            duration_ms=duration_ms,
            triggered=decision.triggered,
            encrypted=encrypted,
            open_risk=open_risk,
            channel=channel,
            client_id=client_id,
            client_ip=client_ip,
            mode=mode,
        )

        # 自动持久化到 DB（fail-open，不影响扫描结论）
        outcome.persist_to_db()

        return outcome

    def health(self) -> Dict[str, Any]:
        """各引擎健康状态。"""
        clamav = self.clamav_engine.health()
        yara = self.yara_engine.health()
        return {
            "clamav": clamav,
            "yara": yara,
            "document": {"available": True},
            "structure": {"available": True},
            "binary": {"available": True},
        }