Merge pull request #67 from christianlouis/copilot/add-gmail-api-mail-source

feat: add Gmail API & main.py helper tests; fix flake8 lint errors
This commit is contained in:
Christian Krakau-Louis
2026-03-30 00:26:28 +02:00
committed by GitHub
9 changed files with 2951 additions and 73 deletions
@@ -0,0 +1,40 @@
"""add Gmail API fields to mail_sources
Revision ID: b2c3d4e5f6a7
Revises: a1b2c3d4e5f6
Create Date: 2026-03-29 19:00:00.000000
"""
from typing import Sequence, Union
import sqlalchemy as sa
from alembic import op
# revision identifiers, used by Alembic.
revision: str = "b2c3d4e5f6a7"
down_revision: Union[str, Sequence[str], None] = "a1b2c3d4e5f6"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Add Gmail OAuth2 credential columns to mail_sources."""
op.add_column("mail_sources", sa.Column("gmail_client_id", sa.String(), nullable=True))
op.add_column("mail_sources", sa.Column("gmail_client_secret", sa.Text(), nullable=True))
op.add_column("mail_sources", sa.Column("gmail_access_token", sa.Text(), nullable=True))
op.add_column("mail_sources", sa.Column("gmail_refresh_token", sa.Text(), nullable=True))
op.add_column("mail_sources", sa.Column("gmail_email", sa.String(), nullable=True))
op.add_column(
"mail_sources",
sa.Column("gmail_ingested_ids", sa.Text(), nullable=True, server_default="[]"),
)
def downgrade() -> None:
"""Remove Gmail OAuth2 credential columns from mail_sources."""
op.drop_column("mail_sources", "gmail_ingested_ids")
op.drop_column("mail_sources", "gmail_email")
op.drop_column("mail_sources", "gmail_refresh_token")
op.drop_column("mail_sources", "gmail_access_token")
op.drop_column("mail_sources", "gmail_client_secret")
op.drop_column("mail_sources", "gmail_client_id")
@@ -3,20 +3,22 @@ Mail Sources API endpoints.
Provides CRUD operations for MailSource objects stored in the database, plus
a *test-connection* action that validates the supplied credentials without
persisting anything.
persisting anything. Gmail API sources additionally have OAuth2 helper
endpoints (authorize-url, callback, fetch).
"""
import logging
from datetime import datetime
from typing import Any, Dict, List, Optional
from fastapi import APIRouter, Depends, HTTPException, status
from fastapi import APIRouter, Depends, HTTPException, Request, status
from pydantic import BaseModel
from sqlalchemy.orm import Session
from app.core.database import get_db
from app.core.security import require_admin_auth
from app.models.mail_source import MailSource
from app.services.gmail_client import GmailClient
from app.services.imap_client import IMAPClient
router = APIRouter()
@@ -41,6 +43,9 @@ class MailSourceBase(BaseModel):
folder: str = "INBOX"
polling_interval: int = 60
enabled: bool = True
# Gmail API OAuth2 fields (only relevant when method == GMAIL_API)
gmail_client_id: Optional[str] = None
gmail_client_secret: Optional[str] = None
class MailSourceCreate(MailSourceBase):
@@ -60,6 +65,8 @@ class MailSourceUpdate(BaseModel):
folder: Optional[str] = None
polling_interval: Optional[int] = None
enabled: Optional[bool] = None
gmail_client_id: Optional[str] = None
gmail_client_secret: Optional[str] = None
class MailSourceResponse(MailSourceBase):
@@ -71,6 +78,10 @@ class MailSourceResponse(MailSourceBase):
updated_at: Optional[datetime] = None
# Mask the stored password in responses
password: Optional[str] = None
# Gmail: show the authorised email address but not tokens
gmail_email: Optional[str] = None
# Indicate whether OAuth tokens are present (without exposing them)
gmail_connected: bool = False
class Config:
from_attributes = True
@@ -87,6 +98,13 @@ class TestConnectionRequest(BaseModel):
method: str = "IMAP"
class GmailCallbackRequest(BaseModel):
"""Payload for the Gmail OAuth2 callback endpoint."""
code: str
redirect_uri: str
# ---------------------------------------------------------------------------
# Helper
# ---------------------------------------------------------------------------
@@ -124,6 +142,10 @@ def _source_to_response(source: MailSource) -> MailSourceResponse:
last_checked=source.last_checked,
created_at=source.created_at,
updated_at=source.updated_at,
gmail_client_id=source.gmail_client_id,
gmail_client_secret="**redacted**" if source.gmail_client_secret else None,
gmail_email=source.gmail_email,
gmail_connected=bool(source.gmail_access_token),
)
@@ -160,6 +182,8 @@ async def create_mail_source(
folder=payload.folder,
polling_interval=payload.polling_interval,
enabled=payload.enabled,
gmail_client_id=payload.gmail_client_id,
gmail_client_secret=payload.gmail_client_secret,
)
db.add(source)
db.commit()
@@ -242,6 +266,43 @@ async def test_stored_mail_source(
"""Test the connection for an already-stored mail source using its saved credentials."""
source = _get_source_or_404(source_id, db)
if source.method == "GMAIL_API":
if not source.gmail_access_token:
return {
"success": False,
"message": "Gmail API source is not yet authorised. "
"Use the 'Connect Gmail' button to complete OAuth2 authorisation.",
"timestamp": datetime.now().isoformat(),
}
try:
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 "",
)
# Attempt to list one message to verify the credentials work
service = gmail_client._build_service() # pylint: disable=protected-access
service.users().getProfile(userId="me").execute()
source.last_checked = datetime.utcnow()
db.commit()
return {
"success": True,
"message": f"Gmail API credentials are valid (account: {source.gmail_email or 'unknown'}).",
"timestamp": datetime.now().isoformat(),
}
except Exception as exc: # pylint: disable=broad-exception-caught
logger.error(
"Gmail API test failed for source id=%d: %s",
int(source_id),
_sanitize_for_log(exc),
)
return {
"success": False,
"message": "Gmail API test failed. Check server logs for details.",
"timestamp": datetime.now().isoformat(),
}
if source.method != "IMAP":
return {
"success": False,
@@ -308,3 +369,303 @@ async def test_connection_adhoc(
"available_mailboxes": stats.get("available_mailboxes", []),
"timestamp": datetime.now().isoformat(),
}
# ---------------------------------------------------------------------------
# Gmail API OAuth2 routes
# ---------------------------------------------------------------------------
@router.get("/{source_id}/gmail/authorize-url", response_model=Dict[str, Any])
async def gmail_authorize_url(
source_id: int,
request: Request,
db: Session = Depends(get_db),
_auth: dict = Depends(require_admin_auth),
) -> Dict[str, Any]:
"""
Return a Google OAuth2 authorization URL for the given GMAIL_API source.
The frontend should redirect the user to this URL. After the user
grants access Google redirects back to
``<origin>/mail-sources/<id>/gmail/callback`` with a ``code`` parameter.
"""
source = _get_source_or_404(source_id, db)
if source.method != "GMAIL_API":
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="This endpoint is only available for GMAIL_API sources.",
)
if not source.gmail_client_id:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="gmail_client_id is not configured for this source.",
)
# Build a redirect_uri that points back to this server's callback endpoint
base_url = str(request.base_url).rstrip("/")
redirect_uri = f"{base_url}/api/v1/mail-sources/{source_id}/gmail/callback"
auth_url = GmailClient.build_authorization_url(
client_id=source.gmail_client_id,
redirect_uri=redirect_uri,
state=str(source_id),
)
return {
"authorization_url": auth_url,
"redirect_uri": redirect_uri,
}
@router.get("/{source_id}/gmail/callback")
async def gmail_oauth_callback(
source_id: int,
request: Request,
db: Session = Depends(get_db),
) -> Any:
"""
Handle the Google OAuth2 redirect after the user grants Gmail access.
Exchanges the authorization ``code`` query parameter for access/refresh
tokens and stores them on the MailSource row. This endpoint is called
directly by Google's redirect, so it does not require the usual API key
authentication; it is protected instead by the state/code being
single-use and bound to the source_id in the URL.
"""
from fastapi.responses import HTMLResponse
code = request.query_params.get("code")
error = request.query_params.get("error")
if error or not code:
html = (
"<html><body><p>Gmail authorisation failed: "
f"{error or 'no code received'}. "
"You may close this window.</p></body></html>"
)
return HTMLResponse(content=html, status_code=400)
source = db.query(MailSource).filter(MailSource.id == source_id).first()
if source is None or source.method != "GMAIL_API":
return HTMLResponse(
content="<html><body><p>Mail source not found.</p></body></html>",
status_code=404,
)
base_url = str(request.base_url).rstrip("/")
redirect_uri = f"{base_url}/api/v1/mail-sources/{source_id}/gmail/callback"
try:
token_data = GmailClient.exchange_code_for_tokens(
client_id=source.gmail_client_id or "",
client_secret=source.gmail_client_secret or "",
code=code,
redirect_uri=redirect_uri,
)
except ValueError as exc:
logger.error(
"Gmail token exchange error for source id=%d: %s",
int(source_id),
_sanitize_for_log(exc),
)
html = (
"<html><body><p>Token exchange failed. "
"Please close this window and try again.</p></body></html>"
)
return HTMLResponse(content=html, status_code=400)
access_token = token_data.get("access_token")
refresh_token = token_data.get("refresh_token")
if not access_token:
return HTMLResponse(
content="<html><body><p>No access token returned by Google.</p></body></html>",
status_code=400,
)
gmail_email = GmailClient.get_gmail_email(access_token)
source.gmail_access_token = access_token
if refresh_token:
source.gmail_refresh_token = refresh_token
if gmail_email:
source.gmail_email = gmail_email
source.updated_at = datetime.utcnow()
db.commit()
logger.info(
"Gmail OAuth2 authorisation complete for source id=%d (account=%s)",
int(source_id),
_sanitize_for_log(gmail_email or "unknown"),
)
html = (
"<html><body>"
"<p>✅ Gmail account connected successfully"
f"{(' (' + gmail_email + ')') if gmail_email else ''}. "
"You may close this window.</p>"
"<script>window.close();</script>"
"</body></html>"
)
return HTMLResponse(content=html)
@router.post("/{source_id}/gmail/callback", response_model=MailSourceResponse)
async def gmail_oauth_callback_post(
source_id: int,
payload: GmailCallbackRequest,
db: Session = Depends(get_db),
_auth: dict = Depends(require_admin_auth),
) -> MailSourceResponse:
"""
Exchange an OAuth2 authorization code for tokens (JSON / programmatic flow).
This POST variant is for clients that handle the OAuth2 redirect
themselves and post the code here as JSON. Requires the standard
admin authentication.
"""
source = _get_source_or_404(source_id, db)
if source.method != "GMAIL_API":
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="This endpoint is only available for GMAIL_API sources.",
)
try:
token_data = GmailClient.exchange_code_for_tokens(
client_id=source.gmail_client_id or "",
client_secret=source.gmail_client_secret or "",
code=payload.code,
redirect_uri=payload.redirect_uri,
)
except ValueError as exc:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=str(exc),
) from exc
access_token = token_data.get("access_token")
refresh_token = token_data.get("refresh_token")
if not access_token:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="Google did not return an access token.",
)
gmail_email = GmailClient.get_gmail_email(access_token)
source.gmail_access_token = access_token
if refresh_token:
source.gmail_refresh_token = refresh_token
if gmail_email:
source.gmail_email = gmail_email
source.updated_at = datetime.utcnow()
db.commit()
db.refresh(source)
logger.info(
"Gmail OAuth2 tokens saved for source id=%d (account=%s)",
int(source_id),
_sanitize_for_log(gmail_email or "unknown"),
)
return _source_to_response(source)
@router.post("/{source_id}/gmail/fetch", response_model=Dict[str, Any])
async def gmail_fetch_reports(
source_id: int,
db: Session = Depends(get_db),
_auth: dict = Depends(require_admin_auth),
) -> Dict[str, Any]:
"""
Manually trigger a Gmail DMARC report fetch for the given source.
Searches Gmail for emails matching the DMARC report heuristic, ingests
any attachments not yet seen, and returns a summary.
"""
source = _get_source_or_404(source_id, db)
if source.method != "GMAIL_API":
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="This endpoint is only available for GMAIL_API sources.",
)
if not source.gmail_access_token:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="Gmail account not yet authorised. Complete OAuth2 flow first.",
)
already = GmailClient.load_ingested_ids(source.gmail_ingested_ids)
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 = client.fetch_reports()
# Persist updated ingested IDs and any refreshed tokens
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 = 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()
logger.info(
"Gmail fetch for source id=%d: processed=%d reports_found=%d",
int(source_id),
int(results.get("processed", 0)),
int(results.get("reports_found", 0)),
)
for err in results.get("errors", []):
logger.warning(
"Gmail fetch warning for source id=%d: %s",
int(source_id),
_sanitize_for_log(err),
)
return {
"success": bool(results.get("success", False)),
"processed": int(results.get("processed", 0)),
"reports_found": int(results.get("reports_found", 0)),
"new_domains": [str(d) for d in results.get("new_domains", [])],
"error_count": len(results.get("errors", [])),
"timestamp": datetime.now().isoformat(),
}
@router.delete("/{source_id}/gmail/connection", status_code=status.HTTP_204_NO_CONTENT)
async def gmail_disconnect(
source_id: int,
db: Session = Depends(get_db),
_auth: dict = Depends(require_admin_auth),
) -> None:
"""Revoke / clear the stored Gmail OAuth2 tokens for this source."""
source = _get_source_or_404(source_id, db)
if source.method != "GMAIL_API":
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="This endpoint is only available for GMAIL_API sources.",
)
source.gmail_access_token = None
source.gmail_refresh_token = None
source.gmail_email = None
source.updated_at = datetime.utcnow()
db.commit()
logger.info("Gmail tokens cleared for source id=%d", int(source_id))
+180 -52
View File
@@ -18,6 +18,7 @@ from app.core.database import Base, SessionLocal, engine
from app.core.security import add_api_key, generate_api_key, require_admin_auth
from app.middleware.security import SecurityHeadersMiddleware
from app.models.mail_source import MailSource # noqa: F401 ensure table is registered
from app.services.gmail_client import GmailClient
from app.models.user import User # noqa: F401 ensure User mapper is registered
from app.services.imap_client import IMAPClient
from app.services.report_store import ReportStore
@@ -73,6 +74,66 @@ def _poll_single_imap_source(source: MailSource) -> None:
)
def _poll_single_gmail_source(source: MailSource) -> None:
"""Fetch DMARC reports for a single GMAIL_API mail source."""
global last_check_time # pylint: disable=global-statement
if not source.gmail_access_token:
logger.info(
"Gmail polling (source id=%d): skipped OAuth2 not yet authorised",
source.id,
)
return
already = GmailClient.load_ingested_ids(source.gmail_ingested_ids)
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 = client.fetch_reports()
db = SessionLocal()
try:
src = db.query(MailSource).get(source.id)
if src:
if results.get("new_ingested_ids"):
all_ids = list(dict.fromkeys(already + results["new_ingested_ids"]))
src.gmail_ingested_ids = GmailClient.dump_ingested_ids(all_ids)
refreshed = client.get_refreshed_tokens()
if refreshed:
src.gmail_access_token = refreshed["access_token"]
if "refresh_token" in refreshed:
src.gmail_refresh_token = refreshed["refresh_token"]
src.last_checked = datetime.utcnow()
db.commit()
finally:
db.close()
last_check_time = datetime.now()
if results["success"]:
logger.info(
"Gmail polling (source id=%d): %s emails processed, %s reports found",
source.id,
results["processed"],
results["reports_found"],
)
if results["new_domains"]:
logger.info("New domains found: %s", ", ".join(results["new_domains"]))
else:
logger.error(
"Gmail polling (source id=%d) failed: %s",
source.id,
results.get("error", "Unknown error"),
)
def _poll_all_enabled_sources() -> None:
"""Iterate over all enabled mail sources and poll each one."""
db = SessionLocal()
@@ -88,17 +149,22 @@ def _poll_all_enabled_sources() -> None:
return
for source in enabled_sources:
if source.method != "IMAP":
if source.method == "GMAIL_API":
try:
_poll_single_gmail_source(source)
except Exception as e: # pylint: disable=broad-exception-caught
logger.error("Error polling Gmail source id=%d: %s", source.id, str(e))
elif source.method == "IMAP":
try:
_poll_single_imap_source(source)
except Exception as e: # pylint: disable=broad-exception-caught
logger.error("Error polling mail source id=%d: %s", source.id, str(e))
else:
logger.info(
"Skipping mail source id=%d method=%r (not yet implemented)",
source.id,
source.method,
)
continue
try:
_poll_single_imap_source(source)
except Exception as e: # pylint: disable=broad-exception-caught
logger.error("Error polling mail source id=%d: %s", source.id, str(e))
def _next_sleep_seconds(min_sleep: int = 60) -> int:
@@ -364,6 +430,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)):
@@ -372,8 +545,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:
@@ -390,50 +561,7 @@ async def trigger_imap_poll(auth: dict = Depends(require_admin_auth)):
}
for source in enabled_sources:
if source.method != "IMAP":
results_summary.append(
{
"source_id": source.id,
"name": source.name,
"skipped": True,
"reason": f"method '{source.method}' not yet implemented",
}
)
continue
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.",
}
)
results_summary.append(_poll_source_for_trigger(source, db))
finally:
db.close()
+15 -1
View File
@@ -15,7 +15,7 @@ class MailSource(Base):
- ``IMAP`` standard IMAP4 (over SSL/TLS or STARTTLS)
- ``POP3`` POP3 inbox (stub for future implementation)
- ``GMAIL_API`` Gmail API with OAuth 2.0 (stub for future implementation)
- ``GMAIL_API`` Gmail API with OAuth 2.0
"""
__tablename__ = "mail_sources"
@@ -38,6 +38,20 @@ class MailSource(Base):
use_ssl = Column(Boolean, default=True)
folder = Column(String, default="INBOX")
# Gmail API OAuth2 credentials (used by GMAIL_API method)
# NOTE: tokens and secrets are stored in plaintext in a production
# environment these fields should be encrypted at the application layer
# (e.g. using Fernet/AES) before persisting, the same way IMAP passwords
# should be. Treat database access as equivalent to credential access.
gmail_client_id = Column(String, nullable=True)
gmail_client_secret = Column(Text, nullable=True)
gmail_access_token = Column(Text, nullable=True)
gmail_refresh_token = Column(Text, nullable=True)
# Email address of the authorised Gmail account
gmail_email = Column(String, nullable=True)
# JSON-encoded list of Gmail message IDs that have already been ingested
gmail_ingested_ids = Column(Text, nullable=True, default="[]")
# Polling behaviour
polling_interval = Column(Integer, default=60) # minutes
+380
View File
@@ -0,0 +1,380 @@
"""
Gmail API client for retrieving DMARC reports.
Connects to Gmail via OAuth 2.0, searches for emails that are likely to
contain DMARC aggregate-report attachments, and processes any new ones.
Already-ingested message IDs are tracked so the same email is never
processed twice (no messages are modified or deleted).
"""
import base64
import email
import json
import logging
from typing import Any, Dict, List, Optional
from urllib.parse import urlencode
import httpx
from google.auth.transport.requests import Request
from google.oauth2.credentials import Credentials
from googleapiclient.discovery import build
from googleapiclient.errors import HttpError
from app.services.dmarc_parser import DMARCParser
from app.services.report_store import ReportStore
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# OAuth2 scopes read-only access to Gmail messages is all we need
# ---------------------------------------------------------------------------
GMAIL_SCOPES = [
"https://www.googleapis.com/auth/gmail.readonly",
]
# ---------------------------------------------------------------------------
# Gmail search query used to find emails likely containing DMARC reports.
#
# Strategy:
# • Require at least one attachment whose name ends in .zip, .gz, or .xml
# (the three formats used by virtually every DMARC sender).
# • Additionally require *either* a keyword in the subject that DMARC senders
# use, or an envelope-from that belongs to a well-known DMARC reporting
# address. This keeps false-positive rates low while catching reports
# from providers that don't follow naming conventions perfectly.
# ---------------------------------------------------------------------------
DMARC_GMAIL_QUERY = (
"has:attachment "
"(filename:zip OR filename:gz OR filename:xml) "
"(subject:dmarc OR subject:report OR subject:rua "
'OR subject:"aggregate report" OR subject:"domain report" '
"OR from:dmarc OR from:dmarc-noreply OR from:reports OR from:postmaster)"
)
# How many message results to fetch per API page
_PAGE_SIZE = 100
class GmailClient:
"""
Client for retrieving DMARC reports from a Gmail account via the Gmail API.
OAuth2 tokens are accepted at construction time and auto-refreshed when
expired. The caller is responsible for persisting any refreshed tokens
returned by :meth:`get_refreshed_tokens`.
"""
def __init__(
self,
client_id: str,
client_secret: str,
access_token: str,
refresh_token: str,
already_ingested_ids: Optional[List[str]] = None,
):
self.client_id = client_id
self.client_secret = client_secret
self._initial_access_token = access_token
self.already_ingested_ids: List[str] = list(already_ingested_ids or [])
self.report_store = ReportStore.get_instance()
self.credentials = Credentials(
token=access_token,
refresh_token=refresh_token,
token_uri="https://oauth2.googleapis.com/token",
client_id=client_id,
client_secret=client_secret,
scopes=GMAIL_SCOPES,
)
# ------------------------------------------------------------------
# Public helpers
# ------------------------------------------------------------------
def get_refreshed_tokens(self) -> Optional[Dict[str, str]]:
"""
Return updated tokens if the google-auth library has refreshed them.
Call this after :meth:`fetch_reports` and persist any non-None result
so the next run doesn't need an extra refresh round-trip.
"""
current = self.credentials.token
if current and current != self._initial_access_token:
result: Dict[str, str] = {"access_token": current}
if self.credentials.refresh_token:
result["refresh_token"] = self.credentials.refresh_token
return result
return None
# ------------------------------------------------------------------
# OAuth2 helpers (static / class methods used by the endpoint layer)
# ------------------------------------------------------------------
@staticmethod
def build_authorization_url(
client_id: str,
redirect_uri: str,
state: Optional[str] = None,
) -> str:
"""
Construct the Google OAuth2 authorization URL.
Requests offline access so a refresh token is issued, and forces
the consent screen so the refresh token is always returned even if
the user has authorised this app before.
"""
params: Dict[str, str] = {
"client_id": client_id,
"response_type": "code",
"scope": " ".join(GMAIL_SCOPES),
"redirect_uri": redirect_uri,
"access_type": "offline",
"prompt": "consent",
}
if state:
params["state"] = state
return "https://accounts.google.com/o/oauth2/v2/auth?" + urlencode(params)
@staticmethod
def exchange_code_for_tokens(
client_id: str,
client_secret: str,
code: str,
redirect_uri: str,
) -> Dict[str, Any]:
"""
Synchronously exchange an authorization code for access+refresh tokens.
Returns the raw JSON from Google's token endpoint. The caller
should check for ``access_token`` in the result before using it.
Raises:
ValueError: if Google returns a non-200 response.
"""
resp = httpx.post(
"https://oauth2.googleapis.com/token",
data={
"code": code,
"client_id": client_id,
"client_secret": client_secret,
"redirect_uri": redirect_uri,
"grant_type": "authorization_code",
},
)
if resp.status_code != 200:
raise ValueError(f"Token exchange failed ({resp.status_code}): {resp.text}")
return resp.json()
@staticmethod
def get_gmail_email(access_token: str) -> Optional[str]:
"""
Return the email address associated with an access token.
Uses the OAuth2 userinfo endpoint. Returns None on failure.
"""
try:
resp = httpx.get(
"https://www.googleapis.com/oauth2/v2/userinfo",
headers={"Authorization": f"Bearer {access_token}"},
)
if resp.status_code == 200:
return resp.json().get("email")
except Exception as exc: # pylint: disable=broad-exception-caught
logger.error("Failed to fetch Gmail email address: %s", exc)
return None
# ------------------------------------------------------------------
# Core fetching logic
# ------------------------------------------------------------------
def fetch_reports(self) -> Dict[str, Any]:
"""
Search Gmail for DMARC report emails and ingest any new ones.
Emails that have already been ingested (tracked via
``already_ingested_ids``) are silently skipped. No messages are
modified or deleted.
Returns:
A dict with keys ``success``, ``processed``, ``reports_found``,
``new_domains``, ``errors``, and ``new_ingested_ids`` (the IDs
added in this run so the caller can persist them).
"""
stats: Dict[str, Any] = {
"success": True,
"processed": 0,
"reports_found": 0,
"new_domains": [],
"errors": [],
"new_ingested_ids": [],
}
try:
service = self._build_service()
except Exception as exc: # pylint: disable=broad-exception-caught
logger.error("Gmail API: failed to build service: %s", exc)
return {**stats, "success": False, "error": str(exc)}
try:
message_ids = self._list_dmarc_message_ids(service)
except Exception as exc: # pylint: disable=broad-exception-caught
logger.error("Gmail API: failed to list messages: %s", exc)
return {**stats, "success": False, "error": str(exc)}
domains_before = set(self.report_store.get_domains())
for msg_id in message_ids:
if msg_id in self.already_ingested_ids:
continue
stats["processed"] += 1
found = self._process_message(service, msg_id, stats)
if found > 0:
stats["new_ingested_ids"].append(msg_id)
self.already_ingested_ids.append(msg_id)
else:
# Track it anyway so we don't re-examine it next run
stats["new_ingested_ids"].append(msg_id)
self.already_ingested_ids.append(msg_id)
domains_after = set(self.report_store.get_domains())
stats["new_domains"] = list(domains_after - domains_before)
return stats
# ------------------------------------------------------------------
# Private helpers
# ------------------------------------------------------------------
def _build_service(self):
"""Build (and auto-refresh if needed) the Gmail API service object."""
if self.credentials.expired and self.credentials.refresh_token:
try:
self.credentials.refresh(Request())
except Exception as exc: # pylint: disable=broad-exception-caught
logger.error("Gmail token refresh failed: %s", exc)
raise
return build("gmail", "v1", credentials=self.credentials, cache_discovery=False)
def _list_dmarc_message_ids(self, service) -> List[str]:
"""Return all Gmail message IDs matching the DMARC search query."""
ids: List[str] = []
page_token: Optional[str] = None
while True:
kwargs: Dict[str, Any] = {
"userId": "me",
"q": DMARC_GMAIL_QUERY,
"maxResults": _PAGE_SIZE,
}
if page_token:
kwargs["pageToken"] = page_token
try:
result = service.users().messages().list(**kwargs).execute()
except HttpError as exc:
logger.error("Gmail API list error: %s", exc)
raise
for msg in result.get("messages", []):
ids.append(msg["id"])
page_token = result.get("nextPageToken")
if not page_token:
break
return ids
def _process_message(self, service, msg_id: str, stats: dict) -> int:
"""
Download a Gmail message and process any DMARC-report attachments.
Returns the number of DMARC reports found in this message.
"""
try:
msg_data = (
service.users().messages().get(userId="me", id=msg_id, format="raw").execute()
)
except HttpError as exc:
logger.error("Gmail API: failed to fetch message %s: %s", msg_id, exc)
stats["errors"].append(f"Failed to fetch message {msg_id}")
return 0
raw_bytes = base64.urlsafe_b64decode(msg_data.get("raw", ""))
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
for part in msg.walk():
if part.get_content_disposition() != "attachment":
continue
filename = self._decode_part_filename(part)
if not self._is_dmarc_attachment(filename):
continue
content = part.get_payload(decode=True)
if not content:
continue
try:
parser = DMARCParser()
reports = parser.parse(content, filename)
for report in reports:
self.report_store.add_report(report)
stats["reports_found"] += 1
reports_found += 1
except Exception as exc: # pylint: disable=broad-exception-caught
logger.error("Failed to parse DMARC attachment %s: %s", filename, exc)
stats["errors"].append(f"Failed to parse {filename}: {exc}")
return reports_found
# ------------------------------------------------------------------
# Convenience: load / save ingested IDs from/to the JSON text column
# ------------------------------------------------------------------
@staticmethod
def load_ingested_ids(json_text: Optional[str]) -> List[str]:
"""Deserialise the gmail_ingested_ids text column into a list."""
if not json_text:
return []
try:
return json.loads(json_text)
except (json.JSONDecodeError, TypeError):
return []
@staticmethod
def dump_ingested_ids(ids: List[str]) -> str:
"""Serialise the list of ingested IDs back to a JSON string."""
return json.dumps(ids)
+157 -17
View File
@@ -61,8 +61,8 @@
<tr>
<th>Name</th>
<th>Method</th>
<th>Server</th>
<th>Username</th>
<th>Server / Account</th>
<th>User / Status</th>
<th>Last Checked</th>
<th>Enabled</th>
<th>Actions</th>
@@ -75,8 +75,8 @@
<td>
<span class="badge badge-outline" x-text="source.method"></span>
</td>
<td x-text="source.server ? source.server + ':' + source.port : '—'"></td>
<td x-text="source.username || '—'"></td>
<td x-text="source.method === 'GMAIL_API' ? (source.gmail_email || '—') : (source.server ? source.server + ':' + source.port : '—')"></td>
<td x-text="source.method === 'GMAIL_API' ? (source.gmail_connected ? '✅ Connected' : '⚠️ Not authorised') : (source.username || '—')"></td>
<td x-text="source.last_checked ? new Date(source.last_checked).toLocaleString() : 'Never'"></td>
<td>
<input
@@ -153,7 +153,7 @@
<select x-model="form.method" class="select select-bordered w-full">
<option value="IMAP">IMAP</option>
<option value="POP3">POP3 (coming soon)</option>
<option value="GMAIL_API">Gmail API (coming soon)</option>
<option value="GMAIL_API">Gmail API (OAuth2)</option>
</select>
</div>
@@ -229,6 +229,72 @@
</div>
</template>
<!-- Gmail API fields -->
<template x-if="form.method === 'GMAIL_API'">
<div class="space-y-4">
<div class="alert alert-info text-sm p-3">
<div>
<p class="font-semibold">Gmail API Setup</p>
<p class="mt-1">Enter your Google OAuth2 Client ID and Secret from
<a href="https://console.cloud.google.com/apis/credentials" target="_blank"
class="underline">Google Cloud Console</a>.
Enable the <strong>Gmail API</strong>, create an OAuth2 client (Web application),
and add your callback URL as an authorised redirect URI.
Then save this source and click <strong>Connect Gmail</strong> to authorise access.</p>
</div>
</div>
<div>
<label class="label"><span class="label-text font-medium">Google Client ID <span class="text-error">*</span></span></label>
<input type="text" x-model="form.gmail_client_id"
placeholder="123456789-abc.apps.googleusercontent.com"
class="input input-bordered w-full" />
</div>
<div>
<label class="label">
<span class="label-text font-medium">Google Client Secret <span class="text-error">*</span></span>
<span class="label-text-alt text-muted-foreground" x-show="editingId">Leave blank to keep existing</span>
</label>
<div class="relative">
<input :type="showPassword ? 'text' : 'password'"
x-model="form.gmail_client_secret"
class="input input-bordered w-full pr-10"
placeholder="GOCSPX-…" />
<button type="button"
class="absolute right-2 top-3 text-muted-foreground hover:text-foreground"
x-on:click="showPassword = !showPassword">
<svg x-show="!showPassword" xmlns="http://www.w3.org/2000/svg" width="16" height="16"
viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"
stroke-linecap="round" stroke-linejoin="round">
<path d="M1 12s4-8 11-8 11 8 11 8-4 8-11 8-11-8-11-8z"></path>
<circle cx="12" cy="12" r="3"></circle>
</svg>
<svg x-show="showPassword" xmlns="http://www.w3.org/2000/svg" width="16" height="16"
viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"
stroke-linecap="round" stroke-linejoin="round">
<path d="M17.94 17.94A10.07 10.07 0 0 1 12 20c-7 0-11-8-11-8a18.45 18.45 0 0 1 5.06-5.94"></path>
<path d="M9.9 4.24A9.12 9.12 0 0 1 12 4c7 0 11 8 11 8a18.5 18.5 0 0 1-2.16 3.19"></path>
<line x1="1" y1="1" x2="23" y2="23"></line>
</svg>
</button>
</div>
</div>
<!-- OAuth connection status (when editing) -->
<template x-if="editingId && gmailConnected">
<div class="alert alert-success text-sm p-3">
<span>✅ Connected as <strong x-text="gmailEmail || 'Gmail account'"></strong></span>
</div>
</template>
<template x-if="editingId && !gmailConnected">
<div class="alert alert-warning text-sm p-3">
<span>⚠️ Not yet authorised. Save this source first, then click <strong>Connect Gmail</strong>.</span>
</div>
</template>
</div>
</template>
<!-- Polling interval -->
<div>
<label class="label"><span class="label-text font-medium">Polling Interval (minutes)</span></label>
@@ -252,18 +318,34 @@
<!-- Form actions -->
<div class="flex justify-between pt-2">
<button type="button" class="btn btn-outline btn-sm"
x-on:click="testAdHoc()"
:disabled="isTesting || isSaving">
<span x-show="!isTesting">Test Connection</span>
<span x-show="isTesting" class="flex items-center" x-cloak>
<svg class="animate-spin h-4 w-4 mr-1" xmlns="http://www.w3.org/2000/svg" fill="none" viewBox="0 0 24 24">
<circle class="opacity-25" cx="12" cy="12" r="10" stroke="currentColor" stroke-width="4"></circle>
<path class="opacity-75" fill="currentColor" d="M4 12a8 8 0 018-8V0C5.373 0 0 5.373 0 12h4zm2 5.291A7.962 7.962 0 014 12H0c0 3.042 1.135 5.824 3 7.938l3-2.647z"></path>
</svg>
Testing...
</span>
</button>
<!-- Test Connection (IMAP/POP3) or Connect Gmail button -->
<div class="flex gap-2">
<template x-if="form.method !== 'GMAIL_API'">
<button type="button" class="btn btn-outline btn-sm"
x-on:click="testAdHoc()"
:disabled="isTesting || isSaving">
<span x-show="!isTesting">Test Connection</span>
<span x-show="isTesting" class="flex items-center" x-cloak>
<svg class="animate-spin h-4 w-4 mr-1" xmlns="http://www.w3.org/2000/svg" fill="none" viewBox="0 0 24 24">
<circle class="opacity-25" cx="12" cy="12" r="10" stroke="currentColor" stroke-width="4"></circle>
<path class="opacity-75" fill="currentColor" d="M4 12a8 8 0 018-8V0C5.373 0 0 5.373 0 12h4zm2 5.291A7.962 7.962 0 014 12H0c0 3.042 1.135 5.824 3 7.938l3-2.647z"></path>
</svg>
Testing...
</span>
</button>
</template>
<template x-if="form.method === 'GMAIL_API' && editingId">
<button type="button" class="btn btn-outline btn-sm"
x-on:click="connectGmail()"
:disabled="isTesting || isSaving">
<svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none"
stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="mr-1">
<path d="M18 2h-3a5 5 0 0 0-5 5v3H7v4h3v8h4v-8h3l1-4h-4V7a1 1 0 0 1 1-1h3z"></path>
</svg>
Connect Gmail
</button>
</template>
</div>
<div class="flex gap-2">
<button type="button" class="btn btn-ghost btn-sm" x-on:click="closeForm()">Cancel</button>
<button type="submit" class="btn btn-default btn-sm" :disabled="isSaving">
@@ -310,6 +392,8 @@ function mailSourcesApp() {
testing: {},
feedback: { message: '', type: '' },
testResult: { message: '', success: false },
gmailConnected: false,
gmailEmail: '',
form: {
name: '',
@@ -322,6 +406,8 @@ function mailSourcesApp() {
folder: 'INBOX',
polling_interval: 60,
enabled: true,
gmail_client_id: '',
gmail_client_secret: '',
},
async init() {
@@ -341,6 +427,8 @@ function mailSourcesApp() {
openAddForm() {
this.editingId = null;
this.gmailConnected = false;
this.gmailEmail = '';
this.form = {
name: '',
method: 'IMAP',
@@ -352,6 +440,8 @@ function mailSourcesApp() {
folder: 'INBOX',
polling_interval: 60,
enabled: true,
gmail_client_id: '',
gmail_client_secret: '',
};
this.testResult = { message: '', success: false };
this.showForm = true;
@@ -359,6 +449,8 @@ function mailSourcesApp() {
openEditForm(source) {
this.editingId = source.id;
this.gmailConnected = source.gmail_connected || false;
this.gmailEmail = source.gmail_email || '';
this.form = {
name: source.name,
method: source.method,
@@ -370,6 +462,8 @@ function mailSourcesApp() {
folder: source.folder || 'INBOX',
polling_interval: source.polling_interval || 60,
enabled: source.enabled !== false,
gmail_client_id: source.gmail_client_id || '',
gmail_client_secret: '', // never pre-fill client secret
};
this.testResult = { message: '', success: false };
this.showForm = true;
@@ -381,6 +475,48 @@ function mailSourcesApp() {
this.testResult = { message: '', success: false };
},
async connectGmail() {
if (!this.editingId) return;
try {
const resp = await fetch(
`/api/v1/mail-sources/${this.editingId}/gmail/authorize-url`
);
if (!resp.ok) {
const err = await resp.json();
this.feedback = { message: `Error: ${err.detail || 'Failed to get authorization URL'}`, type: 'error' };
return;
}
const data = await resp.json();
// Open the OAuth2 URL in a new popup window
const popup = window.open(
data.authorization_url,
'gmail_oauth',
'width=600,height=700,scrollbars=yes'
);
// Poll for popup close and refresh source list
const pollInterval = setInterval(async () => {
if (popup && popup.closed) {
clearInterval(pollInterval);
await this.loadSources();
// Refresh the editing form state
const updated = this.sources.find(s => s.id === this.editingId);
if (updated) {
this.gmailConnected = updated.gmail_connected || false;
this.gmailEmail = updated.gmail_email || '';
if (this.gmailConnected) {
this.feedback = {
message: `Gmail connected successfully (${this.gmailEmail}).`,
type: 'success',
};
}
}
}
}, 1000);
} catch (e) {
this.feedback = { message: `Error: ${e.message}`, type: 'error' };
}
},
async saveSource() {
this.isSaving = true;
this.feedback = { message: '', type: '' };
@@ -390,6 +526,10 @@ function mailSourcesApp() {
if (this.editingId && !payload.password) {
delete payload.password;
}
// Don't send empty gmail_client_secret on edit
if (this.editingId && !payload.gmail_client_secret) {
delete payload.gmail_client_secret;
}
const url = this.editingId
? `/api/v1/mail-sources/${this.editingId}`
+638
View File
@@ -0,0 +1,638 @@
"""
Unit tests for app.services.gmail_client.GmailClient.
All external I/O (httpx, google-auth, googleapiclient) is mocked so these
tests never make real network calls.
"""
import base64
import email as email_mod
import json
from email import encoders as email_encoders
from email.mime.multipart import MIMEMultipart
from email.mime.base import MIMEBase
from email.mime.text import MIMEText
from typing import Optional
from unittest.mock import MagicMock, patch
import pytest
from app.services.gmail_client import GmailClient
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _make_client(
access_token: str = "acc",
refresh_token: str = "ref",
already_ingested: Optional[list] = None,
) -> GmailClient:
"""Instantiate a GmailClient with real Credentials mocked out."""
with patch("app.services.gmail_client.Credentials") as mock_creds_class:
mock_creds = MagicMock()
mock_creds.token = access_token
mock_creds.refresh_token = refresh_token
mock_creds.expired = False
mock_creds_class.return_value = mock_creds
client = GmailClient(
client_id="cid",
client_secret="csec",
access_token=access_token,
refresh_token=refresh_token,
already_ingested_ids=already_ingested or [],
)
# Expose the mock so tests can manipulate it
client._mock_creds = mock_creds # type: ignore[attr-defined]
return client
def _make_raw_email(attachments: list) -> bytes:
"""Build a raw MIME email whose attachments are described by *attachments*.
Each item in *attachments* is a dict with keys:
filename: str
content: bytes
disposition: str (default "attachment")
"""
msg = MIMEMultipart()
msg["Subject"] = "DMARC aggregate report"
msg["From"] = "noreply@example.com"
msg["To"] = "user@gmail.com"
msg.attach(MIMEText("See attached DMARC report.", "plain"))
for att in attachments:
part = MIMEBase("application", "octet-stream")
part.set_payload(att["content"])
email_encoders.encode_base64(part)
disposition = att.get("disposition", "attachment")
part.add_header(
"Content-Disposition",
disposition,
filename=att["filename"],
)
msg.attach(part)
return msg.as_bytes()
def _b64_raw(raw_bytes: bytes) -> str:
"""URL-safe base64-encode bytes (as Gmail API returns them)."""
return base64.urlsafe_b64encode(raw_bytes).decode()
# ===========================================================================
# __init__ / basic construction
# ===========================================================================
class TestGmailClientInit:
def test_init_stores_tokens(self):
client = _make_client(access_token="my-acc", refresh_token="my-ref")
assert client._initial_access_token == "my-acc"
assert client.client_id == "cid"
assert client.client_secret == "csec"
def test_init_already_ingested_defaults_to_empty(self):
client = _make_client()
assert client.already_ingested_ids == []
def test_init_already_ingested_is_copied(self):
ids = ["a", "b"]
client = _make_client(already_ingested=ids)
assert client.already_ingested_ids == ["a", "b"]
# Mutating the original should not affect the client
ids.append("c")
assert "c" not in client.already_ingested_ids
# ===========================================================================
# get_refreshed_tokens
# ===========================================================================
class TestGetRefreshedTokens:
def test_returns_none_when_token_unchanged(self):
client = _make_client(access_token="original")
# credentials.token == _initial_access_token → no refresh happened
client._mock_creds.token = "original"
assert client.get_refreshed_tokens() is None
def test_returns_new_access_token_when_changed(self):
client = _make_client(access_token="original")
client._mock_creds.token = "new-token"
client._mock_creds.refresh_token = None
result = client.get_refreshed_tokens()
assert result is not None
assert result["access_token"] == "new-token"
assert "refresh_token" not in result
def test_returns_both_tokens_when_refresh_token_present(self):
client = _make_client(access_token="original")
client._mock_creds.token = "new-token"
client._mock_creds.refresh_token = "new-refresh"
result = client.get_refreshed_tokens()
assert result is not None
assert result["access_token"] == "new-token"
assert result["refresh_token"] == "new-refresh"
def test_returns_none_when_token_is_none(self):
client = _make_client(access_token="original")
client._mock_creds.token = None
assert client.get_refreshed_tokens() is None
# ===========================================================================
# build_authorization_url (already partly covered; extend for completeness)
# ===========================================================================
class TestBuildAuthorizationUrl:
def test_includes_all_required_params(self):
url = GmailClient.build_authorization_url(
client_id="cid",
redirect_uri="https://example.com/cb",
state="99",
)
assert "client_id=cid" in url
assert "response_type=code" in url
assert "gmail.readonly" in url
assert "access_type=offline" in url
assert "prompt=consent" in url
assert "state=99" in url
def test_state_omitted_when_none(self):
url = GmailClient.build_authorization_url(
client_id="cid",
redirect_uri="https://example.com/cb",
)
assert "state=" not in url
# ===========================================================================
# exchange_code_for_tokens
# ===========================================================================
class TestExchangeCodeForTokens:
def test_success_returns_json(self):
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_resp.json.return_value = {"access_token": "acc", "refresh_token": "ref"}
with patch("app.services.gmail_client.httpx.post", return_value=mock_resp):
result = GmailClient.exchange_code_for_tokens(
client_id="cid",
client_secret="csec",
code="auth-code",
redirect_uri="https://example.com/cb",
)
assert result["access_token"] == "acc"
assert result["refresh_token"] == "ref"
def test_non_200_raises_value_error(self):
mock_resp = MagicMock()
mock_resp.status_code = 400
mock_resp.text = '{"error": "invalid_grant"}'
with patch("app.services.gmail_client.httpx.post", return_value=mock_resp):
with pytest.raises(ValueError, match="400"):
GmailClient.exchange_code_for_tokens(
client_id="cid",
client_secret="csec",
code="bad-code",
redirect_uri="https://example.com/cb",
)
# ===========================================================================
# get_gmail_email
# ===========================================================================
class TestGetGmailEmail:
def test_returns_email_on_200(self):
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_resp.json.return_value = {"email": "user@gmail.com"}
with patch("app.services.gmail_client.httpx.get", return_value=mock_resp):
result = GmailClient.get_gmail_email("valid-token")
assert result == "user@gmail.com"
def test_returns_none_on_non_200(self):
mock_resp = MagicMock()
mock_resp.status_code = 401
with patch("app.services.gmail_client.httpx.get", return_value=mock_resp):
result = GmailClient.get_gmail_email("expired-token")
assert result is None
def test_returns_none_on_exception(self):
with patch(
"app.services.gmail_client.httpx.get", side_effect=Exception("network error")
):
result = GmailClient.get_gmail_email("some-token")
assert result is None
# ===========================================================================
# _build_service
# ===========================================================================
class TestBuildService:
def test_returns_service_when_not_expired(self):
client = _make_client()
client._mock_creds.expired = False
mock_service = MagicMock()
with patch("app.services.gmail_client.build", return_value=mock_service):
svc = client._build_service()
assert svc is mock_service
def test_refreshes_token_when_expired(self):
client = _make_client()
client._mock_creds.expired = True
client._mock_creds.refresh_token = "ref"
mock_service = MagicMock()
with patch("app.services.gmail_client.build", return_value=mock_service), patch(
"app.services.gmail_client.Request"
):
svc = client._build_service()
client._mock_creds.refresh.assert_called_once()
assert svc is mock_service
def test_raises_when_refresh_fails(self):
client = _make_client()
client._mock_creds.expired = True
client._mock_creds.refresh_token = "ref"
client._mock_creds.refresh.side_effect = Exception("refresh failed")
with patch("app.services.gmail_client.Request"), patch(
"app.services.gmail_client.build"
):
with pytest.raises(Exception, match="refresh failed"):
client._build_service()
# ===========================================================================
# _list_dmarc_message_ids
# ===========================================================================
class TestListDmarcMessageIds:
def test_returns_empty_when_no_messages(self):
client = _make_client()
service = MagicMock()
service.users.return_value.messages.return_value.list.return_value.execute.return_value = {
"messages": []
}
ids = client._list_dmarc_message_ids(service)
assert ids == []
def test_returns_ids_from_single_page(self):
client = _make_client()
service = MagicMock()
service.users.return_value.messages.return_value.list.return_value.execute.return_value = {
"messages": [{"id": "id1"}, {"id": "id2"}]
}
ids = client._list_dmarc_message_ids(service)
assert ids == ["id1", "id2"]
def test_follows_next_page_token(self):
client = _make_client()
# First page has a nextPageToken; second page has none
service = MagicMock()
execute = service.users.return_value.messages.return_value.list.return_value.execute
execute.side_effect = [
{"messages": [{"id": "id1"}], "nextPageToken": "page2"},
{"messages": [{"id": "id2"}]},
]
ids = client._list_dmarc_message_ids(service)
assert ids == ["id1", "id2"]
# list() should have been called twice
assert service.users.return_value.messages.return_value.list.call_count == 2
def test_raises_on_http_error(self):
from googleapiclient.errors import HttpError
client = _make_client()
service = MagicMock()
execute = service.users.return_value.messages.return_value.list.return_value.execute
fake_error = HttpError(MagicMock(status=403), b"forbidden")
execute.side_effect = fake_error
with pytest.raises(HttpError):
client._list_dmarc_message_ids(service)
# ===========================================================================
# _decode_part_filename
# ===========================================================================
class TestDecodePartFilename:
def test_plain_ascii_filename(self):
part = MagicMock()
part.get_filename.return_value = "report.xml"
assert GmailClient._decode_part_filename(part) == "report.xml"
def test_none_filename_returns_empty(self):
part = MagicMock()
part.get_filename.return_value = None
assert GmailClient._decode_part_filename(part) == ""
def test_rfc2047_encoded_filename(self):
# Build an RFC 2047 encoded filename
encoded = "=?utf-8?b?cmVwb3J0LnhtbA==?=" # base64("report.xml")
part = MagicMock()
part.get_filename.return_value = encoded
result = GmailClient._decode_part_filename(part)
assert result == "report.xml"
# ===========================================================================
# _is_dmarc_attachment
# ===========================================================================
class TestIsDmarcAttachment:
@pytest.mark.parametrize(
"filename",
[
"report.xml",
"report.XML", # case-insensitive
"report.zip",
"report.gz",
"report.gzip",
"Report.ZIP",
],
)
def test_dmarc_extensions_return_true(self, filename):
assert GmailClient._is_dmarc_attachment(filename) is True
@pytest.mark.parametrize(
"filename",
["report.txt", "image.png", "report.pdf", "report.tar", ""],
)
def test_non_dmarc_extensions_return_false(self, filename):
assert GmailClient._is_dmarc_attachment(filename) is False
# ===========================================================================
# _process_message
# ===========================================================================
class TestProcessMessage:
def test_fetches_and_processes_message(self):
"""Happy path: message fetched, attachments processed."""
client = _make_client()
raw_email = _make_raw_email(
[{"filename": "report.xml", "content": b"<xml/>"}]
)
raw_b64 = _b64_raw(raw_email)
service = MagicMock()
service.users.return_value.messages.return_value.get.return_value.execute.return_value = {
"raw": raw_b64
}
stats = {"reports_found": 0, "errors": []}
with patch.object(client, "_process_attachments", return_value=0) as mock_proc:
count = client._process_message(service, "msg1", stats)
mock_proc.assert_called_once()
assert count == 0 # our mock returns 0
def test_http_error_recorded_and_returns_zero(self):
from googleapiclient.errors import HttpError
client = _make_client()
service = MagicMock()
fake_error = HttpError(MagicMock(status=404), b"not found")
service.users.return_value.messages.return_value.get.return_value.execute.side_effect = (
fake_error
)
stats = {"reports_found": 0, "errors": []}
count = client._process_message(service, "bad-id", stats)
assert count == 0
assert len(stats["errors"]) == 1
assert "bad-id" in stats["errors"][0]
# ===========================================================================
# _process_attachments
# ===========================================================================
class TestProcessAttachments:
def test_no_attachments_returns_zero(self):
client = _make_client()
msg = email_mod.message_from_bytes(
b"From: a@b.com\r\nTo: c@d.com\r\n\r\nHello"
)
stats = {"reports_found": 0, "errors": []}
count = client._process_attachments(msg, stats)
assert count == 0
def test_non_dmarc_attachment_skipped(self):
"""An inline or non-DMARC file should not count as a report."""
client = _make_client()
raw = _make_raw_email(
[{"filename": "photo.png", "content": b"\x89PNG"}]
)
msg = email_mod.message_from_bytes(raw)
stats = {"reports_found": 0, "errors": []}
count = client._process_attachments(msg, stats)
assert count == 0
assert stats["reports_found"] == 0
def test_dmarc_xml_attachment_is_parsed(self):
"""A .xml attachment is parsed via DMARCParser and counts as a report."""
client = _make_client()
raw = _make_raw_email(
[{"filename": "report.xml", "content": b"<xml_content/>"}]
)
msg = email_mod.message_from_bytes(raw)
stats = {"reports_found": 0, "errors": []}
mock_report = {"domain": "example.com", "records": []}
with patch("app.services.gmail_client.DMARCParser") as mock_parser_class:
mock_parser = MagicMock()
mock_parser.parse.return_value = [mock_report]
mock_parser_class.return_value = mock_parser
# Also mock report_store.add_report to avoid real persistence
with patch.object(client.report_store, "add_report"):
count = client._process_attachments(msg, stats)
assert count == 1
assert stats["reports_found"] == 1
def test_dmarc_attachment_with_empty_content_skipped(self):
"""A DMARC-named attachment with truly empty payload is skipped gracefully."""
client = _make_client()
# Build an attachment with empty bytes base64 of b"" is b""
raw = _make_raw_email(
[{"filename": "report.zip", "content": b""}]
)
msg = email_mod.message_from_bytes(raw)
stats = {"reports_found": 0, "errors": []}
count = client._process_attachments(msg, stats)
# Empty payload → `get_payload(decode=True)` returns b"" which is
# falsy, so the attachment is skipped
assert count == 0
def test_parse_exception_adds_error_and_continues(self):
"""A parse error should be recorded in stats but not raise."""
client = _make_client()
raw = _make_raw_email(
[
{"filename": "bad.xml", "content": b"corrupt"},
{"filename": "good.xml", "content": b"<xml/>"},
]
)
msg = email_mod.message_from_bytes(raw)
stats = {"reports_found": 0, "errors": []}
good_report = {"domain": "example.com", "records": []}
call_count = 0
def parse_side_effect(content, filename):
nonlocal call_count
call_count += 1
if call_count == 1:
raise ValueError("bad xml")
return [good_report]
with patch("app.services.gmail_client.DMARCParser") as mock_parser_class:
mock_parser = MagicMock()
mock_parser.parse.side_effect = parse_side_effect
mock_parser_class.return_value = mock_parser
with patch.object(client.report_store, "add_report"):
count = client._process_attachments(msg, stats)
assert len(stats["errors"]) == 1
assert "bad.xml" in stats["errors"][0]
assert count == 1 # second attachment still parsed
# ===========================================================================
# fetch_reports
# ===========================================================================
class TestFetchReports:
def test_returns_failure_when_build_service_raises(self):
client = _make_client()
with patch.object(
client, "_build_service", side_effect=Exception("auth error")
):
result = client.fetch_reports()
assert result["success"] is False
assert "auth error" in result.get("error", "")
def test_returns_failure_when_list_messages_raises(self):
client = _make_client()
mock_service = MagicMock()
with patch.object(client, "_build_service", return_value=mock_service), patch.object(
client, "_list_dmarc_message_ids", side_effect=Exception("list error")
):
result = client.fetch_reports()
assert result["success"] is False
def test_returns_success_with_no_messages(self):
client = _make_client()
mock_service = MagicMock()
with patch.object(client, "_build_service", return_value=mock_service), patch.object(
client, "_list_dmarc_message_ids", return_value=[]
):
result = client.fetch_reports()
assert result["success"] is True
assert result["processed"] == 0
def test_skips_already_ingested_messages(self):
client = _make_client(already_ingested=["id1"])
mock_service = MagicMock()
with patch.object(client, "_build_service", return_value=mock_service), patch.object(
client, "_list_dmarc_message_ids", return_value=["id1", "id2"]
), patch.object(client, "_process_message", return_value=0) as mock_proc:
result = client.fetch_reports()
# Only id2 should be processed; id1 is already ingested
assert mock_proc.call_count == 1
call_args = mock_proc.call_args_list[0][0]
assert call_args[1] == "id2"
assert result["processed"] == 1
def test_tracks_new_ingested_ids(self):
client = _make_client()
mock_service = MagicMock()
with patch.object(client, "_build_service", return_value=mock_service), patch.object(
client, "_list_dmarc_message_ids", return_value=["id1", "id2"]
), patch.object(client, "_process_message", return_value=0):
result = client.fetch_reports()
assert "id1" in result["new_ingested_ids"]
assert "id2" in result["new_ingested_ids"]
def test_reports_new_domains(self):
"""fetch_reports should report domains that appear after ingestion."""
from app.services.report_store import ReportStore
client = _make_client()
mock_service = MagicMock()
def _process_side_effect(service, msg_id, stats):
# Simulate adding a domain to the report store
ReportStore.get_instance().add_report(
{
"org_name": "Test Org",
"report_id": "r1",
"begin_date": "2024-01-01",
"end_date": "2024-01-02",
"domain": "newdomain.example",
"records": [],
}
)
stats["reports_found"] += 1
return 1
with patch.object(client, "_build_service", return_value=mock_service), patch.object(
client, "_list_dmarc_message_ids", return_value=["id1"]
), patch.object(client, "_process_message", side_effect=_process_side_effect):
result = client.fetch_reports()
assert "newdomain.example" in result["new_domains"]
# ===========================================================================
# load_ingested_ids / dump_ingested_ids (already tested in TestGmailClientHelpers
# in test_mail_sources.py; add a few edge-cases here)
# ===========================================================================
class TestIngestedIdHelpers:
def test_load_non_list_json_returns_no_error(self):
# Valid JSON but not a list should gracefully not raise
result = GmailClient.load_ingested_ids('{"key": "value"}')
assert result is not None # no crash
def test_dump_preserves_order(self):
ids = ["z", "a", "m"]
dumped = GmailClient.dump_ingested_ids(ids)
assert json.loads(dumped) == ["z", "a", "m"]
File diff suppressed because it is too large Load Diff
+4 -1
View File
@@ -24,4 +24,7 @@ email-validator>=2.0.0
lxml>=4.9.2
zipfile36>=0.1.3
aiosmtplib>=2.0.2
jinja2>=3.1.2
jinja2>=3.1.2
google-auth>=2.0.0
google-api-python-client>=2.0.0
google-auth-httplib2>=0.2.0