com2care_website/apps/pages/contact_uploads.py

import os
import uuid
import zipfile
from pathlib import PurePosixPath

from django.conf import settings
from django.core.exceptions import ValidationError
from django.core.files.storage import FileSystemStorage
from django.core.files.uploadedfile import UploadedFile
from django.utils.text import get_valid_filename
from PIL import Image

ALLOWED_EXTENSIONS = frozenset(
    {
        ".png",
        ".jpg",
        ".jpeg",
        ".gif",
        ".webp",
        ".bmp",
        ".zip",
        ".log",
        ".txt",
    }
)

IMAGE_EXTENSIONS = frozenset({".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp"})
TEXT_EXTENSIONS = frozenset({".log", ".txt"})
ARCHIVE_EXTENSIONS = frozenset({".zip"})

MAX_FILE_SIZE = getattr(settings, "CONTACT_ATTACHMENT_MAX_SIZE", 10 * 1024 * 1024)
MAX_ATTACHMENTS = getattr(settings, "CONTACT_ATTACHMENT_MAX_COUNT", 3)
MAX_ZIP_ENTRIES = 100
MAX_ZIP_UNCOMPRESSED_SIZE = 50 * 1024 * 1024


class ContactAttachmentStorage(FileSystemStorage):
    def __init__(self):
        super().__init__(location=settings.CONTACT_UPLOAD_ROOT)


contact_attachment_storage = ContactAttachmentStorage()


def contact_attachment_upload_to(instance, _filename):
    ext = os.path.splitext(instance.original_filename)[1].lower()
    if ext not in ALLOWED_EXTENSIONS:
        ext = ".bin"
    subdir = str(instance.submission_id) if instance.submission_id else "pending"
    return f"{subdir}/{uuid.uuid4().hex}{ext}"


def _extension(filename):
    return os.path.splitext(filename)[1].lower()


def _read_header(uploaded_file, size=32):
    uploaded_file.seek(0)
    header = uploaded_file.read(size)
    uploaded_file.seek(0)
    return header


def _validate_magic(header, ext):
    if ext == ".png" and not header.startswith(b"\x89PNG\r\n\x1a\n"):
        raise ValidationError("File content does not match its type.")
    if ext in (".jpg", ".jpeg") and not header.startswith(b"\xff\xd8\xff"):
        raise ValidationError("File content does not match its type.")
    if ext == ".gif" and not (
        header.startswith(b"GIF87a") or header.startswith(b"GIF89a")
    ):
        raise ValidationError("File content does not match its type.")
    if ext == ".webp" and not (
        len(header) >= 12 and header[0:4] == b"RIFF" and header[8:12] == b"WEBP"
    ):
        raise ValidationError("File content does not match its type.")
    if ext == ".bmp" and not header.startswith(b"BM"):
        raise ValidationError("File content does not match its type.")
    if ext == ".zip" and not (
        header.startswith(b"PK\x03\x04")
        or header.startswith(b"PK\x05\x06")
        or header.startswith(b"PK\x07\x08")
    ):
        raise ValidationError("File content does not match its type.")


def _validate_image(uploaded_file):
    try:
        uploaded_file.seek(0)
        with Image.open(uploaded_file) as img:
            img.verify()
        uploaded_file.seek(0)
        with Image.open(uploaded_file) as img:
            img.load()
    except Exception as exc:
        raise ValidationError("Invalid or corrupted image file.") from exc
    finally:
        uploaded_file.seek(0)


def _validate_text(uploaded_file):
    uploaded_file.seek(0)
    data = uploaded_file.read()
    uploaded_file.seek(0)
    if b"\x00" in data:
        raise ValidationError("Text files must not contain binary data.")
    try:
        data.decode("utf-8")
    except UnicodeDecodeError:
        try:
            data.decode("latin-1")
        except UnicodeDecodeError as exc:
            raise ValidationError("Text file is not valid UTF-8 or Latin-1.") from exc


def _validate_zip(uploaded_file):
    uploaded_file.seek(0)
    if not zipfile.is_zipfile(uploaded_file):
        raise ValidationError("Invalid ZIP archive.")
    uploaded_file.seek(0)

    total_uncompressed = 0
    entry_count = 0

    with zipfile.ZipFile(uploaded_file, "r") as archive:
        for info in archive.infolist():
            entry_count += 1
            if entry_count > MAX_ZIP_ENTRIES:
                raise ValidationError("ZIP archive contains too many files.")

            name = info.filename
            if name.startswith("/") or ".." in PurePosixPath(name).parts:
                raise ValidationError("ZIP archive contains unsafe paths.")

            if info.flag_bits & 0x1:
                raise ValidationError("Encrypted ZIP archives are not allowed.")

            total_uncompressed += info.file_size
            if total_uncompressed > MAX_ZIP_UNCOMPRESSED_SIZE:
                raise ValidationError("ZIP archive uncompressed size is too large.")

    uploaded_file.seek(0)


def validate_contact_attachment(uploaded_file):
    if uploaded_file.size > MAX_FILE_SIZE:
        max_mb = MAX_FILE_SIZE // (1024 * 1024)
        raise ValidationError(f"File exceeds the maximum size of {max_mb} MB.")

    name = get_valid_filename(os.path.basename(uploaded_file.name))
    if not name:
        raise ValidationError("Invalid file name.")

    ext = _extension(name)
    if ext not in ALLOWED_EXTENSIONS:
        raise ValidationError(
            "File type not allowed. Permitted: images (PNG, JPG, GIF, WebP, BMP), "
            "ZIP, LOG, or TXT."
        )

    header = _read_header(uploaded_file)
    if ext in IMAGE_EXTENSIONS or ext in ARCHIVE_EXTENSIONS:
        _validate_magic(header, ext)

    if ext in IMAGE_EXTENSIONS:
        _validate_image(uploaded_file)
    elif ext in TEXT_EXTENSIONS:
        _validate_text(uploaded_file)
    elif ext in ARCHIVE_EXTENSIONS:
        _validate_zip(uploaded_file)

    return name


def validate_contact_attachments(files):
    if not files:
        return []

    if len(files) > MAX_ATTACHMENTS:
        raise ValidationError(f"You can attach at most {MAX_ATTACHMENTS} files.")

    validated = []
    for uploaded_file in files:
        if not isinstance(uploaded_file, UploadedFile):
            raise ValidationError("Invalid upload.")
        original_name = validate_contact_attachment(uploaded_file)
        validated.append((uploaded_file, original_name))
    return validated