training-software/tests/test_upload.py

"""
Upload defense tests — M5: Security Hardening and Upload Defense.

Verifies that the upload validators correctly:
- Allow files with permitted extensions and matching magic bytes.
- Reject disallowed extensions.
- Reject oversized files.
- Reject files whose magic bytes contradict the declared extension (polyglots,
  disguised executables, etc.).
"""
import io
import pytest
from django.core.exceptions import ValidationError

from api.upload_validators import (
    validate_file_extension,
    validate_file_size,
    validate_file_signature,
    validate_upload,
)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _mock_file(name: str, content: bytes, size: int | None = None):
    """Return a minimal file-like object accepted by the validators."""
    buf = io.BytesIO(content)
    buf.name = name
    buf.size = size if size is not None else len(content)
    return buf


# ---------------------------------------------------------------------------
# Extension validation
# ---------------------------------------------------------------------------

@pytest.mark.upload
class TestFileExtension:
    def test_allowed_pdf(self, settings):
        settings.ALLOWED_UPLOAD_EXTENSIONS = [".pdf"]
        f = _mock_file("report.pdf", b"%PDF-1.4")
        validate_file_extension(f)  # should not raise

    def test_allowed_jpg(self, settings):
        settings.ALLOWED_UPLOAD_EXTENSIONS = [".jpg", ".jpeg"]
        f = _mock_file("photo.jpg", b"\xFF\xD8\xFF\xE0")
        validate_file_extension(f)  # should not raise

    def test_disallowed_exe(self, settings):
        settings.ALLOWED_UPLOAD_EXTENSIONS = [".pdf", ".jpg"]
        f = _mock_file("malware.exe", b"MZ\x00\x00")
        with pytest.raises(ValidationError, match="not permitted"):
            validate_file_extension(f)

    def test_disallowed_php(self, settings):
        settings.ALLOWED_UPLOAD_EXTENSIONS = [".pdf"]
        f = _mock_file("shell.php", b"<?php system($_GET['cmd']); ?>")
        with pytest.raises(ValidationError, match="not permitted"):
            validate_file_extension(f)

    def test_case_insensitive_extension(self, settings):
        settings.ALLOWED_UPLOAD_EXTENSIONS = [".pdf"]
        f = _mock_file("REPORT.PDF", b"%PDF-1.4")
        validate_file_extension(f)  # uppercase should still pass

    def test_double_extension_blocked(self, settings):
        """file.pdf.exe must be rejected because final extension is .exe."""
        settings.ALLOWED_UPLOAD_EXTENSIONS = [".pdf"]
        f = _mock_file("file.pdf.exe", b"MZ\x00\x00")
        with pytest.raises(ValidationError):
            validate_file_extension(f)


# ---------------------------------------------------------------------------
# Size validation
# ---------------------------------------------------------------------------

@pytest.mark.upload
class TestFileSize:
    def test_within_limit(self, settings):
        settings.MAX_UPLOAD_SIZE_BYTES = 10 * 1024 * 1024  # 10 MB
        f = _mock_file("small.pdf", b"%PDF", size=5 * 1024 * 1024)
        validate_file_size(f)  # should not raise

    def test_exactly_at_limit(self, settings):
        limit = 10 * 1024 * 1024
        settings.MAX_UPLOAD_SIZE_BYTES = limit
        f = _mock_file("exact.pdf", b"%PDF", size=limit)
        validate_file_size(f)  # boundary — allowed

    def test_exceeds_limit(self, settings):
        settings.MAX_UPLOAD_SIZE_BYTES = 10 * 1024 * 1024
        f = _mock_file("huge.pdf", b"%PDF", size=11 * 1024 * 1024)
        with pytest.raises(ValidationError, match="exceeds"):
            validate_file_size(f)

    def test_zero_byte_file_passes_size_check(self, settings):
        settings.MAX_UPLOAD_SIZE_BYTES = 10 * 1024 * 1024
        f = _mock_file("empty.txt", b"", size=0)
        validate_file_size(f)  # size check only; signature check is a different concern


# ---------------------------------------------------------------------------
# Magic-byte / signature validation
# ---------------------------------------------------------------------------

@pytest.mark.upload
class TestFileSignature:
    def test_valid_pdf_signature(self):
        f = _mock_file("report.pdf", b"%PDF-1.7 content here")
        validate_file_signature(f)  # should not raise

    def test_valid_png_signature(self):
        f = _mock_file("image.png", b"\x89PNG\r\n\x1a\n" + b"\x00" * 100)
        validate_file_signature(f)  # should not raise

    def test_valid_jpeg_signature(self):
        f = _mock_file("photo.jpg", b"\xFF\xD8\xFF\xE0" + b"\x00" * 100)
        validate_file_signature(f)  # should not raise

    def test_exe_disguised_as_pdf(self):
        """MZ header in a .pdf file should be rejected."""
        f = _mock_file("evil.pdf", b"MZ\x90\x00" + b"\x00" * 100)
        with pytest.raises(ValidationError, match="does not match"):
            validate_file_signature(f)

    def test_php_disguised_as_jpg(self):
        """PHP script in a .jpg file should be rejected."""
        f = _mock_file("shell.jpg", b"<?php echo 'pwned'; ?>" + b"\x00" * 100)
        with pytest.raises(ValidationError, match="does not match"):
            validate_file_signature(f)

    def test_txt_skips_signature_check(self):
        """Plain text files have no reliable magic bytes — skip signature check."""
        f = _mock_file("data.txt", b"Hello, world!")
        validate_file_signature(f)  # no signature check for .txt

    def test_csv_skips_signature_check(self):
        f = _mock_file("export.csv", b"id,name,value\n1,foo,bar\n")
        validate_file_signature(f)  # no signature check for .csv

    def test_docx_valid_zip_signature(self):
        f = _mock_file("document.docx", b"PK\x03\x04" + b"\x00" * 100)
        validate_file_signature(f)  # should not raise

    def test_file_pointer_reset_after_check(self):
        """Signature check must leave the file pointer at position 0."""
        content = b"%PDF-1.7 content here"
        f = _mock_file("report.pdf", content)
        validate_file_signature(f)
        assert f.tell() == 0


# ---------------------------------------------------------------------------
# Combined validate_upload pipeline
# ---------------------------------------------------------------------------

@pytest.mark.upload
class TestValidateUploadPipeline:
    def test_valid_pdf_passes_all_checks(self, settings):
        settings.ALLOWED_UPLOAD_EXTENSIONS = [".pdf"]
        settings.MAX_UPLOAD_SIZE_BYTES = 10 * 1024 * 1024
        f = _mock_file("report.pdf", b"%PDF-1.7 " + b"a" * 100)
        validate_upload(f)  # should not raise

    def test_bad_extension_short_circuits_pipeline(self, settings):
        """Extension check runs first; the pipeline stops there."""
        settings.ALLOWED_UPLOAD_EXTENSIONS = [".pdf"]
        settings.MAX_UPLOAD_SIZE_BYTES = 10 * 1024 * 1024
        f = _mock_file("malware.exe", b"MZ\x90\x00")
        with pytest.raises(ValidationError, match="not permitted"):
            validate_upload(f)

    def test_oversized_file_with_valid_extension(self, settings):
        settings.ALLOWED_UPLOAD_EXTENSIONS = [".pdf"]
        settings.MAX_UPLOAD_SIZE_BYTES = 1 * 1024  # 1 KB limit
        f = _mock_file("big.pdf", b"%PDF-1.7 " + b"a" * 100, size=2 * 1024)
        with pytest.raises(ValidationError, match="exceeds"):
            validate_upload(f)

    def test_wrong_magic_bytes_with_valid_extension_and_size(self, settings):
        settings.ALLOWED_UPLOAD_EXTENSIONS = [".pdf"]
        settings.MAX_UPLOAD_SIZE_BYTES = 10 * 1024 * 1024
        f = _mock_file("disguised.pdf", b"MZ\x90\x00" + b"a" * 100)
        with pytest.raises(ValidationError, match="does not match"):
            validate_upload(f)