mirror of
https://github.com/mblanke/ThreatHunt.git
synced 2026-03-01 05:50:21 -05:00
- Implemented PlaybookManager for creating and managing investigation playbooks with templates. - Added SavedSearches component for managing bookmarked queries and recurring scans. - Introduced TimelineView for visualizing forensic event timelines with zoomable charts. - Enhanced backend processing with auto-queued jobs for dataset uploads and improved database concurrency. - Updated frontend components for better user experience and performance optimizations. - Documented changes in update log for future reference.
218 lines
8.7 KiB
Python
218 lines
8.7 KiB
Python
"""API routes for investigation playbooks."""
|
|
|
|
import logging
|
|
from datetime import datetime, timezone
|
|
|
|
from fastapi import APIRouter, Depends, HTTPException, status
|
|
from pydantic import BaseModel, Field
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.db import get_db
|
|
from app.db.models import Playbook, PlaybookStep
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
router = APIRouter(prefix="/api/playbooks", tags=["playbooks"])
|
|
|
|
|
|
# -- Request / Response schemas ---
|
|
|
|
class StepCreate(BaseModel):
|
|
title: str
|
|
description: str | None = None
|
|
step_type: str = "manual"
|
|
target_route: str | None = None
|
|
|
|
class PlaybookCreate(BaseModel):
|
|
name: str
|
|
description: str | None = None
|
|
hunt_id: str | None = None
|
|
is_template: bool = False
|
|
steps: list[StepCreate] = []
|
|
|
|
class PlaybookUpdate(BaseModel):
|
|
name: str | None = None
|
|
description: str | None = None
|
|
status: str | None = None
|
|
|
|
class StepUpdate(BaseModel):
|
|
is_completed: bool | None = None
|
|
notes: str | None = None
|
|
|
|
|
|
# -- Default investigation templates ---
|
|
|
|
DEFAULT_TEMPLATES = [
|
|
{
|
|
"name": "Standard Threat Hunt",
|
|
"description": "Step-by-step investigation workflow for a typical threat hunting engagement.",
|
|
"steps": [
|
|
{"title": "Upload Artifacts", "description": "Import CSV exports from Velociraptor or other tools", "step_type": "upload", "target_route": "/upload"},
|
|
{"title": "Create Hunt", "description": "Create a new hunt and associate uploaded datasets", "step_type": "action", "target_route": "/hunts"},
|
|
{"title": "AUP Keyword Scan", "description": "Run AUP keyword scanner for policy violations", "step_type": "analysis", "target_route": "/aup"},
|
|
{"title": "Auto-Triage", "description": "Trigger AI triage on all datasets", "step_type": "analysis", "target_route": "/analysis"},
|
|
{"title": "Review Triage Results", "description": "Review flagged rows and risk scores", "step_type": "review", "target_route": "/analysis"},
|
|
{"title": "Enrich IOCs", "description": "Enrich flagged IPs, hashes, and domains via external sources", "step_type": "analysis", "target_route": "/enrichment"},
|
|
{"title": "Host Profiling", "description": "Generate deep host profiles for suspicious hosts", "step_type": "analysis", "target_route": "/analysis"},
|
|
{"title": "Cross-Hunt Correlation", "description": "Identify shared IOCs and patterns across hunts", "step_type": "analysis", "target_route": "/correlation"},
|
|
{"title": "Document Hypotheses", "description": "Record investigation hypotheses with MITRE mappings", "step_type": "action", "target_route": "/hypotheses"},
|
|
{"title": "Generate Report", "description": "Generate final AI-assisted hunt report", "step_type": "action", "target_route": "/analysis"},
|
|
],
|
|
},
|
|
{
|
|
"name": "Incident Response Triage",
|
|
"description": "Fast-track workflow for active incident response.",
|
|
"steps": [
|
|
{"title": "Upload Artifacts", "description": "Import forensic data from affected hosts", "step_type": "upload", "target_route": "/upload"},
|
|
{"title": "Auto-Triage", "description": "Immediate AI triage for threat indicators", "step_type": "analysis", "target_route": "/analysis"},
|
|
{"title": "IOC Extraction", "description": "Extract all IOCs from flagged data", "step_type": "analysis", "target_route": "/analysis"},
|
|
{"title": "Enrich Critical IOCs", "description": "Priority enrichment of high-risk indicators", "step_type": "analysis", "target_route": "/enrichment"},
|
|
{"title": "Network Map", "description": "Visualize host connections and lateral movement", "step_type": "review", "target_route": "/network"},
|
|
{"title": "Generate Situation Report", "description": "Create executive summary for incident command", "step_type": "action", "target_route": "/analysis"},
|
|
],
|
|
},
|
|
]
|
|
|
|
|
|
# -- Routes ---
|
|
|
|
@router.get("")
|
|
async def list_playbooks(
|
|
include_templates: bool = True,
|
|
hunt_id: str | None = None,
|
|
db: AsyncSession = Depends(get_db),
|
|
):
|
|
q = select(Playbook)
|
|
if hunt_id:
|
|
q = q.where(Playbook.hunt_id == hunt_id)
|
|
if not include_templates:
|
|
q = q.where(Playbook.is_template == False)
|
|
q = q.order_by(Playbook.created_at.desc())
|
|
result = await db.execute(q.limit(100))
|
|
playbooks = result.scalars().all()
|
|
|
|
return {"playbooks": [
|
|
{
|
|
"id": p.id, "name": p.name, "description": p.description,
|
|
"is_template": p.is_template, "hunt_id": p.hunt_id,
|
|
"status": p.status,
|
|
"total_steps": len(p.steps),
|
|
"completed_steps": sum(1 for s in p.steps if s.is_completed),
|
|
"created_at": p.created_at.isoformat() if p.created_at else None,
|
|
}
|
|
for p in playbooks
|
|
]}
|
|
|
|
|
|
@router.get("/templates")
|
|
async def get_templates():
|
|
"""Return built-in investigation templates."""
|
|
return {"templates": DEFAULT_TEMPLATES}
|
|
|
|
|
|
@router.post("", status_code=status.HTTP_201_CREATED)
|
|
async def create_playbook(body: PlaybookCreate, db: AsyncSession = Depends(get_db)):
|
|
pb = Playbook(
|
|
name=body.name,
|
|
description=body.description,
|
|
hunt_id=body.hunt_id,
|
|
is_template=body.is_template,
|
|
)
|
|
db.add(pb)
|
|
await db.flush()
|
|
|
|
created_steps = []
|
|
for i, step in enumerate(body.steps):
|
|
s = PlaybookStep(
|
|
playbook_id=pb.id,
|
|
order_index=i,
|
|
title=step.title,
|
|
description=step.description,
|
|
step_type=step.step_type,
|
|
target_route=step.target_route,
|
|
)
|
|
db.add(s)
|
|
created_steps.append(s)
|
|
|
|
await db.flush()
|
|
|
|
return {
|
|
"id": pb.id, "name": pb.name, "description": pb.description,
|
|
"hunt_id": pb.hunt_id, "is_template": pb.is_template,
|
|
"steps": [
|
|
{"id": s.id, "order_index": s.order_index, "title": s.title,
|
|
"description": s.description, "step_type": s.step_type,
|
|
"target_route": s.target_route, "is_completed": False}
|
|
for s in created_steps
|
|
],
|
|
}
|
|
|
|
|
|
@router.get("/{playbook_id}")
|
|
async def get_playbook(playbook_id: str, db: AsyncSession = Depends(get_db)):
|
|
result = await db.execute(select(Playbook).where(Playbook.id == playbook_id))
|
|
pb = result.scalar_one_or_none()
|
|
if not pb:
|
|
raise HTTPException(status_code=404, detail="Playbook not found")
|
|
|
|
return {
|
|
"id": pb.id, "name": pb.name, "description": pb.description,
|
|
"is_template": pb.is_template, "hunt_id": pb.hunt_id,
|
|
"status": pb.status,
|
|
"created_at": pb.created_at.isoformat() if pb.created_at else None,
|
|
"steps": [
|
|
{
|
|
"id": s.id, "order_index": s.order_index, "title": s.title,
|
|
"description": s.description, "step_type": s.step_type,
|
|
"target_route": s.target_route,
|
|
"is_completed": s.is_completed,
|
|
"completed_at": s.completed_at.isoformat() if s.completed_at else None,
|
|
"notes": s.notes,
|
|
}
|
|
for s in sorted(pb.steps, key=lambda x: x.order_index)
|
|
],
|
|
}
|
|
|
|
|
|
@router.put("/{playbook_id}")
|
|
async def update_playbook(playbook_id: str, body: PlaybookUpdate, db: AsyncSession = Depends(get_db)):
|
|
result = await db.execute(select(Playbook).where(Playbook.id == playbook_id))
|
|
pb = result.scalar_one_or_none()
|
|
if not pb:
|
|
raise HTTPException(status_code=404, detail="Playbook not found")
|
|
|
|
if body.name is not None:
|
|
pb.name = body.name
|
|
if body.description is not None:
|
|
pb.description = body.description
|
|
if body.status is not None:
|
|
pb.status = body.status
|
|
return {"status": "updated"}
|
|
|
|
|
|
@router.delete("/{playbook_id}")
|
|
async def delete_playbook(playbook_id: str, db: AsyncSession = Depends(get_db)):
|
|
result = await db.execute(select(Playbook).where(Playbook.id == playbook_id))
|
|
pb = result.scalar_one_or_none()
|
|
if not pb:
|
|
raise HTTPException(status_code=404, detail="Playbook not found")
|
|
await db.delete(pb)
|
|
return {"status": "deleted"}
|
|
|
|
|
|
@router.put("/steps/{step_id}")
|
|
async def update_step(step_id: int, body: StepUpdate, db: AsyncSession = Depends(get_db)):
|
|
result = await db.execute(select(PlaybookStep).where(PlaybookStep.id == step_id))
|
|
step = result.scalar_one_or_none()
|
|
if not step:
|
|
raise HTTPException(status_code=404, detail="Step not found")
|
|
|
|
if body.is_completed is not None:
|
|
step.is_completed = body.is_completed
|
|
step.completed_at = datetime.now(timezone.utc) if body.is_completed else None
|
|
if body.notes is not None:
|
|
step.notes = body.notes
|
|
return {"status": "updated", "is_completed": step.is_completed}
|
|
|