diff --git a/backend/app/main.py b/backend/app/main.py index 667a89e..d2aa102 100644 --- a/backend/app/main.py +++ b/backend/app/main.py @@ -427,6 +427,113 @@ async def health(): return {"status": "ok", "service": "dmarq"} +# --------------------------------------------------------------------------- +# Helpers for the manual trigger-poll endpoint +# --------------------------------------------------------------------------- + + +def _trigger_poll_imap_source(source: MailSource, db) -> dict: + """Poll a single IMAP source and return a result dict for the API response.""" + global last_check_time # pylint: disable=global-statement + + imap_client = IMAPClient( + server=source.server, + port=source.port or 993, + username=source.username, + password=source.password, + delete_emails=False, + ) + results = imap_client.fetch_reports(days=7) + last_check_time = datetime.now() + source.last_checked = datetime.utcnow() + db.commit() + return { + "source_id": source.id, + "name": source.name, + "success": results["success"], + "processed": results.get("processed", 0), + "reports_found": results.get("reports_found", 0), + "new_domains": results.get("new_domains", []), + } + + +def _trigger_poll_gmail_source(source: MailSource, db) -> dict: + """Poll a single GMAIL_API source and return a result dict for the API response.""" + global last_check_time # pylint: disable=global-statement + + already = GmailClient.load_ingested_ids(source.gmail_ingested_ids) + gmail_client = GmailClient( + client_id=source.gmail_client_id or "", + client_secret=source.gmail_client_secret or "", + access_token=source.gmail_access_token, + refresh_token=source.gmail_refresh_token or "", + already_ingested_ids=already, + ) + results = gmail_client.fetch_reports() + last_check_time = datetime.now() + + if results.get("new_ingested_ids"): + all_ids = list(dict.fromkeys(already + results["new_ingested_ids"])) + source.gmail_ingested_ids = GmailClient.dump_ingested_ids(all_ids) + refreshed = gmail_client.get_refreshed_tokens() + if refreshed: + source.gmail_access_token = refreshed["access_token"] + if "refresh_token" in refreshed: + source.gmail_refresh_token = refreshed["refresh_token"] + source.last_checked = datetime.utcnow() + db.commit() + return { + "source_id": source.id, + "name": source.name, + "success": results["success"], + "processed": results.get("processed", 0), + "reports_found": results.get("reports_found", 0), + "new_domains": results.get("new_domains", []), + } + + +def _poll_source_for_trigger(source: MailSource, db) -> dict: + """Dispatch a single mail source for the manual trigger-poll endpoint. + + Returns a result/summary dict that is included in the API response. + """ + if source.method == "GMAIL_API": + if not source.gmail_access_token: + return { + "source_id": source.id, + "name": source.name, + "skipped": True, + "reason": "Gmail account not yet authorised", + } + try: + return _trigger_poll_gmail_source(source, db) + except Exception as e: # pylint: disable=broad-exception-caught + logger.error("Error polling Gmail source id=%d: %s", source.id, str(e)) + return { + "source_id": source.id, + "name": source.name, + "success": False, + "error": "Failed to poll. Check server logs for details.", + } + if source.method == "IMAP": + try: + return _trigger_poll_imap_source(source, db) + except Exception as e: # pylint: disable=broad-exception-caught + logger.error("Error polling mail source id=%d: %s", source.id, str(e)) + return { + "source_id": source.id, + "name": source.name, + "success": False, + "error": "Failed to poll. Check server logs for details.", + } + return { + "source_id": source.id, + "name": source.name, + "skipped": True, + "reason": f"method '{source.method}' not yet implemented", + } + + # API endpoint to manually trigger IMAP polling @app.post("/api/v1/admin/trigger-poll") async def trigger_imap_poll(auth: dict = Depends(require_admin_auth)): @@ -435,8 +542,6 @@ async def trigger_imap_poll(auth: dict = Depends(require_admin_auth)): Security: Requires either X-API-Key header or Bearer token """ - global last_check_time # pylint: disable=global-statement - results_summary = [] db = SessionLocal() try: @@ -453,103 +558,7 @@ async def trigger_imap_poll(auth: dict = Depends(require_admin_auth)): } for source in enabled_sources: - if source.method == "GMAIL_API": - if not source.gmail_access_token: - results_summary.append( - { - "source_id": source.id, - "name": source.name, - "skipped": True, - "reason": "Gmail account not yet authorised", - } - ) - continue - try: - already = GmailClient.load_ingested_ids(source.gmail_ingested_ids) - gmail_client = GmailClient( - client_id=source.gmail_client_id or "", - client_secret=source.gmail_client_secret or "", - access_token=source.gmail_access_token, - refresh_token=source.gmail_refresh_token or "", - already_ingested_ids=already, - ) - results = gmail_client.fetch_reports() - last_check_time = datetime.now() - - if results.get("new_ingested_ids"): - all_ids = list(dict.fromkeys(already + results["new_ingested_ids"])) - source.gmail_ingested_ids = GmailClient.dump_ingested_ids(all_ids) - refreshed = gmail_client.get_refreshed_tokens() - if refreshed: - source.gmail_access_token = refreshed["access_token"] - if "refresh_token" in refreshed: - source.gmail_refresh_token = refreshed["refresh_token"] - source.last_checked = datetime.utcnow() - db.commit() - - results_summary.append( - { - "source_id": source.id, - "name": source.name, - "success": results["success"], - "processed": results.get("processed", 0), - "reports_found": results.get("reports_found", 0), - "new_domains": results.get("new_domains", []), - } - ) - except Exception as e: # pylint: disable=broad-exception-caught - logger.error("Error polling Gmail source id=%d: %s", source.id, str(e)) - results_summary.append( - { - "source_id": source.id, - "name": source.name, - "success": False, - "error": "Failed to poll. Check server logs for details.", - } - ) - elif source.method == "IMAP": - try: - imap_client = IMAPClient( - server=source.server, - port=source.port or 993, - username=source.username, - password=source.password, - delete_emails=False, - ) - results = imap_client.fetch_reports(days=7) - last_check_time = datetime.now() - source.last_checked = datetime.utcnow() - db.commit() - - results_summary.append( - { - "source_id": source.id, - "name": source.name, - "success": results["success"], - "processed": results.get("processed", 0), - "reports_found": results.get("reports_found", 0), - "new_domains": results.get("new_domains", []), - } - ) - except Exception as e: # pylint: disable=broad-exception-caught - logger.error("Error polling mail source id=%d: %s", source.id, str(e)) - results_summary.append( - { - "source_id": source.id, - "name": source.name, - "success": False, - "error": "Failed to poll. Check server logs for details.", - } - ) - else: - results_summary.append( - { - "source_id": source.id, - "name": source.name, - "skipped": True, - "reason": f"method '{source.method}' not yet implemented", - } - ) + results_summary.append(_poll_source_for_trigger(source, db)) finally: db.close() diff --git a/backend/app/services/gmail_client.py b/backend/app/services/gmail_client.py index 0204195..7dfd345 100644 --- a/backend/app/services/gmail_client.py +++ b/backend/app/services/gmail_client.py @@ -306,6 +306,31 @@ class GmailClient: msg = email.message_from_bytes(raw_bytes) return self._process_attachments(msg, stats) + @staticmethod + def _decode_part_filename(part: email.message.Message) -> str: + """Return the decoded filename for a MIME part (handles RFC 2047 encoding).""" + from email.header import decode_header + + raw_name = part.get_filename() or "" + decoded_parts = [] + for fragment, charset in decode_header(raw_name): + if isinstance(fragment, bytes): + decoded_parts.append(fragment.decode(charset or "utf-8", errors="replace")) + else: + decoded_parts.append(fragment) + return "".join(decoded_parts) + + @staticmethod + def _is_dmarc_attachment(filename: str) -> bool: + """Return True if *filename* looks like a DMARC aggregate-report file.""" + lower = filename.lower() + return ( + lower.endswith(".xml") + or lower.endswith(".zip") + or lower.endswith(".gz") + or lower.endswith(".gzip") + ) + def _process_attachments(self, msg: email.message.Message, stats: dict) -> int: """Walk a parsed email message and extract DMARC report attachments.""" reports_found = 0 @@ -314,27 +339,8 @@ class GmailClient: if part.get_content_disposition() != "attachment": continue - filename = part.get_filename() or "" - if hasattr(filename, "encode"): - # Decode RFC 2047-encoded filenames - from email.header import decode_header - - parts = decode_header(filename) - decoded_parts = [] - for raw, charset in parts: - if isinstance(raw, bytes): - decoded_parts.append(raw.decode(charset or "utf-8", errors="replace")) - else: - decoded_parts.append(raw) - filename = "".join(decoded_parts) - - lower = filename.lower() - if not ( - lower.endswith(".xml") - or lower.endswith(".zip") - or lower.endswith(".gz") - or lower.endswith(".gzip") - ): + filename = self._decode_part_filename(part) + if not self._is_dmarc_attachment(filename): continue content = part.get_payload(decode=True) diff --git a/backend/app/tests/test_mail_sources.py b/backend/app/tests/test_mail_sources.py index b51be58..32f736a 100644 --- a/backend/app/tests/test_mail_sources.py +++ b/backend/app/tests/test_mail_sources.py @@ -489,10 +489,6 @@ class TestGmailAPIMailSource: ) source_id = create_resp.json()["id"] - # Inject tokens directly into DB via the DB session - from sqlalchemy.orm import Session - from app.models.mail_source import MailSource as MS - # Use the authed_client's DB override — patch the ORM object instead mock_service = MagicMock() mock_service.users.return_value.getProfile.return_value.execute.return_value = {