"""上传隔离、扫描与审计工具。""" from __future__ import annotations import hashlib import json import mimetypes import os import shutil import subprocess import time import uuid from dataclasses import dataclass, asdict from datetime import datetime from pathlib import Path from typing import Any, Dict, Optional, TYPE_CHECKING from config import ( MAX_UPLOAD_SIZE, UPLOAD_ALLOWED_EXTENSIONS, UPLOAD_CLAMAV_ARGS, UPLOAD_CLAMAV_BIN, UPLOAD_CLAMAV_ENABLED, UPLOAD_CLAMAV_TIMEOUT_SECONDS, UPLOAD_SCAN_LOG_SUBDIR, ) from utils.logger import setup_logger if TYPE_CHECKING: from modules.user_manager import UserWorkspace from werkzeug.datastructures import FileStorage class UploadSecurityError(Exception): """上传被拒绝或扫描失败时抛出的业务异常。""" def __init__(self, message: str, code: str = "upload_error"): super().__init__(message) self.code = code @dataclass class StageRecord: """暂存区中的上传文件快照。""" path: Path filename: str size: int sha256: str mime: Optional[str] @dataclass class ScanReport: """病毒扫描结果。""" status: str engine: str signature: Optional[str] = None message: Optional[str] = None duration_ms: Optional[int] = None @property def is_clean(self) -> bool: return self.status in {"passed", "skipped"} class UploadQuarantineManager: """负责将上传文件落地到隔离区、执行安全扫描并记录审计日志。""" def __init__(self, workspace: "UserWorkspace", *, logger_name: Optional[str] = None): self.workspace = workspace self.quarantine_dir = Path(workspace.quarantine_dir).resolve() self.quarantine_dir.mkdir(parents=True, exist_ok=True) log_rel = Path(UPLOAD_SCAN_LOG_SUBDIR) / f"{workspace.username}.log" self.logger = setup_logger( logger_name or f"upload_guard.{workspace.username}", str(log_rel), ) self.allowed_extensions = tuple(UPLOAD_ALLOWED_EXTENSIONS) self.clamav_enabled = bool(UPLOAD_CLAMAV_ENABLED) self.clamav_bin = self._resolve_clamav_bin(UPLOAD_CLAMAV_BIN) self.clamav_args = list(UPLOAD_CLAMAV_ARGS) self.clamav_timeout = int(max(1, UPLOAD_CLAMAV_TIMEOUT_SECONDS)) def process_upload( self, file_obj: "FileStorage", target_path: Path, *, username: str, source: str, original_name: str, relative_path: str, ) -> Dict[str, Any]: """保存文件到隔离区、执行扫描并将通过的文件同步到工作目录。""" stage = self._stage_file(file_obj, original_name) metadata: Dict[str, Any] = { "upload_id": uuid.uuid4().hex, "username": username, "source": source, "original_name": original_name, "target_path": str(target_path), "target_relative": relative_path, "staged_path": str(stage.path), "size": stage.size, "sha256": stage.sha256, "mime": stage.mime, } scan_report: Optional[ScanReport] = None promoted_path: Optional[Path] = None try: self._enforce_size(stage.size) self._enforce_extension(target_path.name) scan_report = self._scan(stage.path) metadata["scan"] = asdict(scan_report) if not scan_report.is_clean: metadata["scan_failure_reason"] = scan_report.message or scan_report.signature raise UploadSecurityError("安全审核未通过", code="scan_failed") promoted_path = self._promote(stage.path, target_path) metadata["final_path"] = str(promoted_path) self._log_event(True, metadata) return { "final_path": promoted_path, "metadata": metadata, } except UploadSecurityError as exc: metadata["error"] = { "code": exc.code, "message": str(exc), } if scan_report: metadata.setdefault("scan", asdict(scan_report)) self._log_event(False, metadata) raise finally: if stage.path.exists(): try: stage.path.unlink() except OSError: pass # ------------------------------------------------------------------ # 内部实现 # ------------------------------------------------------------------ def _stage_file(self, file_obj: "FileStorage", original_name: str) -> StageRecord: suffix = "".join(Path(original_name or "").suffixes[-2:]) or Path(original_name or "").suffix unique_name = f"{int(time.time())}_{uuid.uuid4().hex}{suffix or ''}" staged_path = self.quarantine_dir / unique_name file_obj.save(staged_path) size = staged_path.stat().st_size sha256 = self._hash_file(staged_path) mime, _ = mimetypes.guess_type(original_name or "") return StageRecord(staged_path, original_name, size, sha256, mime) def _enforce_size(self, size: int): if size > MAX_UPLOAD_SIZE: raise UploadSecurityError( f"文件大小 {size} 超过上限 {MAX_UPLOAD_SIZE} 字节", code="size_exceeded", ) def _enforce_extension(self, filename: str): if not self.allowed_extensions: return lowered = (filename or "").lower() for pattern in self.allowed_extensions: if lowered.endswith(pattern): return raise UploadSecurityError("文件类型不在允许列表中", code="extension_forbidden") def _scan(self, staged_path: Path) -> ScanReport: if not self.clamav_enabled: return ScanReport(status="skipped", engine="clamdscan", message="已跳过病毒扫描") if not self.clamav_bin: raise UploadSecurityError("未找到 ClamAV 扫描器，请检查配置", code="scanner_missing") command = [self.clamav_bin] + self.clamav_args + [str(staged_path)] start = time.time() try: result = subprocess.run( command, capture_output=True, text=True, timeout=self.clamav_timeout, check=False, ) except subprocess.TimeoutExpired: return ScanReport( status="error", engine="clamdscan", message="扫描超时", duration_ms=int((time.time() - start) * 1000), ) except FileNotFoundError as exc: raise UploadSecurityError("ClamAV 扫描器不可用", code="scanner_missing") from exc duration_ms = int((time.time() - start) * 1000) stdout = (result.stdout or "").strip() stderr = (result.stderr or "").strip() if result.returncode == 0: return ScanReport(status="passed", engine="clamdscan", duration_ms=duration_ms) if result.returncode == 1: signature = self._extract_signature(stdout) or self._extract_signature(stderr) message = signature or "检测到可疑内容" return ScanReport( status="failed", engine="clamdscan", signature=signature, message=message, duration_ms=duration_ms, ) message = stderr or stdout or f"扫描失败（退出码 {result.returncode}）" return ScanReport( status="error", engine="clamdscan", message=message, duration_ms=duration_ms, ) def _promote(self, staged_path: Path, target_path: Path) -> Path: target_path.parent.mkdir(parents=True, exist_ok=True) shutil.move(str(staged_path), str(target_path)) return target_path def _hash_file(self, path: Path) -> str: sha256 = hashlib.sha256() with path.open("rb") as fh: for chunk in iter(lambda: fh.read(1024 * 1024), b""): sha256.update(chunk) return sha256.hexdigest() def _extract_signature(self, text: str) -> Optional[str]: if not text: return None for raw_line in text.splitlines(): line = raw_line.strip() if not line or "FOUND" not in line.upper(): continue parts = line.split(":", 1) tail = parts[-1].strip() if tail.upper().endswith("FOUND"): signature = tail[: -len("FOUND")].strip() if signature: return signature return None def _resolve_clamav_bin(self, candidate: str) -> Optional[str]: if not candidate: return None candidate = candidate.strip() if not candidate: return None if os.path.isabs(candidate): return candidate if Path(candidate).exists() else None resolved = shutil.which(candidate) return resolved or None def _log_event(self, accepted: bool, payload: Dict[str, Any]): entry = { "timestamp": datetime.utcnow().isoformat(), "accepted": accepted, **payload, } try: self.logger.info("UPLOAD_AUDIT %s", json.dumps(entry, ensure_ascii=False)) except Exception: self.logger.warning("无法写入上传审计日志：%s", entry) __all__ = [ "UploadQuarantineManager", "UploadSecurityError", ]