File size: 4,378 Bytes
df5d3a2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b2aea8e
df5d3a2
 
 
 
 
 
 
 
b2aea8e
 
 
 
df5d3a2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b2aea8e
df5d3a2
 
 
 
b2aea8e
 
 
df5d3a2
 
 
 
 
 
 
 
 
b2aea8e
df5d3a2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
"""
SOC Dashboard API Server
Thin FastAPI wrapper around the existing OrchestratorCoordinator.
Run: python server.py  β†’  http://localhost:8000
"""

import os
import sys
import json
import asyncio
import uvicorn
from pathlib import Path
from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware

# Ensure project root is on sys.path so "from src.…" imports work
PROJECT_ROOT = Path(__file__).resolve().parent
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from dotenv import load_dotenv
load_dotenv(PROJECT_ROOT / ".env")

if "GOOGLE_API_KEY" in os.environ and "GEMINI_API_KEY" not in os.environ:
    os.environ["GEMINI_API_KEY"] = os.environ["GOOGLE_API_KEY"]

from src.agents.orchestrator.agent import OrchestratorCoordinator

app = FastAPI(title="Security Triage Fleet β€” SOC Dashboard")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# Single orchestrator instance shared across requests
orchestrator = OrchestratorCoordinator()

# ── Serve the dashboard ──────────────────────────────────────────────
@app.get("/", response_class=HTMLResponse)
async def serve_dashboard():
    html_path = PROJECT_ROOT / "dashboard.html"
    return HTMLResponse(content=html_path.read_text(encoding="utf-8"))


# ── Run a full simulation ────────────────────────────────────────────
@app.post("/api/simulate")
async def run_simulation(request: Request):
    """
    Accepts a JSON body with a 'prompt' field (the target context).
    Runs the 3-phase pipeline and returns the simulation report.
    """
    try:
        api_key = request.headers.get("X-API-Key")
        body = await request.json()
        prompt = body.get("prompt", "")
        if not prompt:
            return JSONResponse(
                status_code=400,
                content={"error": "Missing 'prompt' field in request body."}
            )

        if not api_key:
            api_key = body.get("api_key")

        result_json = await orchestrator.invoke(prompt, api_key=api_key)
        return JSONResponse(content=json.loads(result_json))

    except Exception as e:
        return JSONResponse(
            status_code=500,
            content={"error": str(e), "type": type(e).__name__}
        )


# ── HITL Approve / Deny ──────────────────────────────────────────────
@app.post("/api/approve")
async def approve_hitl(request: Request):
    """
    Re-runs the simulation with 'approve' in the prompt so the orchestrator
    triggers the Green Team remediation phase.
    """
    try:
        api_key = request.headers.get("X-API-Key")
        body = await request.json()
        original_prompt = body.get("prompt", "")
        approved = body.get("approved", False)

        if not api_key:
            api_key = body.get("api_key")

        if approved:
            prompt_with_approval = f"{original_prompt}\napprove"
        else:
            # Deny β€” return current state without remediation
            return JSONResponse(content={
                "simulation_outcome": "DENIED",
                "message": "HITL approval denied. Simulation halted at Phase 2."
            })

        result_json = await orchestrator.invoke(prompt_with_approval, api_key=api_key)
        return JSONResponse(content=json.loads(result_json))

    except Exception as e:
        return JSONResponse(
            status_code=500,
            content={"error": str(e), "type": type(e).__name__}
        )


if __name__ == "__main__":
    print("╔══════════════════════════════════════════════╗")
    print("β•‘  Security Triage Fleet β€” SOC Dashboard       β•‘")
    print("β•‘  http://localhost:8000                        β•‘")
    print("β•šβ•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•")
    uvicorn.run(app, host="0.0.0.0", port=8000)