Aegis/backend/app/services/data_source_service.py

"""Data source management service — framework-agnostic query and sync logic.

Provides list, update, sync, and stats. Sync operations commit internally
since they are long-running and self-contained.
"""

from __future__ import annotations

import logging
from datetime import datetime

from sqlalchemy.orm import Session

from app.domain.errors import BusinessRuleViolation, EntityNotFoundError
from app.domain.ports.import_service import get_import_handler
from app.models.data_source import DataSource

logger = logging.getLogger(__name__)


def list_sources(db: Session) -> list[dict]:
    """Return all registered data sources as a list of dicts."""
    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
    ]


def update_source(db: Session, source_id: str, **fields: object) -> None:
    """Update a data source's fields (is_enabled, sync_frequency, config).

    Raises EntityNotFoundError if source does not exist.
    Does not commit; the router handles that.
    """
    ds = db.query(DataSource).filter(DataSource.id == source_id).first()
    if not ds:
        raise EntityNotFoundError("Data source", source_id)

    if "is_enabled" in fields:
        ds.is_enabled = fields["is_enabled"]
    if "sync_frequency" in fields:
        ds.sync_frequency = fields["sync_frequency"]
    if "config" in fields:
        ds.config = fields["config"]


def sync_source(db: Session, source_id: str) -> dict:
    """Trigger sync for a specific data source.

    Raises EntityNotFoundError if source does not exist.
    Raises BusinessRuleViolation if no sync handler is available.
    Commits internally (long-running, self-contained operation).
    Returns dict with message, source, stats.
    """
    ds = db.query(DataSource).filter(DataSource.id == source_id).first()
    if not ds:
        raise EntityNotFoundError("Data source", source_id)

    handler = get_import_handler(ds.name)
    if handler is None:
        raise BusinessRuleViolation(f"No sync handler available for '{ds.name}'")

    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, exc_info=True)
        ds.last_sync_status = "error"
        ds.last_sync_at = datetime.utcnow()
        ds.last_sync_stats = {"error": str(exc)}
        db.commit()
        raise BusinessRuleViolation(
            f"Sync failed for '{ds.display_name}'. Check server logs for details."
        )

    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,
    }


def sync_all_sources(db: Session) -> list[dict]:
    """Trigger sync for all enabled data sources (sequentially).

    Commits internally (long-running, self-contained operation).
    Returns list of result dicts with source, status, stats/detail.
    """
    enabled_sources = (
        db.query(DataSource)
        .filter(DataSource.is_enabled == True)
        .order_by(DataSource.name)
        .all()
    )

    results = []
    for ds in enabled_sources:
        handler = get_import_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, exc_info=True)
            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": "Sync failed. Check server logs for details.",
            })

    return results


def get_source_stats(db: Session, source_id: str) -> dict:
    """Return detailed statistics for a data source.

    Raises EntityNotFoundError if source does not exist.
    """
    ds = db.query(DataSource).filter(DataSource.id == source_id).first()
    if not ds:
        raise EntityNotFoundError("Data source", source_id)

    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,
    }