Christian Krakau-Louis
194 lines
7.3 KiB
Python
194 lines
7.3 KiB
Python
import base64
|
|
import email
|
|
import gzip
|
|
import io
|
|
import zipfile
|
|
from email import encoders as email_encoders
|
|
from email.mime.application import MIMEApplication
|
|
from email.mime.base import MIMEBase
|
|
from email.mime.multipart import MIMEMultipart
|
|
from email.mime.text import MIMEText
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
|
|
from app.models.report import DMARCReport
|
|
from app.services.dmarc_parser import DMARCParser
|
|
from app.services.gmail_client import GmailClient
|
|
from app.services.imap_client import IMAPClient
|
|
from app.tests.test_data import DMARC_COMPATIBILITY_FIXTURES, load_dmarc_fixture
|
|
|
|
|
|
def _fixture_id(fixture: dict) -> str:
|
|
return fixture["id"]
|
|
|
|
|
|
def _fixture_bytes(fixture: dict) -> bytes:
|
|
return load_dmarc_fixture(fixture["filename"]).encode("utf-8")
|
|
|
|
|
|
def _zip_xml(xml_content: bytes, filename: str = "report.xml") -> bytes:
|
|
buf = io.BytesIO()
|
|
with zipfile.ZipFile(buf, "w") as archive:
|
|
archive.writestr(filename, xml_content)
|
|
return buf.getvalue()
|
|
|
|
|
|
def _make_mime_attachment(filename: str, content: bytes, content_type: str) -> bytes:
|
|
msg = MIMEMultipart()
|
|
msg["Subject"] = "DMARC aggregate report"
|
|
msg["From"] = "reports@example.test"
|
|
msg["To"] = "dmarc@example.test"
|
|
msg.attach(MIMEText("DMARC report attached.", "plain"))
|
|
|
|
part = MIMEApplication(content, Name=filename)
|
|
part["Content-Disposition"] = f'attachment; filename="{filename}"'
|
|
part.set_type(content_type)
|
|
msg.attach(part)
|
|
return msg.as_bytes()
|
|
|
|
|
|
def _make_gmail_raw_attachment(filename: str, content: bytes) -> bytes:
|
|
msg = MIMEMultipart()
|
|
msg["Subject"] = "DMARC aggregate report"
|
|
msg["From"] = "reports@example.test"
|
|
msg["To"] = "dmarc@example.test"
|
|
msg.attach(MIMEText("DMARC report attached.", "plain"))
|
|
|
|
part = MIMEBase("application", "zip")
|
|
part.set_payload(content)
|
|
email_encoders.encode_base64(part)
|
|
part.add_header("Content-Disposition", "attachment", filename=filename)
|
|
msg.attach(part)
|
|
return msg.as_bytes()
|
|
|
|
|
|
def _make_imap_client(db_session) -> IMAPClient:
|
|
with patch("app.services.imap_client.get_settings") as mock_settings:
|
|
mock_settings.return_value = MagicMock(
|
|
IMAP_SERVER="imap.example.test",
|
|
IMAP_PORT=993,
|
|
IMAP_USERNAME="dmarc@example.test",
|
|
IMAP_PASSWORD="password",
|
|
)
|
|
return IMAPClient(db=db_session)
|
|
|
|
|
|
def _make_gmail_client(db_session) -> GmailClient:
|
|
with patch("app.services.gmail_client.Credentials") as mock_credentials_class:
|
|
credentials = MagicMock()
|
|
credentials.token = "access-token"
|
|
credentials.refresh_token = "refresh-token"
|
|
credentials.expired = False
|
|
mock_credentials_class.return_value = credentials
|
|
return GmailClient(
|
|
client_id="client-id",
|
|
client_secret="client-secret",
|
|
access_token="access-token",
|
|
refresh_token="refresh-token",
|
|
db=db_session,
|
|
)
|
|
|
|
|
|
def _assert_expected_report(report: dict, fixture: dict) -> None:
|
|
assert report["domain"] == fixture["domain"]
|
|
assert report["report_id"] == fixture["report_id"]
|
|
assert report["variant"] == fixture["variant"]
|
|
assert report["summary"]["total_count"] == fixture["total_count"]
|
|
|
|
for key, value in fixture.get("policy", {}).items():
|
|
assert report["policy"][key] == value
|
|
|
|
|
|
def _assert_persisted_report(db_session, fixture: dict) -> DMARCReport:
|
|
db_session.flush()
|
|
report = db_session.query(DMARCReport).filter_by(report_id=fixture["report_id"]).one()
|
|
assert report.domain.name == fixture["domain"]
|
|
assert report.report_variant == fixture["variant"]
|
|
assert sum(record.count for record in report.records) == fixture["total_count"]
|
|
return report
|
|
|
|
|
|
@pytest.mark.parametrize("fixture", DMARC_COMPATIBILITY_FIXTURES, ids=_fixture_id)
|
|
def test_aggregate_compatibility_fixtures_parse_xml_zip_and_gzip(fixture):
|
|
xml_bytes = _fixture_bytes(fixture)
|
|
|
|
xml_report = DMARCParser.parse_file(xml_bytes, fixture["filename"])
|
|
zip_report = DMARCParser.parse_file(_zip_xml(xml_bytes), fixture["filename"] + ".zip")
|
|
gzip_report = DMARCParser.parse_file(gzip.compress(xml_bytes), fixture["filename"] + ".gz")
|
|
|
|
for report in (xml_report, zip_report, gzip_report):
|
|
_assert_expected_report(report, fixture)
|
|
|
|
|
|
@pytest.mark.parametrize("fixture", DMARC_COMPATIBILITY_FIXTURES, ids=_fixture_id)
|
|
def test_aggregate_compatibility_fixtures_import_via_upload(
|
|
client: TestClient, db_session, fixture
|
|
):
|
|
zip_bytes = _zip_xml(_fixture_bytes(fixture), fixture["filename"])
|
|
|
|
response = client.post(
|
|
"/api/v1/reports/upload",
|
|
files={"file": (fixture["filename"] + ".zip", zip_bytes, "application/zip")},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
_assert_persisted_report(db_session, fixture)
|
|
|
|
reports = client.get(f"/api/v1/domains/{fixture['domain']}/reports")
|
|
assert reports.status_code == 200
|
|
assert reports.json()["reports"][0]["id"] == fixture["report_id"]
|
|
|
|
export = client.get(f"/api/v1/domains/{fixture['domain']}/reports/export")
|
|
assert export.status_code == 200
|
|
assert fixture["report_id"] in export.text
|
|
assert "report_variant" in export.text
|
|
|
|
|
|
@pytest.mark.parametrize("fixture", DMARC_COMPATIBILITY_FIXTURES, ids=_fixture_id)
|
|
def test_aggregate_compatibility_fixtures_import_via_imap(db_session, fixture):
|
|
client = _make_imap_client(db_session)
|
|
zip_bytes = _zip_xml(_fixture_bytes(fixture), fixture["filename"])
|
|
raw_message = _make_mime_attachment(fixture["filename"] + ".zip", zip_bytes, "application/zip")
|
|
msg = email.message_from_bytes(raw_message)
|
|
stats = {"processed": 0, "reports_found": 0, "errors": []}
|
|
|
|
count = client._process_attachments(msg, stats, message_id="imap-" + fixture["id"])
|
|
|
|
assert count == 1
|
|
assert stats["details"][0]["status"] == "imported"
|
|
assert stats["details"][0]["report_id"] == fixture["report_id"]
|
|
_assert_persisted_report(db_session, fixture)
|
|
|
|
|
|
@pytest.mark.parametrize("fixture", DMARC_COMPATIBILITY_FIXTURES, ids=_fixture_id)
|
|
def test_aggregate_compatibility_fixtures_import_via_gmail(db_session, fixture):
|
|
client = _make_gmail_client(db_session)
|
|
zip_bytes = _zip_xml(_fixture_bytes(fixture), fixture["filename"])
|
|
raw_message = _make_gmail_raw_attachment(fixture["filename"] + ".zip", zip_bytes)
|
|
msg = email.message_from_bytes(raw_message)
|
|
stats = {"reports_found": 0, "errors": []}
|
|
|
|
count = client._process_attachments(msg, stats, message_id="gmail-" + fixture["id"])
|
|
|
|
assert count == 1
|
|
assert stats["details"][0]["status"] == "imported"
|
|
assert stats["details"][0]["report_id"] == fixture["report_id"]
|
|
_assert_persisted_report(db_session, fixture)
|
|
|
|
|
|
def test_gmail_fixture_email_shape_matches_api_raw_encoding():
|
|
"""The fixture MIME shape can be decoded from Gmail's raw message payload form."""
|
|
fixture = DMARC_COMPATIBILITY_FIXTURES[2]
|
|
zip_bytes = _zip_xml(_fixture_bytes(fixture), fixture["filename"])
|
|
encoded = base64.urlsafe_b64encode(
|
|
_make_gmail_raw_attachment(fixture["filename"] + ".zip", zip_bytes)
|
|
)
|
|
|
|
decoded = base64.urlsafe_b64decode(encoded)
|
|
msg = email.message_from_bytes(decoded)
|
|
|
|
assert msg["Subject"] == "DMARC aggregate report"
|
|
assert any(part.get_filename() == fixture["filename"] + ".zip" for part in msg.walk())
|