| 1 | from __future__ import annotations |
| 2 | |
| 3 | from datetime import datetime, timezone |
| 4 | |
| 5 | from fastapi import APIRouter, Depends, HTTPException |
| 6 | from sqlalchemy import select |
| 7 | from sqlalchemy.orm import Session |
| 8 | |
| 9 | from app.api.deps import Perm, get_db, require |
| 10 | from app.models import ( |
| 11 | AuditLog, EnrichmentJob, EntityRelationship, OsintEntity, User, |
| 12 | ) |
| 13 | from app.schemas import EnrichRunRequest |
| 14 | from app.services.enrichment import list_enrichers, run_enricher |
| 15 | |
| 16 | router = APIRouter(tags=["enrichment"]) |
| 17 | |
| 18 | |
| 19 | @router.get("/enrichers") |
| 20 | def enrichers(_: User = Depends(require(Perm.READ))): |
| 21 | return list_enrichers() |
| 22 | |
| 23 | |
| 24 | @router.post("/enrichment/run") |
| 25 | def run(payload: EnrichRunRequest, db: Session = Depends(get_db), |
| 26 | user: User = Depends(require(Perm.RUN_ENRICHMENT))): |
| 27 | job = EnrichmentJob(investigation_id=payload.investigation_id, enricher=payload.enricher, |
| 28 | target=payload.target, status="running") |
| 29 | db.add(job) |
| 30 | db.flush() |
| 31 | |
| 32 | result = run_enricher(payload.enricher, payload.target) |
| 33 | |
| 34 | for ent in result.entities: |
| 35 | existing = db.scalar(select(OsintEntity).where( |
| 36 | OsintEntity.kind == ent.kind, OsintEntity.value == ent.value)) |
| 37 | now = datetime.now(timezone.utc) |
| 38 | if existing: |
| 39 | existing.last_seen = now |
| 40 | existing.attributes = {**(existing.attributes or {}), **ent.attributes} |
| 41 | else: |
| 42 | db.add(OsintEntity(kind=ent.kind, value=ent.value, attributes=ent.attributes, |
| 43 | first_seen=now, last_seen=now)) |
| 44 | for rel in result.relationships: |
| 45 | db.add(EntityRelationship(src_type=rel.src_kind, src_id=rel.src_value, |
| 46 | dst_type=rel.dst_kind, dst_id=rel.dst_value, |
| 47 | relation=rel.relation, confidence=rel.confidence, |
| 48 | evidence_ref=payload.enricher)) |
| 49 | |
| 50 | job.status = "done" if not result.error else "error" |
| 51 | job.result = result.to_dict() |
| 52 | job.confidence = result.confidence |
| 53 | job.finished_at = datetime.now(timezone.utc) |
| 54 | db.add(AuditLog(actor_user_id=user.id, action="enrichment.run", |
| 55 | entity_type="enrichment_job", entity_id=job.id, |
| 56 | meta={"enricher": payload.enricher, "target": payload.target, |
| 57 | "error": result.error})) |
| 58 | db.commit() |
| 59 | return {"job_id": job.id, "status": job.status, "result": result.to_dict()} |
| 60 | |
| 61 | |
| 62 | @router.get("/enrichment/jobs/{job_id}") |
| 63 | def get_job(job_id: str, db: Session = Depends(get_db), _: User = Depends(require(Perm.READ))): |
| 64 | job = db.get(EnrichmentJob, job_id) |
| 65 | if not job: |
| 66 | raise HTTPException(404, "job not found") |
| 67 | return {"id": job.id, "enricher": job.enricher, "target": job.target, |
| 68 | "status": job.status, "confidence": job.confidence, "result": job.result} |