"""API routes for IOC enrichment.""" import logging from fastapi import APIRouter, Depends, HTTPException, Query from pydantic import BaseModel, Field from sqlalchemy.ext.asyncio import AsyncSession from app.db import get_db from app.services.enrichment import ( enrichment_engine, IOCType, Verdict, EnrichmentResultData, ) logger = logging.getLogger(__name__) router = APIRouter(prefix="/api/enrichment", tags=["enrichment"]) # ── Models ──────────────────────────────────────────────────────────── class EnrichIOCRequest(BaseModel): ioc_value: str = Field(..., max_length=2048, description="IOC value to enrich") ioc_type: str = Field(..., description="IOC type: ip, domain, hash_md5, hash_sha1, hash_sha256, url") skip_cache: bool = False class EnrichBatchRequest(BaseModel): iocs: list[dict] = Field( ..., description="List of {value, type} pairs", max_length=50, ) class EnrichmentResultResponse(BaseModel): ioc_value: str ioc_type: str source: str verdict: str score: float tags: list[str] = [] country: str = "" asn: str = "" org: str = "" last_seen: str = "" raw_data: dict = {} error: str = "" latency_ms: int = 0 class EnrichIOCResponse(BaseModel): ioc_value: str ioc_type: str results: list[EnrichmentResultResponse] overall_verdict: str overall_score: float class EnrichBatchResponse(BaseModel): results: dict[str, list[EnrichmentResultResponse]] total_enriched: int def _to_response(r: EnrichmentResultData) -> EnrichmentResultResponse: return EnrichmentResultResponse( ioc_value=r.ioc_value, ioc_type=r.ioc_type.value, source=r.source, verdict=r.verdict.value, score=r.score, tags=r.tags, country=r.country, asn=r.asn, org=r.org, last_seen=r.last_seen, raw_data=r.raw_data, error=r.error, latency_ms=r.latency_ms, ) def _compute_overall(results: list[EnrichmentResultData]) -> tuple[str, float]: """Compute overall verdict from multiple provider results.""" if not results: return Verdict.UNKNOWN.value, 0.0 verdicts = [r.verdict for r in results if r.verdict != Verdict.ERROR] if not verdicts: return Verdict.ERROR.value, 0.0 if Verdict.MALICIOUS in verdicts: return Verdict.MALICIOUS.value, max(r.score for r in results) elif Verdict.SUSPICIOUS in verdicts: return Verdict.SUSPICIOUS.value, max(r.score for r in results) elif Verdict.CLEAN in verdicts: return Verdict.CLEAN.value, 0.0 return Verdict.UNKNOWN.value, 0.0 # ── Routes ──────────────────────────────────────────────────────────── @router.post( "/ioc", response_model=EnrichIOCResponse, summary="Enrich a single IOC", description="Query all configured providers for an IOC (IP, hash, domain, URL).", ) async def enrich_ioc( body: EnrichIOCRequest, db: AsyncSession = Depends(get_db), ): try: ioc_type = IOCType(body.ioc_type) except ValueError: raise HTTPException( status_code=400, detail=f"Invalid IOC type: {body.ioc_type}. Valid: {[t.value for t in IOCType]}", ) results = await enrichment_engine.enrich_ioc( body.ioc_value, ioc_type, db=db, skip_cache=body.skip_cache, ) overall_verdict, overall_score = _compute_overall(results) return EnrichIOCResponse( ioc_value=body.ioc_value, ioc_type=body.ioc_type, results=[_to_response(r) for r in results], overall_verdict=overall_verdict, overall_score=overall_score, ) @router.post( "/batch", response_model=EnrichBatchResponse, summary="Enrich a batch of IOCs", description="Enrich up to 50 IOCs at once across all providers.", ) async def enrich_batch( body: EnrichBatchRequest, db: AsyncSession = Depends(get_db), ): iocs = [] for item in body.iocs: try: ioc_type = IOCType(item["type"]) iocs.append((item["value"], ioc_type)) except (KeyError, ValueError): continue if not iocs: raise HTTPException(status_code=400, detail="No valid IOCs provided") all_results = await enrichment_engine.enrich_batch(iocs, db=db) return EnrichBatchResponse( results={ k: [_to_response(r) for r in v] for k, v in all_results.items() }, total_enriched=len(all_results), ) @router.post( "/dataset/{dataset_id}", summary="Auto-enrich IOCs in a dataset", description="Automatically extract and enrich IOCs from a dataset's IOC columns.", ) async def enrich_dataset( dataset_id: str, max_iocs: int = Query(50, ge=1, le=200), db: AsyncSession = Depends(get_db), ): from app.db.repositories.datasets import DatasetRepository repo = DatasetRepository(db) dataset = await repo.get_dataset(dataset_id) if not dataset: raise HTTPException(status_code=404, detail="Dataset not found") if not dataset.ioc_columns: return {"message": "No IOC columns detected in this dataset", "results": {}} rows = await repo.get_rows(dataset_id, limit=1000) row_data = [r.data for r in rows] all_results = await enrichment_engine.enrich_dataset_iocs( rows=row_data, ioc_columns=dataset.ioc_columns, db=db, max_iocs=max_iocs, ) return { "dataset_id": dataset_id, "dataset_name": dataset.name, "ioc_columns": dataset.ioc_columns, "results": { k: [_to_response(r) for r in v] for k, v in all_results.items() }, "total_enriched": len(all_results), } @router.get( "/status", summary="Enrichment engine status", description="Check which providers are configured and available.", ) async def enrichment_status(): return enrichment_engine.status()