103 lines
4.1 KiB
Python
103 lines
4.1 KiB
Python
"""Bounded, dependency-free validation for supported upload formats."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import codecs
|
|
import zipfile
|
|
from pathlib import Path
|
|
|
|
|
|
class UploadContentError(ValueError):
|
|
"""Raised when bytes do not match the declared file extension."""
|
|
|
|
|
|
_OLE_MAGIC = b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1"
|
|
_PNG_MAGIC = b"\x89PNG\r\n\x1a\n"
|
|
_MAX_ZIP_ENTRIES = 10_000
|
|
_MAX_ZIP_UNCOMPRESSED_BYTES = 200 * 1024 * 1024
|
|
_MAX_ZIP_RATIO = 500
|
|
|
|
|
|
def _validate_text(path: Path) -> None:
|
|
with path.open("rb") as handle:
|
|
sample = handle.read(64 * 1024)
|
|
if not sample:
|
|
return
|
|
if sample.startswith((b"\xff\xfe", b"\xfe\xff")):
|
|
try:
|
|
codecs.getincrementaldecoder("utf-16")(errors="strict").decode(sample, final=False)
|
|
except UnicodeDecodeError as exc:
|
|
raise UploadContentError("Text file has an invalid UTF-16 encoding") from exc
|
|
return
|
|
if b"\x00" in sample:
|
|
raise UploadContentError("Text file contains binary NUL bytes")
|
|
control_bytes = sum(byte < 9 or 13 < byte < 32 for byte in sample)
|
|
if control_bytes / len(sample) > 0.10:
|
|
raise UploadContentError("Text file contains too many binary control bytes")
|
|
for encoding in ("utf-8-sig", "gb18030"):
|
|
try:
|
|
codecs.getincrementaldecoder(encoding)(errors="strict").decode(sample, final=False)
|
|
return
|
|
except UnicodeDecodeError:
|
|
continue
|
|
raise UploadContentError("Text file encoding is not supported")
|
|
|
|
|
|
def _validate_ooxml(path: Path, expected_prefix: str) -> None:
|
|
try:
|
|
with zipfile.ZipFile(path) as archive:
|
|
infos = archive.infolist()
|
|
if len(infos) > _MAX_ZIP_ENTRIES:
|
|
raise UploadContentError("Office archive contains too many entries")
|
|
names = {info.filename.replace("\\", "/") for info in infos}
|
|
if "[Content_Types].xml" not in names or not any(name.startswith(expected_prefix) for name in names):
|
|
raise UploadContentError("Office archive content does not match its extension")
|
|
total_uncompressed = sum(max(0, info.file_size) for info in infos)
|
|
total_compressed = sum(max(0, info.compress_size) for info in infos)
|
|
if total_uncompressed > _MAX_ZIP_UNCOMPRESSED_BYTES:
|
|
raise UploadContentError("Office archive expands beyond the safety limit")
|
|
if (
|
|
total_uncompressed > 10 * 1024 * 1024
|
|
and total_uncompressed > max(1, total_compressed) * _MAX_ZIP_RATIO
|
|
):
|
|
raise UploadContentError("Office archive compression ratio is unsafe")
|
|
except (zipfile.BadZipFile, OSError) as exc:
|
|
raise UploadContentError("Office file is not a valid ZIP container") from exc
|
|
|
|
|
|
def validate_upload_content(path: Path, extension: str) -> None:
|
|
"""Check the supported format's signature without reading the whole file."""
|
|
|
|
extension = extension.lower()
|
|
with path.open("rb") as handle:
|
|
header = handle.read(1024)
|
|
if extension in {".txt", ".md", ".csv"}:
|
|
_validate_text(path)
|
|
return
|
|
if extension == ".pdf":
|
|
if b"%PDF-" not in header:
|
|
raise UploadContentError("PDF signature does not match the .pdf extension")
|
|
return
|
|
if extension == ".png":
|
|
if not header.startswith(_PNG_MAGIC):
|
|
raise UploadContentError("PNG signature does not match the .png extension")
|
|
return
|
|
if extension in {".jpg", ".jpeg"}:
|
|
if not header.startswith(b"\xff\xd8\xff"):
|
|
raise UploadContentError("JPEG signature does not match the image extension")
|
|
return
|
|
if extension in {".doc", ".ppt", ".xls"}:
|
|
if not header.startswith(_OLE_MAGIC):
|
|
raise UploadContentError("Legacy Office signature does not match the extension")
|
|
return
|
|
if extension == ".docx":
|
|
_validate_ooxml(path, "word/")
|
|
return
|
|
if extension == ".pptx":
|
|
_validate_ooxml(path, "ppt/")
|
|
return
|
|
if extension == ".xlsx":
|
|
_validate_ooxml(path, "xl/")
|
|
return
|
|
raise UploadContentError(f"Unsupported file extension: {extension or '(none)'}")
|