Share
## https://sploitus.com/exploit?id=10598381-D0D4-53FD-A2A7-155D7E77213A
# ExploitRAG - RAG-based Cybersecurity Chat System
A production-grade FastAPI application for RAG-based ExploitDB analysis with streaming responses, conversation persistence, and intelligent follow-up detection.
## ๐๏ธ Architecture
```
Client โ FastAPI โ RAG Engine โ OpenAI API
โ
PostgreSQL + ChromaDB + Redis
```
### Tech Stack
- **FastAPI** - Async web framework
- **PostgreSQL** - Conversation & message persistence
- **ChromaDB** - Vector database for exploit embeddings
- **OpenAI API** - LLM (GPT-4) and embeddings
- **Redis** - Caching layer
- **SQLAlchemy 2.0** - Async ORM
- **Alembic** - Database migrations
- **JWT** - Authentication
## ๐ Project Structure
```
backend/
โโโ app/
โ โโโ main.py # FastAPI application entry point
โ โโโ core/
โ โ โโโ config.py # Settings and environment variables
โ โ โโโ security.py # JWT and password hashing
โ โ โโโ openai_client.py # OpenAI API wrapper
โ โโโ db/
โ โ โโโ base.py # SQLAlchemy base
โ โ โโโ session.py # Database session management
โ โ โโโ models/
โ โ โโโ user.py # User model
โ โ โโโ conversation.py # Conversation model
โ โ โโโ message.py # Message model (with JSONB context_sources)
โ โโโ api/
โ โ โโโ auth.py # Authentication endpoints
โ โ โโโ chat.py # Chat endpoints (streaming & sync)
โ โ โโโ conversations.py # Conversation CRUD
โ โโโ rag/
โ โ โโโ chroma.py # ChromaDB client
โ โ โโโ retriever.py # RAG retriever with follow-up detection
โ โ โโโ context.py # Context window management
โ โ โโโ followups.py # Follow-up question generator
โ โโโ services/
โ โ โโโ chat_service.py # Chat orchestration service
โ โ โโโ export_service.py # Conversation export
โ โโโ utils/
โ โโโ tokens.py # Token counting utilities
โโโ alembic/ # Database migrations
โโโ docker-compose.yml # Services (PostgreSQL, Redis, ChromaDB)
โโโ requirements.txt
โโโ .env.example
```
## ๐ Quick Start
### 1. Prerequisites
- Python 3.10+
- Docker & Docker Compose
- OpenAI API key
### 2. Setup
```bash
# Clone the repository
cd backend
# Create virtual environment
python -m venv venv
source venv/bin/activate # On Windows: venv\Scripts\activate
# Install dependencies
pip install -r requirements.txt
# Copy environment variables
cp .env.example .env
# Edit .env and add your OpenAI API key
nano .env # or use your preferred editor
```
### 3. Start Services
```bash
# Start PostgreSQL, Redis, and ChromaDB
docker-compose up -d
# Wait for services to be healthy
docker-compose ps
```
### 4. Initialize Database
```bash
# Run migrations
alembic upgrade head
# Or in development mode, the app will auto-create tables
```
### 5. Run the Application
```bash
# Development mode with auto-reload
python app/main.py
# Or with uvicorn directly
uvicorn app.main:app --reload --host 0.0.0.0 --port 8000
```
The API will be available at `http://localhost:8000`
## ๐ API Documentation
Once running, visit:
- **Swagger UI**: http://localhost:8000/docs
- **ReDoc**: http://localhost:8000/redoc
### Key Endpoints
#### Authentication
- `POST /api/auth/register` - Register new user
- `POST /api/auth/login` - Login and get tokens
- `POST /api/auth/refresh` - Refresh access token
- `GET /api/auth/me` - Get current user info
#### Chat
- `POST /api/chat/query` - Chat with streaming (SSE)
- `POST /api/chat/query-sync` - Chat without streaming
#### Conversations
- `GET /api/conversations` - List conversations
- `POST /api/conversations` - Create conversation
- `GET /api/conversations/{id}` - Get conversation with messages
- `DELETE /api/conversations/{id}` - Delete conversation
- `GET /api/conversations/{id}/messages` - Get conversation messages
## ๐ง Configuration
Key environment variables in `.env`:
```bash
# Database
DATABASE_URL=postgresql+asyncpg://postgres:postgres@localhost:5432/exploitrag
# ChromaDB (choose one)
# - http: use the docker-compose Chroma service (recommended for dev)
# - persistent: embedded/offline storage in CHROMA_PERSIST_PATH
# - cloud: managed Chroma Cloud
CHROMA_MODE=http
CHROMA_HOST=localhost
CHROMA_PORT=8001
CHROMA_SSL=False
# OpenAI
OPENAI_API_KEY=your-key-here
OPENAI_MODEL=gpt-4-turbo-preview
# JWT
JWT_SECRET_KEY=your-secret-key-change-in-production
# RAG Configuration
RAG_TOP_K=5 # Number of exploits to retrieve
RAG_SIMILARITY_THRESHOLD=0.7 # Minimum similarity score
RAG_MAX_CONTEXT_TOKENS=3000 # Max tokens for context
RAG_FOLLOWUP_WINDOW=3 # Messages to check for follow-ups
```
## ๐ก Usage Examples
### 1. Register and Login
```bash
# Register
curl -X POST http://localhost:8000/api/auth/register \
-H "Content-Type: application/json" \
-d '{"email":"user@example.com","password":"secure123"}'
# Response: {"access_token":"...", "refresh_token":"...", "token_type":"bearer"}
# Login
curl -X POST http://localhost:8000/api/auth/login \
-H "Content-Type: application/json" \
-d '{"email":"user@example.com","password":"secure123"}'
```
### 2. Chat Query (Streaming)
```bash
# Using curl with SSE
curl -X POST http://localhost:8000/api/chat/query \
-H "Authorization: Bearer YOUR_ACCESS_TOKEN" \
-H "Content-Type: application/json" \
-d '{"message":"What are SQL injection exploits?"}' \
--no-buffer
# Response (SSE stream):
# data: {"type":"content","content":"SQL"}
# data: {"type":"content","content":" injection"}
# ...
# data: {"type":"metadata","conversation_id":1,"message_id":2,"sources":[...],"followups":[...]}
```
### 3. Chat Query (Non-Streaming)
```bash
curl -X POST http://localhost:8000/api/chat/query-sync \
-H "Authorization: Bearer YOUR_ACCESS_TOKEN" \
-H "Content-Type: application/json" \
-d '{"conversation_id":1,"message":"Tell me more about the first exploit"}'
```
### 4. List Conversations
```bash
curl -X GET http://localhost:8000/api/conversations \
-H "Authorization: Bearer YOUR_ACCESS_TOKEN"
```
## ๐ง RAG Pipeline
### How It Works
1. **User Query** โ Received via `/api/chat/query`
2. **Follow-up Detection** โ Check if query references previous context
3. **Vector Retrieval** โ Query ChromaDB for relevant exploits
4. **Context Building** โ Format exploits with token management
5. **LLM Streaming** โ Stream response from OpenAI
6. **Persistence** โ Save messages with metadata (no full exploit text)
7. **Follow-up Generation** โ Suggest next questions
### Follow-up Detection
The system intelligently detects follow-up questions:
- Short queries (โค5 words)
- Pronouns: "it", "this", "that", "these"
- Continuation phrases: "more", "explain", "how", "why"
When detected, it prioritizes exploits from previous messages.
### Context Window Management
- Automatically truncates context to fit within `RAG_MAX_CONTEXT_TOKENS`
- Prioritizes most relevant exploits
- Uses tiktoken for accurate token counting
## ๐๏ธ Database Schema
### User
- `id`, `email`, `hashed_password`, `is_active`, `created_at`, `updated_at`
### Conversation
- `id`, `user_id`, `title`, `metadata`, `created_at`, `updated_at`
### Message
- `id`, `conversation_id`, `role` (user/assistant), `content`
- `context_sources` (JSONB) - Stores exploit metadata only:
```json
{
"sources": [
{
"exploit_id": "EDB-12345",
"title": "SQL Injection in WordPress",
"similarity_score": 0.85,
"metadata": {...}
}
]
}
```
**Important**: Full exploit text is NOT stored in PostgreSQL, only metadata.
## ๐ Security
- **JWT Authentication** - Access and refresh tokens
- **Password Hashing** - bcrypt with salt
- **CORS** - Configurable allowed origins
- **Input Validation** - Pydantic schemas
## ๐ ChromaDB Setup
Before using the chat system, you need to populate ChromaDB with ExploitDB data:
```python
# Example script to populate ChromaDB (create this separately)
from app.rag.chroma import chroma_client
from app.core.openai_client import openai_client
async def populate_exploitdb():
# Load your ExploitDB data
exploits = load_exploitdb_data() # Your data source
for exploit in exploits:
# Create embedding
embedding = await openai_client.create_embedding(exploit['content'])
# Add to ChromaDB
chroma_client.add(
embeddings=[embedding],
documents=[exploit['content']],
metadatas=[{
'title': exploit['title'],
'exploit_id': exploit['id'],
# ... other metadata
}],
ids=[exploit['id']]
)
```
## ๐งช Testing
```bash
# Test health endpoint
curl http://localhost:8000/health
# Test authentication flow
curl -X POST http://localhost:8000/api/auth/register \
-H "Content-Type: application/json" \
-d '{"email":"test@example.com","password":"test123"}'
# Test chat (after getting token)
curl -X POST http://localhost:8000/api/chat/query-sync \
-H "Authorization: Bearer YOUR_TOKEN" \
-H "Content-Type: application/json" \
-d '{"message":"What is a buffer overflow?"}'
```
## ๐ Troubleshooting
### Services not starting
```bash
# Check service status
docker-compose ps
# View logs
docker-compose logs postgres
docker-compose logs chromadb
docker-compose logs redis
```
### Database connection issues
```bash
# Verify PostgreSQL is running
docker-compose exec postgres pg_isready
# Check database exists
docker-compose exec postgres psql -U postgres -l
```
### ChromaDB connection issues
```bash
# Test ChromaDB health
curl http://localhost:8001/api/v1/heartbeat
```
## ๐ Development
### Create a new migration
```bash
# Auto-generate migration from model changes
alembic revision --autogenerate -m "description"
# Apply migration
alembic upgrade head
# Rollback
alembic downgrade -1
```
### Code Structure Guidelines
- **Async everywhere** - All database and API calls are async
- **Dependency injection** - Use FastAPI's `Depends()`
- **Type hints** - All functions have type annotations
- **Pydantic models** - For request/response validation
- **Clean separation** - Models, services, API routes are separate
## ๐ข Production Deployment
1. **Set DEBUG=False** in `.env`
2. **Use strong JWT_SECRET_KEY**
3. **Configure proper CORS origins**
4. **Use production database** (not docker-compose)
5. **Set up Redis persistence**
6. **Use gunicorn** with uvicorn workers:
```bash
gunicorn app.main:app -w 4 -k uvicorn.workers.UvicornWorker
```
7. **Set up reverse proxy** (nginx/Caddy)
8. **Enable HTTPS**
9. **Monitor with logging/metrics**
## ๐ License
This project is for educational and research purposes in cybersecurity.
## ๐ค Contributing
This is a hackathon project. Feel free to extend and improve!
---
**Built with โค๏ธ for ethical cybersecurity research**