talos/backend/app/db.py

import json
import uuid
from datetime import datetime, timezone
from typing import Any

import aiosqlite

from app.config import settings
from app.models import BulkRun, Finding, Phase, Scan, ScanOptions

_db: aiosqlite.Connection | None = None


def _utcnow() -> datetime:
    return datetime.now(timezone.utc)


def _dt_iso(dt: datetime | None) -> str | None:
    if dt is None:
        return None
    return dt.isoformat()


def _parse_dt(s: str | None) -> datetime | None:
    if not s:
        return None
    return datetime.fromisoformat(s)


async def init_db() -> None:
    global _db
    _db = await aiosqlite.connect(settings.db_path)
    _db.row_factory = aiosqlite.Row
    await _db.executescript(
        """
        CREATE TABLE IF NOT EXISTS scans (
            id TEXT PRIMARY KEY,
            target TEXT NOT NULL,
            profile TEXT NOT NULL,
            options_json TEXT NOT NULL,
            status TEXT NOT NULL,
            created_at TEXT NOT NULL,
            completed_at TEXT,
            llm_unavailable INTEGER DEFAULT 0
        );
        CREATE TABLE IF NOT EXISTS phases (
            id TEXT PRIMARY KEY,
            scan_id TEXT NOT NULL,
            name TEXT NOT NULL,
            tool TEXT NOT NULL,
            command TEXT NOT NULL,
            status TEXT NOT NULL,
            started_at TEXT,
            completed_at TEXT,
            raw_output TEXT DEFAULT '',
            exit_code INTEGER,
            sort_order INTEGER NOT NULL,
            FOREIGN KEY (scan_id) REFERENCES scans(id)
        );
        CREATE TABLE IF NOT EXISTS findings (
            id TEXT PRIMARY KEY,
            scan_id TEXT NOT NULL,
            phase_id TEXT NOT NULL,
            severity TEXT NOT NULL,
            title TEXT NOT NULL,
            description TEXT NOT NULL,
            evidence TEXT NOT NULL,
            recommendation TEXT NOT NULL,
            cve_refs_json TEXT DEFAULT '[]',
            confidence TEXT NOT NULL,
            FOREIGN KEY (scan_id) REFERENCES scans(id)
        );
        CREATE TABLE IF NOT EXISTS bulk_runs (
            id TEXT PRIMARY KEY,
            profile TEXT NOT NULL,
            status TEXT NOT NULL,
            created_at TEXT NOT NULL,
            completed_at TEXT
        );
        """
    )
    await _migrate_schema()
    await _db.commit()


async def _migrate_schema() -> None:
    """Add columns/tables introduced after initial deploy."""
    cursor = await _conn().execute("PRAGMA table_info(scans)")
    cols = {row[1] for row in await cursor.fetchall()}
    if "bulk_id" not in cols:
        await _conn().execute("ALTER TABLE scans ADD COLUMN bulk_id TEXT")


async def close_db() -> None:
    global _db
    if _db:
        await _db.close()
        _db = None


def _conn() -> aiosqlite.Connection:
    if _db is None:
        raise RuntimeError("Database not initialized")
    return _db


async def create_scan(
    target: str,
    profile: str,
    options: ScanOptions,
    *,
    bulk_id: str | None = None,
) -> Scan:
    scan_id = str(uuid.uuid4())
    now = _utcnow()
    scan = Scan(
        id=scan_id,
        target=target,
        profile=profile,  # type: ignore[arg-type]
        options=options,
        status="queued",
        created_at=now,
        bulk_id=bulk_id,
    )
    await _conn().execute(
        """INSERT INTO scans (id, target, profile, options_json, status, created_at, bulk_id)
           VALUES (?,?,?,?,?,?,?)""",
        (scan_id, target, profile, options.model_dump_json(), "queued", _dt_iso(now), bulk_id),
    )
    await _conn().commit()
    return scan


async def create_bulk_run(profile: str, total: int) -> BulkRun:
    bulk_id = str(uuid.uuid4())
    now = _utcnow()
    await _conn().execute(
        "INSERT INTO bulk_runs (id, profile, status, created_at) VALUES (?,?,?,?)",
        (bulk_id, profile, "queued", _dt_iso(now)),
    )
    await _conn().commit()
    return BulkRun(
        id=bulk_id,
        profile=profile,  # type: ignore[arg-type]
        status="queued",
        created_at=now,
        total=total,
    )


async def _count_bulk_scans(bulk_id: str) -> dict[str, int]:
    cursor = await _conn().execute(
        """SELECT status, COUNT(*) AS n FROM scans WHERE bulk_id=? GROUP BY status""",
        (bulk_id,),
    )
    rows = await cursor.fetchall()
    counts: dict[str, int] = {}
    for r in rows:
        counts[r["status"]] = int(r["n"])
    total = sum(counts.values())
    terminal = {"complete", "failed", "cancelled"}
    done = sum(counts.get(s, 0) for s in terminal)
    return {
        "total": total,
        "complete": counts.get("complete", 0),
        "failed": counts.get("failed", 0),
        "running": counts.get("running", 0),
        "queued": counts.get("queued", 0),
        "done": done,
    }


async def refresh_bulk_run_status(bulk_id: str) -> BulkRun | None:
    bulk = await get_bulk_run(bulk_id)
    if not bulk:
        return None
    stats = await _count_bulk_scans(bulk_id)
    if stats["total"] == 0:
        return bulk
    if stats["running"] > 0 or (stats["queued"] > 0 and stats["done"] > 0):
        new_status = "running"
    elif stats["queued"] == stats["total"]:
        new_status = "queued"
    elif stats["done"] == stats["total"]:
        new_status = "complete"
    else:
        new_status = "running"
    completed_at = _dt_iso(_utcnow()) if new_status == "complete" else None
    await _conn().execute(
        "UPDATE bulk_runs SET status=?, completed_at=COALESCE(?, completed_at) WHERE id=?",
        (new_status, completed_at, bulk_id),
    )
    await _conn().commit()
    return await get_bulk_run(bulk_id)


async def get_bulk_run(bulk_id: str) -> BulkRun | None:
    row = await (await _conn().execute("SELECT * FROM bulk_runs WHERE id=?", (bulk_id,))).fetchone()
    if not row:
        return None
    stats = await _count_bulk_scans(bulk_id)
    return BulkRun(
        id=row["id"],
        profile=row["profile"],  # type: ignore[arg-type]
        status=row["status"],  # type: ignore[arg-type]
        created_at=_parse_dt(row["created_at"]) or _utcnow(),
        completed_at=_parse_dt(row["completed_at"]),
        total=stats["total"],
        complete=stats["complete"],
        failed=stats["failed"],
        running=stats["running"],
        queued=stats["queued"],
    )


async def list_scans_for_bulk(bulk_id: str) -> list[Scan]:
    cursor = await _conn().execute(
        "SELECT * FROM scans WHERE bulk_id=? ORDER BY created_at ASC",
        (bulk_id,),
    )
    rows = await cursor.fetchall()
    return [_row_to_scan(r) for r in rows]


async def has_active_scan_for_target(target: str) -> bool:
    row = await (
        await _conn().execute(
            "SELECT 1 FROM scans WHERE target=? AND status IN ('queued','running') LIMIT 1",
            (target,),
        )
    ).fetchone()
    return row is not None


def _row_to_scan(row: aiosqlite.Row) -> Scan:
    bulk_id = row["bulk_id"] if "bulk_id" in row.keys() else None
    return Scan(
        id=row["id"],
        target=row["target"],
        profile=row["profile"],  # type: ignore[arg-type]
        options=ScanOptions.model_validate_json(row["options_json"]),
        status=row["status"],  # type: ignore[arg-type]
        created_at=_parse_dt(row["created_at"]) or _utcnow(),
        completed_at=_parse_dt(row["completed_at"]),
        llm_unavailable=bool(row["llm_unavailable"]),
        bulk_id=bulk_id,
    )


async def update_scan_status(
    scan_id: str,
    status: str,
    llm_unavailable: bool | None = None,
) -> None:
    completed = _dt_iso(_utcnow()) if status in ("complete", "failed", "cancelled") else None
    if llm_unavailable is not None:
        await _conn().execute(
            "UPDATE scans SET status=?, completed_at=?, llm_unavailable=? WHERE id=?",
            (status, completed, int(llm_unavailable), scan_id),
        )
    else:
        await _conn().execute(
            "UPDATE scans SET status=?, completed_at=? WHERE id=?",
            (status, completed, scan_id),
        )
    await _conn().commit()


async def get_scan(scan_id: str) -> Scan | None:
    row = await (await _conn().execute("SELECT * FROM scans WHERE id=?", (scan_id,))).fetchone()
    if not row:
        return None
    return _row_to_scan(row)


async def list_scans(limit: int = 50) -> list[Scan]:
    cursor = await _conn().execute(
        "SELECT * FROM scans ORDER BY created_at DESC LIMIT ?", (limit,)
    )
    rows = await cursor.fetchall()
    return [_row_to_scan(r) for r in rows]


async def get_bulk_export_payload(bulk_id: str) -> dict[str, Any] | None:
    bulk = await get_bulk_run(bulk_id)
    if not bulk:
        return None
    scans = await list_scans_for_bulk(bulk_id)
    scan_details: list[dict[str, Any]] = []
    for s in scans:
        phases = await get_phases(s.id)
        findings = await get_findings(s.id)
        scan_details.append(
            {
                "scan": s.model_dump(mode="json"),
                "phases": [p.model_dump(mode="json") for p in phases],
                "findings": [f.model_dump(mode="json") for f in findings],
            }
        )
    return {
        "bulk": bulk.model_dump(mode="json"),
        "scans": scan_details,
    }


async def create_phases(scan_id: str, phases: list[tuple[str, str, str, str]]) -> list[Phase]:
    """phases: (id, name, tool, command)"""
    out: list[Phase] = []
    for i, (pid, name, tool, command) in enumerate(phases):
        ph = Phase(
            id=pid,
            scan_id=scan_id,
            name=name,
            tool=tool,
            command=command,
            status="pending",
        )
        await _conn().execute(
            """INSERT INTO phases (id, scan_id, name, tool, command, status, sort_order)
               VALUES (?,?,?,?,?,?,?)""",
            (pid, scan_id, name, tool, command, "pending", i),
        )
        out.append(ph)
    await _conn().commit()
    return out


async def update_phase(
    phase_id: str,
    status: str,
    raw_output: str | None = None,
    exit_code: int | None = None,
    started: bool = False,
) -> None:
    now = _dt_iso(_utcnow())
    if started:
        await _conn().execute(
            "UPDATE phases SET status=?, started_at=? WHERE id=?",
            (status, now, phase_id),
        )
    elif raw_output is not None:
        await _conn().execute(
            "UPDATE phases SET status=?, completed_at=?, raw_output=?, exit_code=? WHERE id=?",
            (status, now, raw_output, exit_code, phase_id),
        )
    else:
        await _conn().execute("UPDATE phases SET status=? WHERE id=?", (status, phase_id))
    await _conn().commit()


async def append_phase_output(phase_id: str, line: str) -> None:
    await _conn().execute(
        "UPDATE phases SET raw_output = raw_output || ? || char(10) WHERE id=?",
        (line, phase_id),
    )
    await _conn().commit()


async def get_phases(scan_id: str) -> list[Phase]:
    cursor = await _conn().execute(
        "SELECT * FROM phases WHERE scan_id=? ORDER BY sort_order", (scan_id,)
    )
    rows = await cursor.fetchall()
    return [
        Phase(
            id=r["id"],
            scan_id=r["scan_id"],
            name=r["name"],
            tool=r["tool"],
            command=r["command"],
            status=r["status"],  # type: ignore[arg-type]
            started_at=_parse_dt(r["started_at"]),
            completed_at=_parse_dt(r["completed_at"]),
            raw_output=r["raw_output"] or "",
            exit_code=r["exit_code"],
        )
        for r in rows
    ]


async def save_finding(finding: Finding) -> None:
    await _conn().execute(
        """INSERT INTO findings
           (id, scan_id, phase_id, severity, title, description, evidence, recommendation, cve_refs_json, confidence)
           VALUES (?,?,?,?,?,?,?,?,?,?)""",
        (
            finding.id,
            finding.scan_id,
            finding.phase_id,
            finding.severity,
            finding.title,
            finding.description,
            finding.evidence,
            finding.recommendation,
            json.dumps(finding.cve_refs),
            finding.confidence,
        ),
    )
    await _conn().commit()


async def get_findings(scan_id: str) -> list[Finding]:
    cursor = await _conn().execute(
        "SELECT * FROM findings WHERE scan_id=? ORDER BY severity", (scan_id,)
    )
    rows = await cursor.fetchall()
    severity_order = {"critical": 0, "high": 1, "medium": 2, "low": 3, "info": 4}
    findings = [
        Finding(
            id=r["id"],
            scan_id=r["scan_id"],
            phase_id=r["phase_id"],
            severity=r["severity"],  # type: ignore[arg-type]
            title=r["title"],
            description=r["description"],
            evidence=r["evidence"],
            recommendation=r["recommendation"],
            cve_refs=json.loads(r["cve_refs_json"] or "[]"),
            confidence=r["confidence"],  # type: ignore[arg-type]
        )
        for r in rows
    ]
    findings.sort(key=lambda f: severity_order.get(f.severity, 99))
    return findings