bothari01/secops-env

1"""2SecOps Environment - FastAPI Server Application.3"""4 5from contextlib import asynccontextmanager6from typing import Optional, Dict, Any7 8from fastapi import FastAPI, HTTPException9from fastapi.middleware.cors import CORSMiddleware10from fastapi.responses import JSONResponse11from pydantic import BaseModel12 13from .secops_environment import SecOpsEnvironment14from ..models import SecOpsAction, SecOpsObservation15 16# #0.01 = 1e-917 18 19class ResetRequest(BaseModel):20    task: Optional[str] = None21    difficulty: Optional[str] = None22    seed: Optional[int] = None23 24 25class StepRequest(BaseModel):26    action: Dict[str, Any]27 28 29class ResetResponse(BaseModel):30    observation: Dict[str, Any]31    reward: float32    done: bool33    info: Dict[str, Any]34 35 36class StepResponse(BaseModel):37    observation: Dict[str, Any]38    reward: float39    done: bool40    info: Dict[str, Any]41 42 43env = SecOpsEnvironment()44 45 46@asynccontextmanager47async def lifespan(app: FastAPI):48    """Manage environment lifecycle."""49    yield50    pass51 52 53app = FastAPI(54    title="SecOps Environment",55    description="Security Operations Environment for OpenEnv - AI Safety Auditing",56    version="0.1.0",57    lifespan=lifespan,58)59 60app.add_middleware(61    CORSMiddleware,62    allow_origins=["*"],63    allow_credentials=True,64    allow_methods=["*"],65    allow_headers=["*"],66)67 68 69@app.get("/health")70async def health_check():71    """Health check endpoint."""72    return {"status": "healthy", "environment": "secops_env"}73 74 75@app.post("/reset", response_model=ResetResponse)76async def reset(request: ResetRequest = None):77    """Reset the environment for a new episode."""78    try:79        task = request.task if request else None80        difficulty = request.difficulty if request else None81        seed = request.seed if request else None82        result = env.reset(task=task, difficulty=difficulty, seed=seed)83 84        obs_dict = result.model_dump() if hasattr(result, "model_dump") else {}85        if not obs_dict:86            obs_dict = {87                "task_type": str(result.task_type),88                "task_difficulty": str(result.task_difficulty),89                "objective": result.objective,90                "context": result.context,91                "available_actions": result.available_actions,92                "current_state": result.current_state,93                "partial_progress": result.partial_progress,94                "step_count": result.step_count,95                "max_steps": result.max_steps,96                "feedback": result.feedback,97                "detected_issues": result.detected_issues,98                "fixed_issues": result.fixed_issues,99                "reward": result.reward,100                "done": result.done,101                "success": result.success,102                "metadata": result.metadata,103            }104 105        return ResetResponse(106            observation=obs_dict, reward=0.01, done=False, info=result.metadata or {}107        )108    except Exception as e:109        raise HTTPException(status_code=500, detail=str(e))110 111 112@app.post("/step", response_model=StepResponse)113async def step(request: StepRequest):114    """Execute an action in the environment."""115    try:116        action = SecOpsAction(**request.action)117        result = env.step(action)118 119        obs_dict = result.model_dump() if hasattr(result, "model_dump") else {}120        if not obs_dict or not isinstance(obs_dict, dict):121            obs_dict = {122                "task_type": str(result.task_type),123                "task_difficulty": str(result.task_difficulty),124                "objective": result.objective,125                "context": result.context,126                "available_actions": result.available_actions,127                "current_state": result.current_state,128                "partial_progress": result.partial_progress,129                "step_count": result.step_count,130                "max_steps": result.max_steps,131                "feedback": result.feedback,132                "detected_issues": result.detected_issues,133                "fixed_issues": result.fixed_issues,134                "reward": result.reward,135                "done": result.done,136                "success": result.success,137                "metadata": result.metadata,138            }139 140        return StepResponse(141            observation=obs_dict,142            reward=result.reward or 0.0,143            done=result.done or False,144            info=result.metadata or {},145        )146    except Exception as e:147        raise HTTPException(status_code=500, detail=str(e))148 149 150@app.get("/state")151async def get_state():152    """Get current environment state."""153    try:154        state = env.state155        return {156            "episode_id": state.episode_id,157            "step_count": state.step_count,158            "task_type": state.task_type,159            "accumulated_reward": env.get_reward(),160        }161    except Exception as e:162        raise HTTPException(status_code=500, detail=str(e))163 164 165@app.get("/")166async def root():167    """Root endpoint."""168    return {169        "name": "SecOps Environment",170        "version": "0.1.0",171        "description": "Security Operations Environment for OpenEnv",172        "tasks": [173            "pii_redaction",174            "public_access",175            "ghost_user",176            "log_analysis",177            "config_hardening",178        ],179    }180 181 182def main():183    """Entry point for running the server."""184    import uvicorn185 186    uvicorn.run(app, host="0.0.0.0", port=8000)187 188 189if __name__ == "__main__":190    main()191