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:
@@ -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
@@ -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,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
|
||||
|
||||
|
||||
@@ -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)
|
||||
@@ -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}`
|
||||
|
||||
@@ -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
@@ -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
|
||||
Reference in New Issue
Block a user