mirror of
https://github.com/mblanke/ThreatHunt.git
synced 2026-03-01 05:50:21 -05:00
- 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
312 lines
11 KiB
Python
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}
|