CoolFace
Apppublic

ofc01/sentinel-zero

sourceHugging Faceupdated 3mo agoView on Hugging Face
0likes
app.py161 linesDownload Raw Back to root
1import os
2import sys
3import json
4import asyncio
5import threading
6from fastapi import FastAPI, Request
7from fastapi.responses import StreamingResponse, HTMLResponse, RedirectResponse
8from fastapi.staticfiles import StaticFiles
9from fastapi.middleware.cors import CORSMiddleware
10from dotenv import load_dotenv
11
12# Ensure core files are in path
13sys.path.append(os.path.abspath(os.path.dirname(__file__)))
14
15from core.agent import SentinelZeroAgent
16from core.mcp_client import McpClient
17
18load_dotenv()
19
20app = FastAPI(title="Sentinel Zero Security Triage API")
21
22# Enable CORS for local development
23app.add_middleware(
24    CORSMiddleware,
25    allow_origins=["*"],
26    allow_credentials=True,
27    allow_methods=["*"],
28    allow_headers=["*"],
29)
30
31# Helper to push logs from worker thread to asyncio loop
32class SSELoggerCallback:
33    def __init__(self, loop, queue_obj):
34        self.loop = loop
35        self.queue = queue_obj
36
37    def __call__(self, log_entry):
38        self.loop.call_soon_threadsafe(self.queue.put_nowait, log_entry)
39
40@app.get("/api/status")
41def get_status():
42    """Health check endpoint — returns server status, available modes, and registered MCP tools."""
43    from core.mcp_client import McpClient
44    sift_client = McpClient(mode="sift")
45    splunk_client = McpClient(mode="splunk")
46    return {
47        "status": "online",
48        "system": "Sentinel Zero",
49        "version": "1.0.0",
50        "model": "gemini-2.5-flash",
51        "modes": {
52            "splunk": {
53                "description": "Splunk SIEM alert triage mode",
54                "tools": [fn.__name__ for fn in splunk_client.get_gemini_tools()]
55            },
56            "sift": {
57                "description": "SANS SIFT forensic investigation mode",
58                "tools": [fn.__name__ for fn in sift_client.get_gemini_tools()]
59            }
60        },
61        "mcp_server": "sentinel-sift-server (FastMCP)",
62        "evidence_integrity": "SHA-256 verified",
63        "self_correction": "enabled"
64    }
65
66@app.get("/api/alerts")
67def get_alerts():
68    """Returns mock Splunk security alerts for the dashboard."""
69    mock_path = os.path.join("demo_data", "mock_alerts.json")
70    if os.path.exists(mock_path):
71        with open(mock_path, "r") as f:
72            return json.load(f)
73    return []
74
75@app.get("/api/investigate")
76async def investigate(mode: str = "splunk", task: str = "Investigate high-severity security incidents."):
77    # Load default context based on mode
78    if mode == "splunk":
79        alerts_path = os.path.join("demo_data", "mock_alerts.json")
80        if os.path.exists(alerts_path):
81            with open(alerts_path, "r") as f:
82                context = {"splunk_alerts": json.load(f)}
83        else:
84            context = {}
85    else:
86        context = {
87            "forensic_image": "SEC-PROD-SRV01_disk.raw",
88            "memory_dump": "SEC-PROD-SRV01_memory.dmp",
89            "sift_tools": ["fls", "volatility3", "grep"]
90        }
91
92    loop = asyncio.get_running_loop()
93    event_queue = asyncio.Queue()
94
95    def run_agent_thread():
96        mcp = McpClient(mode=mode)
97        agent = SentinelZeroAgent(mcp_client=mcp, mode=mode)
98        agent.logger.callback = SSELoggerCallback(loop, event_queue)
99        
100        try:
101            result = agent.analyze(task, context)
102            loop.call_soon_threadsafe(event_queue.put_nowait, {
103                "timestamp": "",
104                "type": "result",
105                "data": result
106            })
107        except Exception as e:
108            loop.call_soon_threadsafe(event_queue.put_nowait, {
109                "timestamp": "",
110                "type": "error",
111                "message": str(e)
112            })
113
114    # Execute in a daemon thread so FastAPI doesn't block
115    threading.Thread(target=run_agent_thread, daemon=True).start()
116
117    async def event_generator():
118        while True:
119            item = await event_queue.get()
120            yield f"data: {json.dumps(item)}\n\n"
121            if item.get("type") in ["result", "error"]:
122                break
123
124    return StreamingResponse(event_generator(), media_type="text/event-stream")
125
126from fastapi.responses import FileResponse
127
128@app.get("/")
129def get_index():
130    if os.path.exists("frontend/index.html"):
131        return FileResponse("frontend/index.html")
132    return {"status": "Sentinel Zero API running. Frontend index.html not found."}
133
134@app.get("/style.css")
135def get_style():
136    return FileResponse("frontend/style.css")
137
138@app.get("/app.js")
139def get_js():
140    return FileResponse("frontend/app.js")
141
142@app.get("/logo.png")
143def get_logo():
144    return FileResponse("frontend/logo.png")
145
146@app.get("/forensics_art.png")
147def get_art():
148    return FileResponse("frontend/forensics_art.png")
149
150import os
151if os.path.exists("frontend/frames"):
152    app.mount("/frames", StaticFiles(directory="frontend/frames"), name="frames")
153
154if __name__ == "__main__":
155    import uvicorn
156    # Load configuration
157    # FIX: Default port changed from 8000 → 8001 to match .env and run_dashboard.cmd
158    port = int(os.getenv("PORT", 8001))
159    print(f"Starting Sentinel Zero on http://localhost:{port}")
160    uvicorn.run("app:app", host="0.0.0.0", port=port, reload=True)
161