Fix code formatting and linting issues
- Auto-format all Python files with black and isort - Remove unused imports with autoflake - Fix flake8 issues (missing newlines, blank lines, etc.) - Fix nonlocal/global scope issues in main.py - Fix security.py import order (E402) - Remove f-string without placeholders - Add nosec comment for intentional exception handling - Fix test imports to match refactored DMARCParser API Co-authored-by: christianlouis <361235+christianlouis@users.noreply.github.com>
This commit is contained in:
copilot-swe-agent[bot]
@@ -1,11 +1,11 @@
|
||||
import os
|
||||
import zipfile
|
||||
import gzip
|
||||
import io
|
||||
from datetime import datetime
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
import defusedxml.ElementTree as ET
|
||||
import logging
|
||||
import zipfile
|
||||
from datetime import datetime
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import defusedxml.ElementTree as ET
|
||||
|
||||
# Set up logging
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
@@ -16,35 +16,38 @@ MAX_FILE_SIZE = 10 * 1024 * 1024 # 10 MB
|
||||
MAX_UNCOMPRESSED_SIZE = 100 * 1024 * 1024 # 100 MB for zip bomb protection
|
||||
MAX_FILES_IN_ARCHIVE = 10 # Maximum number of files in a zip archive
|
||||
|
||||
|
||||
class DMARCParser:
|
||||
"""
|
||||
Parser for DMARC Aggregate Reports (XML format)
|
||||
"""
|
||||
|
||||
|
||||
@staticmethod
|
||||
def parse_file(file_content: bytes, filename: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Parse a DMARC report file (XML, zip, or gzip) into a dictionary
|
||||
|
||||
|
||||
Args:
|
||||
file_content: The binary content of the file
|
||||
filename: The name of the file (used to determine type)
|
||||
|
||||
|
||||
Returns:
|
||||
Dict containing the parsed report data
|
||||
|
||||
|
||||
Raises:
|
||||
ValueError: If file is invalid, too large, or potentially malicious
|
||||
"""
|
||||
# Security: Check file size
|
||||
if len(file_content) > MAX_FILE_SIZE:
|
||||
raise ValueError(f"File too large. Maximum size is {MAX_FILE_SIZE / (1024*1024):.1f} MB")
|
||||
|
||||
raise ValueError(
|
||||
f"File too large. Maximum size is {MAX_FILE_SIZE / (1024*1024):.1f} MB"
|
||||
)
|
||||
|
||||
# Determine file type and extract XML content
|
||||
xml_content = DMARCParser._extract_xml_content(file_content, filename)
|
||||
if not xml_content:
|
||||
raise ValueError("Could not extract XML content from file")
|
||||
|
||||
|
||||
# Security: Check uncompressed XML size
|
||||
if len(xml_content) > MAX_UNCOMPRESSED_SIZE:
|
||||
raise ValueError(
|
||||
@@ -52,20 +55,20 @@ class DMARCParser:
|
||||
f"Maximum is {MAX_UNCOMPRESSED_SIZE / (1024*1024):.1f} MB. "
|
||||
"Possible zip bomb attack detected."
|
||||
)
|
||||
|
||||
|
||||
# Parse the XML content
|
||||
return DMARCParser._parse_xml(xml_content)
|
||||
|
||||
|
||||
@staticmethod
|
||||
def _extract_xml_content(file_content: bytes, filename: str) -> Optional[bytes]:
|
||||
"""
|
||||
Extract XML content from various file formats (ZIP, GZIP, or plain XML)
|
||||
|
||||
|
||||
Raises:
|
||||
ValueError: If archive contains too many files or is potentially malicious
|
||||
"""
|
||||
# Try to handle as ZIP file
|
||||
if filename.lower().endswith('.zip'):
|
||||
if filename.lower().endswith(".zip"):
|
||||
try:
|
||||
with zipfile.ZipFile(io.BytesIO(file_content)) as z:
|
||||
# Security: Check number of files in archive
|
||||
@@ -75,7 +78,7 @@ class DMARCParser:
|
||||
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:
|
||||
@@ -84,10 +87,10 @@ class DMARCParser:
|
||||
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'):
|
||||
if file_info.filename.lower().endswith(".xml"):
|
||||
# Security: Double-check individual file size
|
||||
if file_info.file_size > MAX_UNCOMPRESSED_SIZE:
|
||||
raise ValueError(
|
||||
@@ -96,20 +99,20 @@ class DMARCParser:
|
||||
return z.read(file_info.filename)
|
||||
except zipfile.BadZipFile:
|
||||
pass
|
||||
|
||||
|
||||
# Try to handle as GZIP file
|
||||
if filename.lower().endswith('.gz') or filename.lower().endswith('.gzip'):
|
||||
if filename.lower().endswith(".gz") or filename.lower().endswith(".gzip"):
|
||||
try:
|
||||
return gzip.decompress(file_content)
|
||||
except gzip.BadGzipFile:
|
||||
pass
|
||||
|
||||
|
||||
# Assume it's plain XML
|
||||
if filename.lower().endswith('.xml'):
|
||||
if filename.lower().endswith(".xml"):
|
||||
return file_content
|
||||
|
||||
|
||||
return None
|
||||
|
||||
|
||||
@staticmethod
|
||||
def _parse_xml(xml_content: bytes) -> Dict[str, Any]:
|
||||
"""
|
||||
@@ -118,14 +121,14 @@ 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:
|
||||
@@ -135,7 +138,7 @@ class DMARCParser:
|
||||
report["end_date"] = datetime.fromtimestamp(end_ts).isoformat()
|
||||
report["begin_timestamp"] = begin_ts
|
||||
report["end_timestamp"] = end_ts
|
||||
|
||||
|
||||
# Parse policy published
|
||||
policy = root.find("policy_published")
|
||||
if policy is not None:
|
||||
@@ -145,84 +148,93 @@ class DMARCParser:
|
||||
"sp": policy.findtext("sp", ""),
|
||||
"pct": policy.findtext("pct", "100"),
|
||||
}
|
||||
|
||||
|
||||
# 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()
|
||||
})
|
||||
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", "")
|
||||
})
|
||||
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)
|
||||
|
||||
|
||||
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")
|
||||
|
||||
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
|
||||
|
||||
|
||||
# Log parse results for debugging
|
||||
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"Sample record - SPF: {records[0].get('spf_result')}, DKIM: {records[0].get('dkim_result')}")
|
||||
|
||||
logger.info(
|
||||
f"Sample record - SPF: {records[0].get('spf_result')}, 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
|
||||
"pass_rate": (passed_count / total_count * 100) if total_count > 0 else 0,
|
||||
}
|
||||
|
||||
|
||||
return report
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error parsing DMARC XML: {str(e)}")
|
||||
raise ValueError(f"Error parsing DMARC XML: {str(e)}")
|
||||
raise ValueError(f"Error parsing DMARC XML: {str(e)}")
|
||||
|
||||
+145
-129
@@ -1,11 +1,9 @@
|
||||
import imaplib
|
||||
import email
|
||||
import os
|
||||
import imaplib
|
||||
import logging
|
||||
import tempfile
|
||||
from email.header import decode_header
|
||||
from typing import List, Dict, Any, Optional, Tuple
|
||||
from datetime import datetime, timedelta
|
||||
from email.header import decode_header
|
||||
from typing import Any, Dict, Tuple
|
||||
|
||||
from app.core.config import get_settings
|
||||
from app.services.dmarc_parser import DMARCParser
|
||||
@@ -14,20 +12,23 @@ from app.services.report_store import ReportStore
|
||||
# Setup logger
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class IMAPClient:
|
||||
"""
|
||||
Client for retrieving DMARC reports from an IMAP mailbox
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
server: str = None,
|
||||
port: int = None,
|
||||
username: str = None,
|
||||
password: str = None,
|
||||
delete_emails: bool = False):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
server: str = None,
|
||||
port: int = None,
|
||||
username: str = None,
|
||||
password: str = None,
|
||||
delete_emails: bool = False,
|
||||
):
|
||||
"""
|
||||
Initialize the IMAP client with credentials
|
||||
|
||||
|
||||
Args:
|
||||
server: IMAP server hostname (if None, uses settings)
|
||||
port: IMAP server port (if None, uses settings)
|
||||
@@ -36,22 +37,22 @@ class IMAPClient:
|
||||
delete_emails: Whether to delete emails after processing (default: False)
|
||||
"""
|
||||
settings = get_settings()
|
||||
|
||||
|
||||
self.server = server or settings.IMAP_SERVER
|
||||
self.port = port or settings.IMAP_PORT
|
||||
self.username = username or settings.IMAP_USERNAME
|
||||
self.password = password or settings.IMAP_PASSWORD
|
||||
self.delete_emails = delete_emails
|
||||
|
||||
|
||||
self.report_store = ReportStore.get_instance()
|
||||
|
||||
|
||||
if not all([self.server, self.username, self.password]):
|
||||
logger.warning("IMAP credentials not fully configured")
|
||||
|
||||
|
||||
def test_connection(self) -> Tuple[bool, str, Dict[str, Any]]:
|
||||
"""
|
||||
Test the IMAP connection and gather basic mailbox statistics
|
||||
|
||||
|
||||
Returns:
|
||||
Tuple of (success, message, stats)
|
||||
- success: Boolean indicating if connection was successful
|
||||
@@ -60,56 +61,58 @@ class IMAPClient:
|
||||
"""
|
||||
if not all([self.server, self.username, self.password]):
|
||||
return False, "IMAP credentials not fully configured", {}
|
||||
|
||||
|
||||
try:
|
||||
# Create IMAP4 connection
|
||||
mail = imaplib.IMAP4_SSL(self.server, self.port)
|
||||
# Login
|
||||
mail.login(self.username, self.password)
|
||||
|
||||
|
||||
# List available mailboxes
|
||||
status, mailbox_list = mail.list()
|
||||
available_mailboxes = []
|
||||
|
||||
if status == 'OK':
|
||||
|
||||
if status == "OK":
|
||||
for mailbox in mailbox_list:
|
||||
if isinstance(mailbox, bytes):
|
||||
try:
|
||||
# Extract mailbox name from response
|
||||
mailbox_str = mailbox.decode('utf-8')
|
||||
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(' '):
|
||||
if mailbox_name.startswith(" "):
|
||||
mailbox_name = mailbox_name[1:]
|
||||
available_mailboxes.append(mailbox_name)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Silently skip mailboxes that can't be parsed
|
||||
# This is expected for some IMAP server responses
|
||||
pass # nosec B110
|
||||
|
||||
# Select inbox and get message count
|
||||
status, data = mail.select('INBOX')
|
||||
status, data = mail.select("INBOX")
|
||||
message_count = 0
|
||||
unread_count = 0
|
||||
|
||||
if status == 'OK':
|
||||
|
||||
if status == "OK":
|
||||
message_count = int(data[0])
|
||||
|
||||
|
||||
# Count unread messages
|
||||
status, data = mail.search(None, 'UNSEEN')
|
||||
if status == 'OK':
|
||||
status, data = mail.search(None, "UNSEEN")
|
||||
if status == "OK":
|
||||
unread_count = len(data[0].split())
|
||||
|
||||
|
||||
# Gather some stats about potential DMARC reports
|
||||
dmarc_count = 0
|
||||
status, data = mail.search(None, 'SUBJECT "DMARC"')
|
||||
if status == 'OK':
|
||||
if status == "OK":
|
||||
dmarc_count = len(data[0].split())
|
||||
|
||||
|
||||
# Close connection
|
||||
mail.close()
|
||||
mail.logout()
|
||||
|
||||
|
||||
stats = {
|
||||
"message_count": message_count,
|
||||
"unread_count": unread_count,
|
||||
@@ -117,127 +120,123 @@ class IMAPClient:
|
||||
"available_mailboxes": available_mailboxes,
|
||||
"server": self.server,
|
||||
"port": self.port,
|
||||
"timestamp": datetime.now().isoformat()
|
||||
"timestamp": datetime.now().isoformat(),
|
||||
}
|
||||
|
||||
|
||||
return True, "Connection successful", stats
|
||||
except Exception as e:
|
||||
logger.error(f"IMAP connection test failed: {str(e)}")
|
||||
return False, f"Connection failed: {str(e)}", {}
|
||||
|
||||
|
||||
def fetch_reports(self, days: int = 7) -> Dict[str, Any]:
|
||||
"""
|
||||
Fetch and process DMARC reports from the configured mailbox
|
||||
|
||||
|
||||
Args:
|
||||
days: Number of days to look back for emails
|
||||
|
||||
|
||||
Returns:
|
||||
Dictionary with stats about processing results
|
||||
"""
|
||||
if not all([self.server, self.username, self.password]):
|
||||
logger.error("IMAP credentials not fully configured")
|
||||
return {
|
||||
"success": False,
|
||||
"error": "IMAP credentials not configured",
|
||||
"processed": 0
|
||||
}
|
||||
|
||||
return {"success": False, "error": "IMAP credentials not configured", "processed": 0}
|
||||
|
||||
stats = {
|
||||
"success": True,
|
||||
"processed": 0,
|
||||
"reports_found": 0,
|
||||
"new_domains": [],
|
||||
"errors": []
|
||||
"errors": [],
|
||||
}
|
||||
|
||||
|
||||
try:
|
||||
# Connect to the mail server
|
||||
mail = imaplib.IMAP4_SSL(self.server, self.port)
|
||||
mail.login(self.username, self.password)
|
||||
mail.select('INBOX')
|
||||
|
||||
mail.select("INBOX")
|
||||
|
||||
# Calculate the date range for search
|
||||
date_since = (datetime.now() - timedelta(days=days)).strftime("%d-%b-%Y")
|
||||
|
||||
|
||||
# Search for all emails containing possible DMARC reports
|
||||
search_criteria = f'(SINCE {date_since})'
|
||||
search_criteria = f"(SINCE {date_since})"
|
||||
status, data = mail.search(None, search_criteria)
|
||||
|
||||
if status != 'OK':
|
||||
|
||||
if status != "OK":
|
||||
logger.error("Error searching mailbox")
|
||||
stats["success"] = False
|
||||
stats["error"] = "Error searching mailbox"
|
||||
mail.logout()
|
||||
return stats
|
||||
|
||||
|
||||
# Get list of email IDs
|
||||
email_ids = data[0].split()
|
||||
|
||||
|
||||
# Track domains before processing to identify new ones
|
||||
domains_before = set(self.report_store.get_domains())
|
||||
|
||||
|
||||
# Process each email
|
||||
for email_id in email_ids:
|
||||
try:
|
||||
# Fetch the email
|
||||
status, msg_data = mail.fetch(email_id, '(RFC822)')
|
||||
|
||||
if status != 'OK':
|
||||
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')
|
||||
|
||||
mail.store(email_id, "+FLAGS", "\\Seen")
|
||||
|
||||
# Delete email if configured
|
||||
if self.delete_emails:
|
||||
mail.store(email_id, '+FLAGS', '\\Deleted')
|
||||
|
||||
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)
|
||||
|
||||
|
||||
# Actually remove emails marked for deletion
|
||||
if self.delete_emails:
|
||||
mail.expunge()
|
||||
|
||||
|
||||
# Logout
|
||||
mail.logout()
|
||||
|
||||
|
||||
# Identify new domains
|
||||
domains_after = set(self.report_store.get_domains())
|
||||
stats["new_domains"] = list(domains_after - domains_before)
|
||||
|
||||
|
||||
return stats
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error fetching DMARC reports: {str(e)}")
|
||||
return {
|
||||
"success": False,
|
||||
"error": f"Error connecting to mailbox: {str(e)}",
|
||||
"processed": 0
|
||||
"success": False,
|
||||
"error": f"Error connecting to mailbox: {str(e)}",
|
||||
"processed": 0,
|
||||
}
|
||||
|
||||
|
||||
def _is_dmarc_report_email(self, msg: email.message.Message) -> bool:
|
||||
"""
|
||||
Check if an email likely contains DMARC reports
|
||||
|
||||
|
||||
Args:
|
||||
msg: Email message object
|
||||
|
||||
|
||||
Returns:
|
||||
True if the email is likely a DMARC report, False otherwise
|
||||
"""
|
||||
@@ -245,43 +244,54 @@ class IMAPClient:
|
||||
subject = ""
|
||||
if "Subject" in msg:
|
||||
subject = self._decode_email_header(msg["Subject"])
|
||||
|
||||
|
||||
# Get email from
|
||||
from_addr = ""
|
||||
if "From" in msg:
|
||||
from_addr = self._decode_email_header(msg["From"])
|
||||
|
||||
|
||||
# Common keywords in DMARC report emails
|
||||
dmarc_keywords = [
|
||||
"dmarc", "aggregate", "report", "rua",
|
||||
"authentication", "domain", "failure"
|
||||
"dmarc",
|
||||
"aggregate",
|
||||
"report",
|
||||
"rua",
|
||||
"authentication",
|
||||
"domain",
|
||||
"failure",
|
||||
]
|
||||
|
||||
|
||||
# Common senders of DMARC reports
|
||||
dmarc_senders = [
|
||||
"noreply@", "dmarc-noreply@", "postmaster@",
|
||||
"microsoft.com", "google.com", "yahoo.com",
|
||||
"hotmail.com", "outlook.com", "mail.ru"
|
||||
"noreply@",
|
||||
"dmarc-noreply@",
|
||||
"postmaster@",
|
||||
"microsoft.com",
|
||||
"google.com",
|
||||
"yahoo.com",
|
||||
"hotmail.com",
|
||||
"outlook.com",
|
||||
"mail.ru",
|
||||
]
|
||||
|
||||
|
||||
# Check if subject contains DMARC keywords
|
||||
if any(keyword in subject.lower() for keyword in dmarc_keywords):
|
||||
return True
|
||||
|
||||
|
||||
# Check if sender matches common DMARC report senders
|
||||
if any(sender in from_addr.lower() for sender in dmarc_senders):
|
||||
return True
|
||||
|
||||
|
||||
# Check for attachments with typical DMARC report filenames
|
||||
return self._has_dmarc_attachments(msg)
|
||||
|
||||
|
||||
def _decode_email_header(self, header: str) -> str:
|
||||
"""
|
||||
Decode an email header that might contain non-ASCII characters
|
||||
|
||||
|
||||
Args:
|
||||
header: Email header string
|
||||
|
||||
|
||||
Returns:
|
||||
Decoded header text
|
||||
"""
|
||||
@@ -289,90 +299,96 @@ class IMAPClient:
|
||||
for text, encoding in decode_header(header):
|
||||
if isinstance(text, bytes):
|
||||
if encoding:
|
||||
decoded_parts.append(text.decode(encoding or 'utf-8', errors='replace'))
|
||||
decoded_parts.append(text.decode(encoding or "utf-8", errors="replace"))
|
||||
else:
|
||||
decoded_parts.append(text.decode('utf-8', errors='replace'))
|
||||
decoded_parts.append(text.decode("utf-8", errors="replace"))
|
||||
else:
|
||||
decoded_parts.append(text)
|
||||
|
||||
|
||||
return " ".join(decoded_parts)
|
||||
|
||||
|
||||
def _has_dmarc_attachments(self, msg: email.message.Message) -> bool:
|
||||
"""
|
||||
Check if the email has attachments that might be DMARC reports
|
||||
|
||||
|
||||
Args:
|
||||
msg: Email message object
|
||||
|
||||
|
||||
Returns:
|
||||
True if the email has potential DMARC report attachments
|
||||
"""
|
||||
for part in msg.walk():
|
||||
content_disposition = part.get_content_disposition()
|
||||
if content_disposition == 'attachment':
|
||||
if content_disposition == "attachment":
|
||||
filename = part.get_filename()
|
||||
if filename:
|
||||
# Decode filename if needed
|
||||
filename = self._decode_email_header(filename)
|
||||
|
||||
|
||||
# Check file extension
|
||||
if (filename.lower().endswith('.xml') or
|
||||
filename.lower().endswith('.zip') or
|
||||
filename.lower().endswith('.gz') or
|
||||
filename.lower().endswith('.gzip')):
|
||||
if (
|
||||
filename.lower().endswith(".xml")
|
||||
or filename.lower().endswith(".zip")
|
||||
or filename.lower().endswith(".gz")
|
||||
or filename.lower().endswith(".gzip")
|
||||
):
|
||||
return True
|
||||
|
||||
|
||||
# Check content type
|
||||
content_type = part.get_content_type()
|
||||
if (content_type == 'application/zip' or
|
||||
content_type == 'application/gzip' or
|
||||
content_type == 'application/x-gzip' or
|
||||
content_type == 'application/xml' or
|
||||
content_type == 'text/xml'):
|
||||
if (
|
||||
content_type == "application/zip"
|
||||
or content_type == "application/gzip"
|
||||
or content_type == "application/x-gzip"
|
||||
or content_type == "application/xml"
|
||||
or content_type == "text/xml"
|
||||
):
|
||||
return True
|
||||
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def _process_attachments(self, msg: email.message.Message) -> int:
|
||||
"""
|
||||
Process email attachments that might be DMARC reports
|
||||
|
||||
|
||||
Args:
|
||||
msg: Email message object
|
||||
|
||||
|
||||
Returns:
|
||||
Number of DMARC reports found and processed
|
||||
"""
|
||||
reports_found = 0
|
||||
|
||||
|
||||
for part in msg.walk():
|
||||
content_disposition = part.get_content_disposition()
|
||||
|
||||
if content_disposition == 'attachment':
|
||||
|
||||
if content_disposition == "attachment":
|
||||
filename = part.get_filename()
|
||||
if filename:
|
||||
# Decode filename if needed
|
||||
filename = self._decode_email_header(filename)
|
||||
|
||||
|
||||
# Check if it's a likely DMARC report file
|
||||
if (filename.lower().endswith('.xml') or
|
||||
filename.lower().endswith('.zip') or
|
||||
filename.lower().endswith('.gz') or
|
||||
filename.lower().endswith('.gzip')):
|
||||
|
||||
if (
|
||||
filename.lower().endswith(".xml")
|
||||
or filename.lower().endswith(".zip")
|
||||
or filename.lower().endswith(".gz")
|
||||
or filename.lower().endswith(".gzip")
|
||||
):
|
||||
|
||||
try:
|
||||
# Get attachment content
|
||||
content = part.get_payload(decode=True)
|
||||
|
||||
|
||||
# Parse the DMARC report
|
||||
report = DMARCParser.parse_file(content, filename)
|
||||
|
||||
|
||||
# Add the report to the store
|
||||
self.report_store.add_report(report)
|
||||
|
||||
|
||||
reports_found += 1
|
||||
logger.info(f"Successfully processed DMARC report: {filename}")
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing attachment {filename}: {str(e)}")
|
||||
|
||||
return reports_found
|
||||
|
||||
return reports_found
|
||||
|
||||
@@ -1,18 +1,18 @@
|
||||
from typing import Dict, List, Any, Optional
|
||||
import threading
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
|
||||
class ReportStore:
|
||||
"""
|
||||
In-memory store for DMARC reports
|
||||
In-memory store for DMARC reports
|
||||
(for Milestone 1, will be replaced with database in Milestone 3)
|
||||
"""
|
||||
|
||||
|
||||
_instance = None
|
||||
_lock = threading.Lock()
|
||||
|
||||
|
||||
@classmethod
|
||||
def get_instance(cls) -> 'ReportStore':
|
||||
def get_instance(cls) -> "ReportStore":
|
||||
"""
|
||||
Get singleton instance of the report store
|
||||
"""
|
||||
@@ -21,7 +21,7 @@ class ReportStore:
|
||||
if cls._instance is None:
|
||||
cls._instance = ReportStore()
|
||||
return cls._instance
|
||||
|
||||
|
||||
def __init__(self):
|
||||
"""
|
||||
Initialize empty report store
|
||||
@@ -32,16 +32,16 @@ class ReportStore:
|
||||
self.domain_summary: Dict[str, Dict[str, Any]] = {}
|
||||
# Domain -> sources (sending IPs)
|
||||
self.domain_sources: Dict[str, Dict[str, Dict[str, Any]]] = {}
|
||||
|
||||
|
||||
def add_report(self, report: Dict[str, Any]) -> None:
|
||||
"""
|
||||
Add a new report to the store
|
||||
|
||||
|
||||
Args:
|
||||
report: Parsed DMARC report from DMARCParser
|
||||
"""
|
||||
domain = report.get("domain", "unknown")
|
||||
|
||||
|
||||
# Initialize data structures if this is a new domain
|
||||
if domain not in self.domain_reports:
|
||||
self.domain_reports[domain] = []
|
||||
@@ -52,21 +52,21 @@ class ReportStore:
|
||||
"reports_processed": 0,
|
||||
}
|
||||
self.domain_sources[domain] = {}
|
||||
|
||||
|
||||
# Add the new report
|
||||
self.domain_reports[domain].append(report)
|
||||
|
||||
|
||||
# Update summary stats for this domain
|
||||
summary = report.get("summary", {})
|
||||
self.domain_summary[domain]["total_count"] += summary.get("total_count", 0)
|
||||
self.domain_summary[domain]["passed_count"] += summary.get("passed_count", 0)
|
||||
self.domain_summary[domain]["failed_count"] += summary.get("failed_count", 0)
|
||||
self.domain_summary[domain]["reports_processed"] += 1
|
||||
|
||||
|
||||
# Set policy from the latest report
|
||||
if "policy" in report:
|
||||
self.domain_summary[domain]["policy"] = report["policy"]
|
||||
|
||||
|
||||
# Update source data
|
||||
report_records = report.get("records", [])
|
||||
for record in report_records:
|
||||
@@ -76,72 +76,71 @@ class ReportStore:
|
||||
"count": 0,
|
||||
"spf_result": "unknown",
|
||||
"dkim_result": "unknown",
|
||||
"disposition": "none"
|
||||
"disposition": "none",
|
||||
}
|
||||
|
||||
|
||||
# Update source counts and results
|
||||
self.domain_sources[domain][source_ip]["count"] += record.get("count", 0)
|
||||
self.domain_sources[domain][source_ip]["spf_result"] = record.get("spf", "unknown")
|
||||
self.domain_sources[domain][source_ip]["spf_result"] = record.get("spf", "unknown")
|
||||
self.domain_sources[domain][source_ip]["dkim_result"] = record.get("dkim", "unknown")
|
||||
self.domain_sources[domain][source_ip]["disposition"] = record.get("disposition", "none")
|
||||
|
||||
self.domain_sources[domain][source_ip]["disposition"] = record.get(
|
||||
"disposition", "none"
|
||||
)
|
||||
|
||||
# Calculate compliance rate (percentage of passing emails)
|
||||
if self.domain_summary[domain]["total_count"] > 0:
|
||||
pass_rate = (
|
||||
self.domain_summary[domain]["passed_count"] /
|
||||
self.domain_summary[domain]["total_count"] * 100
|
||||
self.domain_summary[domain]["passed_count"]
|
||||
/ self.domain_summary[domain]["total_count"]
|
||||
* 100
|
||||
)
|
||||
self.domain_summary[domain]["compliance_rate"] = round(pass_rate, 1)
|
||||
else:
|
||||
self.domain_summary[domain]["compliance_rate"] = 0
|
||||
|
||||
|
||||
def get_domains(self) -> List[str]:
|
||||
"""
|
||||
Get list of all domains with reports
|
||||
"""
|
||||
return list(self.domain_reports.keys())
|
||||
|
||||
|
||||
def get_domain_summary(self, domain: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Get summary statistics for a domain
|
||||
|
||||
|
||||
Args:
|
||||
domain: Domain name
|
||||
|
||||
|
||||
Returns:
|
||||
Dictionary with summary stats or empty dict if domain not found
|
||||
"""
|
||||
return self.domain_summary.get(domain, {})
|
||||
|
||||
|
||||
def get_all_domain_summaries(self) -> Dict[str, Dict[str, Any]]:
|
||||
"""
|
||||
Get summary statistics for all domains
|
||||
|
||||
|
||||
Returns:
|
||||
Dictionary mapping domain names to their summary stats
|
||||
"""
|
||||
return self.domain_summary
|
||||
|
||||
|
||||
def get_domain_reports(self, domain: str, limit: Optional[int] = None) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Get all reports for a domain
|
||||
|
||||
|
||||
Args:
|
||||
domain: Domain name
|
||||
limit: Optional limit on number of reports to return
|
||||
|
||||
|
||||
Returns:
|
||||
List of reports or empty list if domain not found
|
||||
"""
|
||||
reports = self.domain_reports.get(domain, [])
|
||||
|
||||
|
||||
# Sort reports by date (most recent first)
|
||||
sorted_reports = sorted(
|
||||
reports,
|
||||
key=lambda r: r.get("end_date", 0),
|
||||
reverse=True
|
||||
)
|
||||
|
||||
sorted_reports = sorted(reports, key=lambda r: r.get("end_date", 0), reverse=True)
|
||||
|
||||
# Calculate pass rate for each report
|
||||
for report in sorted_reports:
|
||||
total = report.get("summary", {}).get("total_count", 0)
|
||||
@@ -150,39 +149,36 @@ class ReportStore:
|
||||
report["pass_rate"] = round((passed / total) * 100, 1)
|
||||
else:
|
||||
report["pass_rate"] = 0
|
||||
|
||||
|
||||
# Apply limit if provided
|
||||
if limit is not None:
|
||||
return sorted_reports[:limit]
|
||||
return sorted_reports
|
||||
|
||||
|
||||
def get_domain_sources(self, domain: str, days: int = 30) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Get sending sources for a domain
|
||||
|
||||
|
||||
Args:
|
||||
domain: Domain name
|
||||
days: Number of days to look back
|
||||
|
||||
|
||||
Returns:
|
||||
List of source entries or empty list if domain not found
|
||||
"""
|
||||
if domain not in self.domain_sources:
|
||||
return []
|
||||
|
||||
|
||||
# For Milestone 1, we don't filter by date
|
||||
# In a future milestone, we'll add date-based filtering
|
||||
sources = []
|
||||
for ip, data in self.domain_sources[domain].items():
|
||||
source_entry = {
|
||||
"source_ip": ip,
|
||||
**data
|
||||
}
|
||||
source_entry = {"source_ip": ip, **data}
|
||||
sources.append(source_entry)
|
||||
|
||||
|
||||
# Sort sources by count (highest first)
|
||||
return sorted(sources, key=lambda s: s["count"], reverse=True)
|
||||
|
||||
|
||||
def clear(self) -> None:
|
||||
"""
|
||||
Clear all data in the store
|
||||
@@ -190,20 +186,20 @@ class ReportStore:
|
||||
self.domain_reports = {}
|
||||
self.domain_summary = {}
|
||||
self.domain_sources = {}
|
||||
|
||||
|
||||
def delete_domain_with_cleanup(self, domain: str) -> bool:
|
||||
"""
|
||||
Delete a domain and all its associated data
|
||||
|
||||
|
||||
Args:
|
||||
domain: Domain name to delete
|
||||
|
||||
|
||||
Returns:
|
||||
True if domain was deleted, False otherwise
|
||||
"""
|
||||
if domain not in self.domain_reports:
|
||||
return False
|
||||
|
||||
|
||||
try:
|
||||
# Remove all data for this domain
|
||||
self.domain_reports.pop(domain, None)
|
||||
@@ -212,4 +208,4 @@ class ReportStore:
|
||||
return True
|
||||
except Exception:
|
||||
# If any exception occurs during deletion, return False
|
||||
return False
|
||||
return False
|
||||
|
||||
Reference in New Issue
Block a user