ThreatHunt/backend/app/api/routes/notebooks.py

"""API routes for investigation notebooks and playbooks."""

import logging
from typing import Optional

from fastapi import APIRouter, Depends, HTTPException, Query
from pydantic import BaseModel, Field
from sqlalchemy import select, func, desc
from sqlalchemy.ext.asyncio import AsyncSession

from app.db import get_db
from app.db.models import Notebook, PlaybookRun, _new_id, _utcnow
from app.services.playbook import (
    get_builtin_playbooks,
    get_playbook_template,
    validate_notebook_cells,
)

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/api/notebooks", tags=["notebooks"])


# ── Pydantic models ──────────────────────────────────────────────────


class NotebookCreate(BaseModel):
    title: str
    description: Optional[str] = None
    cells: Optional[list[dict]] = None
    hunt_id: Optional[str] = None
    case_id: Optional[str] = None
    tags: Optional[list[str]] = None


class NotebookUpdate(BaseModel):
    title: Optional[str] = None
    description: Optional[str] = None
    cells: Optional[list[dict]] = None
    tags: Optional[list[str]] = None


class CellUpdate(BaseModel):
    """Update a single cell or add a new one."""
    cell_id: str
    cell_type: Optional[str] = None
    source: Optional[str] = None
    output: Optional[str] = None
    metadata: Optional[dict] = None


class PlaybookStart(BaseModel):
    playbook_name: str
    hunt_id: Optional[str] = None
    case_id: Optional[str] = None
    started_by: Optional[str] = None


class StepComplete(BaseModel):
    notes: Optional[str] = None
    status: str = "completed"  # completed | skipped


# ── Helpers ───────────────────────────────────────────────────────────


def _notebook_to_dict(nb: Notebook) -> dict:
    return {
        "id": nb.id,
        "title": nb.title,
        "description": nb.description,
        "cells": nb.cells or [],
        "hunt_id": nb.hunt_id,
        "case_id": nb.case_id,
        "owner_id": nb.owner_id,
        "tags": nb.tags or [],
        "cell_count": len(nb.cells or []),
        "created_at": nb.created_at.isoformat() if nb.created_at else None,
        "updated_at": nb.updated_at.isoformat() if nb.updated_at else None,
    }


def _run_to_dict(run: PlaybookRun) -> dict:
    return {
        "id": run.id,
        "playbook_name": run.playbook_name,
        "status": run.status,
        "current_step": run.current_step,
        "total_steps": run.total_steps,
        "step_results": run.step_results or [],
        "hunt_id": run.hunt_id,
        "case_id": run.case_id,
        "started_by": run.started_by,
        "created_at": run.created_at.isoformat() if run.created_at else None,
        "updated_at": run.updated_at.isoformat() if run.updated_at else None,
        "completed_at": run.completed_at.isoformat() if run.completed_at else None,
    }


# ── Notebook CRUD ─────────────────────────────────────────────────────


@router.get("", summary="List notebooks")
async def list_notebooks(
    hunt_id: str | None = Query(None),
    limit: int = Query(50, ge=1, le=200),
    offset: int = Query(0, ge=0),
    db: AsyncSession = Depends(get_db),
):
    stmt = select(Notebook)
    count_stmt = select(func.count(Notebook.id))
    if hunt_id:
        stmt = stmt.where(Notebook.hunt_id == hunt_id)
        count_stmt = count_stmt.where(Notebook.hunt_id == hunt_id)

    total = (await db.execute(count_stmt)).scalar() or 0
    results = (await db.execute(
        stmt.order_by(desc(Notebook.updated_at)).offset(offset).limit(limit)
    )).scalars().all()

    return {"notebooks": [_notebook_to_dict(n) for n in results], "total": total}


@router.get("/{notebook_id}", summary="Get notebook")
async def get_notebook(notebook_id: str, db: AsyncSession = Depends(get_db)):
    nb = await db.get(Notebook, notebook_id)
    if not nb:
        raise HTTPException(status_code=404, detail="Notebook not found")
    return _notebook_to_dict(nb)


@router.post("", summary="Create notebook")
async def create_notebook(body: NotebookCreate, db: AsyncSession = Depends(get_db)):
    cells = validate_notebook_cells(body.cells or [])
    if not cells:
        # Start with a default markdown cell
        cells = [{"id": "cell-0", "cell_type": "markdown", "source": "# Investigation Notes\n\nStart documenting your findings here.", "output": None, "metadata": {}}]

    nb = Notebook(
        id=_new_id(),
        title=body.title,
        description=body.description,
        cells=cells,
        hunt_id=body.hunt_id,
        case_id=body.case_id,
        tags=body.tags,
    )
    db.add(nb)
    await db.commit()
    await db.refresh(nb)
    return _notebook_to_dict(nb)


@router.put("/{notebook_id}", summary="Update notebook")
async def update_notebook(
    notebook_id: str, body: NotebookUpdate, db: AsyncSession = Depends(get_db)
):
    nb = await db.get(Notebook, notebook_id)
    if not nb:
        raise HTTPException(status_code=404, detail="Notebook not found")

    if body.title is not None:
        nb.title = body.title
    if body.description is not None:
        nb.description = body.description
    if body.cells is not None:
        nb.cells = validate_notebook_cells(body.cells)
    if body.tags is not None:
        nb.tags = body.tags

    await db.commit()
    await db.refresh(nb)
    return _notebook_to_dict(nb)


@router.post("/{notebook_id}/cells", summary="Add or update a cell")
async def upsert_cell(
    notebook_id: str, body: CellUpdate, db: AsyncSession = Depends(get_db)
):
    nb = await db.get(Notebook, notebook_id)
    if not nb:
        raise HTTPException(status_code=404, detail="Notebook not found")

    cells = list(nb.cells or [])
    found = False
    for i, c in enumerate(cells):
        if c.get("id") == body.cell_id:
            if body.cell_type is not None:
                cells[i]["cell_type"] = body.cell_type
            if body.source is not None:
                cells[i]["source"] = body.source
            if body.output is not None:
                cells[i]["output"] = body.output
            if body.metadata is not None:
                cells[i]["metadata"] = body.metadata
            found = True
            break

    if not found:
        cells.append({
            "id": body.cell_id,
            "cell_type": body.cell_type or "markdown",
            "source": body.source or "",
            "output": body.output,
            "metadata": body.metadata or {},
        })

    nb.cells = cells
    await db.commit()
    await db.refresh(nb)
    return _notebook_to_dict(nb)


@router.delete("/{notebook_id}/cells/{cell_id}", summary="Delete a cell")
async def delete_cell(
    notebook_id: str, cell_id: str, db: AsyncSession = Depends(get_db)
):
    nb = await db.get(Notebook, notebook_id)
    if not nb:
        raise HTTPException(status_code=404, detail="Notebook not found")

    cells = [c for c in (nb.cells or []) if c.get("id") != cell_id]
    nb.cells = cells
    await db.commit()
    return {"ok": True, "remaining_cells": len(cells)}


@router.delete("/{notebook_id}", summary="Delete notebook")
async def delete_notebook(notebook_id: str, db: AsyncSession = Depends(get_db)):
    nb = await db.get(Notebook, notebook_id)
    if not nb:
        raise HTTPException(status_code=404, detail="Notebook not found")
    await db.delete(nb)
    await db.commit()
    return {"ok": True}


# ── Playbooks ─────────────────────────────────────────────────────────


@router.get("/playbooks/templates", summary="List built-in playbook templates")
async def list_playbook_templates():
    templates = get_builtin_playbooks()
    return {
        "templates": [
            {
                "name": t["name"],
                "description": t["description"],
                "category": t["category"],
                "tags": t["tags"],
                "step_count": len(t["steps"]),
            }
            for t in templates
        ]
    }


@router.get("/playbooks/templates/{name}", summary="Get playbook template detail")
async def get_playbook_template_detail(name: str):
    template = get_playbook_template(name)
    if not template:
        raise HTTPException(status_code=404, detail="Playbook template not found")
    return template


@router.post("/playbooks/start", summary="Start a playbook run")
async def start_playbook(body: PlaybookStart, db: AsyncSession = Depends(get_db)):
    template = get_playbook_template(body.playbook_name)
    if not template:
        raise HTTPException(status_code=404, detail="Playbook template not found")

    run = PlaybookRun(
        id=_new_id(),
        playbook_name=body.playbook_name,
        status="in-progress",
        current_step=1,
        total_steps=len(template["steps"]),
        step_results=[],
        hunt_id=body.hunt_id,
        case_id=body.case_id,
        started_by=body.started_by,
    )
    db.add(run)
    await db.commit()
    await db.refresh(run)
    return _run_to_dict(run)


@router.get("/playbooks/runs", summary="List playbook runs")
async def list_playbook_runs(
    status: str | None = Query(None),
    hunt_id: str | None = Query(None),
    limit: int = Query(50, ge=1, le=200),
    db: AsyncSession = Depends(get_db),
):
    stmt = select(PlaybookRun)
    if status:
        stmt = stmt.where(PlaybookRun.status == status)
    if hunt_id:
        stmt = stmt.where(PlaybookRun.hunt_id == hunt_id)

    results = (await db.execute(
        stmt.order_by(desc(PlaybookRun.created_at)).limit(limit)
    )).scalars().all()

    return {"runs": [_run_to_dict(r) for r in results]}


@router.get("/playbooks/runs/{run_id}", summary="Get playbook run detail")
async def get_playbook_run(run_id: str, db: AsyncSession = Depends(get_db)):
    run = await db.get(PlaybookRun, run_id)
    if not run:
        raise HTTPException(status_code=404, detail="Run not found")

    # Also include the template steps
    template = get_playbook_template(run.playbook_name)
    result = _run_to_dict(run)
    result["steps"] = template["steps"] if template else []
    return result


@router.post("/playbooks/runs/{run_id}/complete-step", summary="Complete current playbook step")
async def complete_step(
    run_id: str, body: StepComplete, db: AsyncSession = Depends(get_db)
):
    run = await db.get(PlaybookRun, run_id)
    if not run:
        raise HTTPException(status_code=404, detail="Run not found")
    if run.status != "in-progress":
        raise HTTPException(status_code=400, detail="Run is not in progress")

    step_results = list(run.step_results or [])
    step_results.append({
        "step": run.current_step,
        "status": body.status,
        "notes": body.notes,
        "completed_at": _utcnow().isoformat(),
    })
    run.step_results = step_results

    if run.current_step >= run.total_steps:
        run.status = "completed"
        run.completed_at = _utcnow()
    else:
        run.current_step += 1

    await db.commit()
    await db.refresh(run)
    return _run_to_dict(run)


@router.post("/playbooks/runs/{run_id}/abort", summary="Abort a playbook run")
async def abort_run(run_id: str, db: AsyncSession = Depends(get_db)):
    run = await db.get(PlaybookRun, run_id)
    if not run:
        raise HTTPException(status_code=404, detail="Run not found")
    run.status = "aborted"
    run.completed_at = _utcnow()
    await db.commit()
    return _run_to_dict(run)