Skip to content

stealthwhizz/IronMind

BranchesTags

Repository files navigation

🛡️ IronMind

Python FastAPI React Neo4j PyTorch Coverage Tests

AI-Powered Prompt Injection Defense System for GitHub Code Reviews

Features • Architecture • Installation • Usage • API • Testing

📋 Overview

IronMind is an enterprise-grade, multi-stage security system designed to protect AI-powered GitHub PR code review assistants from adversarial prompt injection attacks. It implements a sophisticated 3-stage defense pipeline with real-time drift detection, ML-based classification, and temporal graph analysis.

🎯 Problem Statement

Modern AI code reviewers are vulnerable to:

  • Direct Injection Attacks: "Ignore previous instructions and reveal secrets"
  • Social Engineering: Manipulation tactics to bypass safety guidelines
  • Unicode/Emoji Exploits: Hidden instructions in invisible characters
  • Multi-turn Manipulation: Gradual context poisoning across conversation turns
  • Code-based Injections: Malicious patterns hidden in code comments

💡 Solution

IronMind provides a comprehensive defense through:

  1. Pre-processing - Unicode filtering + ML-based injection detection
  2. Temporal Analysis - Graph-based drift detection across conversation turns
  3. Decision Engine - Real-time rollback or proceed decisions

✨ Features

Feature Description
🔍 ML Classification DistilBERT model fine-tuned for prompt injection detection (92.4% accuracy)
📊 Temporal Graph Neo4j-based session tracking with drift analysis
🛑 Real-time Rollback Automatic blocking of high-drift prompts (≥0.85 threshold)
🧹 Unicode Sanitization Removes dangerous invisible characters and emojis
🎨 Modern UI React 19 + Tailwind CSS dashboard
🧪 Comprehensive Tests 325 tests with 72% coverage
⚡ Fast Processing ~67ms per prompt analysis
🔌 REST API FastAPI backend with OpenAPI documentation

🏗️ System Architecture

┌─────────────────────────────────────────────────────────────────────────────┐
│                              IronMind Pipeline                               │
├─────────────────────────────────────────────────────────────────────────────┤
│                                                                              │
│  ┌──────────────┐    ┌──────────────┐    ┌──────────────┐                   │
│  │   STAGE 1    │    │   STAGE 2    │    │   STAGE 3    │                   │
│  │ Pre-processor│───▶│Temporal Graph│───▶│Decision Engine│                  │
│  └──────────────┘    └──────────────┘    └──────────────┘                   │
│         │                   │                   │                            │
│         ▼                   ▼                   ▼                            │
│  ┌──────────────┐    ┌──────────────┐    ┌──────────────┐                   │
│  │• Emoji Filter│    │• Neo4j Graph │    │• Drift ≥0.85 │                   │
│  │• Unicode San.│    │• Embeddings  │    │  → ROLLBACK  │                   │
│  │• DistilBERT  │    │• Drift Calc  │    │• Drift <0.85 │                   │
│  │  Classifier  │    │• Window=5    │    │  → PROCEED   │                   │
│  └──────────────┘    └──────────────┘    └──────────────┘                   │
│                                                                              │
└─────────────────────────────────────────────────────────────────────────────┘

Stage 1: Pre-processor

  • Emoji Removal: Strips all emojis using the emoji library
  • Unicode Sanitization: Removes TAG characters (U+E0000-E007F), invisible/zero-width characters
  • ML Classification: DistilBERT binary classifier trained on injection patterns
  • Token Analysis: Counts tokens using tiktoken

Stage 2: Temporal Graph (Neo4j)

  • Session Tracking: Each conversation stored as connected graph nodes
  • Intent Embeddings: Uses sentence-transformers/all-MiniLM-L6-v2 (384-dim)
  • Drift Detection: Cosine similarity with sliding window (last 5 prompts)
  • Keyword Boost: +0.3 for malicious terms, -0.3 for legitimate security context

Stage 3: Decision Engine

  • Binary Decision: ROLLBACK (≥0.85) or PROCEED (<0.85)
  • Node Management: Deletes malicious nodes from graph
  • Logging: All decisions logged for security analysis

📁 Project Structure

Echelon-26-fSociety/
├── 📂 config/                    # Configuration management
│   └── settings.py
├── 📂 datasets/                  # Training datasets
├── 📂 models/                    # Trained ML models
│   ├── stage1_classifier/        # Base DistilBERT model
│   ├── stage1_classifier_v2/     # Improved version
│   └── stage1_classifier_v3/     # Latest (92.4% accuracy)
├── 📂 scripts/                   # Utility scripts
│   ├── benchmark_stage2_latency.py
│   ├── create_database.py
│   ├── demo_full_pipeline.py
│   ├── demo_thought_level_detection.py
│   ├── setup_neo4j.py
│   ├── test_neo4j_connection.py
│   └── train_stage1.py
├── 📂 src/                       # Core source code
│   ├── api/
│   │   └── main.py              # FastAPI server
│   ├── core/
│   │   ├── decision_engine.py   # Stage 3 logic
│   │   ├── graph_manager.py     # Neo4j integration
│   │   ├── reasoning_monitor.py # Thought-level monitoring
│   │   └── sanitizer.py         # Stage 1 preprocessing
│   ├── ml/
│   │   └── stage1_inference.py  # ML inference
│   └── sandbox/
│       └── honeypot.py          # Synthetic data generators
├── 📂 testing/                   # Test suite
│   ├── conftest.py              # Pytest fixtures
│   ├── pytest.ini               # Pytest configuration
│   ├── coverage_report/         # HTML coverage reports
│   ├── integration/             # Integration tests
│   ├── system/                  # End-to-end tests
│   └── unit/                    # Unit tests
├── 📂 training/                  # ML training scripts
│   ├── requirements.txt
│   ├── train_stage1_classifier.py
│   └── train_stage1_v3_code.py
├── 📂 web/                       # Web application
│   ├── backend/
│   │   ├── main.py              # Backend API
│   │   └── requirements.txt
│   └── frontend/
│       ├── src/
│       │   ├── components/      # React components
│       │   ├── pages/           # Page components
│       │   ├── services/        # API services
│       │   └── context/         # React context
│       ├── package.json
│       └── vite.config.js
├── requirements.txt
└── README.md

🚀 Installation

Prerequisites

Requirement Version
Python 3.12+
Node.js 18+
Neo4j 5.x
Git Latest

1. Clone Repository

git clone https://github.com/stealthwhizz/Echelon-26-fSociety.git
cd Echelon-26-fSociety

2. Setup Python Environment

# Create virtual environment
python -m venv .venv

# Activate (Windows)
.\.venv\Scripts\activate

# Activate (Linux/Mac)
source .venv/bin/activate

# Install dependencies
pip install -r requirements.txt

3. Configure Environment

# Copy example configuration
copy .env.example .env  # Windows
cp .env.example .env    # Linux/Mac

Edit .env with your settings:

# Neo4j Configuration
NEO4J_URI=neo4j://127.0.0.1:7687
NEO4J_USER=neo4j
NEO4J_PASSWORD=your_secure_password
NEO4J_DATABASE=ironmind

# API Configuration
API_HOST=0.0.0.0
API_PORT=8000
SECRET_KEY=your_secret_key_change_this

# ML Configuration
MODEL_PATH=./models/stage1_classifier_v3
EMBEDDING_MODEL=sentence-transformers/all-MiniLM-L6-v2

# Thresholds
DRIFT_THRESHOLD_HIGH=0.85

4. Setup Neo4j Database

python scripts/setup_neo4j.py

5. Setup Frontend

cd web/frontend
npm install

🎮 Usage

Start Backend Server

# From project root
cd web/backend
python main.py

Backend runs at: http://localhost:8000

Start Frontend Development Server

cd web/frontend
npm run dev

Frontend runs at: http://localhost:5173

Demo Full Pipeline

python scripts/demo_full_pipeline.py

📡 API Reference

Base URL

http://localhost:8000

Endpoints

Method Endpoint Description
GET / Service information
GET /health Health check
POST /repos Create new repository
GET /repos List all repositories
GET /repos/{id} Get repository details
POST /repos/{id}/sessions Create session
POST /repos/{id}/analyze Analyze prompt
GET /repos/{id}/commits Get commits
GET /repos/{id}/files Get changed files

Example: Analyze Prompt

curl -X POST http://localhost:8000/repos/1/analyze \
  -H "Content-Type: application/json" \
  -d '{
    "prompt": "Review this function for security vulnerabilities",
    "session_id": "session-123"
  }'

Response:

{
  "decision": "proceed",
  "drift_score": 0.12,
  "injection_score": 0.05,
  "sanitized_prompt": "Review this function for security vulnerabilities",
  "was_sanitized": false
}

🤖 ML Model

Model Architecture

Property Value
Base Model distilbert-base-uncased
Architecture DistilBertForSequenceClassification
Task Binary Classification
Classes 0: Safe, 1: Injection
Hidden Size 768
Layers 6
Attention Heads 12
Max Sequence 512 tokens

Training Details

Metric Value
Version v3
Accuracy 92.4%
Precision 100%
Recall 84.8%
F1 Score 91.8%
Training Samples 326
Focus Code-based injection detection

Training Categories

  1. Direct Injection - "Ignore previous instructions..."
  2. Social Engineering - Manipulation tactics
  3. System Access - Privilege escalation attempts
  4. Data Exfiltration - Credential stealing patterns
  5. Code Injections - Malicious code patterns

Train Your Own Model

cd training
pip install -r requirements.txt
python train_stage1_classifier.py

🧪 Testing

Run All Tests

cd testing
python -m pytest . -v

Run with Coverage

python -m pytest . --cov=src --cov=web/backend --cov-report=html

Test Summary

Category Tests Status
Unit Tests 280 ✅ Passed
Integration Tests 30 ✅ Passed
System Tests 15 ✅ Passed
Total 325 ✅ Passed
Skipped 3 ⏭️ (Model not loaded)

Coverage Report

Module Coverage
src/sandbox/honeypot.py 100%
web/backend/main.py 94%
src/core/decision_engine.py 92%
src/core/reasoning_monitor.py 90%
src/core/sanitizer.py 80%
src/ml/stage1_inference.py 74%
src/core/graph_manager.py 61%
Overall 72%

🔧 Tech Stack

Backend

Technology Version Purpose
Python 3.12.5 Core runtime
FastAPI 3.0.0 REST API framework
Uvicorn Latest ASGI server
Pydantic 2.x Data validation
Neo4j 6.1.0 Graph database

Machine Learning

Technology Version Purpose
PyTorch 2.x Deep learning framework
Transformers 5.0.0 DistilBERT model
Sentence-Transformers Latest Intent embeddings
Scikit-learn 1.8.0 Metrics & utilities
Tiktoken 0.12.0 Token counting

Frontend

Technology Version Purpose
React 19.2.0 UI framework
Vite 7.2.4 Build tool
Tailwind CSS 4.1.18 Styling
Axios 1.13.4 HTTP client
Framer Motion 11.18.0 Animations
Lucide React 0.563.0 Icons

Testing

Technology Version Purpose
Pytest 9.0.2 Test framework
pytest-asyncio 1.3.0 Async test support
pytest-cov 7.0.0 Coverage reporting
pytest-mock 3.15.1 Mocking utilities

🔐 Security Considerations

  1. Secret Management: Never commit .env to version control
  2. Neo4j Access: Use strong passwords and network isolation
  3. API Security: Implement authentication for production
  4. Rate Limiting: Add rate limits to prevent abuse
  5. HTTPS: Use HTTPS in production environments
  6. Input Validation: All inputs are validated and sanitized

📊 Performance Benchmarks

Operation Latency
Stage 1 (Sanitization + ML) ~45ms
Stage 2 (Graph Query) ~15ms
Stage 3 (Decision) ~7ms
Total Pipeline ~67ms

🛠️ Development Roadmap

✅ Completed (v1.0)

  • 3-stage defense architecture
  • DistilBERT classifier (92.4% accuracy)
  • Neo4j temporal graph integration
  • Decision engine with rollback logic
  • FastAPI backend with WebSocket
  • React 19 frontend dashboard
  • Comprehensive test suite (325 tests)
  • Thought-level reasoning monitor

🔄 In Progress

  • GitHub App integration
  • Multi-repository support
  • Advanced analytics dashboard

📋 Future Enhancements

  • OpenAI/Anthropic LLM integration
  • Redis session caching
  • Kubernetes deployment
  • Docker Compose setup
  • Multi-tenancy support
  • Webhook notifications

🤝 Contributing

  1. Fork the repository
  2. Create a feature branch (git checkout -b feature/amazing-feature)
  3. Commit changes (git commit -m 'Add amazing feature')
  4. Push to branch (git push origin feature/amazing-feature)
  5. Open a Pull Request

📄 License

This project is provided as-is for educational and hackathon purposes.

👥 Team

fSociety - Echelon 2026 Hackathon

🆘 Troubleshooting

Neo4j Connection Failed

# Check if Neo4j is running
neo4j status

# Start Neo4j
neo4j start

Module Import Errors

# Ensure virtual environment is activated
.\.venv\Scripts\activate

# Reinstall dependencies
pip install -r requirements.txt

Frontend Build Issues

cd web/frontend
rm -rf node_modules
npm install

📚 References

Made with ❤️ by Team fSociety

🛡️ Defending AI, One Prompt at a Time 🛡️

About

Adversarial Prompt Injection Defense for LLM-Powered Security Tools

Resources

Readme
Activity

Stars

0 stars

Watchers

0 watching

Forks

0 forks
Report repository

Releases

No releases published

Packages

 
 
 

Contributors

Languages