Fix flake8 C901 complexity violations and F824 unused global in 5 files

Agent-Logs-Url: https://github.com/christianlouis/dmarq/sessions/e58d59de-a79d-48e1-a4fc-ffb61b868593

Co-authored-by: christianlouis <361235+christianlouis@users.noreply.github.com>
This commit is contained in:
copilot-swe-agent[bot]
2026-03-29 10:58:33 +00:00
parent 7752479444
commit 271b29186a
5 changed files with 312 additions and 257 deletions
+140 -123
View File
@@ -59,6 +59,48 @@ class DMARCParser:
# Parse the XML content
return DMARCParser._parse_xml(xml_content)
@staticmethod
def _extract_from_zip(file_content: bytes) -> Optional[bytes]:
"""Extract the first XML file from a ZIP archive.
Raises:
ValueError: If the archive exceeds size/count security limits.
"""
try:
with zipfile.ZipFile(io.BytesIO(file_content)) as z:
file_list = z.infolist()
# Security: Check number of files in archive
if len(file_list) > MAX_FILES_IN_ARCHIVE:
raise ValueError(
f"ZIP archive contains too many files ({len(file_list)}). "
f"Maximum is {MAX_FILES_IN_ARCHIVE}."
)
# Security: Check for zip bomb by examining compression ratios
total_uncompressed = sum(f.file_size for f in file_list)
if total_uncompressed > MAX_UNCOMPRESSED_SIZE:
raise ValueError(
f"ZIP archive uncompressed size too large "
f"({total_uncompressed / (1024*1024):.1f} MB). "
f"Maximum is {MAX_UNCOMPRESSED_SIZE / (1024*1024):.1f} MB. "
"Possible zip bomb attack detected."
)
# Find the first XML file in the archive
for file_info in file_list:
if file_info.filename.lower().endswith(".xml"):
# Security: Double-check individual file size
if file_info.file_size > MAX_UNCOMPRESSED_SIZE:
raise ValueError(
f"XML file in archive too large "
f"({file_info.file_size / (1024*1024):.1f} MB)"
)
return z.read(file_info.filename)
except zipfile.BadZipFile:
pass
return None
@staticmethod
def _extract_xml_content(file_content: bytes, filename: str) -> Optional[bytes]:
"""
@@ -69,36 +111,9 @@ class DMARCParser:
"""
# Try to handle as ZIP file
if filename.lower().endswith(".zip"):
try:
with zipfile.ZipFile(io.BytesIO(file_content)) as z:
# Security: Check number of files in archive
file_list = z.infolist()
if len(file_list) > MAX_FILES_IN_ARCHIVE:
raise ValueError(
f"ZIP archive contains too many files ({len(file_list)}). "
f"Maximum is {MAX_FILES_IN_ARCHIVE}."
)
# Security: Check for zip bomb by examining compression ratios
total_uncompressed = sum(f.file_size for f in file_list)
if total_uncompressed > MAX_UNCOMPRESSED_SIZE:
raise ValueError(
f"ZIP archive uncompressed size too large ({total_uncompressed / (1024*1024):.1f} MB). "
f"Maximum is {MAX_UNCOMPRESSED_SIZE / (1024*1024):.1f} MB. "
"Possible zip bomb attack detected."
)
# Find the first XML file in the archive
for file_info in file_list:
if file_info.filename.lower().endswith(".xml"):
# Security: Double-check individual file size
if file_info.file_size > MAX_UNCOMPRESSED_SIZE:
raise ValueError(
f"XML file in archive too large ({file_info.file_size / (1024*1024):.1f} MB)"
)
return z.read(file_info.filename)
except zipfile.BadZipFile:
pass
result = DMARCParser._extract_from_zip(file_content)
if result is not None:
return result
# Try to handle as GZIP file
if filename.lower().endswith(".gz") or filename.lower().endswith(".gzip"):
@@ -113,6 +128,87 @@ class DMARCParser:
return None
@staticmethod
def _parse_metadata(root) -> dict:
"""Parse the report_metadata section of a DMARC XML report."""
report: dict = {}
metadata = root.find("report_metadata")
if metadata is not None:
report["report_id"] = metadata.findtext("report_id", "")
report["org_name"] = metadata.findtext("org_name", "")
report["email"] = metadata.findtext("email", "")
date_range = metadata.find("date_range")
if date_range is not None:
begin_ts = int(date_range.findtext("begin", 0))
end_ts = int(date_range.findtext("end", 0))
report["begin_date"] = datetime.fromtimestamp(begin_ts).isoformat()
report["end_date"] = datetime.fromtimestamp(end_ts).isoformat()
report["begin_timestamp"] = begin_ts
report["end_timestamp"] = end_ts
return report
@staticmethod
def _parse_record(record_elem) -> dict:
"""Parse a single <record> element into a dictionary."""
record: dict = {}
row = record_elem.find("row")
if row is not None:
record["source_ip"] = row.findtext("source_ip", "")
record["count"] = int(row.findtext("count", 0))
policy_evaluated = row.find("policy_evaluated")
if policy_evaluated is not None:
record["disposition"] = policy_evaluated.findtext("disposition", "none")
record["dkim_result"] = policy_evaluated.findtext("dkim", "").lower()
record["spf_result"] = policy_evaluated.findtext("spf", "").lower()
identifiers = record_elem.find("identifiers")
if identifiers is not None:
record["header_from"] = identifiers.findtext("header_from", "")
auth_results = record_elem.find("auth_results")
if auth_results is not None:
spf_entries = [
{
"domain": spf.findtext("domain", ""),
"result": spf.findtext("result", "").lower(),
}
for spf in auth_results.findall("spf")
]
if spf_entries:
record["spf"] = spf_entries
dkim_entries = [
{
"domain": dkim.findtext("domain", ""),
"result": dkim.findtext("result", "").lower(),
"selector": dkim.findtext("selector", ""),
}
for dkim in auth_results.findall("dkim")
]
if dkim_entries:
record["dkim"] = dkim_entries
return record
@staticmethod
def _compute_summary(records: list) -> dict:
"""Compute aggregate pass/fail statistics for a list of records."""
total_count = sum(r["count"] for r in records)
passed_count = sum(
r["count"]
for r in records
if r.get("spf_result") == "pass" or r.get("dkim_result") == "pass"
)
failed_count = total_count - passed_count
return {
"total_count": total_count,
"passed_count": passed_count,
"failed_count": failed_count,
"pass_rate": (passed_count / total_count * 100) if total_count > 0 else 0,
}
@staticmethod
def _parse_xml(xml_content: bytes) -> Dict[str, Any]:
"""
@@ -120,24 +216,8 @@ class DMARCParser:
"""
try:
root = ET.fromstring(xml_content)
report = {}
# Parse report metadata
metadata = root.find("report_metadata")
if metadata is not None:
report["report_id"] = metadata.findtext("report_id", "")
report["org_name"] = metadata.findtext("org_name", "")
report["email"] = metadata.findtext("email", "")
# Parse date range
date_range = metadata.find("date_range")
if date_range is not None:
begin_ts = int(date_range.findtext("begin", 0))
end_ts = int(date_range.findtext("end", 0))
report["begin_date"] = datetime.fromtimestamp(begin_ts).isoformat()
report["end_date"] = datetime.fromtimestamp(end_ts).isoformat()
report["begin_timestamp"] = begin_ts
report["end_timestamp"] = end_ts
report = DMARCParser._parse_metadata(root)
# Parse policy published
policy = root.find("policy_published")
@@ -150,89 +230,26 @@ class DMARCParser:
}
# Parse records
records = []
for record_elem in root.findall("record"):
record = {}
# Parse row
row = record_elem.find("row")
if row is not None:
record["source_ip"] = row.findtext("source_ip", "")
record["count"] = int(row.findtext("count", 0))
policy_evaluated = row.find("policy_evaluated")
if policy_evaluated is not None:
record["disposition"] = policy_evaluated.findtext("disposition", "none")
record["dkim_result"] = policy_evaluated.findtext("dkim", "").lower()
record["spf_result"] = policy_evaluated.findtext("spf", "").lower()
# Parse identifiers
identifiers = record_elem.find("identifiers")
if identifiers is not None:
record["header_from"] = identifiers.findtext("header_from", "")
# Parse auth results
auth_results = record_elem.find("auth_results")
if auth_results is not None:
# SPF results
spf_entries = []
for spf in auth_results.findall("spf"):
spf_entries.append(
{
"domain": spf.findtext("domain", ""),
"result": spf.findtext("result", "").lower(),
}
)
if spf_entries:
record["spf"] = spf_entries
# DKIM results
dkim_entries = []
for dkim in auth_results.findall("dkim"):
dkim_entries.append(
{
"domain": dkim.findtext("domain", ""),
"result": dkim.findtext("result", "").lower(),
"selector": dkim.findtext("selector", ""),
}
)
if dkim_entries:
record["dkim"] = dkim_entries
records.append(record)
records = [DMARCParser._parse_record(elem) for elem in root.findall("record")]
report["records"] = records
# Calculate summary stats
total_count = sum(r["count"] for r in records)
# Count records that pass either SPF or DKIM (or both)
passed_count = sum(
r["count"]
for r in records
if r.get("spf_result") == "pass" or r.get("dkim_result") == "pass"
)
failed_count = total_count - passed_count
report["summary"] = DMARCParser._compute_summary(records)
# Log parse results for debugging
total_count = report["summary"]["total_count"]
logger.info(f"Parsed DMARC report for domain: {report.get('domain')}")
logger.info(f"Found {len(records)} record entries with {total_count} total messages")
logger.info(f"Messages passed: {passed_count}, failed: {failed_count}")
if len(records) > 0:
# Log the first record for debugging
logger.info(
f"Found {len(records)} record entries with {total_count} total messages"
)
logger.info(
f"Messages passed: {report['summary']['passed_count']}, "
f"failed: {report['summary']['failed_count']}"
)
if records:
logger.info(
f"Sample record - SPF: {records[0].get('spf_result')}, DKIM: {records[0].get('dkim_result')}"
f"Sample record - SPF: {records[0].get('spf_result')}, "
f"DKIM: {records[0].get('dkim_result')}"
)
report["summary"] = {
"total_count": total_count,
"passed_count": passed_count,
"failed_count": failed_count,
"pass_rate": (passed_count / total_count * 100) if total_count > 0 else 0,
}
return report
except Exception as e:
+52 -52
View File
@@ -49,6 +49,30 @@ class IMAPClient:
if not all([self.server, self.username, self.password]):
logger.warning("IMAP credentials not fully configured")
def _list_mailboxes(self, mailbox_data: list) -> list:
"""Parse the raw IMAP LIST response into a list of mailbox name strings."""
available_mailboxes = []
for mailbox in mailbox_data:
if isinstance(mailbox, bytes):
try:
mailbox_str = mailbox.decode("utf-8")
# Extract the mailbox name (after the last quote)
parts = mailbox_str.split('"')
if len(parts) > 2:
mailbox_name = parts[-1].strip()
if mailbox_name.startswith(" "):
mailbox_name = mailbox_name[1:]
available_mailboxes.append(mailbox_name)
except Exception:
# Silently skip mailboxes that can't be parsed; they are simply
# omitted from the returned list so callers should expect it may
# be incomplete. Some IMAP servers return non-standard list
# responses or use different delimiters/encodings that don't follow
# RFC 3501 (special characters, non-UTF-8 encodings, malformed
# responses). This is expected behaviour and not a critical error.
pass # nosec B110
return available_mailboxes
def test_connection(self) -> Tuple[bool, str, Dict[str, Any]]:
"""
Test the IMAP connection and gather basic mailbox statistics
@@ -70,28 +94,7 @@ class IMAPClient:
# List available mailboxes
status, mailbox_list = mail.list()
available_mailboxes = []
if status == "OK":
for mailbox in mailbox_list:
if isinstance(mailbox, bytes):
try:
# Extract mailbox name from response
mailbox_str = mailbox.decode("utf-8")
# Extract the mailbox name (after the last quote)
parts = mailbox_str.split('"')
if len(parts) > 2:
mailbox_name = parts[-1].strip()
if mailbox_name.startswith(" "):
mailbox_name = mailbox_name[1:]
available_mailboxes.append(mailbox_name)
except Exception:
# Silently skip mailboxes that can't be parsed
# Some IMAP servers return non-standard list responses or
# use different delimiters/encodings that don't follow RFC 3501
# Common cases: special characters, non-UTF8 encodings, malformed responses
# This is expected behavior and not a critical error
pass # nosec B110
available_mailboxes = self._list_mailboxes(mailbox_list) if status == "OK" else []
# Select inbox and get message count
status, data = mail.select("INBOX")
@@ -131,6 +134,32 @@ class IMAPClient:
logger.error(f"IMAP connection test failed: {str(e)}")
return False, f"Connection failed: {str(e)}", {}
def _process_single_email(self, mail, email_id: bytes, stats: dict) -> None:
"""Fetch, parse, and store DMARC attachments from one email message."""
try:
status, msg_data = mail.fetch(email_id, "(RFC822)")
if status != "OK":
logger.error(f"Error fetching email ID {email_id}")
return
raw_email = msg_data[0][1]
msg = email.message_from_bytes(raw_email)
if self._is_dmarc_report_email(msg):
reports_found = self._process_attachments(msg)
stats["reports_found"] += reports_found
# Mark email as read (and optionally delete)
mail.store(email_id, "+FLAGS", "\\Seen")
if self.delete_emails:
mail.store(email_id, "+FLAGS", "\\Deleted")
stats["processed"] += 1
except Exception as e:
error_msg = f"Error processing email ID {email_id}: {str(e)}"
logger.error(error_msg)
stats["errors"].append(error_msg)
def fetch_reports(self, days: int = 7) -> Dict[str, Any]:
"""
Fetch and process DMARC reports from the configured mailbox
@@ -181,36 +210,7 @@ class IMAPClient:
# Process each email
for email_id in email_ids:
try:
# Fetch the email
status, msg_data = mail.fetch(email_id, "(RFC822)")
if status != "OK":
logger.error(f"Error fetching email ID {email_id}")
continue
# Parse the email
raw_email = msg_data[0][1]
msg = email.message_from_bytes(raw_email)
# Check if this email might contain DMARC reports
if self._is_dmarc_report_email(msg):
# Process attachments
reports_found = self._process_attachments(msg)
stats["reports_found"] += reports_found
# Mark email as read
mail.store(email_id, "+FLAGS", "\\Seen")
# Delete email if configured
if self.delete_emails:
mail.store(email_id, "+FLAGS", "\\Deleted")
stats["processed"] += 1
except Exception as e:
error_msg = f"Error processing email ID {email_id}: {str(e)}"
logger.error(error_msg)
stats["errors"].append(error_msg)
self._process_single_email(mail, email_id, stats)
# Actually remove emails marked for deletion
if self.delete_emails: