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:
@@ -35,6 +35,77 @@ ALLOWED_MIME_TYPES = {
|
||||
ALLOWED_EXTENSIONS = {".xml", ".zip", ".gz", ".gzip"}
|
||||
|
||||
|
||||
def _validate_mime_type(file_content: bytes) -> None:
|
||||
"""Validate the MIME type of the uploaded file using python-magic.
|
||||
|
||||
No-ops silently when python-magic is unavailable.
|
||||
Raises HTTPException on a disallowed MIME type.
|
||||
"""
|
||||
if not HAS_MAGIC:
|
||||
logger.debug("MIME type validation skipped (python-magic not available)")
|
||||
return
|
||||
try:
|
||||
mime_type = magic.from_buffer(file_content, mime=True)
|
||||
if mime_type not in ALLOWED_MIME_TYPES:
|
||||
logger.warning(f"Rejected file with MIME type: {mime_type}")
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Invalid file type. File must be XML, ZIP, or GZIP format.",
|
||||
)
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
# If magic fails, log but continue (fallback to extension check)
|
||||
logger.warning(f"MIME type detection failed: {str(e)}")
|
||||
|
||||
|
||||
def _validate_upload_file(file: UploadFile, file_content: bytes) -> None:
|
||||
"""Run all pre-parse validation checks on an uploaded file.
|
||||
|
||||
Raises HTTPException for any validation failure.
|
||||
"""
|
||||
# Security: Validate filename is provided
|
||||
if not file.filename:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST, detail="Filename is required"
|
||||
)
|
||||
|
||||
# Security: Validate file extension
|
||||
file_ext = "." + file.filename.rsplit(".", 1)[-1].lower() if "." in file.filename else ""
|
||||
if file_ext not in ALLOWED_EXTENSIONS:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=f"Invalid file type. Allowed types: {', '.join(ALLOWED_EXTENSIONS)}",
|
||||
)
|
||||
|
||||
# Security: Validate file is not empty
|
||||
if len(file_content) == 0:
|
||||
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="File is empty")
|
||||
|
||||
# Security: Validate MIME type (if python-magic is available)
|
||||
_validate_mime_type(file_content)
|
||||
|
||||
|
||||
def _handle_upload_value_error(filename: str, error_message: str) -> None:
|
||||
"""Translate a parser ValueError into a sanitized HTTPException.
|
||||
|
||||
Always raises — never returns.
|
||||
"""
|
||||
logger.error(f"ValueError processing report {filename}: {error_message}")
|
||||
if "too large" in error_message.lower():
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE, detail="File too large"
|
||||
)
|
||||
elif "zip bomb" in error_message.lower():
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid archive file"
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid report format"
|
||||
)
|
||||
|
||||
|
||||
class UploadResponse(BaseModel):
|
||||
"""Response model for report upload"""
|
||||
|
||||
@@ -89,42 +160,9 @@ async def upload_report(file: UploadFile = File(...)):
|
||||
- Sanitized error messages
|
||||
"""
|
||||
try:
|
||||
# Security: Validate filename is provided
|
||||
if not file.filename:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST, detail="Filename is required"
|
||||
)
|
||||
|
||||
# Security: Validate file extension
|
||||
file_ext = "." + file.filename.rsplit(".", 1)[-1].lower() if "." in file.filename else ""
|
||||
if file_ext not in ALLOWED_EXTENSIONS:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=f"Invalid file type. Allowed types: {', '.join(ALLOWED_EXTENSIONS)}",
|
||||
)
|
||||
|
||||
# Read the file content
|
||||
# Read content first so validators can inspect it
|
||||
file_content = await file.read()
|
||||
|
||||
# Security: Validate file is not empty
|
||||
if len(file_content) == 0:
|
||||
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="File is empty")
|
||||
|
||||
# Security: Validate MIME type using python-magic (if available)
|
||||
if HAS_MAGIC:
|
||||
try:
|
||||
mime_type = magic.from_buffer(file_content, mime=True)
|
||||
if mime_type not in ALLOWED_MIME_TYPES:
|
||||
logger.warning(f"Rejected file with MIME type: {mime_type}")
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Invalid file type. File must be XML, ZIP, or GZIP format.",
|
||||
)
|
||||
except Exception as e:
|
||||
# If magic fails, log but continue (fallback to extension check)
|
||||
logger.warning(f"MIME type detection failed: {str(e)}")
|
||||
else:
|
||||
logger.debug("MIME type validation skipped (python-magic not available)")
|
||||
_validate_upload_file(file, file_content)
|
||||
|
||||
# Parse the report
|
||||
parser = DMARCParser()
|
||||
@@ -141,7 +179,6 @@ async def upload_report(file: UploadFile = File(...)):
|
||||
# Validate domain format (not DNS resolution to avoid external calls)
|
||||
is_valid, error_msg, error_code = validate_domain(domain, check_dns=False)
|
||||
if not is_valid and error_code != DomainValidationError.DNS_RESOLUTION_FAILED:
|
||||
# Allow domains that fail DNS resolution but have valid format
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=f"Invalid domain in report: {error_msg}",
|
||||
@@ -161,26 +198,10 @@ async def upload_report(file: UploadFile = File(...)):
|
||||
)
|
||||
|
||||
except HTTPException:
|
||||
# Re-raise HTTP exceptions as-is
|
||||
raise
|
||||
except ValueError as e:
|
||||
# Security: Sanitize error messages from parser
|
||||
error_message = str(e)
|
||||
# Log full error for debugging
|
||||
logger.error(f"ValueError processing report {file.filename}: {error_message}")
|
||||
# Return sanitized message
|
||||
if "too large" in error_message.lower():
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE, detail="File too large"
|
||||
)
|
||||
elif "zip bomb" in error_message.lower():
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid archive file"
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid report format"
|
||||
)
|
||||
_handle_upload_value_error(file.filename, str(e))
|
||||
except Exception as e:
|
||||
# Security: Don't expose internal errors to client
|
||||
logger.error(f"Unexpected error processing report {file.filename}: {str(e)}")
|
||||
|
||||
@@ -148,7 +148,6 @@ def create_app() -> FastAPI:
|
||||
@app.on_event("shutdown")
|
||||
async def shutdown_event():
|
||||
"""Clean up background tasks on application shutdown"""
|
||||
global background_task
|
||||
if background_task:
|
||||
logger.info("Cancelling IMAP polling background task")
|
||||
background_task.cancel()
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -17,6 +17,45 @@ class DomainValidationError:
|
||||
DNS_RESOLUTION_FAILED = "dns_resolution_failed"
|
||||
|
||||
|
||||
def _validate_domain_characters(
|
||||
domain_name: str,
|
||||
) -> Tuple[bool, Optional[str], Optional[str]]:
|
||||
"""Check a domain name for whitespace and suspicious characters."""
|
||||
if " " in domain_name or "\t" in domain_name or "\n" in domain_name:
|
||||
return (
|
||||
False,
|
||||
"Domain name cannot contain whitespace",
|
||||
DomainValidationError.INVALID_CHARACTERS,
|
||||
)
|
||||
if any(char in domain_name for char in ["<", ">", '"', "'", "\\", "|", ";", "&", "$", "`"]):
|
||||
return (
|
||||
False,
|
||||
"Domain name contains invalid characters",
|
||||
DomainValidationError.INVALID_CHARACTERS,
|
||||
)
|
||||
return True, None, None
|
||||
|
||||
|
||||
def _validate_domain_labels(
|
||||
labels: list,
|
||||
) -> Tuple[bool, Optional[str], Optional[str]]:
|
||||
"""Check each DNS label for length and hyphen-placement rules."""
|
||||
for label in labels:
|
||||
if len(label) > 63:
|
||||
return (
|
||||
False,
|
||||
f"Domain label too long: '{label}' (max 63 characters per label)",
|
||||
DomainValidationError.LABEL_TOO_LONG,
|
||||
)
|
||||
if label.startswith("-") or label.endswith("-"):
|
||||
return (
|
||||
False,
|
||||
f"Domain label cannot start or end with hyphen: '{label}'",
|
||||
DomainValidationError.INVALID_LABEL,
|
||||
)
|
||||
return True, None, None
|
||||
|
||||
|
||||
def validate_domain(
|
||||
domain_name: str, check_dns: bool = True
|
||||
) -> Tuple[bool, Optional[str], Optional[str]]:
|
||||
@@ -41,21 +80,10 @@ def validate_domain(
|
||||
if len(domain_name) > 253:
|
||||
return False, "Domain name too long (max 253 characters)", DomainValidationError.TOO_LONG
|
||||
|
||||
# Security: Check for whitespace
|
||||
if " " in domain_name or "\t" in domain_name or "\n" in domain_name:
|
||||
return (
|
||||
False,
|
||||
"Domain name cannot contain whitespace",
|
||||
DomainValidationError.INVALID_CHARACTERS,
|
||||
)
|
||||
|
||||
# Security: Check for suspicious characters
|
||||
if any(char in domain_name for char in ["<", ">", '"', "'", "\\", "|", ";", "&", "$", "`"]):
|
||||
return (
|
||||
False,
|
||||
"Domain name contains invalid characters",
|
||||
DomainValidationError.INVALID_CHARACTERS,
|
||||
)
|
||||
# Security: Check for whitespace and suspicious characters
|
||||
char_ok, char_msg, char_code = _validate_domain_characters(domain_name)
|
||||
if not char_ok:
|
||||
return False, char_msg, char_code
|
||||
|
||||
# Check domain format with regex
|
||||
# This regex allows domain names with alphanumeric characters, hyphens,
|
||||
@@ -67,19 +95,9 @@ def validate_domain(
|
||||
|
||||
# Security: Check each label length (max 63 characters per label)
|
||||
labels = domain_name.split(".")
|
||||
for label in labels:
|
||||
if len(label) > 63:
|
||||
return (
|
||||
False,
|
||||
f"Domain label too long: '{label}' (max 63 characters per label)",
|
||||
DomainValidationError.LABEL_TOO_LONG,
|
||||
)
|
||||
if label.startswith("-") or label.endswith("-"):
|
||||
return (
|
||||
False,
|
||||
f"Domain label cannot start or end with hyphen: '{label}'",
|
||||
DomainValidationError.INVALID_LABEL,
|
||||
)
|
||||
label_ok, label_msg, label_code = _validate_domain_labels(labels)
|
||||
if not label_ok:
|
||||
return False, label_msg, label_code
|
||||
|
||||
# Check if domain exists by attempting to resolve DNS (optional)
|
||||
if check_dns:
|
||||
|
||||
Reference in New Issue
Block a user