Initial commit

tests/test_attachment_extractor.py

@@ -0,0 +1,62 @@
+import gzip
+import io
+import zipfile
+from email.message import EmailMessage
+from pathlib import Path
+
+import pytest
+
+from app.attachment_extractor import AttachmentExtractionError, extract_dmarc_attachments, extract_payload
+
+
+def _xml() -> bytes:
+    return Path("tests/fixtures/sample_dmarc.xml").read_bytes()
+
+
+def test_gzip_attachment_extraction():
+    gz = gzip.compress(_xml())
+    reports = extract_payload("report.xml.gz", "application/octet-stream", gz, 20)
+
+    assert len(reports) == 1
+    assert reports[0].payload.startswith(b"<?xml")
+    assert len(reports[0].sha256) == 64
+
+
+def test_zip_attachment_extraction_rejects_traversal():
+    buf = io.BytesIO()
+    with zipfile.ZipFile(buf, "w") as archive:
+        archive.writestr("report.xml", _xml())
+        archive.writestr("../evil.xml", _xml())
+
+    with pytest.raises(AttachmentExtractionError, match="unsafe zip path"):
+        extract_payload("reports.zip", "application/zip", buf.getvalue(), 20)
+
+
+def test_zip_attachment_extraction_rejects_nested_archives():
+    buf = io.BytesIO()
+    with zipfile.ZipFile(buf, "w") as archive:
+        archive.writestr("nested.zip", b"not allowed")
+
+    with pytest.raises(AttachmentExtractionError, match="nested archive"):
+        extract_payload("reports.zip", "application/zip", buf.getvalue(), 20)
+
+
+def test_zip_attachment_extraction_caps_reports_per_archive():
+    buf = io.BytesIO()
+    with zipfile.ZipFile(buf, "w") as archive:
+        archive.writestr("one.xml", _xml())
+        archive.writestr("two.xml", _xml())
+
+    with pytest.raises(AttachmentExtractionError, match="archive XML report limit"):
+        extract_payload("reports.zip", "application/zip", buf.getvalue(), 20, max_reports_per_archive=1)
+
+
+def test_message_attachment_detection_with_octet_stream_valid_filename():
+    msg = EmailMessage()
+    msg["Subject"] = "Report domain tukutoi.com"
+    msg.set_content("attached")
+    msg.add_attachment(gzip.compress(_xml()), maintype="application", subtype="octet-stream", filename="report.gz")
+
+    reports = extract_dmarc_attachments(msg, 20)
+
+    assert len(reports) == 1