talos/backend/app/main.py

import asyncio
import json
from collections import defaultdict
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Any

import httpx
import structlog
from fastapi import FastAPI, HTTPException, Request, WebSocket, WebSocketDisconnect
from fastapi.responses import FileResponse, JSONResponse, PlainTextResponse
from fastapi.staticfiles import StaticFiles

from app import db, export, llm, profiles, scanner, tools
from app.config import settings
from app.models import BulkScanCreate, ScanCreate
from app.safety import validate_scan_request

structlog.configure(processors=[structlog.processors.JSONRenderer()])
log = structlog.get_logger(__name__)

def _frontend_dir() -> Path:
    here = Path(__file__).resolve()
    for candidate in (here.parent.parent / "frontend", here.parent.parent.parent / "frontend"):
        if (candidate / "index.html").exists():
            return candidate
    return here.parent.parent.parent / "frontend"


FRONTEND_DIR = _frontend_dir()


@asynccontextmanager
async def lifespan(_app: FastAPI):
    Path(settings.db_path).parent.mkdir(parents=True, exist_ok=True)
    await db.init_db()
    yield
    await db.close_db()


app = FastAPI(title="Talos", lifespan=lifespan)

if FRONTEND_DIR.exists():
    app.mount("/static", StaticFiles(directory=FRONTEND_DIR), name="static")


@app.get("/")
async def index() -> FileResponse:
    return FileResponse(FRONTEND_DIR / "index.html")


@app.get("/api/health")
async def health() -> dict[str, Any]:
    ollama_ok = await llm.ollama_health()
    tool_checks: dict[str, str] = {}
    for name, cmd in [("nmap", ["nmap", "--version"]), ("curl", ["curl", "--version"])]:
        try:
            r = await tools.run_tool(cmd, timeout=10)
            tool_checks[name] = "ok" if r.exit_code == 0 else "error"
        except Exception:
            tool_checks[name] = "error"
    return {
        "ollama": "ok" if ollama_ok else "down",
        "tools": tool_checks,
        "db": "ok",
        "max_concurrent_scans": settings.max_concurrent_scans,
    }


@app.get("/api/profiles")
async def api_profiles() -> list[dict[str, object]]:
    return profiles.get_profiles_payload()


@app.post("/api/scans")
async def create_scan(body: ScanCreate, request: Request) -> dict[str, str]:
    target = validate_scan_request(body.target, body.authorized)
    busy, reason = await scanner.target_busy(target)
    if busy:
        raise HTTPException(status_code=429, detail=reason)
    scan = await db.create_scan(target, body.profile, body.options)
    source_ip = request.client.host if request.client else "unknown"
    await scanner.queue_scan(scan, source_ip=source_ip)
    return {"scan_id": scan.id}


@app.post("/api/scans/bulk")
async def create_bulk_scans(body: BulkScanCreate, request: Request) -> dict[str, Any]:
    source_ip = request.client.host if request.client else "unknown"
    targets = [t.strip() for t in (body.targets or []) if t and t.strip()]
    if not targets:
        raise HTTPException(status_code=400, detail="At least one target is required")

    bulk = await db.create_bulk_run(body.profile, total=0)
    queued: list[dict[str, str]] = []
    rejected: list[dict[str, str]] = []
    seen_targets: set[str] = set()

    for raw in targets:
        try:
            target = validate_scan_request(raw, body.authorized)
            if target in seen_targets:
                rejected.append({"target": target, "reason": "Duplicate target in this bulk list"})
                continue
            seen_targets.add(target)
            busy, reason = await scanner.target_busy(target)
            if busy:
                rejected.append({"target": target, "reason": reason})
                continue
            scan = await db.create_scan(
                target, body.profile, body.options, bulk_id=bulk.id
            )
            await scanner.queue_scan(scan, source_ip=source_ip)
            queued.append({"target": target, "scan_id": scan.id})
        except HTTPException as exc:
            rejected.append({"target": raw, "reason": str(exc.detail)})

    await db.refresh_bulk_run_status(bulk.id)
    bulk = await db.get_bulk_run(bulk.id)
    return {
        "bulk_id": bulk.id,
        "queued": queued,
        "rejected": rejected,
        "bulk": bulk.model_dump(mode="json") if bulk else None,
    }


@app.get("/api/bulks/{bulk_id}")
async def get_bulk_status(bulk_id: str) -> dict[str, Any]:
    bulk = await db.refresh_bulk_run_status(bulk_id)
    if not bulk:
        raise HTTPException(status_code=404, detail="Bulk run not found")
    scans = await db.list_scans_for_bulk(bulk_id)
    return {
        "bulk": bulk.model_dump(mode="json"),
        "scans": [
            {
                "id": s.id,
                "target": s.target,
                "status": s.status,
                "created_at": s.created_at.isoformat(),
            }
            for s in scans
        ],
    }


@app.get("/api/bulks/{bulk_id}/export")
async def export_bulk(bulk_id: str, format: str = "cursor") -> Any:
    payload = await db.get_bulk_export_payload(bulk_id)
    if not payload:
        raise HTTPException(status_code=404, detail="Bulk run not found")
    if format in ("md", "cursor"):
        text = export.render_bulk_markdown(payload, for_cursor=(format == "cursor"))
        return PlainTextResponse(text, media_type="text/markdown")
    if format == "json":
        return JSONResponse(export.render_bulk_json(payload))
    raise HTTPException(status_code=400, detail="format must be cursor, md, or json")


@app.get("/api/scans")
async def list_scans() -> list[dict[str, Any]]:
    scans = await db.list_scans()
    return [
        {
            "id": s.id,
            "target": s.target,
            "profile": s.profile,
            "status": s.status,
            "created_at": s.created_at.isoformat(),
        }
        for s in scans
    ]


@app.get("/api/scans/{scan_id}")
async def get_scan_detail(scan_id: str) -> dict[str, Any]:
    scan = await db.get_scan(scan_id)
    if not scan:
        raise HTTPException(status_code=404, detail="Scan not found")
    phases = await db.get_phases(scan_id)
    findings = await db.get_findings(scan_id)
    return {
        "scan": scan.model_dump(mode="json"),
        "phases": [p.model_dump(mode="json") for p in phases],
        "findings": [f.model_dump(mode="json") for f in findings],
    }


@app.get("/api/scans/{scan_id}/findings")
async def get_scan_findings(scan_id: str) -> dict[str, list[dict[str, Any]]]:
    scan = await db.get_scan(scan_id)
    if not scan:
        raise HTTPException(status_code=404, detail="Scan not found")
    findings = await db.get_findings(scan_id)
    grouped: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for f in findings:
        grouped[f.severity].append(f.model_dump(mode="json"))
    return dict(grouped)


@app.post("/api/scans/{scan_id}/cancel")
async def cancel_scan(scan_id: str) -> dict[str, bool]:
    scan = await db.get_scan(scan_id)
    if not scan:
        raise HTTPException(status_code=404, detail="Scan not found")
    await scanner.cancel_scan(scan_id)
    return {"cancelled": True}


@app.get("/api/scans/{scan_id}/export")
async def export_scan(scan_id: str, format: str = "json") -> Any:
    detail = await get_scan_detail(scan_id)
    if format == "md":
        lines = [
            f"# Talos Report — {detail['scan']['target']}",
            f"Profile: {detail['scan']['profile']}",
            f"Status: {detail['scan']['status']}",
            "",
        ]
        for f in detail["findings"]:
            lines.append(f"## [{f['severity'].upper()}] {f['title']}")
            lines.append(f["description"])
            lines.append(f"**Recommendation:** {f['recommendation']}")
            lines.append("")
        return PlainTextResponse("\n".join(lines), media_type="text/markdown")
    return JSONResponse(detail)


@app.websocket("/ws/scans/{scan_id}")
async def ws_scan(websocket: WebSocket, scan_id: str) -> None:
    scan = await db.get_scan(scan_id)
    if not scan:
        await websocket.close(code=4004)
        return
    await websocket.accept()
    bc = scanner.get_broadcaster(scan_id)
    queue = bc.subscribe()
    try:
        # If the client connects immediately after queueing a scan, the scan task may not
        # have created phase rows yet. Poll briefly so the UI can render the full phase list.
        phases = await db.get_phases(scan_id)
        if not phases and scan.status in ("queued", "running"):
            for _ in range(20):  # ~2s total
                await asyncio.sleep(0.1)
                phases = await db.get_phases(scan_id)
                if phases:
                    break
        for p in phases:
            await websocket.send_json(
                {"type": "phase", "phase_id": p.id, "phase": p.name, "status": p.status}
            )
        findings = await db.get_findings(scan_id)
        for f in findings:
            await websocket.send_json({"type": "finding", "finding": f.model_dump(mode="json")})

        while True:
            try:
                msg = await asyncio.wait_for(queue.get(), timeout=1.0)
                await websocket.send_json(msg)
                if msg.get("type") == "done":
                    break
            except TimeoutError:
                if not scanner.is_scan_active(scan_id):
                    current = await db.get_scan(scan_id)
                    if current and current.status in ("complete", "failed", "cancelled"):
                        await websocket.send_json(
                            {"type": "done", "scan_id": scan_id, "status": current.status}
                        )
                        break
    except WebSocketDisconnect:
        pass
    finally:
        if queue in bc._queues:
            bc._queues.remove(queue)