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