Files
gh-christianlouis-dmarq/backend/app/services/imap_client.py
T
2026-05-22 19:29:49 +02:00

413 lines
15 KiB
Python

import email
import imaplib
import logging
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
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__( # pylint: disable=too-many-positional-arguments,too-many-arguments
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)
username: IMAP username (if None, uses settings)
password: IMAP password (if None, uses settings)
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 _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: # pylint: disable=broad-exception-caught
# 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
Returns:
Tuple of (success, message, stats)
- success: Boolean indicating if connection was successful
- message: String message describing the result
- stats: Dictionary with mailbox statistics (if successful)
"""
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 = self._list_mailboxes(mailbox_list) if status == "OK" else []
# Select inbox and get message count
status, data = mail.select("INBOX")
message_count = 0
unread_count = 0
if status == "OK":
message_count = int(data[0])
# Count unread messages
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":
dmarc_count = len(data[0].split())
# Close connection
mail.close()
mail.logout()
stats = {
"message_count": message_count,
"unread_count": unread_count,
"dmarc_count": dmarc_count,
"available_mailboxes": available_mailboxes,
"server": self.server,
"port": self.port,
"timestamp": datetime.now().isoformat(),
}
return True, "Connection successful", stats
except Exception as e: # pylint: disable=broad-exception-caught
logger.error("IMAP connection test failed: %s", str(e))
return False, "Connection failed. Check server address and credentials.", {}
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("Error fetching email ID %s", 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)
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: # pylint: disable=broad-exception-caught
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
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}
stats = {
"success": True,
"processed": 0,
"reports_found": 0,
"duplicate_reports": 0,
"new_domains": [],
"errors": [],
}
try:
# Connect to the mail server
mail = imaplib.IMAP4_SSL(self.server, self.port)
mail.login(self.username, self.password)
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})"
status, data = mail.search(None, search_criteria)
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:
self._process_single_email(mail, email_id, stats)
# 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: # pylint: disable=broad-exception-caught
logger.error("Error fetching DMARC reports: %s", str(e))
return {
"success": False,
"error": "Error connecting to mailbox. Check server logs for details.",
"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
"""
# Get email subject
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",
]
# Common senders of DMARC reports
dmarc_senders = [
"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
"""
decoded_parts = []
for text, encoding in decode_header(header):
if isinstance(text, bytes):
if encoding:
decoded_parts.append(text.decode(encoding or "utf-8", errors="replace"))
else:
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":
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")
):
return True
# Check content type
content_type = part.get_content_type()
if content_type in (
"application/zip",
"application/gzip",
"application/x-gzip",
"application/xml",
"text/xml",
):
return True
return False
def _process_attachments(self, msg: email.message.Message, stats: dict | None = None) -> 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":
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")
):
try:
# Get attachment content
content = part.get_payload(decode=True)
# Parse the DMARC report
report = DMARCParser.parse_file(content, filename)
domain = report.get("domain", "unknown")
report_id = report.get("report_id", "")
if report_id and self.report_store.has_report(domain, report_id):
logger.info(
"Skipping duplicate DMARC report %s for %s",
report_id,
domain,
)
if stats is not None:
stats["duplicate_reports"] = (
stats.get("duplicate_reports", 0) + 1
)
continue
# Add the report to the store
self.report_store.add_report(report)
reports_found += 1
logger.info("Successfully processed DMARC report: %s", filename)
except Exception as e: # pylint: disable=broad-exception-caught
logger.error("Error processing attachment %s: %s", filename, str(e))
return reports_found