"""Magic-bytes file type validation. Prevents polyglot file attacks where a file is served under a trusted extension (e.g. ``.jpg``) but its actual content is something else (e.g. a PHP script, a ZIP archive, or an HTML file containing JavaScript). ``check_magic_bytes(path, content)`` returns the detected type string and raises ``PolyglotFileError`` when the declared extension does not match the content's magic bytes. Supported type detection (first-bytes signatures): MIDI — MThd header MP3 — ID3 tag or MPEG sync word WebP — RIFF….WEBP container PNG — PNG signature JPEG — JPEG SOI marker ZIP — PK signature (also covers .muse objects if ever zipped) PDF — %PDF header HTML — starts with bool: """Return True if content matches ALL (offset, expected) pairs in *checks*.""" for offset, expected in checks: end = offset + len(expected) if len(content) < end: return False if content[offset:end] != expected: return False return True def _is_midi(content: bytes) -> bool: return _match_any(content, _MIDI_MAGIC) def _is_mp3(content: bytes) -> bool: return ( _match_any(content, _MP3_ID3) or _match_any(content, _MP3_SYNC_1) or _match_any(content, _MP3_SYNC_2) or _match_any(content, _MP3_SYNC_3) ) def _is_webp(content: bytes) -> bool: return ( len(content) >= 12 and content[0:4] == b"RIFF" and content[8:12] == b"WEBP" ) def _is_png(content: bytes) -> bool: return _match_any(content, _PNG_MAGIC) def _is_jpeg(content: bytes) -> bool: return _match_any(content, _JPEG_MAGIC) def _is_zip(content: bytes) -> bool: return _match_any(content, _ZIP_MAGIC) def _is_pdf(content: bytes) -> bool: return _match_any(content, _PDF_MAGIC) def _looks_like_html(content: bytes) -> bool: head = content[:16].lower() return head.startswith(b" bool: return content[:2] == b"#!" # ── Per-extension expected checker ──────────────────────────────────────────── # Maps extension → (display_name, checker_fn) for extensions we care about. # Extensions not in this map are passed through without validation. _EXTENSION_CHECKERS = { ".mid": ("MIDI", _is_midi), ".midi": ("MIDI", _is_midi), ".mp3": ("MP3", _is_mp3), ".webp": ("WebP", _is_webp), ".png": ("PNG", _is_png), ".jpg": ("JPEG", _is_jpeg), ".jpeg": ("JPEG", _is_jpeg), ".zip": ("ZIP", _is_zip), ".pdf": ("PDF", _is_pdf), } # Content that is NEVER acceptable inside any uploaded binary file regardless # of extension — catches the most dangerous polyglot attacks. _FORBIDDEN_CHECKERS: list[_CheckerEntry] = [ ("HTML/JS", _looks_like_html), ("shebang", _is_shebang), ] # Extensions where a shebang (#!) is legitimate and must not be flagged. # Script/text files routinely start with #!/usr/bin/env . _SHEBANG_ALLOWED_EXTENSIONS: frozenset[str] = frozenset({ ".py", ".sh", ".bash", ".zsh", ".rb", ".pl", ".r", ".js", ".ts", ".mjs", ".cjs", ".lua", ".php", ".tcl", ".awk", ".bats", # BATS (Bash Automated Testing System) — starts with #!/usr/bin/env bats }) def check_magic_bytes(path: str, content: bytes) -> str: """Validate that *content* magic bytes match the extension of *path*. Args: path: The file path or filename (used only to extract the extension). content: The raw file bytes. Only the first ~16 bytes are examined. Returns: The detected type string (e.g. ``"MIDI"``, ``"MP3"``, ``"unknown"``). Raises: PolyglotFileError: When the extension implies a known type but the magic bytes do not match, or when the content matches a forbidden pattern (HTML, shebang) regardless of extension. """ if not content: return "empty" ext = Path(path).suffix.lower() # Check for universally forbidden content signatures. for label, checker_fn in _FORBIDDEN_CHECKERS: if checker_fn(content): # Allow HTML if the extension is explicitly .html / .htm if label == "HTML/JS" and ext in (".html", ".htm"): return "HTML" # Allow shebang (#!) for script/text extensions — they legitimately # start with #!/usr/bin/env and are not polyglot attacks. if label == "shebang" and ext in _SHEBANG_ALLOWED_EXTENSIONS: continue raise PolyglotFileError( f"File '{path}' has {ext!r} extension but content looks like " f"{label} — possible polyglot attack." ) # Extension-specific check. if ext not in _EXTENSION_CHECKERS: return "unknown" # no check for this extension — allow through type_name, checker_fn = _EXTENSION_CHECKERS[ext] if not checker_fn(content): raise PolyglotFileError( f"File '{path}' has {ext!r} extension but magic bytes do not match " f"{type_name} format — possible polyglot file." ) return type_name