Aegis/backend/app/routers/data_sources.py

"""Data sources management endpoints (admin only).

Provides a centralized panel for managing all external data sources
(Atomic Red Team, Sigma, LOLBAS, GTFOBins, CALDERA, Elastic, etc.)
including sync triggers, enable/disable toggles, and statistics.
"""

import logging
from datetime import datetime

from fastapi import APIRouter, Depends, HTTPException, status
from sqlalchemy.orm import Session

from app.database import get_db
from app.dependencies.auth import require_role
from app.models.user import User
from app.models.data_source import DataSource
from app.services.audit_service import log_action

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/data-sources", tags=["data-sources"])


# ---------------------------------------------------------------------------
# Sync dispatcher — maps source name → import function
# ---------------------------------------------------------------------------

def _get_sync_handler(source_name: str):
    """Lazily import and return the sync function for *source_name*.

    We import lazily to avoid circular imports and to only load the
    modules that are actually needed.
    """
    handlers = {
        "atomic_red_team": ("app.services.atomic_import_service", "import_atomic_red_team"),
        "sigma": ("app.services.sigma_import_service", "sync"),
        "lolbas": ("app.services.lolbas_import_service", "sync"),
        "gtfobins": ("app.services.lolbas_import_service", "sync_gtfobins"),
        "caldera": ("app.services.caldera_import_service", "sync"),
        "elastic_rules": ("app.services.elastic_import_service", "sync"),
        "mitre_cti": ("app.services.threat_actor_import_service", "sync"),
        "d3fend": ("app.services.d3fend_import_service", "sync"),
    }

    if source_name not in handlers:
        return None

    module_path, func_name = handlers[source_name]
    import importlib
    mod = importlib.import_module(module_path)
    return getattr(mod, func_name)


# ---------------------------------------------------------------------------
# Endpoints
# ---------------------------------------------------------------------------


@router.get("")
def list_data_sources(
    db: Session = Depends(get_db),
    current_user: User = Depends(require_role("admin")),
):
    """List all registered data sources.

    **Requires** the ``admin`` role.
    """
    sources = db.query(DataSource).order_by(DataSource.name).all()
    return [
        {
            "id": str(s.id),
            "name": s.name,
            "display_name": s.display_name,
            "type": s.type,
            "url": s.url,
            "description": s.description,
            "is_enabled": s.is_enabled,
            "last_sync_at": s.last_sync_at.isoformat() if s.last_sync_at else None,
            "last_sync_status": s.last_sync_status,
            "last_sync_stats": s.last_sync_stats,
            "sync_frequency": s.sync_frequency,
            "config": s.config,
            "created_at": s.created_at.isoformat() if s.created_at else None,
        }
        for s in sources
    ]


@router.patch("/{source_id}")
def update_data_source(
    source_id: str,
    body: dict,
    db: Session = Depends(get_db),
    current_user: User = Depends(require_role("admin")),
):
    """Update a data source (enable/disable, change config).

    **Requires** the ``admin`` role.

    Body fields (all optional):
    - ``is_enabled`` (bool)
    - ``sync_frequency`` (str)
    - ``config`` (dict)
    """
    ds = db.query(DataSource).filter(DataSource.id == source_id).first()
    if not ds:
        raise HTTPException(status_code=404, detail="Data source not found")

    if "is_enabled" in body:
        ds.is_enabled = bool(body["is_enabled"])
    if "sync_frequency" in body:
        ds.sync_frequency = body["sync_frequency"]
    if "config" in body:
        ds.config = body["config"]

    db.commit()

    log_action(
        db,
        user_id=current_user.id,
        action="update_data_source",
        entity_type="data_source",
        entity_id=str(ds.id),
        details={"updates": body},
    )

    return {"message": "Data source updated", "id": str(ds.id)}


@router.post("/{source_id}/sync")
def sync_data_source(
    source_id: str,
    db: Session = Depends(get_db),
    current_user: User = Depends(require_role("admin")),
):
    """Trigger sync/import for a specific data source.

    **Requires** the ``admin`` role.
    """
    ds = db.query(DataSource).filter(DataSource.id == source_id).first()
    if not ds:
        raise HTTPException(status_code=404, detail="Data source not found")

    handler = _get_sync_handler(ds.name)
    if handler is None:
        raise HTTPException(
            status_code=400,
            detail=f"No sync handler available for '{ds.name}'",
        )

    # Mark as in_progress
    ds.last_sync_status = "in_progress"
    db.commit()

    try:
        summary = handler(db)
    except Exception as exc:
        logger.error("Sync failed for %s: %s", ds.name, exc)
        ds.last_sync_status = "error"
        ds.last_sync_at = datetime.utcnow()
        ds.last_sync_stats = {"error": str(exc)}
        db.commit()
        raise HTTPException(
            status_code=500,
            detail=f"Sync failed: {str(exc)}",
        )

    # Update DS record (the handler may already have done this,
    # but we ensure it here as well)
    ds.last_sync_at = datetime.utcnow()
    ds.last_sync_status = "success"
    ds.last_sync_stats = summary
    db.commit()

    return {
        "message": f"Sync complete for {ds.display_name}",
        "source": ds.name,
        "stats": summary,
    }


@router.post("/sync-all")
def sync_all_data_sources(
    db: Session = Depends(get_db),
    current_user: User = Depends(require_role("admin")),
):
    """Trigger sync for all enabled data sources (sequentially).

    **Requires** the ``admin`` role.
    """
    enabled_sources = (
        db.query(DataSource)
        .filter(DataSource.is_enabled == True)
        .order_by(DataSource.name)
        .all()
    )

    results = []
    for ds in enabled_sources:
        handler = _get_sync_handler(ds.name)
        if handler is None:
            results.append({
                "source": ds.name,
                "status": "skipped",
                "detail": "No sync handler available",
            })
            continue

        ds.last_sync_status = "in_progress"
        db.commit()

        try:
            summary = handler(db)
            ds.last_sync_at = datetime.utcnow()
            ds.last_sync_status = "success"
            ds.last_sync_stats = summary
            db.commit()
            results.append({
                "source": ds.name,
                "status": "success",
                "stats": summary,
            })
        except Exception as exc:
            logger.error("Sync failed for %s: %s", ds.name, exc)
            ds.last_sync_status = "error"
            ds.last_sync_at = datetime.utcnow()
            ds.last_sync_stats = {"error": str(exc)}
            db.commit()
            results.append({
                "source": ds.name,
                "status": "error",
                "detail": str(exc),
            })

    log_action(
        db,
        user_id=current_user.id,
        action="sync_all_data_sources",
        entity_type="data_source",
        entity_id=None,
        details={"results": results},
    )

    return {"message": "Sync all complete", "results": results}


@router.get("/{source_id}/stats")
def get_data_source_stats(
    source_id: str,
    db: Session = Depends(get_db),
    current_user: User = Depends(require_role("admin")),
):
    """Get detailed statistics for a specific data source.

    **Requires** the ``admin`` role.
    """
    ds = db.query(DataSource).filter(DataSource.id == source_id).first()
    if not ds:
        raise HTTPException(status_code=404, detail="Data source not found")

    # Count items from this source
    from app.models.test_template import TestTemplate
    from app.models.detection_rule import DetectionRule

    template_count = 0
    rule_count = 0

    if ds.type == "attack_procedure":
        template_count = (
            db.query(TestTemplate)
            .filter(TestTemplate.source == ds.name)
            .count()
        )
    elif ds.type == "detection_rule":
        rule_count = (
            db.query(DetectionRule)
            .filter(DetectionRule.source == ds.name)
            .count()
        )

    return {
        "id": str(ds.id),
        "name": ds.name,
        "display_name": ds.display_name,
        "type": ds.type,
        "is_enabled": ds.is_enabled,
        "last_sync_at": ds.last_sync_at.isoformat() if ds.last_sync_at else None,
        "last_sync_status": ds.last_sync_status,
        "last_sync_stats": ds.last_sync_stats,
        "total_templates": template_count,
        "total_rules": rule_count,
    }