gh-christianlouis-dmarq/backend/app/main.py

import asyncio
import logging
import os
from datetime import datetime

from fastapi import Depends, FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates

import app.models.domain  # noqa: F401 – ensure Domain/UserDomain tables are registered
import app.models.mail_source_import  # noqa: F401 – ensure import history table is registered
import app.models.report  # noqa: F401 – ensure DMARCReport/ReportRecord tables are registered
import app.models.setting  # noqa: F401 – ensure Setting table is registered
import app.models.user  # noqa: F401 – ensure User table is registered
from app.api.api_v1.api import api_router
from app.core.config import get_settings
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.auth import AuthRedirectMiddleware
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.services.imap_client import IMAPClient
from app.services.import_history import record_import_attempt
from app.services.report_persistence import hydrate_report_store_from_db
from app.services.report_store import ReportStore

# Set up logging
logger = logging.getLogger(__name__)

settings = get_settings()

# Global variables for background task management
background_task = None
last_check_time = None


def _poll_single_imap_source(source: MailSource) -> None:
    """Fetch DMARC reports for a single IMAP mail source and update its last_checked timestamp."""
    global last_check_time  # pylint: disable=global-statement

    db = SessionLocal()
    try:
        src = db.query(MailSource).get(source.id)
        poll_source = src or source
        imap_client = IMAPClient(
            server=poll_source.server,
            port=poll_source.port or 993,
            username=poll_source.username,
            password=poll_source.password,
            delete_emails=False,
            db=db,
        )
        started_at = datetime.utcnow()
        results = imap_client.fetch_reports(days=9999)
        if src:
            src.last_checked = datetime.utcnow()
            record_import_attempt(db, src, results, started_at=started_at, trigger="scheduled")
            db.commit()
    finally:
        db.close()

    last_check_time = datetime.now()

    if results["success"]:
        logger.info(
            "IMAP 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(
            "IMAP polling (source id=%d) failed: %s",
            source.id,
            results.get("error", "Unknown error"),
        )


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

    db = SessionLocal()
    try:
        src = db.query(MailSource).get(source.id)
        poll_source = src or source
        already = GmailClient.load_ingested_ids(poll_source.gmail_ingested_ids)
        client = GmailClient(
            client_id=poll_source.gmail_client_id or "",
            client_secret=poll_source.gmail_client_secret or "",
            access_token=poll_source.gmail_access_token,
            refresh_token=poll_source.gmail_refresh_token or "",
            already_ingested_ids=already,
            db=db,
        )

        started_at = datetime.utcnow()
        results = client.fetch_reports()
        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()
            record_import_attempt(db, src, results, started_at=started_at, trigger="scheduled")
            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() -> list[MailSource]:
    """Iterate over all enabled mail sources and poll each one."""
    db = SessionLocal()
    try:
        enabled_sources = (
            db.query(MailSource).filter(MailSource.enabled == True).all()  # noqa: E712
        )
    finally:
        db.close()

    if not enabled_sources:
        logger.info("No enabled mail sources configured – polling skipped")
        return enabled_sources

    for source in enabled_sources:
        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,
            )
    return enabled_sources


def _next_sleep_seconds(
    min_sleep: int = 60, enabled_sources: list[MailSource] | None = None
) -> int:
    """Return how many seconds to sleep until the next polling cycle."""
    try:
        if enabled_sources is None:
            db = SessionLocal()
            try:
                enabled_sources = (
                    db.query(MailSource).filter(MailSource.enabled == True).all()  # noqa: E712
                )
            finally:
                db.close()
        intervals = [s.polling_interval or 60 for s in enabled_sources]
        return max(min_sleep, min(intervals, default=3600) * 60)
    except Exception:  # pylint: disable=broad-exception-caught
        return 3600


async def scheduled_imap_polling():
    """Background task for periodically checking IMAP for new DMARC reports"""
    try:
        while True:
            logger.info("Starting scheduled IMAP polling for DMARC reports")
            try:
                enabled_sources = _poll_all_enabled_sources()
            except Exception as e:  # pylint: disable=broad-exception-caught
                logger.error("Error in IMAP polling task: %s", str(e))
                enabled_sources = None

            try:
                await asyncio.sleep(_next_sleep_seconds(enabled_sources=enabled_sources))
            except asyncio.CancelledError:
                raise
            except Exception as e:  # pylint: disable=broad-exception-caught
                logger.error("Error sleeping in IMAP polling task: %s", str(e))
                await asyncio.sleep(3600)

    except asyncio.CancelledError:
        logger.info("IMAP polling task cancelled")


def _migrate_imap_env_vars_to_db() -> None:
    """
    One-time migration: if IMAP_* environment variables are configured and no
    MailSource rows exist yet, create an initial MailSource from those settings.

    This ensures that existing deployments continue to work without manual
    reconfiguration after the upgrade.
    """
    if not all([settings.IMAP_SERVER, settings.IMAP_USERNAME, settings.IMAP_PASSWORD]):
        return

    db = SessionLocal()
    try:
        if db.query(MailSource).first() is not None:
            return  # already migrated or manually configured

        migrated = MailSource(
            name="Default IMAP (migrated from environment)",
            method="IMAP",
            server=settings.IMAP_SERVER,
            port=settings.IMAP_PORT,
            username=settings.IMAP_USERNAME,
            password=settings.IMAP_PASSWORD,
            use_ssl=True,
            folder="INBOX",
            polling_interval=60,
            enabled=True,
        )
        db.add(migrated)
        db.commit()
        logger.info(
            "Migrated IMAP settings from environment variables to "
            "database (MailSource id=%d). "
            "You can now manage this source via the Mail Sources admin UI.",
            migrated.id,
        )
    except Exception as e:  # pylint: disable=broad-exception-caught
        logger.error("Failed to migrate IMAP env vars to database: %s", str(e))
    finally:
        db.close()


def create_app() -> FastAPI:
    """Create and configure the FastAPI application"""
    application = FastAPI(
        title=settings.PROJECT_NAME,
        openapi_url=f"{settings.API_V1_STR}/openapi.json",
        version="0.1.0",
    )

    # Add security headers middleware
    # Determine environment from settings or environment variable
    environment = os.getenv("ENVIRONMENT", "development")
    application.add_middleware(SecurityHeadersMiddleware, environment=environment)

    # Auth redirect middleware – protects HTML pages; must sit outside CORS
    application.add_middleware(AuthRedirectMiddleware)

    # Improved CORS configuration - restrict to specific methods and headers
    if settings.BACKEND_CORS_ORIGINS:
        application.add_middleware(
            CORSMiddleware,
            allow_origins=[str(origin) for origin in settings.BACKEND_CORS_ORIGINS],
            allow_credentials=True,
            # Security: Restrict to only necessary HTTP methods
            allow_methods=["GET", "POST", "PUT", "DELETE", "OPTIONS"],
            # Security: Specify allowed headers instead of wildcard
            allow_headers=[
                "Content-Type",
                "Authorization",
                "X-API-Key",
                "Accept",
                "Origin",
                "X-Requested-With",
            ],
            # Security: Limit exposed headers
            expose_headers=["Content-Length", "X-RateLimit-Limit"],
            max_age=600,  # Cache preflight requests for 10 minutes
        )

    # Include API router
    application.include_router(api_router, prefix=settings.API_V1_STR)

    # Mount static files directory
    application.mount(
        "/static",
        StaticFiles(directory=os.path.join(os.path.dirname(__file__), "static")),
        name="static",
    )

    # Set up event handlers for startup and shutdown
    @application.on_event("startup")
    async def startup_event():
        """Initialize background tasks and security on application startup"""
        global background_task  # pylint: disable=global-statement

        # Ensure all tables exist (no-op if already present)
        Base.metadata.create_all(bind=engine)

        # Warn loudly when authentication is completely disabled
        if settings.AUTH_DISABLED:
            logger.warning(
                "%s\n"
                "⚠️  AUTH_DISABLED=true — authentication is turned OFF.\n"
                "All requests have unrestricted admin access.\n"
                "Do NOT expose this instance directly to the internet.\n"
                "%s",
                "=" * 80,
                "=" * 80,
            )

        # Load or generate the admin API key
        if settings.ADMIN_API_KEY:
            api_key = settings.ADMIN_API_KEY
            add_api_key(api_key)
            logger.info(
                "Admin API key loaded from ADMIN_API_KEY environment variable "
                "(length: %d chars).",
                len(api_key),
            )
        else:
            api_key = generate_api_key()
            add_api_key(api_key)
            logger.warning(
                "%s\nIMPORTANT: Admin API Key Generated\n"
                "Key length: %d chars. Full key stored securely in memory.\n"
                "Set ADMIN_API_KEY in your environment to use a fixed key across restarts.\n"
                "Use this key in the X-API-Key header for admin endpoints.\n%s",
                "=" * 80,
                len(api_key),
                "=" * 80,
            )

        # One-time migration: if IMAP_* env vars are set and no mail sources exist,
        # create an initial MailSource from those settings so existing deployments
        # continue to work without manual reconfiguration.
        _migrate_imap_env_vars_to_db()

        # Start background polling task (iterates over DB-enabled mail sources)
        logger.info("Starting IMAP polling background task")
        background_task = asyncio.create_task(scheduled_imap_polling())

    @application.on_event("shutdown")
    async def shutdown_event():
        """Clean up background tasks on application shutdown"""
        if background_task:
            logger.info("Cancelling IMAP polling background task")
            background_task.cancel()
            try:
                await background_task
            except asyncio.CancelledError:
                pass

    return application


app = create_app()  # noqa: F811 – intentional rebind; `app` package imported above for side-effects

# Initialize Jinja2 templates
templates_dir = os.path.join(os.path.dirname(__file__), "templates")
templates = Jinja2Templates(directory=templates_dir)


@app.get("/", response_class=HTMLResponse)
async def index(request: Request):
    return templates.TemplateResponse(request, "index.html")


@app.get("/favicon.ico", include_in_schema=False)
async def favicon():
    return FileResponse(os.path.join(os.path.dirname(__file__), "static", "img", "favicon.ico"))


# Individual page routes
@app.get("/dashboard", response_class=HTMLResponse)
async def dashboard(request: Request):
    return templates.TemplateResponse(request, "index.html")


@app.get("/login", response_class=HTMLResponse)
async def login(request: Request, next: str = "/"):
    return templates.TemplateResponse(
        request,
        "login.html",
        {
            "app_name": settings.PROJECT_NAME,
            "logto_configured": settings.logto_configured,
            "auth_disabled": settings.AUTH_DISABLED,
            "next": next,
        },
    )


@app.get("/setup", response_class=HTMLResponse)
async def setup(request: Request):
    return templates.TemplateResponse(
        request,
        "setup.html",
        {
            "app_name": settings.PROJECT_NAME,
            "logto_configured": settings.logto_configured,
        },
    )


@app.get("/domains", response_class=HTMLResponse)
async def domains(request: Request):
    return templates.TemplateResponse(request, "domains.html")


@app.get("/domain/{domain_id}", response_class=HTMLResponse)
async def domain_details(request: Request, domain_id: str):
    """View detailed reports for a specific domain"""
    store = ReportStore.get_instance()
    db = SessionLocal()
    try:
        hydrate_report_store_from_db(db, store)
    finally:
        db.close()
    known_domains = store.get_domains()

    if domain_id not in known_domains:
        # Domain not found, redirect to domains list
        return templates.TemplateResponse(
            request, "domains.html", {"error": f"Domain {domain_id} not found"}
        )

    domain_summary = store.get_domain_summary(domain_id)

    return templates.TemplateResponse(
        request,
        "domain_details.html",
        {
            "domain_id": domain_id,
            "domain": {
                "name": domain_id,
                "description": "",  # Add description if available
                "policy": domain_summary.get("policy", "unknown"),
            },
        },
    )


@app.get("/domains/{domain_id}", response_class=HTMLResponse)
async def domain_details_plural(request: Request, domain_id: str):
    """View detailed reports for a specific domain (plural /domains/ path alias)"""
    return await domain_details(request, domain_id)


@app.get("/reports", response_class=HTMLResponse)
async def reports(request: Request):
    return templates.TemplateResponse(request, "reports.html")


@app.get("/reports/{report_id}", response_class=HTMLResponse)
async def report_detail(request: Request, report_id: str):
    """View detailed information for a specific DMARC report"""
    return templates.TemplateResponse(request, "report_detail.html", {"report_id": report_id})


@app.get("/settings", response_class=HTMLResponse)
async def settings_page(request: Request):
    return templates.TemplateResponse(request, "settings.html")


@app.get("/profile", response_class=HTMLResponse)
async def profile_page(request: Request):
    return templates.TemplateResponse(
        request,
        "profile.html",
        {
            "app_name": settings.PROJECT_NAME,
            "logto_configured": settings.logto_configured,
            "auth_disabled": settings.AUTH_DISABLED,
        },
    )


@app.get("/mail-sources", response_class=HTMLResponse)
async def mail_sources_page(request: Request):
    return templates.TemplateResponse(request, "mail_sources.html")


@app.get("/upload", response_class=HTMLResponse)
async def upload_page(request: Request):
    return templates.TemplateResponse(request, "upload.html")


@app.get("/health", status_code=200, tags=["health"])
@app.get("/healthz", status_code=200, tags=["health"], include_in_schema=False)
async def health():
    """Root-level health check endpoint for Kubernetes liveness/readiness probes."""
    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,
        db=db,
    )
    started_at = datetime.utcnow()
    results = imap_client.fetch_reports(days=7)
    last_check_time = datetime.now()
    source.last_checked = datetime.utcnow()
    record_import_attempt(db, source, results, started_at=started_at, trigger="manual")
    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,
        db=db,
    )
    started_at = datetime.utcnow()
    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()
    record_import_attempt(db, source, results, started_at=started_at, trigger="manual")
    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)):
    """
    Manually trigger IMAP polling for all enabled mail sources (admin only).

    Security: Requires either X-API-Key header or Bearer token
    """
    results_summary = []
    db = SessionLocal()
    try:
        enabled_sources = (
            db.query(MailSource).filter(MailSource.enabled == True).all()  # noqa: E712
        )

        if not enabled_sources:
            return {
                "success": True,
                "message": "No enabled mail sources configured.",
                "sources_polled": 0,
                "authenticated_by": auth.get("auth_type"),
            }

        for source in enabled_sources:
            results_summary.append(_poll_source_for_trigger(source, db))
    finally:
        db.close()

    return {
        "success": all(r.get("success", True) for r in results_summary),
        "timestamp": last_check_time.isoformat() if last_check_time else None,
        "sources": results_summary,
        "authenticated_by": auth.get("auth_type"),
    }


# API endpoint to check status of IMAP polling (public, read-only)
@app.get("/api/v1/poll-status")
async def get_poll_status():
    """
    Get the status of IMAP polling (read-only, no authentication required).
    """
    return {
        "is_running": background_task is not None and not background_task.done(),
        "last_check": last_check_time.isoformat() if last_check_time else None,
    }