Files
ThreatHunt/backend/app/api/routes/annotations.py
mblanke 9b98ab9614 feat: interactive network map, IOC highlighting, AUP hunt selector, type filters
- NetworkMap: hunt-scoped force-directed graph with click-to-inspect popover
- NetworkMap: zoom/pan (wheel, drag, buttons), viewport transform
- NetworkMap: clickable IP/Host/Domain/URL legend chips to filter node types
- NetworkMap: brighter colors, 20% smaller nodes
- DatasetViewer: IOC columns highlighted with colored headers + cell tinting
- AUPScanner: hunt dropdown replacing dataset checkboxes, auto-select all
- Rename 'Social Media (Personal)' theme to 'Social Media' with DB migration
- Fix /api/hunts timeout: Dataset.rows lazy='noload' (was selectin cascade)
- Add OS column mapping to normalizer
- Full backend services, DB models, alembic migrations, new routes
- New components: Dashboard, HuntManager, FileUpload, NetworkMap, etc.
- Docker Compose deployment with nginx reverse proxy
2026-02-19 15:41:15 -05:00

312 lines
11 KiB
Python

"""API routes for annotations and hypotheses."""
import logging
from fastapi import APIRouter, Depends, HTTPException, Query
from pydantic import BaseModel, Field
from sqlalchemy import select, func
from sqlalchemy.ext.asyncio import AsyncSession
from app.db import get_db
from app.db.models import Annotation, Hypothesis
logger = logging.getLogger(__name__)
router = APIRouter(tags=["annotations"])
# ── Annotation models ─────────────────────────────────────────────────
class AnnotationCreate(BaseModel):
row_id: int | None = None
dataset_id: str | None = None
text: str = Field(..., max_length=2000)
severity: str = Field(default="info") # info|low|medium|high|critical
tag: str | None = None # suspicious|benign|needs-review
highlight_color: str | None = None
class AnnotationUpdate(BaseModel):
text: str | None = None
severity: str | None = None
tag: str | None = None
highlight_color: str | None = None
class AnnotationResponse(BaseModel):
id: str
row_id: int | None
dataset_id: str | None
author_id: str | None
text: str
severity: str
tag: str | None
highlight_color: str | None
created_at: str
updated_at: str
class AnnotationListResponse(BaseModel):
annotations: list[AnnotationResponse]
total: int
# ── Hypothesis models ─────────────────────────────────────────────────
class HypothesisCreate(BaseModel):
hunt_id: str | None = None
title: str = Field(..., max_length=256)
description: str | None = None
mitre_technique: str | None = None
status: str = Field(default="draft")
class HypothesisUpdate(BaseModel):
title: str | None = None
description: str | None = None
mitre_technique: str | None = None
status: str | None = None # draft|active|confirmed|rejected
evidence_row_ids: list[int] | None = None
evidence_notes: str | None = None
class HypothesisResponse(BaseModel):
id: str
hunt_id: str | None
title: str
description: str | None
mitre_technique: str | None
status: str
evidence_row_ids: list | None
evidence_notes: str | None
created_at: str
updated_at: str
class HypothesisListResponse(BaseModel):
hypotheses: list[HypothesisResponse]
total: int
# ── Annotation routes ─────────────────────────────────────────────────
ann_router = APIRouter(prefix="/api/annotations")
@ann_router.post("", response_model=AnnotationResponse, summary="Create annotation")
async def create_annotation(body: AnnotationCreate, db: AsyncSession = Depends(get_db)):
ann = Annotation(
row_id=body.row_id,
dataset_id=body.dataset_id,
text=body.text,
severity=body.severity,
tag=body.tag,
highlight_color=body.highlight_color,
)
db.add(ann)
await db.flush()
return AnnotationResponse(
id=ann.id, row_id=ann.row_id, dataset_id=ann.dataset_id,
author_id=ann.author_id, text=ann.text, severity=ann.severity,
tag=ann.tag, highlight_color=ann.highlight_color,
created_at=ann.created_at.isoformat(), updated_at=ann.updated_at.isoformat(),
)
@ann_router.get("", response_model=AnnotationListResponse, summary="List annotations")
async def list_annotations(
dataset_id: str | None = Query(None),
row_id: int | None = Query(None),
tag: str | None = Query(None),
severity: str | None = Query(None),
limit: int = Query(100, ge=1, le=1000),
offset: int = Query(0, ge=0),
db: AsyncSession = Depends(get_db),
):
stmt = select(Annotation).order_by(Annotation.created_at.desc())
if dataset_id:
stmt = stmt.where(Annotation.dataset_id == dataset_id)
if row_id:
stmt = stmt.where(Annotation.row_id == row_id)
if tag:
stmt = stmt.where(Annotation.tag == tag)
if severity:
stmt = stmt.where(Annotation.severity == severity)
stmt = stmt.limit(limit).offset(offset)
result = await db.execute(stmt)
annotations = result.scalars().all()
count_stmt = select(func.count(Annotation.id))
if dataset_id:
count_stmt = count_stmt.where(Annotation.dataset_id == dataset_id)
total = (await db.execute(count_stmt)).scalar_one()
return AnnotationListResponse(
annotations=[
AnnotationResponse(
id=a.id, row_id=a.row_id, dataset_id=a.dataset_id,
author_id=a.author_id, text=a.text, severity=a.severity,
tag=a.tag, highlight_color=a.highlight_color,
created_at=a.created_at.isoformat(), updated_at=a.updated_at.isoformat(),
)
for a in annotations
],
total=total,
)
@ann_router.put("/{annotation_id}", response_model=AnnotationResponse, summary="Update annotation")
async def update_annotation(
annotation_id: str, body: AnnotationUpdate, db: AsyncSession = Depends(get_db)
):
result = await db.execute(select(Annotation).where(Annotation.id == annotation_id))
ann = result.scalar_one_or_none()
if not ann:
raise HTTPException(status_code=404, detail="Annotation not found")
if body.text is not None:
ann.text = body.text
if body.severity is not None:
ann.severity = body.severity
if body.tag is not None:
ann.tag = body.tag
if body.highlight_color is not None:
ann.highlight_color = body.highlight_color
await db.flush()
return AnnotationResponse(
id=ann.id, row_id=ann.row_id, dataset_id=ann.dataset_id,
author_id=ann.author_id, text=ann.text, severity=ann.severity,
tag=ann.tag, highlight_color=ann.highlight_color,
created_at=ann.created_at.isoformat(), updated_at=ann.updated_at.isoformat(),
)
@ann_router.delete("/{annotation_id}", summary="Delete annotation")
async def delete_annotation(annotation_id: str, db: AsyncSession = Depends(get_db)):
result = await db.execute(select(Annotation).where(Annotation.id == annotation_id))
ann = result.scalar_one_or_none()
if not ann:
raise HTTPException(status_code=404, detail="Annotation not found")
await db.delete(ann)
return {"message": "Annotation deleted", "id": annotation_id}
# ── Hypothesis routes ─────────────────────────────────────────────────
hyp_router = APIRouter(prefix="/api/hypotheses")
@hyp_router.post("", response_model=HypothesisResponse, summary="Create hypothesis")
async def create_hypothesis(body: HypothesisCreate, db: AsyncSession = Depends(get_db)):
hyp = Hypothesis(
hunt_id=body.hunt_id,
title=body.title,
description=body.description,
mitre_technique=body.mitre_technique,
status=body.status,
)
db.add(hyp)
await db.flush()
return HypothesisResponse(
id=hyp.id, hunt_id=hyp.hunt_id, title=hyp.title,
description=hyp.description, mitre_technique=hyp.mitre_technique,
status=hyp.status, evidence_row_ids=hyp.evidence_row_ids,
evidence_notes=hyp.evidence_notes,
created_at=hyp.created_at.isoformat(), updated_at=hyp.updated_at.isoformat(),
)
@hyp_router.get("", response_model=HypothesisListResponse, summary="List hypotheses")
async def list_hypotheses(
hunt_id: str | None = Query(None),
status: str | None = Query(None),
limit: int = Query(100, ge=1, le=1000),
offset: int = Query(0, ge=0),
db: AsyncSession = Depends(get_db),
):
stmt = select(Hypothesis).order_by(Hypothesis.updated_at.desc())
if hunt_id:
stmt = stmt.where(Hypothesis.hunt_id == hunt_id)
if status:
stmt = stmt.where(Hypothesis.status == status)
stmt = stmt.limit(limit).offset(offset)
result = await db.execute(stmt)
hyps = result.scalars().all()
count_stmt = select(func.count(Hypothesis.id))
if hunt_id:
count_stmt = count_stmt.where(Hypothesis.hunt_id == hunt_id)
total = (await db.execute(count_stmt)).scalar_one()
return HypothesisListResponse(
hypotheses=[
HypothesisResponse(
id=h.id, hunt_id=h.hunt_id, title=h.title,
description=h.description, mitre_technique=h.mitre_technique,
status=h.status, evidence_row_ids=h.evidence_row_ids,
evidence_notes=h.evidence_notes,
created_at=h.created_at.isoformat(), updated_at=h.updated_at.isoformat(),
)
for h in hyps
],
total=total,
)
@hyp_router.get("/{hypothesis_id}", response_model=HypothesisResponse, summary="Get hypothesis")
async def get_hypothesis(hypothesis_id: str, db: AsyncSession = Depends(get_db)):
result = await db.execute(select(Hypothesis).where(Hypothesis.id == hypothesis_id))
hyp = result.scalar_one_or_none()
if not hyp:
raise HTTPException(status_code=404, detail="Hypothesis not found")
return HypothesisResponse(
id=hyp.id, hunt_id=hyp.hunt_id, title=hyp.title,
description=hyp.description, mitre_technique=hyp.mitre_technique,
status=hyp.status, evidence_row_ids=hyp.evidence_row_ids,
evidence_notes=hyp.evidence_notes,
created_at=hyp.created_at.isoformat(), updated_at=hyp.updated_at.isoformat(),
)
@hyp_router.put("/{hypothesis_id}", response_model=HypothesisResponse, summary="Update hypothesis")
async def update_hypothesis(
hypothesis_id: str, body: HypothesisUpdate, db: AsyncSession = Depends(get_db)
):
result = await db.execute(select(Hypothesis).where(Hypothesis.id == hypothesis_id))
hyp = result.scalar_one_or_none()
if not hyp:
raise HTTPException(status_code=404, detail="Hypothesis not found")
if body.title is not None:
hyp.title = body.title
if body.description is not None:
hyp.description = body.description
if body.mitre_technique is not None:
hyp.mitre_technique = body.mitre_technique
if body.status is not None:
hyp.status = body.status
if body.evidence_row_ids is not None:
hyp.evidence_row_ids = body.evidence_row_ids
if body.evidence_notes is not None:
hyp.evidence_notes = body.evidence_notes
await db.flush()
return HypothesisResponse(
id=hyp.id, hunt_id=hyp.hunt_id, title=hyp.title,
description=hyp.description, mitre_technique=hyp.mitre_technique,
status=hyp.status, evidence_row_ids=hyp.evidence_row_ids,
evidence_notes=hyp.evidence_notes,
created_at=hyp.created_at.isoformat(), updated_at=hyp.updated_at.isoformat(),
)
@hyp_router.delete("/{hypothesis_id}", summary="Delete hypothesis")
async def delete_hypothesis(hypothesis_id: str, db: AsyncSession = Depends(get_db)):
result = await db.execute(select(Hypothesis).where(Hypothesis.id == hypothesis_id))
hyp = result.scalar_one_or_none()
if not hyp:
raise HTTPException(status_code=404, detail="Hypothesis not found")
await db.delete(hyp)
return {"message": "Hypothesis deleted", "id": hypothesis_id}