|
| 1 | +"""文件类型校验：扩展名 / MIME 白名单 + magic bytes 防伪造。""" |
| 2 | + |
| 3 | +from __future__ import annotations |
| 4 | + |
| 5 | +from dataclasses import dataclass |
| 6 | +from pathlib import Path |
| 7 | +from typing import Optional |
| 8 | + |
| 9 | +from fastapi import HTTPException, UploadFile |
| 10 | +from fnmatch import fnmatchcase |
| 11 | + |
| 12 | +from core.settings import settings |
| 13 | + |
| 14 | + |
| 15 | +@dataclass(frozen=True) |
| 16 | +class FileKind: |
| 17 | + name: str |
| 18 | + extensions: tuple[str, ...] |
| 19 | + mimes: tuple[str, ...] |
| 20 | + signatures: tuple[bytes, ...] |
| 21 | + |
| 22 | + |
| 23 | +FILE_KINDS: tuple[FileKind, ...] = ( |
| 24 | + FileKind("png", (".png",), ("image/png",), (b"\x89PNG\r\n\x1a\n",)), |
| 25 | + FileKind("jpg", (".jpg", ".jpeg"), ("image/jpeg",), (b"\xff\xd8\xff",)), |
| 26 | + FileKind("gif", (".gif",), ("image/gif",), (b"GIF87a", b"GIF89a")), |
| 27 | + FileKind("webp", (".webp",), ("image/webp",), ()), |
| 28 | + FileKind("bmp", (".bmp",), ("image/bmp", "image/x-ms-bmp"), (b"BM",)), |
| 29 | + FileKind("pdf", (".pdf",), ("application/pdf",), (b"%PDF",)), |
| 30 | + FileKind( |
| 31 | + "zip", |
| 32 | + (".zip", ".docx", ".xlsx", ".pptx", ".apk", ".jar"), |
| 33 | + ( |
| 34 | + "application/zip", |
| 35 | + "application/x-zip-compressed", |
| 36 | + "application/vnd.openxmlformats-officedocument.wordprocessingml.document", |
| 37 | + "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", |
| 38 | + "application/vnd.openxmlformats-officedocument.presentationml.presentation", |
| 39 | + "application/java-archive", |
| 40 | + "application/vnd.android.package-archive", |
| 41 | + ), |
| 42 | + (b"PK\x03\x04", b"PK\x05\x06", b"PK\x07\x08"), |
| 43 | + ), |
| 44 | + FileKind( |
| 45 | + "rar", |
| 46 | + (".rar",), |
| 47 | + ("application/x-rar-compressed", "application/vnd.rar"), |
| 48 | + (b"Rar!\x1a\x07\x00", b"Rar!\x1a\x07\x01\x00"), |
| 49 | + ), |
| 50 | + FileKind("7z", (".7z",), ("application/x-7z-compressed",), (b"7z\xbc\xaf'\x1c",)), |
| 51 | + FileKind("gz", (".gz", ".tgz"), ("application/gzip", "application/x-gzip"), (b"\x1f\x8b",)), |
| 52 | + FileKind( |
| 53 | + "mp3", |
| 54 | + (".mp3",), |
| 55 | + ("audio/mpeg",), |
| 56 | + (b"ID3", b"\xff\xfb", b"\xff\xf3", b"\xff\xf2"), |
| 57 | + ), |
| 58 | + FileKind("mp4", (".mp4", ".m4a", ".mov"), ("video/mp4", "audio/mp4", "video/quicktime"), ()), |
| 59 | + FileKind( |
| 60 | + "exe", |
| 61 | + (".exe", ".dll", ".sys"), |
| 62 | + ("application/x-msdownload", "application/x-dosexec"), |
| 63 | + (b"MZ",), |
| 64 | + ), |
| 65 | + FileKind("elf", (".elf", ".so", ".o"), ("application/x-executable",), (b"\x7fELF",)), |
| 66 | +) |
| 67 | + |
| 68 | +KNOWN_EXTENSIONS = {ext for kind in FILE_KINDS for ext in kind.extensions} |
| 69 | + |
| 70 | + |
| 71 | +def normalize_allowed_file_types() -> list[str]: |
| 72 | + raw_value = settings.allowed_file_types |
| 73 | + if isinstance(raw_value, str): |
| 74 | + values = [item.strip() for item in raw_value.split(",")] |
| 75 | + elif isinstance(raw_value, (list, tuple, set)): |
| 76 | + values = [str(item).strip() for item in raw_value] |
| 77 | + else: |
| 78 | + values = [] |
| 79 | + normalized = [item.lower() for item in values if item] |
| 80 | + return normalized or ["*"] |
| 81 | + |
| 82 | + |
| 83 | +def _extension_of(file_name: str) -> str: |
| 84 | + return Path(file_name or "").suffix.lower() |
| 85 | + |
| 86 | + |
| 87 | +def _match_allow_rule(rule: str, file_name: str, content_type: str) -> bool: |
| 88 | + if rule in {"*", "*/*"}: |
| 89 | + return True |
| 90 | + if "/" in rule: |
| 91 | + return fnmatchcase(content_type, rule) |
| 92 | + extension = rule if rule.startswith(".") else f".{rule}" |
| 93 | + return file_name.endswith(extension) |
| 94 | + |
| 95 | + |
| 96 | +def is_type_allowed(file_name: str, content_type: Optional[str] = None) -> bool: |
| 97 | + allowed = normalize_allowed_file_types() |
| 98 | + if any(rule in {"*", "*/*"} for rule in allowed): |
| 99 | + return True |
| 100 | + normalized_name = (file_name or "").strip().lower() |
| 101 | + normalized_content_type = (content_type or "").strip().lower() |
| 102 | + return any( |
| 103 | + _match_allow_rule(rule, normalized_name, normalized_content_type) |
| 104 | + for rule in allowed |
| 105 | + ) |
| 106 | + |
| 107 | + |
| 108 | +def detect_file_kind(header: bytes) -> Optional[FileKind]: |
| 109 | + if not header: |
| 110 | + return None |
| 111 | + |
| 112 | + if len(header) >= 12 and header.startswith(b"RIFF") and header[8:12] == b"WEBP": |
| 113 | + return next(kind for kind in FILE_KINDS if kind.name == "webp") |
| 114 | + |
| 115 | + if len(header) >= 12 and header[4:8] == b"ftyp": |
| 116 | + return next(kind for kind in FILE_KINDS if kind.name == "mp4") |
| 117 | + |
| 118 | + candidates: list[tuple[int, FileKind]] = [] |
| 119 | + for kind in FILE_KINDS: |
| 120 | + for signature in kind.signatures: |
| 121 | + if signature and header.startswith(signature): |
| 122 | + candidates.append((len(signature), kind)) |
| 123 | + if not candidates: |
| 124 | + return None |
| 125 | + candidates.sort(key=lambda item: item[0], reverse=True) |
| 126 | + return candidates[0][1] |
| 127 | + |
| 128 | + |
| 129 | +def validate_file_type(file_name: str, content_type: Optional[str] = None) -> None: |
| 130 | + if not is_type_allowed(file_name, content_type): |
| 131 | + raise HTTPException(status_code=403, detail="不允许上传该类型文件") |
| 132 | + |
| 133 | + |
| 134 | +def _kind_matches_ext(kind: FileKind, extension: str) -> bool: |
| 135 | + return extension in kind.extensions |
| 136 | + |
| 137 | + |
| 138 | +def _kind_matches_mime(kind: FileKind, mime: str) -> bool: |
| 139 | + return mime in kind.mimes |
| 140 | + |
| 141 | + |
| 142 | +def validate_file_magic( |
| 143 | + file_name: str, |
| 144 | + content_type: Optional[str], |
| 145 | + header: Optional[bytes], |
| 146 | +) -> None: |
| 147 | + """白名单通过后，用 magic bytes 防止扩展名 / MIME 伪造。""" |
| 148 | + validate_file_type(file_name, content_type) |
| 149 | + if not header: |
| 150 | + return |
| 151 | + |
| 152 | + claimed_ext = _extension_of(file_name) |
| 153 | + claimed_mime = (content_type or "").strip().lower() |
| 154 | + detected = detect_file_kind(header) |
| 155 | + |
| 156 | + if claimed_ext in KNOWN_EXTENSIONS: |
| 157 | + if detected is None or not _kind_matches_ext(detected, claimed_ext): |
| 158 | + raise HTTPException( |
| 159 | + status_code=403, |
| 160 | + detail="文件内容与扩展名不匹配，疑似伪造类型", |
| 161 | + ) |
| 162 | + |
| 163 | + if claimed_mime and any(claimed_mime in kind.mimes for kind in FILE_KINDS): |
| 164 | + if detected is None or not _kind_matches_mime(detected, claimed_mime): |
| 165 | + raise HTTPException( |
| 166 | + status_code=403, |
| 167 | + detail="文件内容与 Content-Type 不匹配，疑似伪造类型", |
| 168 | + ) |
| 169 | + |
| 170 | + allowed = normalize_allowed_file_types() |
| 171 | + if detected and not any(rule in {"*", "*/*"} for rule in allowed): |
| 172 | + allowed_by_detected = any( |
| 173 | + is_type_allowed(f"file{ext}", mime) |
| 174 | + for ext in detected.extensions |
| 175 | + for mime in (detected.mimes or (None,)) |
| 176 | + ) |
| 177 | + if not allowed_by_detected: |
| 178 | + raise HTTPException(status_code=403, detail="不允许上传该类型文件") |
| 179 | + |
| 180 | + |
| 181 | +async def read_upload_header(file: UploadFile, size: int = 64) -> bytes: |
| 182 | + await file.seek(0) |
| 183 | + header = await file.read(size) |
| 184 | + await file.seek(0) |
| 185 | + return header or b"" |
| 186 | + |
| 187 | + |
| 188 | +async def validate_upload_file(file: UploadFile) -> None: |
| 189 | + header = await read_upload_header(file) |
| 190 | + validate_file_magic(file.filename or "", file.content_type, header) |
| 191 | + |
| 192 | + |
| 193 | +def validate_header_bytes( |
| 194 | + file_name: str, content_type: Optional[str], header: bytes |
| 195 | +) -> None: |
| 196 | + validate_file_magic(file_name, content_type, header) |
0 commit comments