"""FastAPI server exposing the Adaptive AI Firewall environment. Endpoints: POST /reset — Start a new episode POST /step — Multi-session step (batch actions) POST /step_single — Single-session step (Gymnasium-compatible) GET /state — Current environment state GET /tools — List available tool names POST /tool/{name} — Call a specific tool GET /health — Health check GET /stats — Current episode statistics """ from __future__ import annotations import csv import json import os from pathlib import Path from typing import Any from fastapi import FastAPI, HTTPException from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import HTMLResponse from dotenv import load_dotenv from server.embedded_dashboard_images import IMAGES from server.firewall_environment import FirewallEnvironment, ACTIONS from models import ( HealthResponse, NetworkStatsResponse, ResetRequest, StateResponse, StepRequest, StepResponse, StepSingleRequest, StepSingleResponse, ToolRequest, ToolsListResponse, ) load_dotenv() def _clean_env_value(value: str) -> str: return value.strip().strip("`").strip().strip("'").strip('"').strip() def _resolve_api_key(value: str | None) -> str: return _clean_env_value(value or os.getenv("HF_TOKEN") or "") def _resolve_model(value: str | None) -> str: return _clean_env_value(value or os.getenv("MODEL_NAME") or "") def _resolve_base_url(value: str | None) -> str: return _clean_env_value( value or os.getenv("API_BASE_URL") or "" ) ROOT_DIR = Path(__file__).resolve().parent.parent PERFORMANCE_MATRIX_PATH = ROOT_DIR / "output" / "performance_matrix.csv" GRAPH_SPECS = [ ( "Training Loss", "01_training_loss.png", "Loss trend across self-play rounds.", ), ( "Reward Analysis", "02_reward_analysis.png", "Raw score versus difficulty-normalized reward.", ), ( "Elo Progression", "03_elo_progression.png", "Agent Elo compared with adaptive difficulty Elo.", ), ( "Win Rate", "04_win_rate.png", "Rolling win rate and Elo delta per round.", ), ( "Detection vs FP", "05_detection_fp_rate.png", "Detection, false-positive rate, and efficiency.", ), ( "Difficulty Curve", "06_difficulty_progression.png", "Adaptive curriculum difficulty progression.", ), ] def _load_performance_matrix(limit: int = 10) -> tuple[list[str], list[dict[str, str]]]: """Load a preview of the performance matrix CSV for dashboard rendering.""" if not PERFORMANCE_MATRIX_PATH.exists(): return [], [] with PERFORMANCE_MATRIX_PATH.open("r", encoding="utf-8", newline="") as handle: reader = csv.DictReader(handle) rows = list(reader) headers = [ "Round", "Raw_Score", "Abs_Training_Loss", "Detection_Rate", "FP_Rate", "Efficiency", "Agent_Elo", "Difficulty_Elo", ] preview = [{key: row.get(key, "") for key in headers} for row in rows[:limit]] return headers, preview def _build_graph_cards() -> str: cards: list[str] = [] for title, filename, description in GRAPH_SPECS: image_src = IMAGES.get(filename, "") media = ( f'{title}' if image_src else '
Image unavailable
' ) cards.append( f"""

{title}

{description}

{media}
""" ) return "\n".join(cards) def _build_table_html() -> str: headers, rows = _load_performance_matrix() if not headers or not rows: return '
No performance matrix data available.
' header_html = "".join(f"{header}" for header in headers) body_rows = [] for row in rows: body_rows.append( "" + "".join(f"{row.get(header, '')}" for header in headers) + "" ) body_html = "\n".join(body_rows) return ( '
' + header_html + "" + body_html + "
" ) def build_playground_html() -> str: """Render the main dashboard HTML.""" html = """ Adaptive Firewall Dashboard

Adaptive Firewall Dashboard

Monitor the firewall environment, step through actions, inspect live state, and review the training visuals from the current performance matrix in one place.

Running on Hugging Face Space

Playground

Click Reset to start a new episode, then use Step to apply a single action.

Ready
Action space: 6 discrete actions
0: ALLOW
1: BLOCK
2: INSPECT
3: SANDBOX
4: RATE_LIMIT
5: QUARANTINE
{}

Live Episode

Core state values update after every API action.

Task easy Current difficulty level
Step Count 0 Episode steps processed
Queue Length 0 Sessions waiting for action
Budget 0 Remaining budget
Total Reward 0 Accumulated environment reward
Focus Session - Current session in focus

Firewall Summary

The dashboard below uses the generated training artifacts and the latest API responses.

Detection Rate - Updated from /stats
FP Rate - Updated from /stats
Efficiency - Updated from /stats

Performance Matrix

Preview of the tracked performance matrix used for training analysis.

__TABLE_HTML__

Performance Visuals

Embedded graphs generated from the available dashboard images so the Space can display them reliably.

__GRAPH_CARDS__
""" return html.replace("__GRAPH_CARDS__", _build_graph_cards()).replace("__TABLE_HTML__", _build_table_html()) env = FirewallEnvironment(seed=42) env.reset(task="easy") app = FastAPI( title="Adaptive AI Firewall OpenEnv", version="0.2.0", description="RL environment for adaptive firewall decision making on encrypted traffic.", ) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) @app.get("/", response_class=HTMLResponse) def root() -> HTMLResponse: """Redirect root to the playground UI.""" return HTMLResponse(content=build_playground_html()) @app.get("/health", response_model=HealthResponse) def health() -> HealthResponse: return HealthResponse(status="ok", version="0.2.0") @app.post("/reset", response_model=StateResponse) def reset(request: ResetRequest = ResetRequest()) -> StateResponse: try: state = env.reset(task=request.task, seed=request.seed) return StateResponse(**state) except ValueError as e: raise HTTPException(status_code=400, detail=str(e)) from e @app.post("/step", response_model=StepResponse) def step(request: StepRequest = StepRequest()) -> StepResponse: result = env.step(action_map=request.actions) return StepResponse(**result) @app.post("/step_single", response_model=StepSingleResponse) def step_single(request: StepSingleRequest = None) -> StepSingleResponse: if request is None: raise HTTPException(status_code=422, detail="Body is required for /step_single") try: result = env.step_single(action=request.action) return StepSingleResponse(**result) except ValueError as e: raise HTTPException(status_code=400, detail=str(e)) from e @app.get("/state", response_model=StateResponse) def state() -> StateResponse: return StateResponse(**env.state()) @app.get("/stats", response_model=NetworkStatsResponse) def stats() -> NetworkStatsResponse: return NetworkStatsResponse(**env.get_network_stats()) @app.get("/tools", response_model=ToolsListResponse) def list_tools() -> ToolsListResponse: return ToolsListResponse(tools=env.list_tools()) @app.get("/web", response_class=HTMLResponse) def web_interface() -> HTMLResponse: return HTMLResponse(content=build_playground_html()) @app.get("/schema") def schema() -> Any: return { "observation_space": { "type": "Box", "shape": [22], "low": 0.0, "high": 1.0, }, "action_space": { "type": "Discrete", "n": 6, "actions": ACTIONS, }, } @app.post("/tool/{name}") def call_tool(name: str, request: ToolRequest) -> Any: try: if name == "evaluate_session": return env.evaluate_session(request.kwargs["session_id"]) if name == "take_action": reward, record = env.take_action( session_id=request.kwargs["session_id"], action=int(request.kwargs["action"]), ) return {"reward": reward, "record": record} if name == "get_network_stats": return env.get_network_stats() if name == "get_threat_intelligence": return env.get_threat_intelligence() raise HTTPException(status_code=404, detail=f"unknown tool: {name}") except KeyError as exc: raise HTTPException(status_code=400, detail=f"missing key: {exc}") from exc except ValueError as exc: raise HTTPException(status_code=400, detail=str(exc)) from exc def main() -> None: import uvicorn uvicorn.run(app, host="0.0.0.0", port=7860) if __name__ == "__main__": main()