upload-guard 0.1.0__py3-none-any.whl

upload_guard/__init__.py

@@ -0,0 +1,41 @@
+"""upload_guard: pure-Python upload validation.
+
+Quick start::
+
+    from upload_guard import UploadGuard, UploadRejected
+
+    guard = UploadGuard(allowed=["image/*", "application/pdf"], max_size=10 * 1024 * 1024)
+    try:
+        result = guard.check(upload)          # bytes, path, file object, FastAPI/Django/Flask upload
+    except UploadRejected as exc:
+        print(exc.codes)                      # e.g. ['extension.mismatch']
+    else:
+        print(result.mime, result.safe_filename, result.sanitized)
+"""
+from ._version import __version__
+from .exceptions import UnsupportedSource, UploadGuardError, UploadRejected
+from .extension import TypeRule, canonical_extension, extension_matches, normalize_extension, split_filename
+from .guard import Policy, UploadGuard, check, detect_type, scan
+from .security.archive import ArchiveLimits, inspect_archive, inspect_compressed, inspect_zip
+from .security.filename import DANGEROUS_EXTENSIONS, check_filename, sanitize_filename
+from .security.polyglot import check_trailing, scan_embedded
+from .security.svg import SvgPolicy, SvgResult, sanitize_svg, scan_svg
+from .sniff import HEADER_SIZE, detect, detect_stream
+from .types import ArchiveReport, DetectedType, Finding, ScanResult, Severity
+
+__all__ = [
+    "__version__",
+    # core
+    "UploadGuard", "Policy", "check", "scan", "detect_type",
+    "ScanResult", "Finding", "Severity", "DetectedType", "ArchiveReport",
+    "UploadGuardError", "UploadRejected", "UnsupportedSource",
+    # sniffing
+    "detect", "detect_stream", "HEADER_SIZE",
+    # extensions
+    "TypeRule", "canonical_extension", "normalize_extension", "extension_matches", "split_filename",
+    # security
+    "check_filename", "sanitize_filename", "DANGEROUS_EXTENSIONS",
+    "SvgPolicy", "SvgResult", "scan_svg", "sanitize_svg",
+    "ArchiveLimits", "inspect_archive", "inspect_zip", "inspect_compressed",
+    "scan_embedded", "check_trailing",
+]

upload_guard/_version.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

upload_guard/adapters.py

@@ -0,0 +1,214 @@
+"""Turn *anything upload-shaped* into a seekable binary stream plus metadata.
+
+Supported inputs (duck-typed, no framework imports):
+
+* ``bytes`` / ``bytearray`` / ``memoryview``
+* ``str`` / ``os.PathLike`` paths
+* any binary file object (``io.BytesIO``, open files, ``tempfile`` objects)
+* Starlette / FastAPI ``UploadFile`` (``.file``, ``.filename``, ``.content_type``, ``.size``)
+* Django ``UploadedFile`` (``.name``, ``.size``, ``.content_type``, ``.read``/``.seek``)
+* Werkzeug / Flask ``FileStorage`` (``.stream``, ``.filename``, ``.mimetype``)
+
+The stream position is captured on entry and restored on :meth:`FileSource.close`,
+so frameworks can still save the file afterwards.
+"""
+from __future__ import annotations
+
+import io
+import os
+import tempfile
+from typing import Any, BinaryIO, Optional
+
+from .exceptions import UnsupportedSource
+
+_SPOOL_MAX = 8 * 1024 * 1024
+
+
+class FileSource:
+    """A seekable binary stream with optional ``filename`` / ``content_type`` / ``size``."""
+
+    def __init__(
+        self,
+        stream: BinaryIO,
+        *,
+        filename: Optional[str] = None,
+        content_type: Optional[str] = None,
+        size: Optional[int] = None,
+        owns_stream: bool = False,
+        oversized: bool = False,
+    ) -> None:
+        self.stream = stream
+        self.filename = filename
+        self.content_type = content_type
+        self._size = size
+        self.owns_stream = owns_stream
+        self.oversized = oversized  # non-seekable input exceeded the read cap
+        try:
+            self._start_pos = stream.tell()
+        except (OSError, ValueError, AttributeError):
+            self._start_pos = 0
+
+    # -- construction -------------------------------------------------------------
+    @classmethod
+    def from_any(
+        cls,
+        obj: Any,
+        *,
+        filename: Optional[str] = None,
+        content_type: Optional[str] = None,
+        max_read: Optional[int] = None,
+    ) -> "FileSource":
+        if isinstance(obj, FileSource):
+            if filename is not None:
+                obj.filename = filename
+            if content_type is not None:
+                obj.content_type = content_type
+            return obj
+
+        if isinstance(obj, (bytes, bytearray, memoryview)):
+            data = bytes(obj)
+            return cls(io.BytesIO(data), filename=filename, content_type=content_type, size=len(data), owns_stream=True)
+
+        if isinstance(obj, (str, os.PathLike)):
+            path = os.fspath(obj)
+            fh = open(path, "rb")  # noqa: SIM115 - closed by FileSource.close()
+            size = os.fstat(fh.fileno()).st_size
+            return cls(fh, filename=filename or os.path.basename(path), content_type=content_type, size=size, owns_stream=True)
+
+        # Starlette / FastAPI UploadFile
+        inner = getattr(obj, "file", None)
+        if inner is not None and hasattr(inner, "read") and hasattr(obj, "filename"):
+            return cls(
+                _ensure_seekable(inner, max_read)[0],
+                filename=filename or _str_or_none(getattr(obj, "filename", None)),
+                content_type=content_type or _str_or_none(getattr(obj, "content_type", None)),
+                size=_int_or_none(getattr(obj, "size", None)),
+            )
+
+        # Werkzeug FileStorage
+        inner = getattr(obj, "stream", None)
+        if inner is not None and hasattr(inner, "read") and hasattr(obj, "filename"):
+            stream, owns, oversized = _ensure_seekable(inner, max_read)
+            return cls(
+                stream,
+                filename=filename or _str_or_none(getattr(obj, "filename", None)),
+                content_type=content_type or _str_or_none(getattr(obj, "mimetype", None) or getattr(obj, "content_type", None)),
+                size=_int_or_none(getattr(obj, "content_length", None)) or None,
+                owns_stream=owns,
+                oversized=oversized,
+            )
+
+        # Django UploadedFile (proxies read/seek to .file) or any file-like object
+        if hasattr(obj, "read"):
+            django_name = getattr(obj, "name", None) if hasattr(obj, "chunks") else None
+            stream, owns, oversized = _ensure_seekable(obj, max_read)
+            name = filename
+            if name is None:
+                raw_name = django_name if django_name is not None else getattr(obj, "name", None)
+                if isinstance(raw_name, str):
+                    name = os.path.basename(raw_name) if not hasattr(obj, "chunks") else raw_name
+            ctype = content_type or _str_or_none(getattr(obj, "content_type", None))
+            size = _int_or_none(getattr(obj, "size", None)) if hasattr(obj, "chunks") else None
+            return cls(stream, filename=name, content_type=ctype, size=size, owns_stream=owns, oversized=oversized)
+
+        raise UnsupportedSource("cannot read uploads from %r" % type(obj).__name__)
+
+    # -- reading helpers -------------------------------------------------------------
+    @property
+    def size(self) -> Optional[int]:
+        if self._size is None:
+            try:
+                pos = self.stream.tell()
+                self.stream.seek(0, os.SEEK_END)
+                self._size = self.stream.tell()
+                self.stream.seek(pos)
+            except (OSError, ValueError, AttributeError):
+                return None
+        return self._size
+
+    def read_at(self, offset: int, n: int) -> bytes:
+        pos = self.stream.tell()
+        try:
+            self.stream.seek(offset)
+            return self.stream.read(n) or b""
+        finally:
+            self.stream.seek(pos)
+
+    def read_head(self, n: int) -> bytes:
+        return self.read_at(0, n)
+
+    def read_all(self, cap: Optional[int] = None) -> Optional[bytes]:
+        """Whole content, or None when it exceeds ``cap`` bytes."""
+        if cap is not None and self.size is not None and self.size > cap:
+            return None
+        pos = self.stream.tell()
+        try:
+            self.stream.seek(0)
+            if cap is None:
+                return self.stream.read() or b""
+            data = self.stream.read(cap + 1) or b""
+            return None if len(data) > cap else data
+        finally:
+            self.stream.seek(pos)
+
+    # -- lifecycle -------------------------------------------------------------------
+    def restore(self) -> None:
+        try:
+            self.stream.seek(self._start_pos)
+        except (OSError, ValueError, AttributeError):
+            pass
+
+    def close(self) -> None:
+        if self.owns_stream:
+            try:
+                self.stream.close()
+            except Exception:  # noqa: BLE001
+                pass
+        else:
+            self.restore()
+
+    def __enter__(self) -> "FileSource":
+        return self
+
+    def __exit__(self, *exc: object) -> None:
+        self.close()
+
+
+def _str_or_none(value: Any) -> Optional[str]:
+    return value if isinstance(value, str) and value else None
+
+
+def _int_or_none(value: Any) -> Optional[int]:
+    return value if isinstance(value, int) and not isinstance(value, bool) and value >= 0 else None
+
+
+def _is_seekable(stream: Any) -> bool:
+    try:
+        if hasattr(stream, "seekable") and not stream.seekable():
+            return False
+        pos = stream.tell()
+        stream.seek(pos)
+        return True
+    except (OSError, ValueError, AttributeError, io.UnsupportedOperation):
+        return False
+
+
+def _ensure_seekable(stream: Any, max_read: Optional[int]):
+    """Return ``(seekable_stream, owns, oversized)``; spools non-seekable input to a temp file."""
+    if _is_seekable(stream):
+        return stream, False, False
+    spool = tempfile.SpooledTemporaryFile(max_size=_SPOOL_MAX)
+    remaining = None if max_read is None else max_read + 1
+    oversized = False
+    while True:
+        chunk = stream.read(64 * 1024 if remaining is None else min(64 * 1024, remaining))
+        if not chunk:
+            break
+        spool.write(chunk)
+        if remaining is not None:
+            remaining -= len(chunk)
+            if remaining <= 0:
+                oversized = True
+                break
+    spool.seek(0)
+    return spool, True, oversized

upload_guard/cli.py

@@ -0,0 +1,97 @@
+"""``upload_guard`` command line interface."""
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from typing import List, Optional
+
+from ._version import __version__
+from .guard import UploadGuard, detect_type
+from .types import Severity
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(prog="upload_guard", description="Validate uploaded files: magic-number sniffing, extension checks, SVG/zip-bomb/path-traversal protection.")
+    parser.add_argument("--version", action="version", version="upload_guard %s" % __version__)
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    scan = sub.add_parser("scan", help="run the full policy against one or more files")
+    scan.add_argument("files", nargs="+")
+    scan.add_argument("--allow", action="append", default=None, metavar="RULE", help='allowed type rule, e.g. "image/*", "application/pdf", ".docx" (repeatable)')
+    scan.add_argument("--block", action="append", default=None, metavar="RULE", help="blocked type rule (repeatable; default blocks executables and scripts)")
+    scan.add_argument("--max-size", type=int, default=None, metavar="BYTES")
+    scan.add_argument("--threshold", choices=[s.name for s in Severity], default="MEDIUM", help="minimum severity that rejects (default MEDIUM)")
+    scan.add_argument("--no-sanitize", action="store_true", help="report SVG problems instead of sanitising")
+    scan.add_argument("--lenient-extension", action="store_true", help="do not require the extension to match the detected type")
+    scan.add_argument("--json", action="store_true", help="machine readable output")
+    scan.add_argument("--write-sanitized", metavar="PATH", help="write the sanitised SVG (single file only)")
+
+    det = sub.add_parser("detect", help="print the detected type only")
+    det.add_argument("files", nargs="+")
+    det.add_argument("--json", action="store_true")
+    return parser
+
+
+def _print_result(path: str, result) -> None:
+    status = "OK  " if result.ok else "FAIL"
+    print("%s %s  ->  %s (%s)%s" % (status, path, result.detected.mime, result.detected.description, "" if result.size is None else "  %d bytes" % result.size))
+    for f in result.findings:
+        tag = "fixed" if f.remediated else f.severity.name.lower()
+        print("      [%-8s] %-28s %s" % (tag, f.code, f.message))
+    if result.archive:
+        a = result.archive
+        print("      archive: %s, %d entries, %d -> %d bytes%s" % (a.format, a.entry_count, a.total_compressed, a.total_uncompressed, " (truncated)" if a.truncated else ""))
+
+
+def main(argv: Optional[List[str]] = None) -> int:
+    parser = _build_parser()
+    args = parser.parse_args(argv)
+
+    if args.command == "detect":
+        out = []
+        for path in args.files:
+            try:
+                t = detect_type(path)
+            except OSError as exc:
+                print("%s: %s" % (path, exc), file=sys.stderr)
+                return 2
+            out.append({"file": path, **t.to_dict()})
+            if not args.json:
+                print("%s: %s (%s)" % (path, t.mime, t.description))
+        if args.json:
+            print(json.dumps(out, indent=2))
+        return 0
+
+    guard = UploadGuard(
+        allowed=args.allow,
+        blocked=args.block if args.block is not None else ("category:executable", "category:script"),
+        max_size=args.max_size,
+        reject_threshold=Severity[args.threshold],
+        sanitize_svg=not args.no_sanitize,
+        strict_extension=not args.lenient_extension,
+    )
+    exit_code = 0
+    results = []
+    for path in args.files:
+        try:
+            result = guard.scan(path)
+        except OSError as exc:
+            print("%s: %s" % (path, exc), file=sys.stderr)
+            exit_code = 2
+            continue
+        results.append({"file": path, **result.to_dict()})
+        if not result.ok:
+            exit_code = max(exit_code, 1)
+        if not args.json:
+            _print_result(path, result)
+        if args.write_sanitized and result.sanitized is not None and len(args.files) == 1:
+            with open(args.write_sanitized, "wb") as fh:
+                fh.write(result.sanitized)
+    if args.json:
+        print(json.dumps(results, indent=2))
+    return exit_code
+
+
+if __name__ == "__main__":  # pragma: no cover
+    sys.exit(main())

upload_guard/exceptions.py

@@ -0,0 +1,33 @@
+"""Exception hierarchy."""
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List
+
+if TYPE_CHECKING:  # pragma: no cover
+    from .types import Finding, ScanResult
+
+
+class UploadGuardError(Exception):
+    """Base class for all upload_guard errors."""
+
+
+class UnsupportedSource(UploadGuardError, TypeError):
+    """The object passed in is not something upload_guard knows how to read."""
+
+
+class UploadRejected(UploadGuardError, ValueError):
+    """Raised by :meth:`UploadGuard.check` when the upload fails the policy.
+
+    ``result`` holds the full :class:`ScanResult`; ``findings`` is a shortcut
+    to the findings that caused the rejection.
+    """
+
+    def __init__(self, result: "ScanResult") -> None:
+        self.result = result
+        self.findings: List["Finding"] = result.errors
+        message = "; ".join(f"[{f.code}] {f.message}" for f in self.findings) or "upload rejected"
+        super().__init__(message)
+
+    @property
+    def codes(self) -> List[str]:
+        return [f.code for f in self.findings]

upload_guard/extension.py

@@ -0,0 +1,179 @@
+"""Extension handling: normalisation, aliases, declared-vs-detected comparison
+and allow-list rules such as ``"image/*"``, ``"application/pdf"`` or ``".docx"``."""
+from __future__ import annotations
+
+import os
+import re
+from typing import Dict, Iterable, List, Optional, Sequence, Tuple, Union
+
+from .types import DetectedType
+
+# Alternative spellings that should be treated as the same extension.
+EXTENSION_ALIASES: Dict[str, str] = {
+    "jpeg": "jpg", "jpe": "jpg", "jfif": "jpg", "pjpeg": "jpg",
+    "tiff": "tif",
+    "htm": "html", "xhtml": "html",
+    "yml": "yaml",
+    "mpeg": "mpg", "mpe": "mpg",
+    "midi": "mid",
+    "markdown": "md",
+    "text": "txt",
+    "gzip": "gz",
+    "sqlite3": "sqlite",
+    "aiff": "aif",
+    "wave": "wav",
+    "heif": "heic",
+    "tar.gz": "tgz",
+}
+
+_EXT_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_\-+~]{0,15}$")
+
+
+def normalize_extension(ext: Optional[str]) -> str:
+    """``".JPG "`` → ``"jpg"``. Returns ``""`` for empty/None."""
+    if not ext:
+        return ""
+    return ext.strip().lstrip(".").strip().lower()
+
+
+def canonical_extension(ext: Optional[str]) -> str:
+    norm = normalize_extension(ext)
+    return EXTENSION_ALIASES.get(norm, norm)
+
+
+def split_filename(filename: Optional[str]) -> Tuple[str, Optional[str]]:
+    """Return ``(stem, extension)`` using the *last* dot; extension is lower-case without the dot.
+
+    Trailing dots/spaces are ignored (``"a.txt. "`` → ``("a", "txt")``). A leading
+    dot alone (``".bashrc"``) is not an extension.
+    """
+    if not filename:
+        return "", None
+    name = os.path.basename(filename.replace("\\", "/")).rstrip(". ")
+    if "." not in name:
+        return name, None
+    stem, _, ext = name.rpartition(".")
+    if not stem:  # ".bashrc"
+        return name, None
+    ext = ext.lower()
+    if not _EXT_RE.match(ext):
+        return name, None
+    return stem, ext
+
+
+def all_extensions(filename: Optional[str]) -> List[str]:
+    """Every dotted suffix: ``"a.tar.gz"`` → ``["tar", "gz"]``."""
+    if not filename:
+        return []
+    name = os.path.basename(filename.replace("\\", "/")).rstrip(". ").lstrip(".")
+    parts = name.split(".")
+    if len(parts) < 2:
+        return []
+    return [p.lower() for p in parts[1:] if _EXT_RE.match(p)]
+
+
+def extension_matches(ext: Optional[str], detected: DetectedType) -> bool:
+    """True when ``ext`` is one of the extensions that ``detected`` may legitimately use."""
+    canon = canonical_extension(ext)
+    if not canon:
+        return False
+    return canon in {canonical_extension(e) for e in detected.extensions}
+
+
+# ---------------------------------------------------------------------------
+# Registry of every known type (for reverse look-ups and double-extension checks)
+# ---------------------------------------------------------------------------
+def _all_known_types() -> List[DetectedType]:
+    from .sniff import containers, signatures
+
+    seen: Dict[Tuple[str, Tuple[str, ...]], DetectedType] = {}
+    for module in (signatures, containers):
+        for value in vars(module).values():
+            if isinstance(value, DetectedType):
+                seen.setdefault((value.mime, value.extensions), value)
+    return list(seen.values())
+
+
+_KNOWN_TYPES: Optional[List[DetectedType]] = None
+
+
+def known_types() -> List[DetectedType]:
+    global _KNOWN_TYPES
+    if _KNOWN_TYPES is None:
+        _KNOWN_TYPES = _all_known_types()
+    return _KNOWN_TYPES
+
+
+def known_extensions() -> frozenset:
+    return frozenset(canonical_extension(e) for t in known_types() for e in t.extensions)
+
+
+def types_for_extension(ext: Optional[str]) -> List[DetectedType]:
+    canon = canonical_extension(ext)
+    if not canon:
+        return []
+    return [t for t in known_types() if canon in {canonical_extension(e) for e in t.extensions}]
+
+
+def mime_for_extension(ext: Optional[str]) -> Optional[str]:
+    types = types_for_extension(ext)
+    return types[0].mime if types else None
+
+
+# ---------------------------------------------------------------------------
+# Allow-list rules
+# ---------------------------------------------------------------------------
+_MIME_RE = re.compile(r"^[a-z0-9!#$&^_.+-]+/(\*|[a-z0-9!#$&^_.+-]+)$", re.IGNORECASE)
+
+
+class TypeRule:
+    """One allow/deny rule: a MIME type, a MIME wildcard (``image/*``), an extension
+    (``.pdf`` / ``pdf``) or a category (``category:image``)."""
+
+    __slots__ = ("kind", "value", "raw")
+
+    def __init__(self, raw: str) -> None:
+        self.raw = raw
+        text = raw.strip().lower()
+        if not text:
+            raise ValueError("empty type rule")
+        if text.startswith("category:"):
+            self.kind, self.value = "category", text.split(":", 1)[1].strip()
+        elif "/" in text:
+            if not _MIME_RE.match(text):
+                raise ValueError("invalid MIME type rule: %r" % raw)
+            if text.endswith("/*"):
+                self.kind, self.value = "mime_prefix", text[:-1]  # "image/"
+            else:
+                self.kind, self.value = "mime", text
+        else:
+            self.kind, self.value = "ext", canonical_extension(text)
+            if not self.value:
+                raise ValueError("invalid extension rule: %r" % raw)
+
+    def matches(self, detected: DetectedType) -> bool:
+        if self.kind == "mime":
+            return detected.mime.lower() == self.value
+        if self.kind == "mime_prefix":
+            return detected.mime.lower().startswith(self.value)
+        if self.kind == "category":
+            return detected.category == self.value
+        return any(canonical_extension(e) == self.value for e in detected.extensions)
+
+    def __repr__(self) -> str:  # pragma: no cover - debugging aid
+        return "TypeRule(%r)" % self.raw
+
+
+RuleInput = Union[str, TypeRule, Iterable[Union[str, TypeRule]], None]
+
+
+def compile_rules(rules: RuleInput) -> List[TypeRule]:
+    if rules is None:
+        return []
+    if isinstance(rules, (str, TypeRule)):
+        rules = [rules]
+    return [r if isinstance(r, TypeRule) else TypeRule(r) for r in rules]
+
+
+def any_rule_matches(rules: Sequence[TypeRule], detected: DetectedType) -> bool:
+    return any(r.matches(detected) for r in rules)