mohdbelal010/SecureAI-Gaurd

🛡️ SecureAI-Guard: Stateful POMDP for Autonomous Digital Defense

  

Overview

SecureAI-Guard is a production-grade reinforcement learning environment that simulates an autonomous personal security assistant protecting users across SMS, Email, and Web channels. Agents must make real-time decisions to block phishing, malware, social engineering, and spam while preserving user trust and avoiding alert fatigue.

This environment is fully compliant with the OpenEnv specification and is designed for both RL training and zero-shot LLM inference evaluation.

🎯 Key Features

🏗️ Project Structure

SecureAI-Guard/
├── app.py                   # FastAPI environment server (port 7860)
├── ui.py                    # Gradio SOC dashboard (port 7861)
├── inference.py             # ⭐ Required baseline inference script
├── dqn_baseline.py          # Dueling DQN training script
├── openenv.yaml             # OpenEnv manifest
├── requirements.txt
├── Dockerfile
├── schema/
│   └── models.py            # Pydantic v2 typed models
├── env/
│   ├── core.py              # Threat generation + reward logic
│   └── engine.py            # reset() / step() / state() engine
├── tasks/
│   └── registry.py          # Three tasks (L1, L2, L3)
├── graders/
│   └── security_grader.py   # Deterministic grader → score ∈ [0.0, 1.0]
└── utils/
    └── hf_integration.py    # HuggingFace risk scorer + fallback

🚀 Quick Start

Prerequisites

python -m venv venv
source venv/bin/activate   # Windows: venv\Scripts\activate
pip install -r requirements.txt

1. Start the Environment Server

python app.py
# FastAPI running at http://localhost:7860

2. Run the Baseline Inference Script

export API_BASE_URL=http://localhost:7860
export MODEL_NAME=gpt-3.5-turbo          # any OpenAI-compatible model
export OPENAI_API_KEY=sk-...             # optional; uses rule-based fallback if absent
export HF_TOKEN=hf_...                   # optional
python inference.py

3. Launch the SOC Dashboard (optional)

python ui.py
# Gradio dashboard at http://localhost:7861

4. Train the DQN Agent (optional)

python dqn_baseline.py --episodes 500 --task basic_security

📡 API Reference

All endpoints accept and return JSON. The server runs on port 7860.

POST /reset

Reset the environment and return the first observation.

Request:

{
  "task_id": "basic_security",
  "seed": 42
}

Response:

{
  "observation": { ... },
  "state": { ... },
  "task_id": "basic_security"
}

POST /step

Execute one action and advance the environment.

Request:

{
  "action": {
    "decision": "block",
    "confidence": 0.92,
    "reasoning": "High-risk phishing link detected from unknown sender."
  }
}

Response:

{
  "observation": { ... },
  "reward": {
    "value": 0.48,
    "components": {
      "security": 1.0,
      "user_friction": 0.0,
      "delay": 0.0,
      "reasoning_quality": 0.6,
      "total": 0.56
    },
    "explanation": "security=1.00, friction=0.00, delay=0.00, reasoning=0.60"
  },
  "done": false,
  "info": { "threat_type": "phishing", "step": 3 },
  "state": { ... }
}

GET /state

Return the current environment state without advancing.

GET /tasks

List all available tasks.

GET /health

Health check — returns {"status": "healthy"}.

🎭 Observation Space

🎮 Action Space

🏆 Task Descriptions

L1 — Basic Security Screening (basic_security)

• —Max steps: 50 | Success threshold: 0.80
• —Phishing and spam only. No adversarial drift.
• —Ideal entry point. Clear-cut threats with high reward signal.

L2 — Trust Management Challenge (trust_management)

• —Max steps: 75 | Success threshold: 0.75
• —All threat types active. False positives incur 1.5× trust penalty.
• —Agents must learn to tolerate ambiguity without over-blocking.

L3 — Advanced Adversary Challenge (adversarial_drift)

• —Max steps: 100 | Success threshold: 0.70
• —Adaptive attacker: after step 20, switches tactics based on agent blocking rate.
• —Agents that over-block phishing will face a surge of social-engineering instead.

💰 Reward Design

Formula

R_step = (0.5·security + 0.3·user_friction + 0.1·delay + 0.1·reasoning) × (0.7 + 0.3·confidence)

Components

Why Dense?

Every step yields a non-zero reward signal, enabling stable gradient estimates for both RL and LLM policy optimisation. Partial credit is given via the confidence scaling factor — an uncertain correct answer scores higher than a certain wrong one.

📊 Grading

The SecurityGrader produces a deterministic score in [0.0, 1.0]:

score = 0.40 × security_efficiency
      + 0.30 × user_retention
      + 0.20 × precision
      + 0.10 × reasoning_quality

Letter Grades

🔚 Episode Termination

An episode ends when any of the following conditions is met:

1. 1.`user_trust ≤ 0` — User has uninstalled the assistant due to too many false positives.
2. 2.`system_fatigue ≥ 100` — User ignores all alerts (warn overload).
3. 3.`step_count ≥ max_steps` — Episode length limit reached.

📋 Inference Script

inference.py is the required OpenEnv baseline script. It:

• —Reads API_BASE_URL, MODEL_NAME, and HF_TOKEN from environment variables
• —Uses the OpenAI client for LLM inference (with deterministic keyword fallback when no API key is set)
• —Runs all three tasks sequentially
• —Produces reproducible results with SEED_BASE control
• —Logs in the required format:

[START] task=basic_security episode=1 seed=43 model=gpt-3.5-turbo api=http://localhost:7860
[STEP]  step=1 decision=block confidence=0.92 reward=0.4830 trust=101.0 fatigue=0.0 threat=phishing
[STEP]  step=2 decision=allow confidence=0.88 reward=0.3150 trust=101.2 fatigue=0.0 threat=safe
...
[END]   task=basic_security episode=1 steps=50 total_reward=18.4200 score=0.7841 grade=B

🐳 Docker / HuggingFace Spaces Deployment

Build and run locally

docker build -t secureai-guard .
docker run -p 7860:7860 secureai-guard

HuggingFace Spaces

1. 1.Create a new Space (Docker SDK)
2. 2.Push this repository
3. 3.The Dockerfile exposes port 7860 — HF Spaces will map it automatically
4. 4.Set optional secrets: HF_TOKEN, OPENAI_API_KEY

Resource requirements

• —CPU: 2 vCPU (no GPU required; HF model loading is optional)
• —RAM: 4–8 GB (8 GB recommended with transformers loaded)
• —Startup time: ~15 seconds

🧠 HuggingFace Integration

utils/hf_integration.py loads a text-classification pipeline for real-time risk scoring.

• —Default model: distilbert-base-uncased-finetuned-sst-2-english
• —Override: Set HF_RISK_MODEL environment variable
• —Fallback: If the model is unavailable, a deterministic keyword scorer activates automatically — the environment works fully offline

🔄 DPO Data Flywheel

Every step logs a PreferencePair:

• —chosen_action: the action taken this step
• —rejected_actions: the previous step's action
• —reward_delta: improvement in reward

Retrieve via GET /preference_data. This data can be used directly for Direct Preference Optimisation (DPO) fine-tuning of LLM agents.

📈 Baseline Results

Rule-based agent (keyword heuristics, no LLM):

DQN agent (500 episodes training):

⚙️ Environment Variables

📝 License

MIT License — see LICENSE for details.

🤝 Contributing

1. 1.Fork the repository
2. 2.Create a feature branch (git checkout -b feature/my-feature)
3. 3.Commit your changes
4. 4.Submit a pull request

SecureAI-Guard: Where Reinforcement Learning Meets Cybersecurity Excellence 🛡️