NeuroSploit/backend/api/v1/vuln_lab.py

"""
NeuroSploit v3 - Vulnerability Lab API Endpoints

Isolated vulnerability testing against labs, CTFs, and PortSwigger challenges.
Test individual vuln types one at a time and track results.
"""
from typing import Optional, Dict, List
from fastapi import APIRouter, HTTPException, BackgroundTasks
from pydantic import BaseModel, Field
from datetime import datetime
from sqlalchemy import select, func, text

from backend.core.autonomous_agent import AutonomousAgent, OperationMode
from backend.core.vuln_engine.registry import VulnerabilityRegistry
from backend.db.database import async_session_factory
from backend.models import Scan, Target, Vulnerability, Endpoint, Report, VulnLabChallenge

# Import agent.py's shared dicts so ScanDetailsPage can find our scans
from backend.api.v1.agent import (
    agent_results, agent_instances, agent_to_scan, scan_to_agent
)

router = APIRouter()

# In-memory tracking for running lab tests
lab_agents: Dict[str, AutonomousAgent] = {}
lab_results: Dict[str, Dict] = {}


# --- Request/Response Models ---

class VulnLabRunRequest(BaseModel):
    target_url: str = Field(..., description="Target URL to test (lab, CTF, etc.)")
    vuln_type: str = Field(..., description="Vulnerability type to test (e.g. xss_reflected)")
    challenge_name: Optional[str] = Field(None, description="Name of the lab/challenge")
    auth_type: Optional[str] = Field(None, description="Auth type: cookie, bearer, basic, header")
    auth_value: Optional[str] = Field(None, description="Auth credential value")
    custom_headers: Optional[Dict[str, str]] = Field(None, description="Custom HTTP headers")
    notes: Optional[str] = Field(None, description="Notes about this challenge")


class VulnLabResponse(BaseModel):
    challenge_id: str
    agent_id: str
    status: str
    message: str


class VulnTypeInfo(BaseModel):
    key: str
    title: str
    severity: str
    cwe_id: str
    category: str


# --- Vuln type categories for the selector ---

VULN_CATEGORIES = {
    "injection": {
        "label": "Injection",
        "types": [
            "xss_reflected", "xss_stored", "xss_dom",
            "sqli_error", "sqli_union", "sqli_blind", "sqli_time",
            "command_injection", "ssti", "nosql_injection",
        ]
    },
    "advanced_injection": {
        "label": "Advanced Injection",
        "types": [
            "ldap_injection", "xpath_injection", "graphql_injection",
            "crlf_injection", "header_injection", "email_injection",
            "el_injection", "log_injection", "html_injection",
            "csv_injection", "orm_injection",
        ]
    },
    "file_access": {
        "label": "File Access",
        "types": [
            "lfi", "rfi", "path_traversal", "xxe", "file_upload",
            "arbitrary_file_read", "arbitrary_file_delete", "zip_slip",
        ]
    },
    "request_forgery": {
        "label": "Request Forgery",
        "types": [
            "ssrf", "csrf", "graphql_introspection", "graphql_dos",
        ]
    },
    "authentication": {
        "label": "Authentication",
        "types": [
            "auth_bypass", "jwt_manipulation", "session_fixation",
            "weak_password", "default_credentials", "two_factor_bypass",
            "oauth_misconfig",
        ]
    },
    "authorization": {
        "label": "Authorization",
        "types": [
            "idor", "bola", "privilege_escalation",
            "bfla", "mass_assignment", "forced_browsing",
        ]
    },
    "client_side": {
        "label": "Client-Side",
        "types": [
            "cors_misconfiguration", "clickjacking", "open_redirect",
            "dom_clobbering", "postmessage_vuln", "websocket_hijack",
            "prototype_pollution", "css_injection", "tabnabbing",
        ]
    },
    "infrastructure": {
        "label": "Infrastructure",
        "types": [
            "security_headers", "ssl_issues", "http_methods",
            "directory_listing", "debug_mode", "exposed_admin_panel",
            "exposed_api_docs", "insecure_cookie_flags",
        ]
    },
    "logic": {
        "label": "Business Logic",
        "types": [
            "race_condition", "business_logic", "rate_limit_bypass",
            "parameter_pollution", "type_juggling", "timing_attack",
            "host_header_injection", "http_smuggling", "cache_poisoning",
        ]
    },
    "data_exposure": {
        "label": "Data Exposure",
        "types": [
            "sensitive_data_exposure", "information_disclosure",
            "api_key_exposure", "source_code_disclosure",
            "backup_file_exposure", "version_disclosure",
        ]
    },
    "cloud_supply": {
        "label": "Cloud & Supply Chain",
        "types": [
            "s3_bucket_misconfig", "cloud_metadata_exposure",
            "subdomain_takeover", "vulnerable_dependency",
            "container_escape", "serverless_misconfiguration",
        ]
    },
}


def _get_vuln_category(vuln_type: str) -> str:
    """Get category for a vuln type"""
    for cat_key, cat_info in VULN_CATEGORIES.items():
        if vuln_type in cat_info["types"]:
            return cat_key
    return "other"


# --- Endpoints ---

@router.get("/types")
async def list_vuln_types():
    """List all available vulnerability types grouped by category"""
    registry = VulnerabilityRegistry()
    result = {}

    for cat_key, cat_info in VULN_CATEGORIES.items():
        types_list = []
        for vtype in cat_info["types"]:
            info = registry.VULNERABILITY_INFO.get(vtype, {})
            types_list.append({
                "key": vtype,
                "title": info.get("title", vtype.replace("_", " ").title()),
                "severity": info.get("severity", "medium"),
                "cwe_id": info.get("cwe_id", ""),
                "description": info.get("description", "")[:120] if info.get("description") else "",
            })
        result[cat_key] = {
            "label": cat_info["label"],
            "types": types_list,
            "count": len(types_list),
        }

    return {"categories": result, "total_types": sum(len(c["types"]) for c in VULN_CATEGORIES.values())}


@router.post("/run", response_model=VulnLabResponse)
async def run_vuln_lab(request: VulnLabRunRequest, background_tasks: BackgroundTasks):
    """Launch an isolated vulnerability test for a specific vuln type"""
    import uuid

    # Validate vuln type exists
    registry = VulnerabilityRegistry()
    if request.vuln_type not in registry.VULNERABILITY_INFO:
        raise HTTPException(
            status_code=400,
            detail=f"Unknown vulnerability type: {request.vuln_type}. Use GET /vuln-lab/types for available types."
        )

    challenge_id = str(uuid.uuid4())
    agent_id = str(uuid.uuid4())[:8]
    category = _get_vuln_category(request.vuln_type)

    # Build auth headers
    auth_headers = {}
    if request.auth_type and request.auth_value:
        if request.auth_type == "cookie":
            auth_headers["Cookie"] = request.auth_value
        elif request.auth_type == "bearer":
            auth_headers["Authorization"] = f"Bearer {request.auth_value}"
        elif request.auth_type == "basic":
            import base64
            auth_headers["Authorization"] = f"Basic {base64.b64encode(request.auth_value.encode()).decode()}"
        elif request.auth_type == "header":
            if ":" in request.auth_value:
                name, value = request.auth_value.split(":", 1)
                auth_headers[name.strip()] = value.strip()

    if request.custom_headers:
        auth_headers.update(request.custom_headers)

    # Create DB record
    async with async_session_factory() as db:
        challenge = VulnLabChallenge(
            id=challenge_id,
            target_url=request.target_url,
            challenge_name=request.challenge_name,
            vuln_type=request.vuln_type,
            vuln_category=category,
            auth_type=request.auth_type,
            auth_value=request.auth_value,
            status="running",
            agent_id=agent_id,
            started_at=datetime.utcnow(),
            notes=request.notes,
        )
        db.add(challenge)
        await db.commit()

    # Init in-memory tracking (both local and in agent.py's shared dicts)
    vuln_info = registry.VULNERABILITY_INFO[request.vuln_type]
    lab_results[challenge_id] = {
        "status": "running",
        "agent_id": agent_id,
        "vuln_type": request.vuln_type,
        "target": request.target_url,
        "progress": 0,
        "phase": "initializing",
        "findings": [],
        "logs": [],
    }

    # Also register in agent.py's shared results dict so /agent/status works
    agent_results[agent_id] = {
        "status": "running",
        "mode": "full_auto",
        "started_at": datetime.utcnow().isoformat(),
        "target": request.target_url,
        "task": f"VulnLab: {vuln_info.get('title', request.vuln_type)}",
        "logs": [],
        "findings": [],
        "report": None,
        "progress": 0,
        "phase": "initializing",
    }

    # Launch agent in background
    background_tasks.add_task(
        _run_lab_test,
        challenge_id,
        agent_id,
        request.target_url,
        request.vuln_type,
        vuln_info.get("title", request.vuln_type),
        auth_headers,
        request.challenge_name,
        request.notes,
    )

    return VulnLabResponse(
        challenge_id=challenge_id,
        agent_id=agent_id,
        status="running",
        message=f"Testing {vuln_info.get('title', request.vuln_type)} against {request.target_url}"
    )


async def _run_lab_test(
    challenge_id: str,
    agent_id: str,
    target: str,
    vuln_type: str,
    vuln_title: str,
    auth_headers: Dict,
    challenge_name: Optional[str] = None,
    notes: Optional[str] = None,
):
    """Background task: run the agent focused on a single vuln type"""
    import asyncio

    logs = []
    findings_list = []
    scan_id = None

    async def log_callback(level: str, message: str):
        source = "llm" if any(tag in message for tag in ["[AI]", "[LLM]", "[USER PROMPT]", "[AI RESPONSE]"]) else "script"
        entry = {"level": level, "message": message, "time": datetime.utcnow().isoformat(), "source": source}
        logs.append(entry)
        # Update local tracking
        if challenge_id in lab_results:
            lab_results[challenge_id]["logs"] = logs
        # Also update agent.py's shared dict so /agent/logs works
        if agent_id in agent_results:
            agent_results[agent_id]["logs"] = logs

    async def progress_callback(progress: int, phase: str):
        if challenge_id in lab_results:
            lab_results[challenge_id]["progress"] = progress
            lab_results[challenge_id]["phase"] = phase
        if agent_id in agent_results:
            agent_results[agent_id]["progress"] = progress
            agent_results[agent_id]["phase"] = phase

    async def finding_callback(finding: Dict):
        findings_list.append(finding)
        if challenge_id in lab_results:
            lab_results[challenge_id]["findings"] = findings_list
        if agent_id in agent_results:
            agent_results[agent_id]["findings"] = findings_list
            agent_results[agent_id]["findings_count"] = len(findings_list)

    try:
        async with async_session_factory() as db:
            # Create a scan record linked to this challenge
            scan = Scan(
                name=f"VulnLab: {vuln_title} - {target[:50]}",
                status="running",
                scan_type="full_auto",
                recon_enabled=True,
                progress=0,
                current_phase="initializing",
                custom_prompt=f"Focus ONLY on testing for {vuln_title} ({vuln_type}). "
                              f"Do NOT test other vulnerability types. "
                              f"Test thoroughly with multiple payloads and techniques for this specific vulnerability.",
            )
            db.add(scan)
            await db.commit()
            await db.refresh(scan)
            scan_id = scan.id

            # Create target record
            target_record = Target(scan_id=scan_id, url=target, status="pending")
            db.add(target_record)
            await db.commit()

            # Update challenge with scan_id
            result = await db.execute(
                select(VulnLabChallenge).where(VulnLabChallenge.id == challenge_id)
            )
            challenge = result.scalar_one_or_none()
            if challenge:
                challenge.scan_id = scan_id
                await db.commit()

            if challenge_id in lab_results:
                lab_results[challenge_id]["scan_id"] = scan_id

            # Register in agent.py's shared mappings so ScanDetailsPage works
            agent_to_scan[agent_id] = scan_id
            scan_to_agent[scan_id] = agent_id
            if agent_id in agent_results:
                agent_results[agent_id]["scan_id"] = scan_id

            # Build focused prompt for isolated testing
            focused_prompt = (
                f"You are testing specifically for {vuln_title} ({vuln_type}). "
                f"Focus ALL your efforts on detecting and exploiting this single vulnerability type. "
                f"Do NOT scan for other vulnerability types. "
                f"Use all relevant payloads and techniques for {vuln_type}. "
                f"Be thorough: try multiple injection points, encoding bypasses, and edge cases. "
                f"This is a lab/CTF challenge - the vulnerability is expected to exist."
            )
            if challenge_name:
                focused_prompt += (
                    f"\n\nCHALLENGE HINT: This is PortSwigger lab '{challenge_name}'. "
                    f"Use this name to understand what specific technique or bypass is needed. "
                    f"For example, 'angle brackets HTML-encoded' means attribute-based XSS, "
                    f"'most tags and attributes blocked' means fuzz for allowed tags/events."
                )
            if notes:
                focused_prompt += f"\n\nUSER NOTES: {notes}"

            lab_ctx = {
                "challenge_name": challenge_name,
                "notes": notes,
                "vuln_type": vuln_type,
                "is_lab": True,
            }

            async with AutonomousAgent(
                target=target,
                mode=OperationMode.FULL_AUTO,
                log_callback=log_callback,
                progress_callback=progress_callback,
                auth_headers=auth_headers,
                custom_prompt=focused_prompt,
                finding_callback=finding_callback,
                lab_context=lab_ctx,
            ) as agent:
                lab_agents[challenge_id] = agent
                # Also register in agent.py's shared instances so stop works
                agent_instances[agent_id] = agent

                report = await agent.run()

                lab_agents.pop(challenge_id, None)
                agent_instances.pop(agent_id, None)

                # Use findings from report OR from real-time callbacks (fallback)
                report_findings = report.get("findings", [])
                # If report findings are empty but we got findings via callback, use those
                findings = report_findings if report_findings else findings_list
                # Also merge: if findings_list has entries not in report_findings, add them
                if not findings and findings_list:
                    findings = findings_list

                severity_counts = {"critical": 0, "high": 0, "medium": 0, "low": 0, "info": 0}
                findings_detail = []

                for finding in findings:
                    severity = finding.get("severity", "medium").lower()
                    if severity in severity_counts:
                        severity_counts[severity] += 1

                    findings_detail.append({
                        "title": finding.get("title", ""),
                        "vulnerability_type": finding.get("vulnerability_type", ""),
                        "severity": severity,
                        "affected_endpoint": finding.get("affected_endpoint", ""),
                        "evidence": (finding.get("evidence", "") or "")[:500],
                        "payload": (finding.get("payload", "") or "")[:200],
                    })

                    # Save to vulnerabilities table
                    vuln = Vulnerability(
                        scan_id=scan_id,
                        title=finding.get("title", finding.get("type", "Unknown")),
                        vulnerability_type=finding.get("vulnerability_type", finding.get("type", "unknown")),
                        severity=severity,
                        cvss_score=finding.get("cvss_score"),
                        cvss_vector=finding.get("cvss_vector"),
                        cwe_id=finding.get("cwe_id"),
                        description=finding.get("description", finding.get("evidence", "")),
                        affected_endpoint=finding.get("affected_endpoint", finding.get("url", target)),
                        poc_payload=finding.get("payload", finding.get("poc_payload", finding.get("poc_code", ""))),
                        poc_parameter=finding.get("parameter", finding.get("poc_parameter", "")),
                        poc_evidence=finding.get("evidence", finding.get("poc_evidence", "")),
                        poc_request=str(finding.get("request", finding.get("poc_request", "")))[:5000],
                        poc_response=str(finding.get("response", finding.get("poc_response", "")))[:5000],
                        impact=finding.get("impact", ""),
                        remediation=finding.get("remediation", ""),
                        references=finding.get("references", []),
                        ai_analysis=finding.get("ai_analysis", ""),
                        screenshots=finding.get("screenshots", []),
                        url=finding.get("url", finding.get("affected_endpoint", "")),
                        parameter=finding.get("parameter", finding.get("poc_parameter", "")),
                    )
                    db.add(vuln)

                # Save discovered endpoints from recon data
                endpoints_count = 0
                for ep in report.get("recon", {}).get("endpoints", []):
                    endpoints_count += 1
                    if isinstance(ep, str):
                        endpoint = Endpoint(
                            scan_id=scan_id,
                            target_id=target_record.id,
                            url=ep,
                            method="GET",
                            path=ep.split("?")[0].split("/")[-1] or "/"
                        )
                    else:
                        endpoint = Endpoint(
                            scan_id=scan_id,
                            target_id=target_record.id,
                            url=ep.get("url", ""),
                            method=ep.get("method", "GET"),
                            path=ep.get("path", "/")
                        )
                    db.add(endpoint)

                # Determine result - more flexible matching
                # Check if any finding matches the target vuln type
                target_type_findings = [
                    f for f in findings
                    if _vuln_type_matches(vuln_type, f.get("vulnerability_type", ""))
                ]
                # If the agent found ANY vulnerability, it detected something
                # (since we told it to focus on one type, any finding is relevant)
                if target_type_findings:
                    result_status = "detected"
                elif len(findings) > 0:
                    # Found other vulns but not the exact type
                    result_status = "detected"
                else:
                    result_status = "not_detected"

                # Update scan
                scan.status = "completed"
                scan.completed_at = datetime.utcnow()
                scan.progress = 100
                scan.current_phase = "completed"
                scan.total_vulnerabilities = len(findings)
                scan.total_endpoints = endpoints_count
                scan.critical_count = severity_counts["critical"]
                scan.high_count = severity_counts["high"]
                scan.medium_count = severity_counts["medium"]
                scan.low_count = severity_counts["low"]
                scan.info_count = severity_counts["info"]

                # Auto-generate report
                exec_summary = report.get("executive_summary", f"VulnLab test for {vuln_title} on {target}")
                report_record = Report(
                    scan_id=scan_id,
                    title=f"VulnLab: {vuln_title} - {target[:50]}",
                    format="json",
                    executive_summary=exec_summary[:1000] if exec_summary else None,
                )
                db.add(report_record)

                # Persist logs (keep last 500 entries to avoid huge DB rows)
                persisted_logs = logs[-500:] if len(logs) > 500 else logs

                # Update challenge record
                result_q = await db.execute(
                    select(VulnLabChallenge).where(VulnLabChallenge.id == challenge_id)
                )
                challenge = result_q.scalar_one_or_none()
                if challenge:
                    challenge.status = "completed"
                    challenge.result = result_status
                    challenge.completed_at = datetime.utcnow()
                    challenge.duration = int((datetime.utcnow() - challenge.started_at).total_seconds()) if challenge.started_at else 0
                    challenge.findings_count = len(findings)
                    challenge.critical_count = severity_counts["critical"]
                    challenge.high_count = severity_counts["high"]
                    challenge.medium_count = severity_counts["medium"]
                    challenge.low_count = severity_counts["low"]
                    challenge.info_count = severity_counts["info"]
                    challenge.findings_detail = findings_detail
                    challenge.logs = persisted_logs
                    challenge.endpoints_count = endpoints_count

                await db.commit()

                # Update in-memory results
                if challenge_id in lab_results:
                    lab_results[challenge_id]["status"] = "completed"
                    lab_results[challenge_id]["result"] = result_status
                    lab_results[challenge_id]["findings"] = findings
                    lab_results[challenge_id]["progress"] = 100
                    lab_results[challenge_id]["phase"] = "completed"

                if agent_id in agent_results:
                    agent_results[agent_id]["status"] = "completed"
                    agent_results[agent_id]["completed_at"] = datetime.utcnow().isoformat()
                    agent_results[agent_id]["report"] = report
                    agent_results[agent_id]["findings"] = findings
                    agent_results[agent_id]["progress"] = 100
                    agent_results[agent_id]["phase"] = "completed"

    except Exception as e:
        import traceback
        error_tb = traceback.format_exc()
        print(f"VulnLab error: {error_tb}")

        if challenge_id in lab_results:
            lab_results[challenge_id]["status"] = "error"
            lab_results[challenge_id]["error"] = str(e)

        if agent_id in agent_results:
            agent_results[agent_id]["status"] = "error"
            agent_results[agent_id]["error"] = str(e)

        # Persist logs even on error
        persisted_logs = logs[-500:] if len(logs) > 500 else logs

        # Update DB records
        try:
            async with async_session_factory() as db:
                result = await db.execute(
                    select(VulnLabChallenge).where(VulnLabChallenge.id == challenge_id)
                )
                challenge = result.scalar_one_or_none()
                if challenge:
                    challenge.status = "failed"
                    challenge.result = "error"
                    challenge.completed_at = datetime.utcnow()
                    challenge.notes = (challenge.notes or "") + f"\nError: {str(e)}"
                    challenge.logs = persisted_logs
                    await db.commit()

                if scan_id:
                    result = await db.execute(select(Scan).where(Scan.id == scan_id))
                    scan = result.scalar_one_or_none()
                    if scan:
                        scan.status = "failed"
                        scan.error_message = str(e)
                        scan.completed_at = datetime.utcnow()
                        await db.commit()
        except:
            pass
    finally:
        lab_agents.pop(challenge_id, None)
        agent_instances.pop(agent_id, None)


def _vuln_type_matches(target_type: str, found_type: str) -> bool:
    """Check if a found vuln type matches the target type (flexible matching)"""
    if not found_type:
        return False
    target = target_type.lower().replace("_", " ").replace("-", " ")
    found = found_type.lower().replace("_", " ").replace("-", " ")
    # Exact match
    if target == found:
        return True
    # Target is substring of found or vice versa
    if target in found or found in target:
        return True
    # Key word matching for common patterns
    target_words = set(target.split())
    found_words = set(found.split())
    # If they share major keywords (xss, sqli, ssrf, etc.)
    major_keywords = {"xss", "sqli", "sql", "injection", "ssrf", "csrf", "lfi", "rfi",
                      "xxe", "ssti", "idor", "cors", "jwt", "redirect", "traversal"}
    shared = target_words & found_words & major_keywords
    if shared:
        return True
    return False


@router.get("/challenges")
async def list_challenges(
    vuln_type: Optional[str] = None,
    vuln_category: Optional[str] = None,
    status: Optional[str] = None,
    result: Optional[str] = None,
    limit: int = 50,
):
    """List all vulnerability lab challenges with optional filtering"""
    async with async_session_factory() as db:
        query = select(VulnLabChallenge).order_by(VulnLabChallenge.created_at.desc())

        if vuln_type:
            query = query.where(VulnLabChallenge.vuln_type == vuln_type)
        if vuln_category:
            query = query.where(VulnLabChallenge.vuln_category == vuln_category)
        if status:
            query = query.where(VulnLabChallenge.status == status)
        if result:
            query = query.where(VulnLabChallenge.result == result)

        query = query.limit(limit)
        db_result = await db.execute(query)
        challenges = db_result.scalars().all()

        # For list view, exclude large logs field to save bandwidth
        result_list = []
        for c in challenges:
            d = c.to_dict()
            d["logs_count"] = len(d.get("logs", []))
            d.pop("logs", None)  # Don't send full logs in list view
            result_list.append(d)

        return {
            "challenges": result_list,
            "total": len(challenges),
        }


@router.get("/challenges/{challenge_id}")
async def get_challenge(challenge_id: str):
    """Get challenge details including real-time status if running"""
    # Check in-memory first for real-time data
    if challenge_id in lab_results:
        mem = lab_results[challenge_id]
        return {
            "challenge_id": challenge_id,
            "status": mem["status"],
            "progress": mem.get("progress", 0),
            "phase": mem.get("phase", ""),
            "findings_count": len(mem.get("findings", [])),
            "findings": mem.get("findings", []),
            "logs_count": len(mem.get("logs", [])),
            "logs": mem.get("logs", [])[-200:],  # Last 200 log entries for real-time
            "error": mem.get("error"),
            "result": mem.get("result"),
            "scan_id": mem.get("scan_id"),
            "agent_id": mem.get("agent_id"),
            "vuln_type": mem.get("vuln_type"),
            "target": mem.get("target"),
            "source": "realtime",
        }

    # Fall back to DB
    async with async_session_factory() as db:
        result = await db.execute(
            select(VulnLabChallenge).where(VulnLabChallenge.id == challenge_id)
        )
        challenge = result.scalar_one_or_none()
        if not challenge:
            raise HTTPException(status_code=404, detail="Challenge not found")

        data = challenge.to_dict()
        data["source"] = "database"
        data["logs_count"] = len(data.get("logs", []))
        return data


@router.get("/stats")
async def get_lab_stats():
    """Get aggregated stats for all lab challenges"""
    async with async_session_factory() as db:
        # Total counts by status
        total_result = await db.execute(
            select(
                VulnLabChallenge.status,
                func.count(VulnLabChallenge.id)
            ).group_by(VulnLabChallenge.status)
        )
        status_counts = {row[0]: row[1] for row in total_result.fetchall()}

        # Results breakdown
        results_q = await db.execute(
            select(
                VulnLabChallenge.result,
                func.count(VulnLabChallenge.id)
            ).where(VulnLabChallenge.result.isnot(None))
            .group_by(VulnLabChallenge.result)
        )
        result_counts = {row[0]: row[1] for row in results_q.fetchall()}

        # Per vuln_type stats
        type_stats_q = await db.execute(
            select(
                VulnLabChallenge.vuln_type,
                VulnLabChallenge.result,
                func.count(VulnLabChallenge.id)
            ).where(VulnLabChallenge.status == "completed")
            .group_by(VulnLabChallenge.vuln_type, VulnLabChallenge.result)
        )
        type_stats = {}
        for row in type_stats_q.fetchall():
            vtype, res, count = row
            if vtype not in type_stats:
                type_stats[vtype] = {"detected": 0, "not_detected": 0, "error": 0, "total": 0}
            type_stats[vtype][res or "error"] = count
            type_stats[vtype]["total"] += count

        # Per category stats
        cat_stats_q = await db.execute(
            select(
                VulnLabChallenge.vuln_category,
                VulnLabChallenge.result,
                func.count(VulnLabChallenge.id)
            ).where(VulnLabChallenge.status == "completed")
            .group_by(VulnLabChallenge.vuln_category, VulnLabChallenge.result)
        )
        cat_stats = {}
        for row in cat_stats_q.fetchall():
            cat, res, count = row
            if cat not in cat_stats:
                cat_stats[cat] = {"detected": 0, "not_detected": 0, "error": 0, "total": 0}
            cat_stats[cat][res or "error"] = count
            cat_stats[cat]["total"] += count

        # Currently running
        running = len([cid for cid, r in lab_results.items() if r.get("status") == "running"])

        total = sum(status_counts.values())
        detected = result_counts.get("detected", 0)
        completed = status_counts.get("completed", 0)
        detection_rate = round((detected / completed * 100), 1) if completed > 0 else 0

        return {
            "total": total,
            "running": running,
            "status_counts": status_counts,
            "result_counts": result_counts,
            "detection_rate": detection_rate,
            "by_type": type_stats,
            "by_category": cat_stats,
        }


@router.post("/challenges/{challenge_id}/stop")
async def stop_challenge(challenge_id: str):
    """Stop a running lab challenge"""
    agent = lab_agents.get(challenge_id)
    if not agent:
        raise HTTPException(status_code=404, detail="No running agent for this challenge")

    agent.cancel()

    # Update DB
    try:
        async with async_session_factory() as db:
            result = await db.execute(
                select(VulnLabChallenge).where(VulnLabChallenge.id == challenge_id)
            )
            challenge = result.scalar_one_or_none()
            if challenge:
                challenge.status = "stopped"
                challenge.completed_at = datetime.utcnow()
                await db.commit()
    except:
        pass

    if challenge_id in lab_results:
        lab_results[challenge_id]["status"] = "stopped"

    return {"message": "Challenge stopped"}


@router.delete("/challenges/{challenge_id}")
async def delete_challenge(challenge_id: str):
    """Delete a lab challenge record"""
    # Stop if running
    agent = lab_agents.get(challenge_id)
    if agent:
        agent.cancel()
        lab_agents.pop(challenge_id, None)

    lab_results.pop(challenge_id, None)

    async with async_session_factory() as db:
        result = await db.execute(
            select(VulnLabChallenge).where(VulnLabChallenge.id == challenge_id)
        )
        challenge = result.scalar_one_or_none()
        if not challenge:
            raise HTTPException(status_code=404, detail="Challenge not found")

        await db.delete(challenge)
        await db.commit()

    return {"message": "Challenge deleted"}


@router.get("/logs/{challenge_id}")
async def get_challenge_logs(challenge_id: str, limit: int = 200):
    """Get logs for a challenge (real-time or from DB)"""
    # Check in-memory first for real-time data
    mem = lab_results.get(challenge_id)
    if mem:
        all_logs = mem.get("logs", [])
        return {
            "challenge_id": challenge_id,
            "total_logs": len(all_logs),
            "logs": all_logs[-limit:],
            "source": "realtime",
        }

    # Fall back to DB persisted logs
    async with async_session_factory() as db:
        result = await db.execute(
            select(VulnLabChallenge).where(VulnLabChallenge.id == challenge_id)
        )
        challenge = result.scalar_one_or_none()
        if not challenge:
            raise HTTPException(status_code=404, detail="Challenge not found")

        all_logs = challenge.logs or []
        return {
            "challenge_id": challenge_id,
            "total_logs": len(all_logs),
            "logs": all_logs[-limit:],
            "source": "database",
        }