mirror of
https://github.com/mblanke/ThreatHunt.git
synced 2026-03-01 14:00:20 -05:00
Compare commits
3 Commits
copilot/im
...
Claude-Ite
| Author | SHA1 | Date | |
|---|---|---|---|
| 037191f981 | |||
| 3c7e9b9eee | |||
| b398f6624c |
3
.env
Normal file
3
.env
Normal file
@@ -0,0 +1,3 @@
|
||||
DATABASE_URL=postgresql://admin:secure_password_123@database:5432/threat_hunter
|
||||
SECRET_KEY=your-very-secret-key-change-in-production
|
||||
FLASK_ENV=production
|
||||
67
.gitignore
vendored
67
.gitignore
vendored
@@ -1,67 +0,0 @@
|
||||
# Python
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
*.so
|
||||
.Python
|
||||
env/
|
||||
venv/
|
||||
ENV/
|
||||
build/
|
||||
develop-eggs/
|
||||
dist/
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
wheels/
|
||||
*.egg-info/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
|
||||
# Node
|
||||
node_modules/
|
||||
npm-debug.log*
|
||||
yarn-debug.log*
|
||||
yarn-error.log*
|
||||
.pnp
|
||||
.pnp.js
|
||||
|
||||
# Testing
|
||||
.coverage
|
||||
.pytest_cache/
|
||||
htmlcov/
|
||||
|
||||
# IDE
|
||||
.vscode/
|
||||
.idea/
|
||||
*.swp
|
||||
*.swo
|
||||
*~
|
||||
|
||||
# Environment
|
||||
.env
|
||||
.env.local
|
||||
.env.development.local
|
||||
.env.test.local
|
||||
.env.production.local
|
||||
|
||||
# OS
|
||||
.DS_Store
|
||||
Thumbs.db
|
||||
|
||||
# Database
|
||||
*.db
|
||||
*.sqlite
|
||||
*.sqlite3
|
||||
|
||||
# Logs
|
||||
*.log
|
||||
logs/
|
||||
|
||||
# Docker
|
||||
*.pid
|
||||
|
||||
5
.vscode/settings.json
vendored
Normal file
5
.vscode/settings.json
vendored
Normal file
@@ -0,0 +1,5 @@
|
||||
{
|
||||
"cSpell.words": [
|
||||
"jsonify"
|
||||
]
|
||||
}
|
||||
472
ARCHITECTURE.md
472
ARCHITECTURE.md
@@ -1,472 +0,0 @@
|
||||
# Architecture Documentation
|
||||
|
||||
This document describes the architecture and design decisions for VelociCompanion.
|
||||
|
||||
## System Overview
|
||||
|
||||
VelociCompanion is a multi-tenant, cloud-native threat hunting companion designed to work with Velociraptor. It provides secure authentication, data isolation, and role-based access control.
|
||||
|
||||
```
|
||||
┌─────────────┐ ┌─────────────┐ ┌──────────────┐
|
||||
│ │ │ │ │ │
|
||||
│ Frontend │────▶│ Backend │────▶│ PostgreSQL │
|
||||
│ (React) │ │ (FastAPI) │ │ Database │
|
||||
│ │ │ │ │ │
|
||||
└─────────────┘ └─────────────┘ └──────────────┘
|
||||
│
|
||||
▼
|
||||
┌─────────────┐
|
||||
│ │
|
||||
│ Velociraptor│
|
||||
│ Servers │
|
||||
│ │
|
||||
└─────────────┘
|
||||
```
|
||||
|
||||
## Technology Stack
|
||||
|
||||
### Backend
|
||||
- **FastAPI**: Modern, fast web framework for building APIs
|
||||
- **SQLAlchemy**: SQL toolkit and ORM
|
||||
- **PostgreSQL**: Relational database
|
||||
- **Alembic**: Database migration tool
|
||||
- **Python-Jose**: JWT token handling
|
||||
- **Passlib**: Password hashing with bcrypt
|
||||
|
||||
### Frontend
|
||||
- **React**: UI library
|
||||
- **TypeScript**: Type-safe JavaScript
|
||||
- **Axios**: HTTP client
|
||||
- **React Router**: Client-side routing
|
||||
|
||||
### Infrastructure
|
||||
- **Docker**: Containerization
|
||||
- **Docker Compose**: Multi-container orchestration
|
||||
|
||||
## Core Components
|
||||
|
||||
### 1. Authentication System
|
||||
|
||||
#### JWT Token Flow
|
||||
```
|
||||
1. User submits credentials (username/password)
|
||||
2. Backend verifies credentials
|
||||
3. Backend generates JWT token with:
|
||||
- user_id (sub)
|
||||
- tenant_id
|
||||
- role
|
||||
- expiration time
|
||||
4. Frontend stores token in localStorage
|
||||
5. All subsequent requests include token in Authorization header
|
||||
6. Backend validates token and extracts user context
|
||||
```
|
||||
|
||||
#### Password Security
|
||||
- Passwords are hashed using bcrypt with automatic salt generation
|
||||
- Password hashes are never exposed in API responses
|
||||
- Plaintext passwords are never logged or stored
|
||||
|
||||
#### Token Security
|
||||
- Tokens expire after 30 minutes (configurable)
|
||||
- Tokens are signed with HS256 algorithm
|
||||
- Secret key must be at least 32 characters
|
||||
|
||||
### 2. Multi-Tenancy
|
||||
|
||||
#### Data Isolation
|
||||
Every database query is automatically scoped to the user's tenant:
|
||||
|
||||
```python
|
||||
# Example: Listing hosts
|
||||
hosts = db.query(Host).filter(Host.tenant_id == current_user.tenant_id).all()
|
||||
```
|
||||
|
||||
#### Tenant Creation
|
||||
- Default tenant is created automatically on first user registration
|
||||
- Admin users can create additional tenants
|
||||
- Users are assigned to exactly one tenant
|
||||
|
||||
#### Cross-Tenant Access
|
||||
- Regular users: Can only access data in their tenant
|
||||
- Admin users: Can access all data in their tenant
|
||||
- Super-admin (future): Could access multiple tenants
|
||||
|
||||
### 3. Role-Based Access Control (RBAC)
|
||||
|
||||
#### Roles
|
||||
- **user**: Standard user with read/write access to their tenant's data
|
||||
- **admin**: Elevated privileges within their tenant
|
||||
- Can manage users in their tenant
|
||||
- Can create/modify/delete resources
|
||||
- Can view all data in their tenant
|
||||
|
||||
#### Permission Enforcement
|
||||
```python
|
||||
# Endpoint requiring admin role
|
||||
@router.get("/users")
|
||||
async def list_users(
|
||||
current_user: User = Depends(require_role(["admin"]))
|
||||
):
|
||||
# Only admins can access this
|
||||
pass
|
||||
```
|
||||
|
||||
### 4. Database Schema
|
||||
|
||||
#### Core Tables
|
||||
|
||||
**tenants**
|
||||
- id (PK)
|
||||
- name (unique)
|
||||
- description
|
||||
- created_at
|
||||
|
||||
**users**
|
||||
- id (PK)
|
||||
- username (unique)
|
||||
- password_hash
|
||||
- role
|
||||
- tenant_id (FK → tenants)
|
||||
- is_active
|
||||
- created_at
|
||||
|
||||
**hosts**
|
||||
- id (PK)
|
||||
- hostname
|
||||
- ip_address
|
||||
- os
|
||||
- tenant_id (FK → tenants)
|
||||
- host_metadata (JSON)
|
||||
- created_at
|
||||
- last_seen
|
||||
|
||||
**cases**
|
||||
- id (PK)
|
||||
- title
|
||||
- description
|
||||
- status (open, closed, investigating)
|
||||
- severity (low, medium, high, critical)
|
||||
- tenant_id (FK → tenants)
|
||||
- created_at
|
||||
- updated_at
|
||||
|
||||
**artifacts**
|
||||
- id (PK)
|
||||
- artifact_type (hash, ip, domain, email, etc.)
|
||||
- value
|
||||
- description
|
||||
- case_id (FK → cases)
|
||||
- artifact_metadata (JSON)
|
||||
- created_at
|
||||
|
||||
#### Relationships
|
||||
```
|
||||
tenants (1) ──< (N) users
|
||||
tenants (1) ──< (N) hosts
|
||||
tenants (1) ──< (N) cases
|
||||
cases (1) ──< (N) artifacts
|
||||
```
|
||||
|
||||
### 5. API Design
|
||||
|
||||
#### RESTful Principles
|
||||
- Resources are nouns (users, hosts, cases)
|
||||
- HTTP methods represent actions (GET, POST, PUT, DELETE)
|
||||
- Proper status codes (200, 201, 401, 403, 404)
|
||||
|
||||
#### Authentication
|
||||
All endpoints except `/auth/register` and `/auth/login` require authentication.
|
||||
|
||||
```
|
||||
Authorization: Bearer <jwt_token>
|
||||
```
|
||||
|
||||
#### Response Format
|
||||
Success:
|
||||
```json
|
||||
{
|
||||
"id": 1,
|
||||
"username": "john",
|
||||
"role": "user",
|
||||
"tenant_id": 1
|
||||
}
|
||||
```
|
||||
|
||||
Error:
|
||||
```json
|
||||
{
|
||||
"detail": "User not found"
|
||||
}
|
||||
```
|
||||
|
||||
### 6. Frontend Architecture
|
||||
|
||||
#### Component Structure
|
||||
```
|
||||
src/
|
||||
├── components/ # Reusable UI components
|
||||
│ └── PrivateRoute.tsx
|
||||
├── context/ # React Context providers
|
||||
│ └── AuthContext.tsx
|
||||
├── pages/ # Page components
|
||||
│ ├── Login.tsx
|
||||
│ └── Dashboard.tsx
|
||||
├── utils/ # Utilities
|
||||
│ └── api.ts # API client
|
||||
├── App.tsx # Main app component
|
||||
└── index.tsx # Entry point
|
||||
```
|
||||
|
||||
#### State Management
|
||||
- **AuthContext**: Global authentication state
|
||||
- Current user
|
||||
- Login/logout functions
|
||||
- Loading state
|
||||
- Authentication status
|
||||
|
||||
#### Routing
|
||||
```
|
||||
/login → Login page (public)
|
||||
/ → Dashboard (protected)
|
||||
/* → Redirect to / (protected)
|
||||
```
|
||||
|
||||
### 7. Security Architecture
|
||||
|
||||
#### Authentication Flow
|
||||
1. Frontend sends credentials to `/api/auth/login`
|
||||
2. Backend validates and returns JWT token
|
||||
3. Frontend stores token in localStorage
|
||||
4. Token included in all API requests
|
||||
5. Backend validates token on each request
|
||||
|
||||
#### Authorization Flow
|
||||
1. Extract JWT from Authorization header
|
||||
2. Verify token signature and expiration
|
||||
3. Extract user_id from token payload
|
||||
4. Load user from database
|
||||
5. Check user's role for endpoint access
|
||||
6. Apply tenant scoping to queries
|
||||
|
||||
#### Security Headers
|
||||
```python
|
||||
# CORS configuration
|
||||
allow_origins=["http://localhost:3000"]
|
||||
allow_credentials=True
|
||||
allow_methods=["*"]
|
||||
allow_headers=["*"]
|
||||
```
|
||||
|
||||
## Data Flow Examples
|
||||
|
||||
### User Registration
|
||||
```
|
||||
1. POST /api/auth/register
|
||||
{username: "john", password: "pass123"}
|
||||
2. Backend hashes password
|
||||
3. Create default tenant if needed
|
||||
4. Create user record
|
||||
5. Return user object (without password_hash)
|
||||
```
|
||||
|
||||
### Host Ingestion
|
||||
```
|
||||
1. Velociraptor sends data to POST /api/ingestion/ingest
|
||||
- Must include valid JWT token
|
||||
2. Extract tenant_id from current user
|
||||
3. Find or create host with hostname
|
||||
4. Update host metadata
|
||||
5. Return success response
|
||||
```
|
||||
|
||||
### Listing Resources
|
||||
```
|
||||
1. GET /api/hosts with Authorization header
|
||||
2. Validate JWT token
|
||||
3. Extract tenant_id from user
|
||||
4. Query: SELECT * FROM hosts WHERE tenant_id = ?
|
||||
5. Return filtered results
|
||||
```
|
||||
|
||||
## Deployment Architecture
|
||||
|
||||
### Development
|
||||
```
|
||||
┌──────────────────────────────────────┐
|
||||
│ Docker Compose │
|
||||
├──────────────────────────────────────┤
|
||||
│ Frontend:3000 Backend:8000 DB:5432│
|
||||
└──────────────────────────────────────┘
|
||||
```
|
||||
|
||||
### Production (Recommended)
|
||||
```
|
||||
┌─────────────┐ ┌─────────────┐
|
||||
│ Nginx/ │ │ Frontend │
|
||||
│ Traefik │────▶│ (Static) │
|
||||
│ (HTTPS) │ └─────────────┘
|
||||
└──────┬──────┘
|
||||
│
|
||||
▼
|
||||
┌─────────────┐ ┌──────────────┐
|
||||
│ Backend │ │ PostgreSQL │
|
||||
│ (Multiple │────▶│ (Managed) │
|
||||
│ instances) │ └──────────────┘
|
||||
└─────────────┘
|
||||
```
|
||||
|
||||
## Performance Considerations
|
||||
|
||||
### Database Indexing
|
||||
- Primary keys on all tables
|
||||
- Unique index on usernames
|
||||
- Index on tenant_id columns for fast filtering
|
||||
- Index on hostname for host lookups
|
||||
|
||||
### Query Optimization
|
||||
- Always filter by tenant_id early in queries
|
||||
- Use pagination for large result sets (skip/limit)
|
||||
- Lazy load relationships when not needed
|
||||
|
||||
### Caching (Future)
|
||||
- Cache tenant information
|
||||
- Cache user profiles
|
||||
- Cache frequently accessed hosts
|
||||
|
||||
## Monitoring & Logging
|
||||
|
||||
### Health Checks
|
||||
```
|
||||
GET /health → {"status": "healthy"}
|
||||
```
|
||||
|
||||
### Logging
|
||||
- Request logging via Uvicorn
|
||||
- Error tracking in application logs
|
||||
- Database query logging (development only)
|
||||
|
||||
### Metrics (Future)
|
||||
- Request count per endpoint
|
||||
- Authentication success/failure rate
|
||||
- Database query performance
|
||||
- Active user count
|
||||
|
||||
## Migration Strategy
|
||||
|
||||
### Database Migrations
|
||||
```bash
|
||||
# Create migration
|
||||
alembic revision --autogenerate -m "Description"
|
||||
|
||||
# Apply migration
|
||||
alembic upgrade head
|
||||
|
||||
# Rollback
|
||||
alembic downgrade -1
|
||||
```
|
||||
|
||||
### Schema Evolution
|
||||
1. Create migration for schema changes
|
||||
2. Test migration in development
|
||||
3. Apply to staging environment
|
||||
4. Verify data integrity
|
||||
5. Apply to production during maintenance window
|
||||
|
||||
## Testing Strategy
|
||||
|
||||
### Unit Tests (Future)
|
||||
- Test individual functions
|
||||
- Mock database connections
|
||||
- Test password hashing
|
||||
- Test JWT token creation/verification
|
||||
|
||||
### Integration Tests (Future)
|
||||
- Test API endpoints
|
||||
- Test authentication flow
|
||||
- Test multi-tenancy isolation
|
||||
- Test RBAC enforcement
|
||||
|
||||
### Manual Testing
|
||||
- Use test_api.sh script
|
||||
- Use FastAPI's /docs interface
|
||||
- Test frontend authentication flow
|
||||
|
||||
## Future Enhancements
|
||||
|
||||
### Phase 2
|
||||
- Refresh tokens for longer sessions
|
||||
- Password reset functionality
|
||||
- Email verification
|
||||
- Two-factor authentication (2FA)
|
||||
|
||||
### Phase 3
|
||||
- Audit logging
|
||||
- Advanced search and filtering
|
||||
- Real-time notifications
|
||||
- Velociraptor direct integration
|
||||
|
||||
### Phase 4
|
||||
- Machine learning for threat detection
|
||||
- Automated playbooks
|
||||
- Integration with SIEM systems
|
||||
- Advanced reporting and analytics
|
||||
|
||||
## Troubleshooting Guide
|
||||
|
||||
### Common Issues
|
||||
|
||||
**Token Expired**
|
||||
- Tokens expire after 30 minutes
|
||||
- User must login again
|
||||
- Consider implementing refresh tokens
|
||||
|
||||
**Permission Denied**
|
||||
- User lacks required role
|
||||
- Check user's role in database
|
||||
- Verify endpoint requires correct role
|
||||
|
||||
**Data Not Visible**
|
||||
- Check tenant_id of user
|
||||
- Verify data belongs to correct tenant
|
||||
- Ensure tenant_id is being applied to queries
|
||||
|
||||
**Database Connection Failed**
|
||||
- Check DATABASE_URL environment variable
|
||||
- Verify PostgreSQL is running
|
||||
- Check network connectivity
|
||||
|
||||
## Development Guidelines
|
||||
|
||||
### Adding New Endpoints
|
||||
|
||||
1. Create route in `app/api/routes/`
|
||||
2. Add authentication dependency
|
||||
3. Apply tenant scoping to queries
|
||||
4. Add role check if needed
|
||||
5. Create Pydantic schemas
|
||||
6. Update router registration in main.py
|
||||
7. Test with /docs interface
|
||||
|
||||
### Adding New Models
|
||||
|
||||
1. Create model in `app/models/`
|
||||
2. Add tenant_id foreign key
|
||||
3. Create migration
|
||||
4. Create Pydantic schemas
|
||||
5. Create CRUD routes
|
||||
6. Apply tenant scoping
|
||||
|
||||
### Code Style
|
||||
|
||||
- Follow PEP 8 for Python
|
||||
- Use type hints
|
||||
- Write docstrings for functions
|
||||
- Keep functions small and focused
|
||||
- Use meaningful variable names
|
||||
|
||||
## References
|
||||
|
||||
- [FastAPI Documentation](https://fastapi.tiangolo.com/)
|
||||
- [SQLAlchemy Documentation](https://docs.sqlalchemy.org/)
|
||||
- [JWT RFC](https://tools.ietf.org/html/rfc7519)
|
||||
- [OAuth 2.0 RFC](https://tools.ietf.org/html/rfc6749)
|
||||
@@ -1,311 +0,0 @@
|
||||
# Deployment Checklist
|
||||
|
||||
Use this checklist to deploy VelociCompanion to production.
|
||||
|
||||
## Pre-Deployment
|
||||
|
||||
### Security Review
|
||||
- [ ] Generate new SECRET_KEY (minimum 32 characters, cryptographically random)
|
||||
- [ ] Update DATABASE_URL with production credentials
|
||||
- [ ] Use strong database password (not default postgres/postgres)
|
||||
- [ ] Review CORS settings in `backend/app/main.py`
|
||||
- [ ] Enable HTTPS/TLS for all communications
|
||||
- [ ] Configure firewall rules
|
||||
- [ ] Set up VPN or IP whitelist for database access
|
||||
|
||||
### Configuration
|
||||
- [ ] Create production `.env` file
|
||||
- [ ] Set ACCESS_TOKEN_EXPIRE_MINUTES appropriately (30 minutes recommended)
|
||||
- [ ] Configure frontend REACT_APP_API_URL
|
||||
- [ ] Review all environment variables
|
||||
- [ ] Set up backup strategy for database
|
||||
|
||||
### Infrastructure
|
||||
- [ ] Provision database server or use managed service (RDS, Cloud SQL, etc.)
|
||||
- [ ] Set up load balancer for backend
|
||||
- [ ] Configure CDN for frontend static files
|
||||
- [ ] Set up monitoring and alerting
|
||||
- [ ] Configure log aggregation
|
||||
- [ ] Set up automated backups
|
||||
|
||||
## Deployment Steps
|
||||
|
||||
### 1. Database Setup
|
||||
|
||||
```bash
|
||||
# Create production database
|
||||
createdb velocicompanion_prod
|
||||
|
||||
# Set environment variable
|
||||
export DATABASE_URL="postgresql://user:pass@host:5432/velocicompanion_prod"
|
||||
|
||||
# Run migrations
|
||||
cd backend
|
||||
alembic upgrade head
|
||||
```
|
||||
|
||||
### 2. Backend Deployment
|
||||
|
||||
```bash
|
||||
# Build production image
|
||||
docker build -t velocicompanion-backend:latest ./backend
|
||||
|
||||
# Or deploy with docker-compose
|
||||
docker-compose -f docker-compose.prod.yml up -d backend
|
||||
```
|
||||
|
||||
### 3. Frontend Deployment
|
||||
|
||||
```bash
|
||||
# Build production bundle
|
||||
cd frontend
|
||||
npm install
|
||||
npm run build
|
||||
|
||||
# Deploy build/ directory to CDN or web server
|
||||
# Update API URL in environment
|
||||
```
|
||||
|
||||
### 4. Create Initial Admin User
|
||||
|
||||
```bash
|
||||
# Register first admin user via API
|
||||
curl -X POST https://your-domain.com/api/auth/register \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"username": "admin",
|
||||
"password": "STRONG_PASSWORD_HERE",
|
||||
"role": "admin"
|
||||
}'
|
||||
```
|
||||
|
||||
### 5. Verify Deployment
|
||||
|
||||
```bash
|
||||
# Check health endpoint
|
||||
curl https://your-domain.com/health
|
||||
|
||||
# Expected: {"status":"healthy"}
|
||||
|
||||
# Test authentication
|
||||
curl -X POST https://your-domain.com/api/auth/login \
|
||||
-d "username=admin&password=YOUR_PASSWORD"
|
||||
|
||||
# Should return JWT token
|
||||
```
|
||||
|
||||
## Post-Deployment
|
||||
|
||||
### Monitoring Setup
|
||||
- [ ] Configure application monitoring (e.g., Prometheus, Datadog)
|
||||
- [ ] Set up uptime monitoring (e.g., Pingdom, UptimeRobot)
|
||||
- [ ] Configure error tracking (e.g., Sentry)
|
||||
- [ ] Set up log analysis (e.g., ELK stack, CloudWatch)
|
||||
- [ ] Create dashboards for key metrics
|
||||
|
||||
### Alerts
|
||||
- [ ] High error rate alert
|
||||
- [ ] Slow response time alert
|
||||
- [ ] Database connection issues
|
||||
- [ ] High CPU/memory usage
|
||||
- [ ] Failed authentication attempts
|
||||
- [ ] Disk space low
|
||||
|
||||
### Backup Verification
|
||||
- [ ] Verify automated backups are running
|
||||
- [ ] Test backup restoration process
|
||||
- [ ] Document backup/restore procedures
|
||||
- [ ] Set up backup retention policy
|
||||
|
||||
### Security
|
||||
- [ ] Run security scan
|
||||
- [ ] Review access logs
|
||||
- [ ] Enable rate limiting
|
||||
- [ ] Set up intrusion detection
|
||||
- [ ] Configure SSL certificate auto-renewal
|
||||
|
||||
### Documentation
|
||||
- [ ] Update production endpoints in documentation
|
||||
- [ ] Document deployment process
|
||||
- [ ] Create runbook for common issues
|
||||
- [ ] Train operations team
|
||||
- [ ] Update architecture diagrams
|
||||
|
||||
## Production Environment Variables
|
||||
|
||||
### Backend (.env)
|
||||
```bash
|
||||
DATABASE_URL=postgresql://user:strongpass@db-host:5432/velocicompanion
|
||||
SECRET_KEY=your-32-plus-character-secret-key-here-make-it-random
|
||||
ACCESS_TOKEN_EXPIRE_MINUTES=30
|
||||
ALGORITHM=HS256
|
||||
```
|
||||
|
||||
### Frontend
|
||||
```bash
|
||||
REACT_APP_API_URL=https://api.your-domain.com
|
||||
```
|
||||
|
||||
## Load Balancer Configuration
|
||||
|
||||
### Backend
|
||||
```nginx
|
||||
upstream backend {
|
||||
server backend1:8000;
|
||||
server backend2:8000;
|
||||
server backend3:8000;
|
||||
}
|
||||
|
||||
server {
|
||||
listen 443 ssl;
|
||||
server_name api.your-domain.com;
|
||||
|
||||
ssl_certificate /path/to/cert.pem;
|
||||
ssl_certificate_key /path/to/key.pem;
|
||||
|
||||
location / {
|
||||
proxy_pass http://backend;
|
||||
proxy_set_header Host $host;
|
||||
proxy_set_header X-Real-IP $remote_addr;
|
||||
proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
|
||||
proxy_set_header X-Forwarded-Proto $scheme;
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### Frontend
|
||||
```nginx
|
||||
server {
|
||||
listen 443 ssl;
|
||||
server_name your-domain.com;
|
||||
|
||||
ssl_certificate /path/to/cert.pem;
|
||||
ssl_certificate_key /path/to/key.pem;
|
||||
|
||||
root /var/www/velocicompanion/build;
|
||||
index index.html;
|
||||
|
||||
location / {
|
||||
try_files $uri /index.html;
|
||||
}
|
||||
|
||||
# Cache static assets
|
||||
location ~* \.(js|css|png|jpg|jpeg|gif|ico|svg)$ {
|
||||
expires 1y;
|
||||
add_header Cache-Control "public, immutable";
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Scaling Considerations
|
||||
|
||||
### Horizontal Scaling
|
||||
- Run multiple backend instances behind load balancer
|
||||
- Use managed PostgreSQL with read replicas
|
||||
- Serve frontend from CDN
|
||||
- Implement caching layer (Redis)
|
||||
|
||||
### Vertical Scaling
|
||||
- **Database**: 4GB RAM minimum, 8GB+ for production
|
||||
- **Backend**: 2GB RAM per instance, 2+ CPU cores
|
||||
- **Frontend**: Static files, minimal resources
|
||||
|
||||
### Performance Optimization
|
||||
- [ ] Enable database connection pooling
|
||||
- [ ] Add Redis cache for sessions
|
||||
- [ ] Implement request rate limiting
|
||||
- [ ] Optimize database queries
|
||||
- [ ] Add database indexes
|
||||
- [ ] Enable GZIP compression
|
||||
|
||||
## Disaster Recovery
|
||||
|
||||
### Backup Strategy
|
||||
- **Database**: Daily full backups, hourly incremental
|
||||
- **Files**: Daily backup of configuration files
|
||||
- **Retention**: 30 days of backups
|
||||
- **Off-site**: Copy backups to different region
|
||||
|
||||
### Recovery Procedures
|
||||
1. Restore database from latest backup
|
||||
2. Deploy latest application version
|
||||
3. Run database migrations if needed
|
||||
4. Verify system functionality
|
||||
5. Update DNS if needed
|
||||
|
||||
### RTO/RPO
|
||||
- **RTO** (Recovery Time Objective): 4 hours
|
||||
- **RPO** (Recovery Point Objective): 1 hour
|
||||
|
||||
## Maintenance
|
||||
|
||||
### Regular Tasks
|
||||
- [ ] Review logs weekly
|
||||
- [ ] Update dependencies monthly
|
||||
- [ ] Security patches: Apply within 7 days
|
||||
- [ ] Database optimization quarterly
|
||||
- [ ] Review and rotate access credentials quarterly
|
||||
|
||||
### Update Process
|
||||
1. Test updates in staging environment
|
||||
2. Schedule maintenance window
|
||||
3. Notify users of planned downtime
|
||||
4. Create backup before update
|
||||
5. Deploy updates
|
||||
6. Run smoke tests
|
||||
7. Monitor for issues
|
||||
|
||||
## Rollback Plan
|
||||
|
||||
If deployment fails:
|
||||
|
||||
1. **Immediate**
|
||||
```bash
|
||||
# Rollback to previous version
|
||||
docker-compose down
|
||||
git checkout <previous-tag>
|
||||
docker-compose up -d
|
||||
```
|
||||
|
||||
2. **Database Rollback**
|
||||
```bash
|
||||
# Rollback migration
|
||||
alembic downgrade -1
|
||||
```
|
||||
|
||||
3. **Verify**
|
||||
- Check health endpoint
|
||||
- Test critical paths
|
||||
- Review error logs
|
||||
|
||||
## Support Contacts
|
||||
|
||||
- **Technical Lead**: [Contact Info]
|
||||
- **Database Admin**: [Contact Info]
|
||||
- **Security Team**: [Contact Info]
|
||||
- **On-Call**: [Rotation Schedule]
|
||||
|
||||
## Success Criteria
|
||||
|
||||
- [ ] All services running and healthy
|
||||
- [ ] Users can login successfully
|
||||
- [ ] API response times < 500ms
|
||||
- [ ] Error rate < 1%
|
||||
- [ ] Database queries optimized
|
||||
- [ ] Backups running successfully
|
||||
- [ ] Monitoring and alerts active
|
||||
- [ ] Documentation updated
|
||||
- [ ] Team trained on operations
|
||||
|
||||
## Sign-Off
|
||||
|
||||
- [ ] Technical Lead Approval
|
||||
- [ ] Security Team Approval
|
||||
- [ ] Operations Team Approval
|
||||
- [ ] Product Owner Approval
|
||||
|
||||
---
|
||||
|
||||
**Deployment Date**: _______________
|
||||
**Deployed By**: _______________
|
||||
**Sign-Off**: _______________
|
||||
@@ -1,357 +0,0 @@
|
||||
# Implementation Summary: Phase 1 - Core Infrastructure & Auth
|
||||
|
||||
## Overview
|
||||
|
||||
This document summarizes the complete implementation of Phase 1 for VelociCompanion, a multi-tenant threat hunting companion for Velociraptor. All acceptance criteria have been met.
|
||||
|
||||
## What Was Built
|
||||
|
||||
### 🎯 Complete Backend API (FastAPI)
|
||||
|
||||
#### Core Infrastructure
|
||||
- ✅ FastAPI application with 22 routes
|
||||
- ✅ PostgreSQL database integration via SQLAlchemy
|
||||
- ✅ Alembic database migrations configured
|
||||
- ✅ Docker containerization with health checks
|
||||
- ✅ Environment-based configuration
|
||||
|
||||
#### Authentication System
|
||||
- ✅ JWT token-based authentication using python-jose
|
||||
- ✅ Password hashing with bcrypt (passlib)
|
||||
- ✅ OAuth2 password flow for API compatibility
|
||||
- ✅ Token expiration and validation
|
||||
- ✅ Secure credential handling
|
||||
|
||||
#### Database Models (5 tables)
|
||||
1. **tenants** - Multi-tenant organization data
|
||||
2. **users** - User accounts with roles
|
||||
3. **hosts** - Monitored systems
|
||||
4. **cases** - Threat hunting investigations
|
||||
5. **artifacts** - IOCs and evidence
|
||||
|
||||
#### API Endpoints (22 routes)
|
||||
|
||||
**Authentication (`/api/auth`)**
|
||||
- `POST /register` - Create new user account
|
||||
- `POST /login` - Authenticate and receive JWT
|
||||
- `GET /me` - Get current user profile
|
||||
- `PUT /me` - Update user profile
|
||||
|
||||
**User Management (`/api/users`)** - Admin only
|
||||
- `GET /` - List users in tenant
|
||||
- `GET /{user_id}` - Get user details
|
||||
- `PUT /{user_id}` - Update user
|
||||
- `DELETE /{user_id}` - Deactivate user
|
||||
|
||||
**Tenants (`/api/tenants`)**
|
||||
- `GET /` - List accessible tenants
|
||||
- `POST /` - Create tenant (admin)
|
||||
- `GET /{tenant_id}` - Get tenant details
|
||||
|
||||
**Hosts (`/api/hosts`)**
|
||||
- `GET /` - List hosts (tenant-scoped)
|
||||
- `POST /` - Create host
|
||||
- `GET /{host_id}` - Get host details
|
||||
|
||||
**Ingestion (`/api/ingestion`)**
|
||||
- `POST /ingest` - Ingest Velociraptor data
|
||||
|
||||
**VirusTotal (`/api/vt`)**
|
||||
- `POST /lookup` - Hash reputation lookup
|
||||
|
||||
#### Security Features
|
||||
- ✅ Role-based access control (user, admin)
|
||||
- ✅ Multi-tenant data isolation
|
||||
- ✅ Automatic tenant scoping on all queries
|
||||
- ✅ Password strength enforcement
|
||||
- ✅ Protected routes with authentication
|
||||
- ✅ 0 security vulnerabilities (CodeQL verified)
|
||||
|
||||
### 🎨 Complete Frontend (React + TypeScript)
|
||||
|
||||
#### Core Components
|
||||
- ✅ React 18 with TypeScript
|
||||
- ✅ React Router for navigation
|
||||
- ✅ Axios for API communication
|
||||
- ✅ Context API for state management
|
||||
|
||||
#### Pages
|
||||
1. **Login Page** - Full authentication form
|
||||
2. **Dashboard** - Protected home page with user info
|
||||
3. **Private Routes** - Authentication-protected routing
|
||||
|
||||
#### Features
|
||||
- ✅ JWT token storage in localStorage
|
||||
- ✅ Automatic token inclusion in API requests
|
||||
- ✅ 401 error handling with auto-redirect
|
||||
- ✅ Loading states during authentication
|
||||
- ✅ Clean, responsive UI design
|
||||
|
||||
### 📦 Infrastructure & DevOps
|
||||
|
||||
#### Docker Configuration
|
||||
- ✅ Multi-container Docker Compose setup
|
||||
- ✅ PostgreSQL with health checks
|
||||
- ✅ Backend with automatic migrations
|
||||
- ✅ Frontend with hot reload
|
||||
- ✅ Volume mounts for persistence
|
||||
|
||||
#### Documentation
|
||||
1. **README.md** - Project overview and features
|
||||
2. **QUICKSTART.md** - Step-by-step setup guide
|
||||
3. **ARCHITECTURE.md** - System design and technical details
|
||||
4. **IMPLEMENTATION_SUMMARY.md** - This document
|
||||
|
||||
#### Testing & Validation
|
||||
- ✅ `test_api.sh` - Automated API testing script
|
||||
- ✅ Manual testing procedures documented
|
||||
- ✅ OpenAPI/Swagger documentation at `/docs`
|
||||
- ✅ Health check endpoint
|
||||
|
||||
## File Structure
|
||||
|
||||
```
|
||||
ThreatHunt/
|
||||
├── backend/
|
||||
│ ├── alembic/ # Database migrations
|
||||
│ │ ├── versions/
|
||||
│ │ │ └── f82b3092d056_initial_migration.py
|
||||
│ │ └── env.py
|
||||
│ ├── app/
|
||||
│ │ ├── api/routes/ # API endpoints
|
||||
│ │ │ ├── auth.py # Authentication
|
||||
│ │ │ ├── users.py # User management
|
||||
│ │ │ ├── tenants.py # Tenant management
|
||||
│ │ │ ├── hosts.py # Host management
|
||||
│ │ │ ├── ingestion.py # Data ingestion
|
||||
│ │ │ └── vt.py # VirusTotal
|
||||
│ │ ├── core/ # Core functionality
|
||||
│ │ │ ├── config.py # Settings
|
||||
│ │ │ ├── database.py # DB connection
|
||||
│ │ │ ├── security.py # JWT & passwords
|
||||
│ │ │ └── deps.py # FastAPI dependencies
|
||||
│ │ ├── models/ # SQLAlchemy models
|
||||
│ │ │ ├── user.py
|
||||
│ │ │ ├── tenant.py
|
||||
│ │ │ ├── host.py
|
||||
│ │ │ ├── case.py
|
||||
│ │ │ └── artifact.py
|
||||
│ │ ├── schemas/ # Pydantic schemas
|
||||
│ │ │ ├── auth.py
|
||||
│ │ │ └── user.py
|
||||
│ │ └── main.py # FastAPI app
|
||||
│ ├── requirements.txt # Python dependencies
|
||||
│ ├── Dockerfile
|
||||
│ └── .env.example
|
||||
├── frontend/
|
||||
│ ├── src/
|
||||
│ │ ├── components/
|
||||
│ │ │ └── PrivateRoute.tsx # Auth wrapper
|
||||
│ │ ├── context/
|
||||
│ │ │ └── AuthContext.tsx # Auth state
|
||||
│ │ ├── pages/
|
||||
│ │ │ ├── Login.tsx # Login form
|
||||
│ │ │ └── Dashboard.tsx # Home page
|
||||
│ │ ├── utils/
|
||||
│ │ │ └── api.ts # API client
|
||||
│ │ ├── App.tsx # Main component
|
||||
│ │ └── index.tsx # Entry point
|
||||
│ ├── public/
|
||||
│ │ └── index.html
|
||||
│ ├── package.json
|
||||
│ ├── tsconfig.json
|
||||
│ └── Dockerfile
|
||||
├── docker-compose.yml # Container orchestration
|
||||
├── test_api.sh # API test script
|
||||
├── .gitignore
|
||||
├── README.md
|
||||
├── QUICKSTART.md
|
||||
├── ARCHITECTURE.md
|
||||
└── IMPLEMENTATION_SUMMARY.md
|
||||
```
|
||||
|
||||
## Acceptance Criteria Status
|
||||
|
||||
| Criterion | Status | Evidence |
|
||||
|-----------|--------|----------|
|
||||
| Users can register with username/password | ✅ PASS | `POST /api/auth/register` endpoint |
|
||||
| Users can login and receive JWT token | ✅ PASS | `POST /api/auth/login` returns JWT |
|
||||
| Protected routes require valid JWT | ✅ PASS | All routes use `get_current_user` dependency |
|
||||
| Users can only access data within their tenant | ✅ PASS | All queries filtered by `tenant_id` |
|
||||
| Admin users can manage other users | ✅ PASS | `/api/users` routes with `require_role(["admin"])` |
|
||||
| Alembic migrations are set up and working | ✅ PASS | Initial migration created and tested |
|
||||
| Frontend has basic login flow | ✅ PASS | Login page with AuthContext integration |
|
||||
| All existing functionality continues to work | ✅ PASS | All routes require auth, tenant scoping applied |
|
||||
|
||||
## Technical Achievements
|
||||
|
||||
### Security
|
||||
- **Zero vulnerabilities** detected by CodeQL scanner
|
||||
- Modern cryptographic practices (bcrypt, HS256)
|
||||
- Secure token handling and storage
|
||||
- Protection against common attacks (SQL injection, XSS)
|
||||
|
||||
### Code Quality
|
||||
- **Type safety** with TypeScript and Python type hints
|
||||
- **Clean architecture** with separation of concerns
|
||||
- **RESTful API design** following best practices
|
||||
- **Comprehensive documentation** for developers
|
||||
|
||||
### Performance
|
||||
- **Database indexing** on key columns
|
||||
- **Efficient queries** with proper filtering
|
||||
- **Fast authentication** with JWT (stateless)
|
||||
- **Health checks** for monitoring
|
||||
|
||||
## How to Use
|
||||
|
||||
### Quick Start
|
||||
```bash
|
||||
# 1. Start services
|
||||
docker-compose up -d
|
||||
|
||||
# 2. Register a user
|
||||
curl -X POST http://localhost:8000/api/auth/register \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"username": "admin", "password": "admin123", "role": "admin"}'
|
||||
|
||||
# 3. Login via frontend
|
||||
open http://localhost:3000
|
||||
|
||||
# 4. Or login via API
|
||||
curl -X POST http://localhost:8000/api/auth/login \
|
||||
-d "username=admin&password=admin123"
|
||||
|
||||
# 5. Test all endpoints
|
||||
./test_api.sh
|
||||
```
|
||||
|
||||
### API Documentation
|
||||
Interactive API docs available at:
|
||||
- Swagger UI: http://localhost:8000/docs
|
||||
- ReDoc: http://localhost:8000/redoc
|
||||
|
||||
## What's Next (Future Phases)
|
||||
|
||||
### Phase 2 - Enhanced Authentication
|
||||
- Refresh tokens for longer sessions
|
||||
- Password reset functionality
|
||||
- Two-factor authentication (2FA)
|
||||
- Session management
|
||||
- Audit logging
|
||||
|
||||
### Phase 3 - Advanced Features
|
||||
- Real-time notifications
|
||||
- WebSocket support
|
||||
- Advanced search and filtering
|
||||
- Report generation
|
||||
- Case collaboration features
|
||||
|
||||
### Phase 4 - Integrations
|
||||
- Direct Velociraptor integration
|
||||
- SIEM system connectors
|
||||
- Threat intelligence feeds
|
||||
- Automated response playbooks
|
||||
- ML-based threat detection
|
||||
|
||||
## Migration from Development to Production
|
||||
|
||||
### Before Going Live
|
||||
|
||||
1. **Security Hardening**
|
||||
- Generate secure SECRET_KEY (32+ chars)
|
||||
- Use strong database passwords
|
||||
- Enable HTTPS/TLS
|
||||
- Configure proper CORS origins
|
||||
- Review and restrict network access
|
||||
|
||||
2. **Database**
|
||||
- Use managed PostgreSQL service
|
||||
- Configure backups
|
||||
- Set up replication
|
||||
- Monitor performance
|
||||
|
||||
3. **Application**
|
||||
- Set up load balancer
|
||||
- Deploy multiple backend instances
|
||||
- Configure logging aggregation
|
||||
- Set up monitoring and alerts
|
||||
|
||||
4. **Frontend**
|
||||
- Build production bundle
|
||||
- Serve via CDN
|
||||
- Enable caching
|
||||
- Minify assets
|
||||
|
||||
## Support & Maintenance
|
||||
|
||||
### Logs
|
||||
```bash
|
||||
# View all logs
|
||||
docker-compose logs -f
|
||||
|
||||
# Backend logs
|
||||
docker-compose logs -f backend
|
||||
|
||||
# Database logs
|
||||
docker-compose logs -f db
|
||||
```
|
||||
|
||||
### Database Migrations
|
||||
```bash
|
||||
# Create migration
|
||||
cd backend
|
||||
alembic revision --autogenerate -m "Description"
|
||||
|
||||
# Apply migrations
|
||||
alembic upgrade head
|
||||
|
||||
# Rollback
|
||||
alembic downgrade -1
|
||||
```
|
||||
|
||||
### Troubleshooting
|
||||
See QUICKSTART.md for common issues and solutions.
|
||||
|
||||
## Metrics
|
||||
|
||||
### Code Statistics
|
||||
- **Backend**: 29 Python files, ~2,000 lines
|
||||
- **Frontend**: 8 TypeScript/TSX files, ~800 lines
|
||||
- **Infrastructure**: 3 Dockerfiles, 1 docker-compose.yml
|
||||
- **Documentation**: 4 comprehensive guides
|
||||
- **Total**: ~50 files across the stack
|
||||
|
||||
### Features Delivered
|
||||
- 22 API endpoints
|
||||
- 5 database models
|
||||
- 1 database migration
|
||||
- 2 frontend pages
|
||||
- 4 React components/contexts
|
||||
- 100% authentication coverage
|
||||
- 100% tenant isolation
|
||||
- 0 security vulnerabilities
|
||||
|
||||
## Conclusion
|
||||
|
||||
Phase 1 of VelociCompanion has been successfully completed with all acceptance criteria met. The system provides a solid foundation for multi-tenant threat hunting operations with:
|
||||
|
||||
- ✅ **Secure authentication** with JWT tokens
|
||||
- ✅ **Complete data isolation** between tenants
|
||||
- ✅ **Role-based access control** for permissions
|
||||
- ✅ **Modern tech stack** (FastAPI, React, PostgreSQL)
|
||||
- ✅ **Production-ready infrastructure** with Docker
|
||||
- ✅ **Comprehensive documentation** for users and developers
|
||||
|
||||
The system is ready for:
|
||||
1. Integration with Velociraptor servers
|
||||
2. Deployment to staging/production environments
|
||||
3. User acceptance testing
|
||||
4. Development of Phase 2 features
|
||||
|
||||
## Credits
|
||||
|
||||
Implemented by: GitHub Copilot
|
||||
Repository: https://github.com/mblanke/ThreatHunt
|
||||
Date: December 2025
|
||||
Version: 0.1.0
|
||||
@@ -1,578 +0,0 @@
|
||||
# Phase 5: Distributed LLM Routing Architecture
|
||||
|
||||
## Overview
|
||||
|
||||
Phase 5 introduces a sophisticated distributed Large Language Model (LLM) routing system that intelligently classifies tasks and routes them to specialized models across multiple GPU nodes (GB10 devices). This architecture enables efficient utilization of computational resources and optimal model selection based on task requirements.
|
||||
|
||||
## Architecture Components
|
||||
|
||||
The system consists of four containerized components that work together to provide intelligent, scalable LLM processing:
|
||||
|
||||
### 1. Router Agent (LLM Classifier + Policy Engine)
|
||||
|
||||
**Module**: `app/core/llm_router.py`
|
||||
|
||||
The Router Agent is responsible for:
|
||||
- **Request Classification**: Analyzes incoming requests to determine the task type
|
||||
- **Model Selection**: Routes requests to the most appropriate specialized model
|
||||
- **Policy Enforcement**: Applies routing rules based on configured policies
|
||||
|
||||
**Task Types & Model Routing:**
|
||||
|
||||
| Task Type | Model | Use Case |
|
||||
|-----------|-------|----------|
|
||||
| `general_reasoning` | DeepSeek | Complex analysis and reasoning |
|
||||
| `multilingual` | Qwen72 / Aya | Translation and multilingual tasks |
|
||||
| `structured_parsing` | Phi-4 | Structured data extraction |
|
||||
| `rule_generation` | Qwen-Coder | Code and rule generation |
|
||||
| `adversarial_reasoning` | LLaMA 3.1 | Threat and adversarial analysis |
|
||||
| `classification` | Granite Guardian | Pure classification tasks |
|
||||
|
||||
**Classification Logic:**
|
||||
```python
|
||||
from app.core.llm_router import get_llm_router
|
||||
|
||||
router = get_llm_router()
|
||||
routing_decision = router.route_request({
|
||||
"prompt": "Analyze this threat...",
|
||||
"task_hints": ["threat", "adversary"]
|
||||
})
|
||||
# Routes to LLaMA 3.1 for adversarial reasoning
|
||||
```
|
||||
|
||||
### 2. Job Scheduler (GPU Load Balancer)
|
||||
|
||||
**Module**: `app/core/job_scheduler.py`
|
||||
|
||||
The Job Scheduler manages:
|
||||
- **Node Selection**: Determines which GB10 device is available
|
||||
- **Resource Monitoring**: Tracks GPU VRAM and compute utilization
|
||||
- **Parallelization Decisions**: Determines if jobs should be distributed
|
||||
- **Serial Chaining**: Handles multi-step reasoning workflows
|
||||
|
||||
**GPU Node Configuration:**
|
||||
|
||||
**GB10 Node 1** (`gb10-node-1:8001`)
|
||||
- **Total VRAM**: 80 GB
|
||||
- **Models Loaded**: DeepSeek, Qwen72
|
||||
- **Primary Use**: General reasoning and multilingual tasks
|
||||
|
||||
**GB10 Node 2** (`gb10-node-2:8001`)
|
||||
- **Total VRAM**: 80 GB
|
||||
- **Models Loaded**: Phi-4, Qwen-Coder, LLaMA 3.1, Granite Guardian
|
||||
- **Primary Use**: Specialized tasks (parsing, coding, classification, threat analysis)
|
||||
|
||||
**Scheduling Strategies:**
|
||||
|
||||
1. **Single Node Execution**
|
||||
- Default for simple requests
|
||||
- Selected based on lowest compute utilization
|
||||
- Requires sufficient VRAM for model
|
||||
|
||||
2. **Parallel Execution**
|
||||
- Distributes work across multiple nodes
|
||||
- Used for batch processing or high-priority jobs
|
||||
- Automatic load balancing
|
||||
|
||||
3. **Serial Chaining**
|
||||
- Multi-step dependent operations
|
||||
- Sequential execution with context passing
|
||||
- Used for complex reasoning workflows
|
||||
|
||||
4. **Queued Execution**
|
||||
- When all nodes are at capacity
|
||||
- Priority-based queue management
|
||||
- Automatic dispatch when resources available
|
||||
|
||||
**Example Usage:**
|
||||
```python
|
||||
from app.core.job_scheduler import get_job_scheduler, Job
|
||||
|
||||
scheduler = get_job_scheduler()
|
||||
job = Job(
|
||||
job_id="threat_analysis_001",
|
||||
model="llama31",
|
||||
priority=1,
|
||||
estimated_vram_gb=10,
|
||||
requires_parallel=False,
|
||||
requires_chaining=False,
|
||||
payload={"prompt": "..."}
|
||||
)
|
||||
|
||||
scheduling_decision = await scheduler.schedule_job(job)
|
||||
# Returns node assignment and execution mode
|
||||
```
|
||||
|
||||
### 3. LLM Pool (OpenAI-Compatible Endpoints)
|
||||
|
||||
**Module**: `app/core/llm_pool.py`
|
||||
|
||||
The LLM Pool provides:
|
||||
- **Unified Interface**: OpenAI-compatible API for all models
|
||||
- **Endpoint Management**: Tracks availability and health
|
||||
- **Parallel Execution**: Simultaneous multi-model requests
|
||||
- **Error Handling**: Graceful fallback on failures
|
||||
|
||||
**Available Endpoints:**
|
||||
|
||||
| Model | Endpoint | Node | Specialization |
|
||||
|-------|----------|------|----------------|
|
||||
| DeepSeek | `http://gb10-node-1:8001/deepseek` | Node 1 | General reasoning |
|
||||
| Qwen72 | `http://gb10-node-1:8001/qwen72` | Node 1 | Multilingual |
|
||||
| Phi-4 | `http://gb10-node-2:8001/phi4` | Node 2 | Structured parsing |
|
||||
| Qwen-Coder | `http://gb10-node-2:8001/qwen-coder` | Node 2 | Code generation |
|
||||
| LLaMA 3.1 | `http://gb10-node-2:8001/llama31` | Node 2 | Adversarial reasoning |
|
||||
| Granite Guardian | `http://gb10-node-2:8001/granite-guardian` | Node 2 | Classification |
|
||||
|
||||
**Example Usage:**
|
||||
```python
|
||||
from app.core.llm_pool import get_llm_pool
|
||||
|
||||
pool = get_llm_pool()
|
||||
|
||||
# Single model call
|
||||
result = await pool.call_model(
|
||||
model_name="llama31",
|
||||
prompt="Analyze this threat pattern...",
|
||||
parameters={"temperature": 0.7, "max_tokens": 2048}
|
||||
)
|
||||
|
||||
# Multiple models in parallel
|
||||
results = await pool.call_multiple_models(
|
||||
model_names=["llama31", "deepseek"],
|
||||
prompt="Complex threat analysis...",
|
||||
parameters={"temperature": 0.7}
|
||||
)
|
||||
```
|
||||
|
||||
### 4. Merger Agent (Result Synthesizer)
|
||||
|
||||
**Module**: `app/core/merger_agent.py`
|
||||
|
||||
The Merger Agent provides:
|
||||
- **Result Combination**: Intelligently merges outputs from multiple models
|
||||
- **Strategy Selection**: Multiple merging strategies for different use cases
|
||||
- **Quality Assessment**: Evaluates and ranks responses
|
||||
- **Consensus Building**: Determines agreement across models
|
||||
|
||||
**Merging Strategies:**
|
||||
|
||||
1. **Consensus** (`MergeStrategy.CONSENSUS`)
|
||||
- Takes majority vote for classifications
|
||||
- Selects most common response
|
||||
- Best for: Classification tasks, binary decisions
|
||||
|
||||
2. **Weighted** (`MergeStrategy.WEIGHTED`)
|
||||
- Weights results by confidence scores
|
||||
- Selects highest confidence response
|
||||
- Best for: When models provide confidence scores
|
||||
|
||||
3. **Concatenate** (`MergeStrategy.CONCATENATE`)
|
||||
- Combines all responses sequentially
|
||||
- Preserves all information
|
||||
- Best for: Comprehensive analysis requiring multiple perspectives
|
||||
|
||||
4. **Best Quality** (`MergeStrategy.BEST_QUALITY`)
|
||||
- Selects highest quality response based on metrics
|
||||
- Considers length, completeness, formatting
|
||||
- Best for: Text generation, detailed explanations
|
||||
|
||||
5. **Ensemble** (`MergeStrategy.ENSEMBLE`)
|
||||
- Synthesizes insights from all models
|
||||
- Creates comprehensive summary
|
||||
- Best for: Complex analysis requiring synthesis
|
||||
|
||||
**Example Usage:**
|
||||
```python
|
||||
from app.core.merger_agent import get_merger_agent, MergeStrategy
|
||||
|
||||
merger = get_merger_agent()
|
||||
|
||||
# Multiple model results
|
||||
results = [
|
||||
{"model": "llama31", "response": "...", "confidence": 0.9},
|
||||
{"model": "deepseek", "response": "...", "confidence": 0.85}
|
||||
]
|
||||
|
||||
# Merge with consensus strategy
|
||||
merged = merger.merge_results(results, strategy=MergeStrategy.CONSENSUS)
|
||||
```
|
||||
|
||||
## API Endpoints
|
||||
|
||||
### Process LLM Request
|
||||
```http
|
||||
POST /api/llm/process
|
||||
```
|
||||
|
||||
Processes a request through the complete routing system.
|
||||
|
||||
**Request Body:**
|
||||
```json
|
||||
{
|
||||
"prompt": "Analyze this threat pattern for indicators of compromise",
|
||||
"task_hints": ["threat", "adversary"],
|
||||
"requires_parallel": false,
|
||||
"requires_chaining": false,
|
||||
"parameters": {
|
||||
"temperature": 0.7,
|
||||
"max_tokens": 2048
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
**Response:**
|
||||
```json
|
||||
{
|
||||
"job_id": "job_123_4567",
|
||||
"status": "completed",
|
||||
"routing": {
|
||||
"task_type": "adversarial_reasoning",
|
||||
"model": "llama31",
|
||||
"endpoint": "llama31",
|
||||
"priority": 1
|
||||
},
|
||||
"scheduling": {
|
||||
"job_id": "job_123_4567",
|
||||
"execution_mode": "single",
|
||||
"node": {
|
||||
"node_id": "gb10-node-2",
|
||||
"endpoint": "http://gb10-node-2:8001/llama31"
|
||||
}
|
||||
},
|
||||
"result": {
|
||||
"choices": [...]
|
||||
},
|
||||
"execution_mode": "single"
|
||||
}
|
||||
```
|
||||
|
||||
### List Available Models
|
||||
```http
|
||||
GET /api/llm/models
|
||||
```
|
||||
|
||||
Returns all available LLM models in the pool.
|
||||
|
||||
**Response:**
|
||||
```json
|
||||
{
|
||||
"models": [
|
||||
{
|
||||
"model_name": "deepseek",
|
||||
"node_id": "gb10-node-1",
|
||||
"endpoint_url": "http://gb10-node-1:8001/deepseek",
|
||||
"is_available": true
|
||||
},
|
||||
...
|
||||
],
|
||||
"total": 6
|
||||
}
|
||||
```
|
||||
|
||||
### List GPU Nodes
|
||||
```http
|
||||
GET /api/llm/nodes
|
||||
```
|
||||
|
||||
Returns status of all GPU nodes.
|
||||
|
||||
**Response:**
|
||||
```json
|
||||
{
|
||||
"nodes": [
|
||||
{
|
||||
"node_id": "gb10-node-1",
|
||||
"hostname": "gb10-node-1",
|
||||
"vram_total_gb": 80,
|
||||
"vram_used_gb": 25,
|
||||
"vram_available_gb": 55,
|
||||
"compute_utilization": 0.35,
|
||||
"status": "available",
|
||||
"models_loaded": ["deepseek", "qwen72"]
|
||||
},
|
||||
...
|
||||
],
|
||||
"available_count": 2
|
||||
}
|
||||
```
|
||||
|
||||
### Update Node Status (Admin Only)
|
||||
```http
|
||||
POST /api/llm/nodes/status
|
||||
```
|
||||
|
||||
Updates GPU node status metrics.
|
||||
|
||||
**Request Body:**
|
||||
```json
|
||||
{
|
||||
"node_id": "gb10-node-1",
|
||||
"vram_used_gb": 30,
|
||||
"compute_utilization": 0.45,
|
||||
"status": "available"
|
||||
}
|
||||
```
|
||||
|
||||
### Get Routing Rules
|
||||
```http
|
||||
GET /api/llm/routing/rules
|
||||
```
|
||||
|
||||
Returns current routing rules for task classification.
|
||||
|
||||
### Test Classification
|
||||
```http
|
||||
POST /api/llm/test-classification
|
||||
```
|
||||
|
||||
Tests task classification without executing the request.
|
||||
|
||||
## Usage Examples
|
||||
|
||||
### Example 1: Threat Analysis with Adversarial Reasoning
|
||||
|
||||
```python
|
||||
import httpx
|
||||
|
||||
async def analyze_threat():
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.post(
|
||||
"http://localhost:8000/api/llm/process",
|
||||
headers={"Authorization": f"Bearer {token}"},
|
||||
json={
|
||||
"prompt": "Analyze this suspicious PowerShell script for malicious intent...",
|
||||
"task_hints": ["threat", "adversary", "malicious"],
|
||||
"parameters": {"temperature": 0.3} # Lower temp for analysis
|
||||
}
|
||||
)
|
||||
result = response.json()
|
||||
print(f"Model used: {result['routing']['model']}")
|
||||
print(f"Analysis: {result['result']}")
|
||||
```
|
||||
|
||||
### Example 2: Code Generation for YARA Rules
|
||||
|
||||
```python
|
||||
async def generate_yara_rule():
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.post(
|
||||
"http://localhost:8000/api/llm/process",
|
||||
headers={"Authorization": f"Bearer {token}"},
|
||||
json={
|
||||
"prompt": "Generate a YARA rule to detect this malware family...",
|
||||
"task_hints": ["code", "rule", "generate"],
|
||||
"parameters": {"temperature": 0.5}
|
||||
}
|
||||
)
|
||||
result = response.json()
|
||||
# Routes to Qwen-Coder automatically
|
||||
print(f"Generated rule: {result['result']}")
|
||||
```
|
||||
|
||||
### Example 3: Parallel Processing for Batch Analysis
|
||||
|
||||
```python
|
||||
async def batch_analysis():
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.post(
|
||||
"http://localhost:8000/api/llm/process",
|
||||
headers={"Authorization": f"Bearer {token}"},
|
||||
json={
|
||||
"prompt": "Analyze these 50 log entries for anomalies...",
|
||||
"task_hints": ["classify", "anomaly"],
|
||||
"requires_parallel": True,
|
||||
"batch_size": 50
|
||||
}
|
||||
)
|
||||
result = response.json()
|
||||
# Automatically parallelized across both nodes
|
||||
print(f"Execution mode: {result['execution_mode']}")
|
||||
```
|
||||
|
||||
### Example 4: Serial Chaining for Multi-Step Analysis
|
||||
|
||||
```python
|
||||
async def chained_analysis():
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.post(
|
||||
"http://localhost:8000/api/llm/process",
|
||||
headers={"Authorization": f"Bearer {token}"},
|
||||
json={
|
||||
"prompt": "First extract IOCs, then classify threats, finally generate response plan",
|
||||
"task_hints": ["parse", "classify", "generate"],
|
||||
"requires_chaining": True,
|
||||
"operations": ["extract", "classify", "generate"]
|
||||
}
|
||||
)
|
||||
result = response.json()
|
||||
# Executed serially with context passing
|
||||
print(f"Chain result: {result['result']}")
|
||||
```
|
||||
|
||||
## Integration with Existing Features
|
||||
|
||||
### Integration with Threat Intelligence (Phase 4)
|
||||
|
||||
The distributed LLM system enhances threat intelligence analysis:
|
||||
|
||||
```python
|
||||
from app.core.threat_intel import get_threat_analyzer
|
||||
from app.core.llm_pool import get_llm_pool
|
||||
|
||||
async def enhanced_threat_analysis(host_id):
|
||||
# Step 1: Traditional ML analysis
|
||||
analyzer = get_threat_analyzer()
|
||||
ml_result = analyzer.analyze_host(host_data)
|
||||
|
||||
# Step 2: LLM-based deep analysis if score is concerning
|
||||
if ml_result["score"] > 0.6:
|
||||
pool = get_llm_pool()
|
||||
llm_result = await pool.call_model(
|
||||
"llama31",
|
||||
f"Deep analysis of threat with score {ml_result['score']}: {host_data}",
|
||||
{"temperature": 0.3}
|
||||
)
|
||||
|
||||
return {
|
||||
"ml_analysis": ml_result,
|
||||
"llm_analysis": llm_result,
|
||||
"recommendation": "quarantine" if ml_result["score"] > 0.8 else "investigate"
|
||||
}
|
||||
```
|
||||
|
||||
### Integration with Automated Playbooks (Phase 4)
|
||||
|
||||
LLM routing can trigger automated responses:
|
||||
|
||||
```python
|
||||
from app.core.playbook_engine import get_playbook_engine
|
||||
|
||||
async def llm_triggered_playbook(threat_analysis):
|
||||
if threat_analysis["result"]["severity"] == "critical":
|
||||
engine = get_playbook_engine()
|
||||
await engine.execute_playbook(
|
||||
playbook={
|
||||
"actions": [
|
||||
{"type": "isolate_host", "params": {"host_id": host_id}},
|
||||
{"type": "send_notification", "params": {"message": "Critical threat detected"}},
|
||||
{"type": "create_case", "params": {"title": "Auto-generated from LLM analysis"}}
|
||||
]
|
||||
},
|
||||
context=threat_analysis
|
||||
)
|
||||
```
|
||||
|
||||
## Deployment
|
||||
|
||||
### Docker Compose Configuration
|
||||
|
||||
Add LLM node services to `docker-compose.yml`:
|
||||
|
||||
```yaml
|
||||
services:
|
||||
# Existing services...
|
||||
|
||||
llm-node-1:
|
||||
image: vllm/vllm-openai:latest
|
||||
ports:
|
||||
- "8001:8001"
|
||||
environment:
|
||||
- NVIDIA_VISIBLE_DEVICES=0,1
|
||||
volumes:
|
||||
- ./models:/models
|
||||
command: >
|
||||
--model /models/deepseek
|
||||
--host 0.0.0.0
|
||||
--port 8001
|
||||
deploy:
|
||||
resources:
|
||||
reservations:
|
||||
devices:
|
||||
- driver: nvidia
|
||||
count: 2
|
||||
capabilities: [gpu]
|
||||
|
||||
llm-node-2:
|
||||
image: vllm/vllm-openai:latest
|
||||
ports:
|
||||
- "8002:8001"
|
||||
environment:
|
||||
- NVIDIA_VISIBLE_DEVICES=2,3
|
||||
volumes:
|
||||
- ./models:/models
|
||||
command: >
|
||||
--model /models/phi4
|
||||
--host 0.0.0.0
|
||||
--port 8001
|
||||
deploy:
|
||||
resources:
|
||||
reservations:
|
||||
devices:
|
||||
- driver: nvidia
|
||||
count: 2
|
||||
capabilities: [gpu]
|
||||
```
|
||||
|
||||
### Environment Variables
|
||||
|
||||
Add to `.env`:
|
||||
|
||||
```bash
|
||||
# Phase 5: LLM Configuration
|
||||
LLM_NODE_1_URL=http://gb10-node-1:8001
|
||||
LLM_NODE_2_URL=http://gb10-node-2:8001
|
||||
LLM_ENABLE_PARALLEL=true
|
||||
LLM_MAX_PARALLEL_JOBS=4
|
||||
LLM_DEFAULT_TIMEOUT=60
|
||||
```
|
||||
|
||||
## Performance Considerations
|
||||
|
||||
### Resource Allocation
|
||||
|
||||
- **DeepSeek**: ~40GB VRAM (high priority)
|
||||
- **Qwen72**: ~35GB VRAM (medium priority)
|
||||
- **Phi-4**: ~15GB VRAM (fast inference)
|
||||
- **Qwen-Coder**: ~20GB VRAM
|
||||
- **LLaMA 3.1**: ~25GB VRAM
|
||||
- **Granite Guardian**: ~10GB VRAM (classification only)
|
||||
|
||||
### Load Balancing
|
||||
|
||||
The scheduler automatically:
|
||||
- Monitors VRAM usage on each node
|
||||
- Tracks compute utilization (0.0-1.0)
|
||||
- Routes requests to less loaded nodes
|
||||
- Queues jobs when capacity is reached
|
||||
|
||||
### Optimization Tips
|
||||
|
||||
1. **Use task_hints**: Helps router select optimal model faster
|
||||
2. **Enable parallelization**: For batch jobs over 10 items
|
||||
3. **Monitor node status**: Use `/api/llm/nodes` endpoint
|
||||
4. **Set appropriate temperatures**: Lower (0.3) for analysis, higher (0.7) for generation
|
||||
5. **Leverage caching**: Repeated prompts hit cache layer
|
||||
|
||||
## Security
|
||||
|
||||
- All LLM endpoints require authentication
|
||||
- Admin-only node status updates
|
||||
- Tenant isolation maintained
|
||||
- Audit logging for all LLM requests
|
||||
- Rate limiting per user/tenant
|
||||
|
||||
## Future Enhancements
|
||||
|
||||
- [ ] Model fine-tuning pipeline
|
||||
- [ ] Custom model deployment
|
||||
- [ ] Advanced caching layer
|
||||
- [ ] Multi-region deployment
|
||||
- [ ] Real-time model swapping
|
||||
- [ ] Automated model selection via meta-learning
|
||||
- [ ] Integration with external model APIs (OpenAI, Anthropic)
|
||||
- [ ] Cost tracking and optimization
|
||||
|
||||
## Conclusion
|
||||
|
||||
Phase 5 provides a production-ready distributed LLM routing architecture that intelligently manages computational resources while optimizing for task-specific model selection. The system integrates seamlessly with existing threat hunting capabilities to provide enhanced analysis and automated decision-making.
|
||||
@@ -1,380 +0,0 @@
|
||||
# Phases 2, 3, and 4 Implementation Complete
|
||||
|
||||
All requested phases have been successfully implemented and are ready for use.
|
||||
|
||||
## Overview
|
||||
|
||||
VelociCompanion v1.0.0 is now a complete, production-ready multi-tenant threat hunting platform with:
|
||||
- Advanced authentication (2FA, refresh tokens, password reset)
|
||||
- Real-time notifications via WebSocket
|
||||
- Direct Velociraptor integration
|
||||
- ML-powered threat detection
|
||||
- Automated response playbooks
|
||||
- Advanced reporting capabilities
|
||||
|
||||
## Phase 2: Enhanced Authentication ✅
|
||||
|
||||
### Implemented Features
|
||||
|
||||
#### Refresh Tokens
|
||||
- 30-day expiration refresh tokens
|
||||
- Secure token generation with `secrets.token_urlsafe()`
|
||||
- Revocation support
|
||||
- **Endpoint**: `POST /api/auth/refresh`
|
||||
|
||||
#### Two-Factor Authentication (2FA)
|
||||
- TOTP-based 2FA using pyotp
|
||||
- QR code generation for authenticator apps
|
||||
- **Endpoints**:
|
||||
- `POST /api/auth/2fa/setup` - Generate secret and QR code
|
||||
- `POST /api/auth/2fa/verify` - Enable 2FA with code verification
|
||||
- `POST /api/auth/2fa/disable` - Disable 2FA (requires code)
|
||||
- Integrated into login flow
|
||||
|
||||
#### Password Reset
|
||||
- Secure token-based password reset
|
||||
- 1-hour token expiration
|
||||
- **Endpoints**:
|
||||
- `POST /api/auth/password-reset/request` - Request reset (email)
|
||||
- `POST /api/auth/password-reset/confirm` - Confirm with token
|
||||
|
||||
#### Email Verification
|
||||
- Email field added to User model
|
||||
- `email_verified` flag for future verification flow
|
||||
- Ready for email verification implementation
|
||||
|
||||
#### Audit Logging
|
||||
- Comprehensive audit trail for all actions
|
||||
- Tracks: user_id, tenant_id, action, resource_type, resource_id, IP, user agent
|
||||
- **Endpoints**:
|
||||
- `GET /api/audit` - List audit logs (admin only)
|
||||
- `GET /api/audit/{id}` - Get specific audit log
|
||||
- Filterable by action, resource type, date range
|
||||
|
||||
### Database Changes
|
||||
- `refresh_tokens` table
|
||||
- `password_reset_tokens` table
|
||||
- `audit_logs` table
|
||||
- User model: added `email`, `email_verified`, `totp_secret`, `totp_enabled`
|
||||
|
||||
## Phase 3: Advanced Features ✅
|
||||
|
||||
### Implemented Features
|
||||
|
||||
#### Advanced Search & Filtering
|
||||
- Enhanced `GET /api/hosts` endpoint with:
|
||||
- Hostname filtering (ILIKE pattern matching)
|
||||
- IP address filtering
|
||||
- OS filtering
|
||||
- Dynamic sorting (any field, asc/desc)
|
||||
- Pagination support
|
||||
|
||||
#### Real-time Notifications
|
||||
- WebSocket-based real-time notifications
|
||||
- Persistent notification storage
|
||||
- **Endpoints**:
|
||||
- `WS /api/notifications/ws` - WebSocket connection
|
||||
- `GET /api/notifications` - List notifications
|
||||
- `PUT /api/notifications/{id}` - Mark as read
|
||||
- `POST /api/notifications/mark-all-read` - Mark all read
|
||||
- Filter by read/unread status
|
||||
- Automatic push to connected clients
|
||||
|
||||
#### Velociraptor Integration
|
||||
- Complete Velociraptor API client (async with httpx)
|
||||
- **Configuration**: `POST /api/velociraptor/config`
|
||||
- **Client Management**:
|
||||
- `GET /api/velociraptor/clients` - List clients
|
||||
- `GET /api/velociraptor/clients/{id}` - Get client info
|
||||
- **Artifact Collection**:
|
||||
- `POST /api/velociraptor/collect` - Collect artifact from client
|
||||
- **Hunt Management**:
|
||||
- `POST /api/velociraptor/hunts` - Create hunt
|
||||
- `GET /api/velociraptor/hunts/{id}/results` - Get hunt results
|
||||
- Per-tenant configuration storage
|
||||
|
||||
### Database Changes
|
||||
- `notifications` table
|
||||
|
||||
## Phase 4: Intelligence & Automation ✅
|
||||
|
||||
### Implemented Features
|
||||
|
||||
#### Machine Learning & Threat Intelligence
|
||||
- `ThreatAnalyzer` class for ML-based threat detection
|
||||
- Host threat analysis with scoring (0.0-1.0)
|
||||
- Artifact threat analysis
|
||||
- Anomaly detection capabilities
|
||||
- Threat classification (benign, low, medium, high, critical)
|
||||
- **Endpoints**:
|
||||
- `POST /api/threat-intel/analyze/host/{id}` - Analyze host
|
||||
- `POST /api/threat-intel/analyze/artifact/{id}` - Analyze artifact
|
||||
- `GET /api/threat-intel/scores` - List threat scores (filterable)
|
||||
- Stores results in database with confidence scores and indicators
|
||||
|
||||
#### Automated Playbooks
|
||||
- `PlaybookEngine` for executing automated responses
|
||||
- Supported actions:
|
||||
- `send_notification` - Send notification to user
|
||||
- `create_case` - Auto-create investigation case
|
||||
- `isolate_host` - Isolate compromised host
|
||||
- `collect_artifact` - Trigger artifact collection
|
||||
- `block_ip` - Block malicious IP
|
||||
- `send_email` - Send email alert
|
||||
- **Endpoints**:
|
||||
- `GET /api/playbooks` - List playbooks
|
||||
- `POST /api/playbooks` - Create playbook (admin)
|
||||
- `GET /api/playbooks/{id}` - Get playbook
|
||||
- `POST /api/playbooks/{id}/execute` - Execute playbook
|
||||
- `GET /api/playbooks/{id}/executions` - List execution history
|
||||
- Trigger types: manual, scheduled, event-based
|
||||
- Execution tracking with status and results
|
||||
|
||||
#### Advanced Reporting
|
||||
- Report template system
|
||||
- Multiple format support (PDF, HTML, JSON)
|
||||
- **Endpoints**:
|
||||
- `GET /api/reports/templates` - List templates
|
||||
- `POST /api/reports/templates` - Create template
|
||||
- `POST /api/reports/generate` - Generate report
|
||||
- `GET /api/reports` - List generated reports
|
||||
- `GET /api/reports/{id}` - Get specific report
|
||||
- Template types: case_summary, host_analysis, threat_report
|
||||
- Async report generation with status tracking
|
||||
|
||||
#### SIEM Integration (Foundation)
|
||||
- Architecture ready for SIEM connectors
|
||||
- Audit logs can be forwarded to SIEM
|
||||
- Threat scores exportable to SIEM
|
||||
- Webhook/API structure supports integration
|
||||
- Ready for Splunk, Elastic, etc. connectors
|
||||
|
||||
### Database Changes
|
||||
- `playbooks` table
|
||||
- `playbook_executions` table
|
||||
- `threat_scores` table
|
||||
- `report_templates` table
|
||||
- `reports` table
|
||||
|
||||
## API Statistics
|
||||
|
||||
### Total Endpoints: 70+
|
||||
|
||||
**By Category:**
|
||||
- Authentication & Users: 13 endpoints
|
||||
- Core Resources: 12 endpoints
|
||||
- Integrations: 15 endpoints
|
||||
- Intelligence & Automation: 20+ endpoints
|
||||
- Health & Info: 2 endpoints
|
||||
|
||||
### Authentication Required
|
||||
All endpoints except:
|
||||
- `POST /api/auth/register`
|
||||
- `POST /api/auth/login`
|
||||
- `POST /api/auth/password-reset/request`
|
||||
- `GET /health`
|
||||
- `GET /`
|
||||
|
||||
### Admin-Only Endpoints
|
||||
- User management (`/api/users`)
|
||||
- Tenant creation
|
||||
- Audit log viewing
|
||||
- Playbook creation
|
||||
- Velociraptor hunt creation
|
||||
|
||||
## Security Features
|
||||
|
||||
### Enhanced Security
|
||||
- ✅ TOTP 2FA implementation
|
||||
- ✅ Refresh token rotation
|
||||
- ✅ Password reset with secure tokens
|
||||
- ✅ Comprehensive audit logging
|
||||
- ✅ IP and user agent tracking
|
||||
- ✅ WebSocket authentication
|
||||
- ✅ Multi-tenant isolation (all phases)
|
||||
- ✅ Role-based access control (all endpoints)
|
||||
|
||||
### CodeQL Verification
|
||||
- All phases passed CodeQL security scan
|
||||
- 0 vulnerabilities detected
|
||||
- Best practices followed
|
||||
|
||||
## Database Schema
|
||||
|
||||
### Total Tables: 15
|
||||
|
||||
**Phase 1 (5 tables)**
|
||||
- tenants, users, hosts, cases, artifacts
|
||||
|
||||
**Phase 2 (3 tables)**
|
||||
- refresh_tokens, password_reset_tokens, audit_logs
|
||||
|
||||
**Phase 3 (1 table)**
|
||||
- notifications
|
||||
|
||||
**Phase 4 (6 tables)**
|
||||
- playbooks, playbook_executions, threat_scores, report_templates, reports
|
||||
|
||||
### Migrations
|
||||
All 4 migrations created and tested:
|
||||
1. `f82b3092d056_initial_migration.py`
|
||||
2. `a1b2c3d4e5f6_add_phase_2_tables.py`
|
||||
3. `b2c3d4e5f6g7_add_phase_3_tables.py`
|
||||
4. `c3d4e5f6g7h8_add_phase_4_tables.py`
|
||||
|
||||
## Dependencies Added
|
||||
|
||||
```
|
||||
pyotp==2.9.0 # TOTP 2FA
|
||||
qrcode[pil]==7.4.2 # QR code generation
|
||||
websockets==12.0 # WebSocket support
|
||||
httpx==0.26.0 # Async HTTP client
|
||||
email-validator==2.1.0 # Email validation
|
||||
```
|
||||
|
||||
## Usage Examples
|
||||
|
||||
### Phase 2: 2FA Setup
|
||||
```python
|
||||
# 1. Setup 2FA
|
||||
POST /api/auth/2fa/setup
|
||||
Response: {"secret": "...", "qr_code_uri": "otpauth://..."}
|
||||
|
||||
# 2. Verify and enable
|
||||
POST /api/auth/2fa/verify
|
||||
Body: {"code": "123456"}
|
||||
|
||||
# 3. Login with 2FA
|
||||
POST /api/auth/login
|
||||
Form: username=user&password=pass&scope=123456
|
||||
```
|
||||
|
||||
### Phase 3: Real-time Notifications
|
||||
```javascript
|
||||
// Frontend WebSocket connection
|
||||
const ws = new WebSocket('ws://localhost:8000/api/notifications/ws');
|
||||
ws.send(JSON.stringify({token: 'jwt_token_here'}));
|
||||
|
||||
ws.onmessage = (event) => {
|
||||
const notification = JSON.parse(event.data);
|
||||
// Display notification
|
||||
};
|
||||
```
|
||||
|
||||
### Phase 3: Velociraptor Integration
|
||||
```python
|
||||
# Configure Velociraptor
|
||||
POST /api/velociraptor/config
|
||||
Body: {"base_url": "https://veloci.example.com", "api_key": "..."}
|
||||
|
||||
# Collect artifact
|
||||
POST /api/velociraptor/collect
|
||||
Body: {
|
||||
"client_id": "C.abc123",
|
||||
"artifact_name": "Windows.System.Pslist"
|
||||
}
|
||||
```
|
||||
|
||||
### Phase 4: Threat Analysis
|
||||
```python
|
||||
# Analyze a host
|
||||
POST /api/threat-intel/analyze/host/123
|
||||
Response: {
|
||||
"score": 0.7,
|
||||
"confidence": 0.8,
|
||||
"threat_type": "high",
|
||||
"indicators": [...]
|
||||
}
|
||||
```
|
||||
|
||||
### Phase 4: Automated Playbook
|
||||
```python
|
||||
# Create playbook
|
||||
POST /api/playbooks
|
||||
Body: {
|
||||
"name": "Isolate High-Risk Host",
|
||||
"trigger_type": "event",
|
||||
"actions": [
|
||||
{"type": "send_notification", "params": {"message": "High risk detected"}},
|
||||
{"type": "isolate_host", "params": {"host_id": "${host_id}"}},
|
||||
{"type": "create_case", "params": {"title": "Auto-generated case"}}
|
||||
]
|
||||
}
|
||||
|
||||
# Execute playbook
|
||||
POST /api/playbooks/1/execute
|
||||
```
|
||||
|
||||
## Testing
|
||||
|
||||
### Manual Testing
|
||||
All endpoints have been tested with:
|
||||
- Authentication flows
|
||||
- Multi-tenancy isolation
|
||||
- Role-based access control
|
||||
- Error handling
|
||||
|
||||
### API Documentation
|
||||
Interactive API docs available at:
|
||||
- Swagger UI: `http://localhost:8000/docs`
|
||||
- ReDoc: `http://localhost:8000/redoc`
|
||||
|
||||
## Deployment Notes
|
||||
|
||||
### Environment Variables
|
||||
Add to `.env`:
|
||||
```bash
|
||||
# Phase 2
|
||||
REFRESH_TOKEN_EXPIRE_DAYS=30
|
||||
SMTP_HOST=localhost
|
||||
SMTP_PORT=587
|
||||
SMTP_USER=
|
||||
SMTP_PASSWORD=
|
||||
FROM_EMAIL=noreply@velocicompanion.com
|
||||
|
||||
# Phase 3
|
||||
WS_ENABLED=true
|
||||
```
|
||||
|
||||
### Database Migrations
|
||||
```bash
|
||||
# Run all migrations
|
||||
cd backend
|
||||
alembic upgrade head
|
||||
|
||||
# Or manually in order
|
||||
alembic upgrade f82b3092d056 # Phase 1
|
||||
alembic upgrade a1b2c3d4e5f6 # Phase 2
|
||||
alembic upgrade b2c3d4e5f6g7 # Phase 3
|
||||
alembic upgrade c3d4e5f6g7h8 # Phase 4
|
||||
```
|
||||
|
||||
## What's Next
|
||||
|
||||
The system is now feature-complete with all requested phases implemented:
|
||||
|
||||
✅ **Phase 1**: Core Infrastructure & Auth
|
||||
✅ **Phase 2**: Enhanced Authentication
|
||||
✅ **Phase 3**: Advanced Features
|
||||
✅ **Phase 4**: Intelligence & Automation
|
||||
|
||||
**Version: 1.0.0 - Production Ready**
|
||||
|
||||
### Future Enhancements (Optional)
|
||||
- Email service integration for password reset
|
||||
- Advanced ML models for threat detection
|
||||
- Additional SIEM connectors (Splunk, Elastic, etc.)
|
||||
- Mobile app for notifications
|
||||
- Advanced playbook conditions and branching
|
||||
- Scheduled playbook triggers
|
||||
- Custom dashboard widgets
|
||||
- Export/import for playbooks and reports
|
||||
- Multi-language support
|
||||
|
||||
## Support
|
||||
|
||||
For issues or questions:
|
||||
- Check API documentation at `/docs`
|
||||
- Review ARCHITECTURE.md for technical details
|
||||
- See QUICKSTART.md for setup instructions
|
||||
- Consult DEPLOYMENT_CHECKLIST.md for production deployment
|
||||
263
QUICKSTART.md
263
QUICKSTART.md
@@ -1,263 +0,0 @@
|
||||
# Quick Start Guide
|
||||
|
||||
This guide will help you get VelociCompanion up and running in minutes.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
- Docker and Docker Compose installed
|
||||
- 8GB RAM minimum
|
||||
- Ports 3000, 5432, and 8000 available
|
||||
|
||||
## Step 1: Start the Application
|
||||
|
||||
```bash
|
||||
# Clone the repository
|
||||
git clone https://github.com/mblanke/ThreatHunt.git
|
||||
cd ThreatHunt
|
||||
|
||||
# Start all services
|
||||
docker-compose up -d
|
||||
|
||||
# Check service status
|
||||
docker-compose ps
|
||||
```
|
||||
|
||||
Expected output:
|
||||
```
|
||||
NAME COMMAND SERVICE STATUS PORTS
|
||||
threathunt-backend-1 "sh -c 'alembic upgr…" backend running 0.0.0.0:8000->8000/tcp
|
||||
threathunt-db-1 "docker-entrypoint.s…" db running 0.0.0.0:5432->5432/tcp
|
||||
threathunt-frontend-1 "docker-entrypoint.s…" frontend running 0.0.0.0:3000->3000/tcp
|
||||
```
|
||||
|
||||
## Step 2: Verify Backend is Running
|
||||
|
||||
```bash
|
||||
# Check backend health
|
||||
curl http://localhost:8000/health
|
||||
|
||||
# Expected response:
|
||||
# {"status":"healthy"}
|
||||
|
||||
# View API documentation
|
||||
open http://localhost:8000/docs
|
||||
```
|
||||
|
||||
## Step 3: Access the Frontend
|
||||
|
||||
Open your browser and navigate to:
|
||||
```
|
||||
http://localhost:3000
|
||||
```
|
||||
|
||||
You should see the VelociCompanion login page.
|
||||
|
||||
## Step 4: Create Your First User
|
||||
|
||||
### Option A: Via API (using curl)
|
||||
|
||||
```bash
|
||||
# Register a new user
|
||||
curl -X POST http://localhost:8000/api/auth/register \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"username": "admin",
|
||||
"password": "admin123",
|
||||
"role": "admin"
|
||||
}'
|
||||
|
||||
# Login to get a token
|
||||
curl -X POST http://localhost:8000/api/auth/login \
|
||||
-H "Content-Type: application/x-www-form-urlencoded" \
|
||||
-d "username=admin&password=admin123"
|
||||
```
|
||||
|
||||
### Option B: Via Frontend
|
||||
|
||||
1. The first time you access the app, you'll need to register via API first (as shown above)
|
||||
2. Then login through the frontend at http://localhost:3000/login
|
||||
|
||||
## Step 5: Explore the API
|
||||
|
||||
Use the interactive API documentation at:
|
||||
```
|
||||
http://localhost:8000/docs
|
||||
```
|
||||
|
||||
Click "Authorize" and enter your token in the format:
|
||||
```
|
||||
Bearer YOUR_TOKEN_HERE
|
||||
```
|
||||
|
||||
## Step 6: Test the API
|
||||
|
||||
Run the test script to verify all endpoints:
|
||||
|
||||
```bash
|
||||
./test_api.sh
|
||||
```
|
||||
|
||||
Expected output:
|
||||
```
|
||||
===================================
|
||||
VelociCompanion API Test Script
|
||||
===================================
|
||||
|
||||
1. Testing health endpoint...
|
||||
✓ Health check passed
|
||||
|
||||
2. Registering a new user...
|
||||
✓ User registration successful
|
||||
|
||||
3. Logging in...
|
||||
✓ Login successful
|
||||
|
||||
4. Getting current user profile...
|
||||
✓ Profile retrieval successful
|
||||
|
||||
5. Listing tenants...
|
||||
✓ Tenants list retrieved
|
||||
|
||||
6. Listing hosts...
|
||||
Hosts: []
|
||||
|
||||
7. Testing authentication protection...
|
||||
✓ Authentication protection working
|
||||
|
||||
===================================
|
||||
API Testing Complete!
|
||||
===================================
|
||||
```
|
||||
|
||||
## Common Operations
|
||||
|
||||
### Create a Host
|
||||
|
||||
```bash
|
||||
# Get your token from login
|
||||
TOKEN="your_token_here"
|
||||
|
||||
# Create a host
|
||||
curl -X POST http://localhost:8000/api/hosts \
|
||||
-H "Authorization: Bearer $TOKEN" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"hostname": "workstation-01",
|
||||
"ip_address": "192.168.1.100",
|
||||
"os": "Windows 10"
|
||||
}'
|
||||
```
|
||||
|
||||
### List Hosts
|
||||
|
||||
```bash
|
||||
curl -X GET http://localhost:8000/api/hosts \
|
||||
-H "Authorization: Bearer $TOKEN"
|
||||
```
|
||||
|
||||
### Ingest Data
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8000/api/ingestion/ingest \
|
||||
-H "Authorization: Bearer $TOKEN" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"hostname": "server-01",
|
||||
"data": {
|
||||
"artifact": "Windows.System.TaskScheduler",
|
||||
"results": [...]
|
||||
}
|
||||
}'
|
||||
```
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Database Connection Issues
|
||||
|
||||
```bash
|
||||
# Check if database is running
|
||||
docker-compose logs db
|
||||
|
||||
# Restart database
|
||||
docker-compose restart db
|
||||
```
|
||||
|
||||
### Backend Not Starting
|
||||
|
||||
```bash
|
||||
# Check backend logs
|
||||
docker-compose logs backend
|
||||
|
||||
# Common issues:
|
||||
# - Database not ready: Wait a few seconds and check logs
|
||||
# - Port 8000 in use: Stop other services using that port
|
||||
```
|
||||
|
||||
### Frontend Not Loading
|
||||
|
||||
```bash
|
||||
# Check frontend logs
|
||||
docker-compose logs frontend
|
||||
|
||||
# Rebuild frontend if needed
|
||||
docker-compose build frontend
|
||||
docker-compose up -d frontend
|
||||
```
|
||||
|
||||
### Reset Everything
|
||||
|
||||
```bash
|
||||
# Stop and remove all containers and volumes
|
||||
docker-compose down -v
|
||||
|
||||
# Start fresh
|
||||
docker-compose up -d
|
||||
```
|
||||
|
||||
## Next Steps
|
||||
|
||||
1. **Create Additional Users**: Use the `/api/auth/register` endpoint
|
||||
2. **Set Up Tenants**: Create tenants via `/api/tenants` (admin only)
|
||||
3. **Integrate with Velociraptor**: Configure Velociraptor to send data to `/api/ingestion/ingest`
|
||||
4. **Explore Cases**: Create and manage threat hunting cases
|
||||
5. **Configure VirusTotal**: Set up VirusTotal API integration for hash lookups
|
||||
|
||||
## Security Considerations
|
||||
|
||||
⚠️ **Before deploying to production:**
|
||||
|
||||
1. Change the `SECRET_KEY` in docker-compose.yml or .env file
|
||||
- Must be at least 32 characters
|
||||
- Use a cryptographically random string
|
||||
|
||||
2. Use strong passwords for the database
|
||||
|
||||
3. Enable HTTPS/TLS for API and frontend
|
||||
|
||||
4. Configure proper firewall rules
|
||||
|
||||
5. Review and update CORS settings in `backend/app/main.py`
|
||||
|
||||
## Development Mode
|
||||
|
||||
To run in development mode with hot reload:
|
||||
|
||||
```bash
|
||||
# Backend
|
||||
cd backend
|
||||
python -m venv venv
|
||||
source venv/bin/activate
|
||||
pip install -r requirements.txt
|
||||
uvicorn app.main:app --reload
|
||||
|
||||
# Frontend (in another terminal)
|
||||
cd frontend
|
||||
npm install
|
||||
npm start
|
||||
```
|
||||
|
||||
## Support
|
||||
|
||||
- Documentation: See [README.md](README.md)
|
||||
- API Docs: http://localhost:8000/docs
|
||||
- Issues: GitHub Issues
|
||||
219
README.md
219
README.md
@@ -1,216 +1,69 @@
|
||||
# VelociCompanion
|
||||
<<<<<<< Updated upstream
|
||||
# ThreatHunt
|
||||
=======
|
||||
# Cyber Threat Hunter
|
||||
|
||||
A multi-tenant threat hunting companion for Velociraptor with JWT authentication and role-based access control.
|
||||
A modern web application for threat hunting and security analysis, built with React frontend and Flask backend.
|
||||
|
||||
## Features
|
||||
|
||||
- **JWT Authentication**: Secure token-based authentication system
|
||||
- **Multi-Tenancy**: Complete data isolation between tenants
|
||||
- **Role-Based Access Control**: Admin and user roles with different permissions
|
||||
- **RESTful API**: FastAPI backend with automatic OpenAPI documentation
|
||||
- **React Frontend**: Modern TypeScript React application with authentication
|
||||
- **Database Migrations**: Alembic for database schema management
|
||||
- **Docker Support**: Complete Docker Compose setup for easy deployment
|
||||
- **Security Tools Detection**: Identify running security tools (AV, EDR, VPN)
|
||||
- **CSV Processing**: Upload and analyze security logs
|
||||
- **Baseline Analysis**: System baseline comparison
|
||||
- **Network Analysis**: Network traffic and connection analysis
|
||||
- **VirusTotal Integration**: File and URL reputation checking
|
||||
|
||||
## Project Structure
|
||||
## Architecture
|
||||
|
||||
```
|
||||
ThreatHunt/
|
||||
├── backend/
|
||||
│ ├── alembic/ # Database migrations
|
||||
│ ├── app/
|
||||
│ │ ├── api/routes/ # API endpoints
|
||||
│ │ │ ├── auth.py # Authentication routes
|
||||
│ │ │ ├── users.py # User management
|
||||
│ │ │ ├── tenants.py # Tenant management
|
||||
│ │ │ ├── hosts.py # Host management
|
||||
│ │ │ ├── ingestion.py # Data ingestion
|
||||
│ │ │ └── vt.py # VirusTotal integration
|
||||
│ │ ├── core/ # Core functionality
|
||||
│ │ │ ├── config.py # Configuration
|
||||
│ │ │ ├── database.py # Database setup
|
||||
│ │ │ ├── security.py # Password hashing, JWT
|
||||
│ │ │ └── deps.py # FastAPI dependencies
|
||||
│ │ ├── models/ # SQLAlchemy models
|
||||
│ │ └── schemas/ # Pydantic schemas
|
||||
│ ├── requirements.txt
|
||||
│ └── Dockerfile
|
||||
├── frontend/
|
||||
│ ├── public/
|
||||
│ ├── src/
|
||||
│ │ ├── components/ # React components
|
||||
│ │ ├── context/ # Auth context
|
||||
│ │ ├── pages/ # Page components
|
||||
│ │ ├── utils/ # API utilities
|
||||
│ │ ├── App.tsx
|
||||
│ │ └── index.tsx
|
||||
│ ├── package.json
|
||||
│ └── Dockerfile
|
||||
└── docker-compose.yml
|
||||
|
||||
├── frontend/ # React application
|
||||
├── backend/ # Flask API server
|
||||
├── uploaded/ # File upload storage
|
||||
└── output/ # Analysis results
|
||||
```
|
||||
|
||||
## Getting Started
|
||||
## Quick Start
|
||||
|
||||
### Prerequisites
|
||||
|
||||
- Docker and Docker Compose
|
||||
- Python 3.11+ (for local development)
|
||||
- Node.js 18+ (for local development)
|
||||
|
||||
### Quick Start with Docker
|
||||
|
||||
1. Clone the repository:
|
||||
```bash
|
||||
git clone https://github.com/mblanke/ThreatHunt.git
|
||||
cd ThreatHunt
|
||||
```
|
||||
|
||||
2. Start all services:
|
||||
```bash
|
||||
docker-compose up -d
|
||||
```
|
||||
|
||||
3. Access the application:
|
||||
- Frontend: http://localhost:3000
|
||||
- Backend API: http://localhost:8000
|
||||
- API Documentation: http://localhost:8000/docs
|
||||
|
||||
### Local Development
|
||||
|
||||
#### Backend
|
||||
### Backend Setup
|
||||
|
||||
```bash
|
||||
cd backend
|
||||
python -m venv venv
|
||||
source venv/bin/activate # On Windows: venv\Scripts\activate
|
||||
pip install -r requirements.txt
|
||||
|
||||
# Set up environment variables
|
||||
cp .env.example .env
|
||||
# Edit .env with your settings
|
||||
|
||||
# Run migrations
|
||||
alembic upgrade head
|
||||
|
||||
# Start development server
|
||||
uvicorn app.main:app --reload
|
||||
chmod +x setup_backend.sh
|
||||
./setup_backend.sh
|
||||
source venv/bin/activate
|
||||
python app.py
|
||||
```
|
||||
|
||||
#### Frontend
|
||||
### Frontend Setup
|
||||
|
||||
```bash
|
||||
cd frontend
|
||||
npm install
|
||||
npm start
|
||||
npm run dev
|
||||
```
|
||||
|
||||
## API Endpoints
|
||||
|
||||
### Authentication
|
||||
- `POST /api/auth/register` - Register a new user
|
||||
- `POST /api/auth/login` - Login and receive JWT token
|
||||
- `GET /api/auth/me` - Get current user profile
|
||||
- `PUT /api/auth/me` - Update current user profile
|
||||
- `GET /` - Serve React app
|
||||
- `GET /api/health` - Health check
|
||||
- `POST /api/upload` - File upload
|
||||
- `GET /api/analysis/<id>` - Get analysis results
|
||||
|
||||
### User Management (Admin only)
|
||||
- `GET /api/users` - List all users in tenant
|
||||
- `GET /api/users/{user_id}` - Get user by ID
|
||||
- `PUT /api/users/{user_id}` - Update user
|
||||
- `DELETE /api/users/{user_id}` - Deactivate user
|
||||
## Security Considerations
|
||||
|
||||
### Tenants
|
||||
- `GET /api/tenants` - List tenants
|
||||
- `POST /api/tenants` - Create tenant (admin)
|
||||
- `GET /api/tenants/{tenant_id}` - Get tenant by ID
|
||||
|
||||
### Hosts
|
||||
- `GET /api/hosts` - List hosts (scoped to tenant)
|
||||
- `POST /api/hosts` - Create host
|
||||
- `GET /api/hosts/{host_id}` - Get host by ID
|
||||
|
||||
### Ingestion
|
||||
- `POST /api/ingestion/ingest` - Ingest data from Velociraptor
|
||||
|
||||
### VirusTotal
|
||||
- `POST /api/vt/lookup` - Lookup hash in VirusTotal
|
||||
|
||||
## Authentication Flow
|
||||
|
||||
1. User registers or logs in via `/api/auth/login`
|
||||
2. Backend returns JWT token with user_id, tenant_id, and role
|
||||
3. Frontend stores token in localStorage
|
||||
4. All subsequent API requests include token in Authorization header
|
||||
5. Backend validates token and enforces tenant scoping
|
||||
|
||||
## Multi-Tenancy
|
||||
|
||||
- All data is scoped to tenant_id
|
||||
- Users can only access data within their tenant
|
||||
- Admin users have elevated permissions within their tenant
|
||||
- Cross-tenant access requires explicit permissions
|
||||
|
||||
## Database Migrations
|
||||
|
||||
Create a new migration:
|
||||
```bash
|
||||
cd backend
|
||||
alembic revision --autogenerate -m "Description of changes"
|
||||
```
|
||||
|
||||
Apply migrations:
|
||||
```bash
|
||||
alembic upgrade head
|
||||
```
|
||||
|
||||
Rollback migrations:
|
||||
```bash
|
||||
alembic downgrade -1
|
||||
```
|
||||
|
||||
## Environment Variables
|
||||
|
||||
### Backend
|
||||
- `DATABASE_URL` - PostgreSQL connection string
|
||||
- `SECRET_KEY` - Secret key for JWT signing (min 32 characters)
|
||||
- `ACCESS_TOKEN_EXPIRE_MINUTES` - JWT token expiration time (default: 30)
|
||||
|
||||
### Frontend
|
||||
- `REACT_APP_API_URL` - Backend API URL (default: http://localhost:8000)
|
||||
|
||||
## Security
|
||||
|
||||
- Passwords are hashed using bcrypt
|
||||
- JWT tokens include expiration time
|
||||
- All API endpoints (except login/register) require authentication
|
||||
- Role-based access control for admin operations
|
||||
- Data isolation through tenant scoping
|
||||
|
||||
## Testing
|
||||
|
||||
### Backend
|
||||
```bash
|
||||
cd backend
|
||||
pytest
|
||||
```
|
||||
|
||||
### Frontend
|
||||
```bash
|
||||
cd frontend
|
||||
npm test
|
||||
```
|
||||
- File upload validation
|
||||
- Input sanitization
|
||||
- Rate limiting
|
||||
- CORS configuration
|
||||
|
||||
## Contributing
|
||||
|
||||
1. Fork the repository
|
||||
2. Create a feature branch
|
||||
3. Make your changes
|
||||
4. Submit a pull request
|
||||
2. Create feature branch
|
||||
3. Submit pull request
|
||||
|
||||
## License
|
||||
|
||||
[Your License Here]
|
||||
|
||||
## Support
|
||||
|
||||
For issues and questions, please open an issue on GitHub.
|
||||
MIT License
|
||||
>>>>>>> Stashed changes
|
||||
|
||||
@@ -1,232 +0,0 @@
|
||||
# Validation Report
|
||||
|
||||
**Date**: 2025-12-09
|
||||
**Version**: 1.0.0
|
||||
**Status**: ✅ ALL CHECKS PASSED
|
||||
|
||||
## Summary
|
||||
|
||||
Comprehensive error checking and validation has been performed on all components of the VelociCompanion threat hunting platform.
|
||||
|
||||
## Python Backend Validation
|
||||
|
||||
### ✅ Syntax Check
|
||||
- All Python files compile successfully
|
||||
- No syntax errors found in 53 files
|
||||
|
||||
### ✅ Import Validation
|
||||
- All core modules import correctly
|
||||
- All 12 model classes verified
|
||||
- All schema modules working
|
||||
- All 12 route modules operational
|
||||
- All engine modules (Velociraptor, ThreatAnalyzer, PlaybookEngine) functional
|
||||
|
||||
### ✅ FastAPI Application
|
||||
- Application loads successfully
|
||||
- 53 routes registered correctly
|
||||
- Version 1.0.0 confirmed
|
||||
- All route tags properly assigned
|
||||
|
||||
### ✅ API Endpoints Registered
|
||||
**Authentication** (10 endpoints)
|
||||
- POST /api/auth/register
|
||||
- POST /api/auth/login
|
||||
- POST /api/auth/refresh
|
||||
- GET /api/auth/me
|
||||
- PUT /api/auth/me
|
||||
- POST /api/auth/2fa/setup
|
||||
- POST /api/auth/2fa/verify
|
||||
- POST /api/auth/2fa/disable
|
||||
- POST /api/auth/password-reset/request
|
||||
- POST /api/auth/password-reset/confirm
|
||||
|
||||
**Users** (4 endpoints)
|
||||
- GET /api/users/
|
||||
- GET /api/users/{user_id}
|
||||
- PUT /api/users/{user_id}
|
||||
- DELETE /api/users/{user_id}
|
||||
|
||||
**Tenants** (3 endpoints)
|
||||
- GET /api/tenants/
|
||||
- POST /api/tenants/
|
||||
- GET /api/tenants/{tenant_id}
|
||||
|
||||
**Hosts** (3 endpoints)
|
||||
- GET /api/hosts/
|
||||
- POST /api/hosts/
|
||||
- GET /api/hosts/{host_id}
|
||||
|
||||
**Audit Logs** (2 endpoints)
|
||||
- GET /api/audit/
|
||||
- GET /api/audit/{log_id}
|
||||
|
||||
**Notifications** (3 endpoints)
|
||||
- GET /api/notifications/
|
||||
- PUT /api/notifications/{notification_id}
|
||||
- POST /api/notifications/mark-all-read
|
||||
|
||||
**Velociraptor** (6 endpoints)
|
||||
- POST /api/velociraptor/config
|
||||
- GET /api/velociraptor/clients
|
||||
- GET /api/velociraptor/clients/{client_id}
|
||||
- POST /api/velociraptor/collect
|
||||
- POST /api/velociraptor/hunts
|
||||
- GET /api/velociraptor/hunts/{hunt_id}/results
|
||||
|
||||
**Playbooks** (5 endpoints)
|
||||
- GET /api/playbooks/
|
||||
- POST /api/playbooks/
|
||||
- GET /api/playbooks/{playbook_id}
|
||||
- POST /api/playbooks/{playbook_id}/execute
|
||||
- GET /api/playbooks/{playbook_id}/executions
|
||||
|
||||
**Threat Intelligence** (3 endpoints)
|
||||
- POST /api/threat-intel/analyze/host/{host_id}
|
||||
- POST /api/threat-intel/analyze/artifact/{artifact_id}
|
||||
- GET /api/threat-intel/scores
|
||||
|
||||
**Reports** (5 endpoints)
|
||||
- GET /api/reports/templates
|
||||
- POST /api/reports/templates
|
||||
- POST /api/reports/generate
|
||||
- GET /api/reports/
|
||||
- GET /api/reports/{report_id}
|
||||
|
||||
**Other** (4 endpoints)
|
||||
- POST /api/ingestion/ingest
|
||||
- POST /api/vt/lookup
|
||||
- GET /
|
||||
- GET /health
|
||||
|
||||
**Total**: 53 routes successfully registered
|
||||
|
||||
## Frontend Validation
|
||||
|
||||
### ✅ TypeScript Files
|
||||
- All 8 TypeScript/TSX files validated
|
||||
- Import statements correct
|
||||
- Component hierarchy verified
|
||||
|
||||
### ✅ File Structure
|
||||
```
|
||||
src/
|
||||
├── App.tsx ✓
|
||||
├── index.tsx ✓
|
||||
├── react-app-env.d.ts ✓
|
||||
├── components/
|
||||
│ └── PrivateRoute.tsx ✓
|
||||
├── context/
|
||||
│ └── AuthContext.tsx ✓
|
||||
├── pages/
|
||||
│ ├── Login.tsx ✓
|
||||
│ └── Dashboard.tsx ✓
|
||||
└── utils/
|
||||
└── api.ts ✓
|
||||
```
|
||||
|
||||
### ✅ Configuration Files
|
||||
- package.json: Valid JSON ✓
|
||||
- tsconfig.json: Present ✓
|
||||
- Dockerfile: Present ✓
|
||||
|
||||
## Database Validation
|
||||
|
||||
### ✅ Migration Chain
|
||||
Correct migration dependency chain:
|
||||
1. f82b3092d056 (Phase 1 - Initial) → None
|
||||
2. a1b2c3d4e5f6 (Phase 2) → f82b3092d056
|
||||
3. b2c3d4e5f6g7 (Phase 3) → a1b2c3d4e5f6
|
||||
4. c3d4e5f6g7h8 (Phase 4) → b2c3d4e5f6g7
|
||||
|
||||
### ✅ Database Models
|
||||
All 15 tables defined:
|
||||
- Phase 1: tenants, users, hosts, cases, artifacts
|
||||
- Phase 2: refresh_tokens, password_reset_tokens, audit_logs
|
||||
- Phase 3: notifications
|
||||
- Phase 4: playbooks, playbook_executions, threat_scores, report_templates, reports
|
||||
|
||||
## Infrastructure Validation
|
||||
|
||||
### ✅ Docker Compose
|
||||
- PostgreSQL service configured ✓
|
||||
- Backend service with migrations ✓
|
||||
- Frontend service configured ✓
|
||||
- Health checks enabled ✓
|
||||
- Volume mounts correct ✓
|
||||
|
||||
### ✅ Configuration Files
|
||||
- alembic.ini: Valid ✓
|
||||
- requirements.txt: Valid (email-validator updated to 2.1.2) ✓
|
||||
- .env.example: Present ✓
|
||||
|
||||
## Documentation Validation
|
||||
|
||||
### ✅ Documentation Files Present
|
||||
- README.md ✓
|
||||
- QUICKSTART.md ✓
|
||||
- ARCHITECTURE.md ✓
|
||||
- DEPLOYMENT_CHECKLIST.md ✓
|
||||
- IMPLEMENTATION_SUMMARY.md ✓
|
||||
- PHASES_COMPLETE.md ✓
|
||||
|
||||
### ✅ Internal Links
|
||||
- All markdown cross-references validated
|
||||
- File references correct
|
||||
|
||||
### ✅ Scripts
|
||||
- test_api.sh: Valid bash syntax ✓
|
||||
|
||||
## Dependencies
|
||||
|
||||
### ✅ Python Dependencies
|
||||
All required packages specified:
|
||||
- fastapi==0.109.0
|
||||
- uvicorn[standard]==0.27.0
|
||||
- sqlalchemy==2.0.25
|
||||
- psycopg2-binary==2.9.9
|
||||
- python-jose[cryptography]==3.3.0
|
||||
- passlib[bcrypt]==1.7.4
|
||||
- python-multipart==0.0.6
|
||||
- alembic==1.13.1
|
||||
- pydantic==2.5.3
|
||||
- pydantic-settings==2.1.0
|
||||
- pyotp==2.9.0
|
||||
- qrcode[pil]==7.4.2
|
||||
- websockets==12.0
|
||||
- httpx==0.26.0
|
||||
- email-validator==2.1.2 (updated from 2.1.0)
|
||||
|
||||
### ✅ Node Dependencies
|
||||
- React 18.2.0
|
||||
- TypeScript 5.3.3
|
||||
- React Router 6.21.0
|
||||
- Axios 1.6.2
|
||||
|
||||
## Security
|
||||
|
||||
### ✅ Security Checks
|
||||
- No hardcoded credentials in code
|
||||
- Environment variables used for secrets
|
||||
- JWT tokens properly secured
|
||||
- Password hashing with bcrypt
|
||||
- 0 vulnerabilities reported by CodeQL
|
||||
|
||||
## Issues Fixed
|
||||
|
||||
1. **email-validator version**: Updated from 2.1.0 to 2.1.2 to avoid yanked version warning
|
||||
|
||||
## Conclusion
|
||||
|
||||
✅ **All validation checks passed successfully**
|
||||
|
||||
The VelociCompanion platform is fully functional with:
|
||||
- 53 API endpoints operational
|
||||
- 15 database tables with correct relationships
|
||||
- 4 complete migration files
|
||||
- All imports and dependencies resolved
|
||||
- Frontend components properly structured
|
||||
- Docker infrastructure configured
|
||||
- Comprehensive documentation
|
||||
|
||||
**Status**: Production Ready
|
||||
**Recommended Action**: Deploy to staging for integration testing
|
||||
9
backend/.env
Normal file
9
backend/.env
Normal file
@@ -0,0 +1,9 @@
|
||||
FLASK_ENV=development
|
||||
FLASK_DEBUG=True
|
||||
SECRET_KEY=development-secret-key-change-in-production
|
||||
MAX_CONTENT_LENGTH=104857600
|
||||
UPLOAD_FOLDER=uploaded
|
||||
OUTPUT_FOLDER=output
|
||||
VIRUSTOTAL_API_KEY=
|
||||
DATABASE_URL=sqlite:///threat_hunter.db
|
||||
REDIS_URL=redis://localhost:6379/0
|
||||
@@ -1,3 +0,0 @@
|
||||
DATABASE_URL=postgresql://postgres:postgres@db:5432/velocicompanion
|
||||
SECRET_KEY=your-secret-key-change-in-production-min-32-chars-long
|
||||
ACCESS_TOKEN_EXPIRE_MINUTES=30
|
||||
|
||||
@@ -2,12 +2,22 @@ FROM python:3.11-slim
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# Install dependencies
|
||||
# Install system dependencies
|
||||
RUN apt-get update && apt-get install -y \
|
||||
gcc \
|
||||
postgresql-client \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Copy requirements first for better caching
|
||||
COPY requirements.txt .
|
||||
RUN pip install --no-cache-dir -r requirements.txt
|
||||
|
||||
# Copy application code
|
||||
COPY . .
|
||||
|
||||
# Run migrations and start server
|
||||
CMD ["sh", "-c", "alembic upgrade head && uvicorn app.main:app --host 0.0.0.0 --port 8000 --reload"]
|
||||
# Create directories
|
||||
RUN mkdir -p uploads output
|
||||
|
||||
EXPOSE 5000
|
||||
|
||||
CMD ["python", "app.py"]
|
||||
|
||||
24
backend/Dockerfile.prod
Normal file
24
backend/Dockerfile.prod
Normal file
@@ -0,0 +1,24 @@
|
||||
FROM python:3.11-slim
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# Install system dependencies
|
||||
RUN apt-get update && apt-get install -y \
|
||||
gcc \
|
||||
postgresql-client \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Create non-root user
|
||||
RUN useradd --create-home --shell /bin/bash app
|
||||
|
||||
COPY requirements.txt .
|
||||
RUN pip install --no-cache-dir -r requirements.txt
|
||||
|
||||
COPY . .
|
||||
RUN chown -R app:app /app
|
||||
|
||||
USER app
|
||||
|
||||
EXPOSE 5000
|
||||
|
||||
CMD ["gunicorn", "--bind", "0.0.0.0:5000", "--workers", "4", "app:app"]
|
||||
@@ -1,147 +0,0 @@
|
||||
# A generic, single database configuration.
|
||||
|
||||
[alembic]
|
||||
# path to migration scripts.
|
||||
# this is typically a path given in POSIX (e.g. forward slashes)
|
||||
# format, relative to the token %(here)s which refers to the location of this
|
||||
# ini file
|
||||
script_location = %(here)s/alembic
|
||||
|
||||
# template used to generate migration file names; The default value is %%(rev)s_%%(slug)s
|
||||
# Uncomment the line below if you want the files to be prepended with date and time
|
||||
# see https://alembic.sqlalchemy.org/en/latest/tutorial.html#editing-the-ini-file
|
||||
# for all available tokens
|
||||
# file_template = %%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s
|
||||
|
||||
# sys.path path, will be prepended to sys.path if present.
|
||||
# defaults to the current working directory. for multiple paths, the path separator
|
||||
# is defined by "path_separator" below.
|
||||
prepend_sys_path = .
|
||||
|
||||
|
||||
# timezone to use when rendering the date within the migration file
|
||||
# as well as the filename.
|
||||
# If specified, requires the tzdata library which can be installed by adding
|
||||
# `alembic[tz]` to the pip requirements.
|
||||
# string value is passed to ZoneInfo()
|
||||
# leave blank for localtime
|
||||
# timezone =
|
||||
|
||||
# max length of characters to apply to the "slug" field
|
||||
# truncate_slug_length = 40
|
||||
|
||||
# set to 'true' to run the environment during
|
||||
# the 'revision' command, regardless of autogenerate
|
||||
# revision_environment = false
|
||||
|
||||
# set to 'true' to allow .pyc and .pyo files without
|
||||
# a source .py file to be detected as revisions in the
|
||||
# versions/ directory
|
||||
# sourceless = false
|
||||
|
||||
# version location specification; This defaults
|
||||
# to <script_location>/versions. When using multiple version
|
||||
# directories, initial revisions must be specified with --version-path.
|
||||
# The path separator used here should be the separator specified by "path_separator"
|
||||
# below.
|
||||
# version_locations = %(here)s/bar:%(here)s/bat:%(here)s/alembic/versions
|
||||
|
||||
# path_separator; This indicates what character is used to split lists of file
|
||||
# paths, including version_locations and prepend_sys_path within configparser
|
||||
# files such as alembic.ini.
|
||||
# The default rendered in new alembic.ini files is "os", which uses os.pathsep
|
||||
# to provide os-dependent path splitting.
|
||||
#
|
||||
# Note that in order to support legacy alembic.ini files, this default does NOT
|
||||
# take place if path_separator is not present in alembic.ini. If this
|
||||
# option is omitted entirely, fallback logic is as follows:
|
||||
#
|
||||
# 1. Parsing of the version_locations option falls back to using the legacy
|
||||
# "version_path_separator" key, which if absent then falls back to the legacy
|
||||
# behavior of splitting on spaces and/or commas.
|
||||
# 2. Parsing of the prepend_sys_path option falls back to the legacy
|
||||
# behavior of splitting on spaces, commas, or colons.
|
||||
#
|
||||
# Valid values for path_separator are:
|
||||
#
|
||||
# path_separator = :
|
||||
# path_separator = ;
|
||||
# path_separator = space
|
||||
# path_separator = newline
|
||||
#
|
||||
# Use os.pathsep. Default configuration used for new projects.
|
||||
path_separator = os
|
||||
|
||||
# set to 'true' to search source files recursively
|
||||
# in each "version_locations" directory
|
||||
# new in Alembic version 1.10
|
||||
# recursive_version_locations = false
|
||||
|
||||
# the output encoding used when revision files
|
||||
# are written from script.py.mako
|
||||
# output_encoding = utf-8
|
||||
|
||||
# database URL. This is consumed by the user-maintained env.py script only.
|
||||
# other means of configuring database URLs may be customized within the env.py
|
||||
# file.
|
||||
# sqlalchemy.url is configured in env.py from app.core.config
|
||||
|
||||
|
||||
[post_write_hooks]
|
||||
# post_write_hooks defines scripts or Python functions that are run
|
||||
# on newly generated revision scripts. See the documentation for further
|
||||
# detail and examples
|
||||
|
||||
# format using "black" - use the console_scripts runner, against the "black" entrypoint
|
||||
# hooks = black
|
||||
# black.type = console_scripts
|
||||
# black.entrypoint = black
|
||||
# black.options = -l 79 REVISION_SCRIPT_FILENAME
|
||||
|
||||
# lint with attempts to fix using "ruff" - use the module runner, against the "ruff" module
|
||||
# hooks = ruff
|
||||
# ruff.type = module
|
||||
# ruff.module = ruff
|
||||
# ruff.options = check --fix REVISION_SCRIPT_FILENAME
|
||||
|
||||
# Alternatively, use the exec runner to execute a binary found on your PATH
|
||||
# hooks = ruff
|
||||
# ruff.type = exec
|
||||
# ruff.executable = ruff
|
||||
# ruff.options = check --fix REVISION_SCRIPT_FILENAME
|
||||
|
||||
# Logging configuration. This is also consumed by the user-maintained
|
||||
# env.py script only.
|
||||
[loggers]
|
||||
keys = root,sqlalchemy,alembic
|
||||
|
||||
[handlers]
|
||||
keys = console
|
||||
|
||||
[formatters]
|
||||
keys = generic
|
||||
|
||||
[logger_root]
|
||||
level = WARNING
|
||||
handlers = console
|
||||
qualname =
|
||||
|
||||
[logger_sqlalchemy]
|
||||
level = WARNING
|
||||
handlers =
|
||||
qualname = sqlalchemy.engine
|
||||
|
||||
[logger_alembic]
|
||||
level = INFO
|
||||
handlers =
|
||||
qualname = alembic
|
||||
|
||||
[handler_console]
|
||||
class = StreamHandler
|
||||
args = (sys.stderr,)
|
||||
level = NOTSET
|
||||
formatter = generic
|
||||
|
||||
[formatter_generic]
|
||||
format = %(levelname)-5.5s [%(name)s] %(message)s
|
||||
datefmt = %H:%M:%S
|
||||
@@ -1 +0,0 @@
|
||||
Generic single-database configuration.
|
||||
@@ -1,95 +0,0 @@
|
||||
from logging.config import fileConfig
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from sqlalchemy import engine_from_config
|
||||
from sqlalchemy import pool
|
||||
|
||||
from alembic import context
|
||||
|
||||
# Add app directory to Python path
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||||
|
||||
# Import models and database
|
||||
from app.core.database import Base
|
||||
from app.core.config import settings
|
||||
|
||||
# Import all models to ensure they're registered with Base
|
||||
from app.models.tenant import Tenant
|
||||
from app.models.user import User
|
||||
from app.models.host import Host
|
||||
from app.models.case import Case
|
||||
from app.models.artifact import Artifact
|
||||
|
||||
# this is the Alembic Config object, which provides
|
||||
# access to the values within the .ini file in use.
|
||||
config = context.config
|
||||
|
||||
# Set the database URL from settings
|
||||
config.set_main_option("sqlalchemy.url", settings.database_url)
|
||||
|
||||
# Interpret the config file for Python logging.
|
||||
# This line sets up loggers basically.
|
||||
if config.config_file_name is not None:
|
||||
fileConfig(config.config_file_name)
|
||||
|
||||
# add your model's MetaData object here
|
||||
# for 'autogenerate' support
|
||||
target_metadata = Base.metadata
|
||||
|
||||
# other values from the config, defined by the needs of env.py,
|
||||
# can be acquired:
|
||||
# my_important_option = config.get_main_option("my_important_option")
|
||||
# ... etc.
|
||||
|
||||
|
||||
def run_migrations_offline() -> None:
|
||||
"""Run migrations in 'offline' mode.
|
||||
|
||||
This configures the context with just a URL
|
||||
and not an Engine, though an Engine is acceptable
|
||||
here as well. By skipping the Engine creation
|
||||
we don't even need a DBAPI to be available.
|
||||
|
||||
Calls to context.execute() here emit the given string to the
|
||||
script output.
|
||||
|
||||
"""
|
||||
url = config.get_main_option("sqlalchemy.url")
|
||||
context.configure(
|
||||
url=url,
|
||||
target_metadata=target_metadata,
|
||||
literal_binds=True,
|
||||
dialect_opts={"paramstyle": "named"},
|
||||
)
|
||||
|
||||
with context.begin_transaction():
|
||||
context.run_migrations()
|
||||
|
||||
|
||||
def run_migrations_online() -> None:
|
||||
"""Run migrations in 'online' mode.
|
||||
|
||||
In this scenario we need to create an Engine
|
||||
and associate a connection with the context.
|
||||
|
||||
"""
|
||||
connectable = engine_from_config(
|
||||
config.get_section(config.config_ini_section, {}),
|
||||
prefix="sqlalchemy.",
|
||||
poolclass=pool.NullPool,
|
||||
)
|
||||
|
||||
with connectable.connect() as connection:
|
||||
context.configure(
|
||||
connection=connection, target_metadata=target_metadata
|
||||
)
|
||||
|
||||
with context.begin_transaction():
|
||||
context.run_migrations()
|
||||
|
||||
|
||||
if context.is_offline_mode():
|
||||
run_migrations_offline()
|
||||
else:
|
||||
run_migrations_online()
|
||||
@@ -1,28 +0,0 @@
|
||||
"""${message}
|
||||
|
||||
Revision ID: ${up_revision}
|
||||
Revises: ${down_revision | comma,n}
|
||||
Create Date: ${create_date}
|
||||
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
${imports if imports else ""}
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = ${repr(up_revision)}
|
||||
down_revision: Union[str, Sequence[str], None] = ${repr(down_revision)}
|
||||
branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)}
|
||||
depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)}
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Upgrade schema."""
|
||||
${upgrades if upgrades else "pass"}
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Downgrade schema."""
|
||||
${downgrades if downgrades else "pass"}
|
||||
@@ -1,105 +0,0 @@
|
||||
"""Add Phase 2 tables
|
||||
|
||||
Revision ID: a1b2c3d4e5f6
|
||||
Revises: f82b3092d056
|
||||
Create Date: 2025-12-09 17:28:20.000000
|
||||
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = 'a1b2c3d4e5f6'
|
||||
down_revision: Union[str, Sequence[str], None] = 'f82b3092d056'
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Upgrade schema for Phase 2."""
|
||||
|
||||
# Add new fields to users table
|
||||
op.add_column('users', sa.Column('email', sa.String(), nullable=True))
|
||||
op.add_column('users', sa.Column('email_verified', sa.Boolean(), nullable=False, server_default='false'))
|
||||
op.add_column('users', sa.Column('totp_secret', sa.String(), nullable=True))
|
||||
op.add_column('users', sa.Column('totp_enabled', sa.Boolean(), nullable=False, server_default='false'))
|
||||
op.create_index(op.f('ix_users_email'), 'users', ['email'], unique=True)
|
||||
|
||||
# Create refresh_tokens table
|
||||
op.create_table(
|
||||
'refresh_tokens',
|
||||
sa.Column('id', sa.Integer(), nullable=False),
|
||||
sa.Column('token', sa.String(), nullable=False),
|
||||
sa.Column('user_id', sa.Integer(), nullable=False),
|
||||
sa.Column('expires_at', sa.DateTime(), nullable=False),
|
||||
sa.Column('is_revoked', sa.Boolean(), nullable=False),
|
||||
sa.Column('created_at', sa.DateTime(), nullable=True),
|
||||
sa.ForeignKeyConstraint(['user_id'], ['users.id'], ),
|
||||
sa.PrimaryKeyConstraint('id')
|
||||
)
|
||||
op.create_index(op.f('ix_refresh_tokens_id'), 'refresh_tokens', ['id'], unique=False)
|
||||
op.create_index(op.f('ix_refresh_tokens_token'), 'refresh_tokens', ['token'], unique=True)
|
||||
|
||||
# Create password_reset_tokens table
|
||||
op.create_table(
|
||||
'password_reset_tokens',
|
||||
sa.Column('id', sa.Integer(), nullable=False),
|
||||
sa.Column('token', sa.String(), nullable=False),
|
||||
sa.Column('user_id', sa.Integer(), nullable=False),
|
||||
sa.Column('expires_at', sa.DateTime(), nullable=False),
|
||||
sa.Column('is_used', sa.Boolean(), nullable=False),
|
||||
sa.Column('created_at', sa.DateTime(), nullable=True),
|
||||
sa.ForeignKeyConstraint(['user_id'], ['users.id'], ),
|
||||
sa.PrimaryKeyConstraint('id')
|
||||
)
|
||||
op.create_index(op.f('ix_password_reset_tokens_id'), 'password_reset_tokens', ['id'], unique=False)
|
||||
op.create_index(op.f('ix_password_reset_tokens_token'), 'password_reset_tokens', ['token'], unique=True)
|
||||
|
||||
# Create audit_logs table
|
||||
op.create_table(
|
||||
'audit_logs',
|
||||
sa.Column('id', sa.Integer(), nullable=False),
|
||||
sa.Column('user_id', sa.Integer(), nullable=True),
|
||||
sa.Column('tenant_id', sa.Integer(), nullable=False),
|
||||
sa.Column('action', sa.String(), nullable=False),
|
||||
sa.Column('resource_type', sa.String(), nullable=False),
|
||||
sa.Column('resource_id', sa.Integer(), nullable=True),
|
||||
sa.Column('details', sa.JSON(), nullable=True),
|
||||
sa.Column('ip_address', sa.String(), nullable=True),
|
||||
sa.Column('user_agent', sa.String(), nullable=True),
|
||||
sa.Column('created_at', sa.DateTime(), nullable=True),
|
||||
sa.ForeignKeyConstraint(['tenant_id'], ['tenants.id'], ),
|
||||
sa.ForeignKeyConstraint(['user_id'], ['users.id'], ),
|
||||
sa.PrimaryKeyConstraint('id')
|
||||
)
|
||||
op.create_index(op.f('ix_audit_logs_id'), 'audit_logs', ['id'], unique=False)
|
||||
op.create_index(op.f('ix_audit_logs_created_at'), 'audit_logs', ['created_at'], unique=False)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Downgrade schema for Phase 2."""
|
||||
|
||||
# Drop audit_logs table
|
||||
op.drop_index(op.f('ix_audit_logs_created_at'), table_name='audit_logs')
|
||||
op.drop_index(op.f('ix_audit_logs_id'), table_name='audit_logs')
|
||||
op.drop_table('audit_logs')
|
||||
|
||||
# Drop password_reset_tokens table
|
||||
op.drop_index(op.f('ix_password_reset_tokens_token'), table_name='password_reset_tokens')
|
||||
op.drop_index(op.f('ix_password_reset_tokens_id'), table_name='password_reset_tokens')
|
||||
op.drop_table('password_reset_tokens')
|
||||
|
||||
# Drop refresh_tokens table
|
||||
op.drop_index(op.f('ix_refresh_tokens_token'), table_name='refresh_tokens')
|
||||
op.drop_index(op.f('ix_refresh_tokens_id'), table_name='refresh_tokens')
|
||||
op.drop_table('refresh_tokens')
|
||||
|
||||
# Remove new fields from users table
|
||||
op.drop_index(op.f('ix_users_email'), table_name='users')
|
||||
op.drop_column('users', 'totp_enabled')
|
||||
op.drop_column('users', 'totp_secret')
|
||||
op.drop_column('users', 'email_verified')
|
||||
op.drop_column('users', 'email')
|
||||
@@ -1,50 +0,0 @@
|
||||
"""Add Phase 3 tables
|
||||
|
||||
Revision ID: b2c3d4e5f6g7
|
||||
Revises: a1b2c3d4e5f6
|
||||
Create Date: 2025-12-09 17:30:00.000000
|
||||
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = 'b2c3d4e5f6g7'
|
||||
down_revision: Union[str, Sequence[str], None] = 'a1b2c3d4e5f6'
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Upgrade schema for Phase 3."""
|
||||
|
||||
# Create notifications table
|
||||
op.create_table(
|
||||
'notifications',
|
||||
sa.Column('id', sa.Integer(), nullable=False),
|
||||
sa.Column('user_id', sa.Integer(), nullable=False),
|
||||
sa.Column('tenant_id', sa.Integer(), nullable=False),
|
||||
sa.Column('title', sa.String(), nullable=False),
|
||||
sa.Column('message', sa.Text(), nullable=False),
|
||||
sa.Column('notification_type', sa.String(), nullable=False),
|
||||
sa.Column('is_read', sa.Boolean(), nullable=False),
|
||||
sa.Column('link', sa.String(), nullable=True),
|
||||
sa.Column('created_at', sa.DateTime(), nullable=True),
|
||||
sa.ForeignKeyConstraint(['tenant_id'], ['tenants.id'], ),
|
||||
sa.ForeignKeyConstraint(['user_id'], ['users.id'], ),
|
||||
sa.PrimaryKeyConstraint('id')
|
||||
)
|
||||
op.create_index(op.f('ix_notifications_id'), 'notifications', ['id'], unique=False)
|
||||
op.create_index(op.f('ix_notifications_created_at'), 'notifications', ['created_at'], unique=False)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Downgrade schema for Phase 3."""
|
||||
|
||||
# Drop notifications table
|
||||
op.drop_index(op.f('ix_notifications_created_at'), table_name='notifications')
|
||||
op.drop_index(op.f('ix_notifications_id'), table_name='notifications')
|
||||
op.drop_table('notifications')
|
||||
@@ -1,152 +0,0 @@
|
||||
"""Add Phase 4 tables
|
||||
|
||||
Revision ID: c3d4e5f6g7h8
|
||||
Revises: b2c3d4e5f6g7
|
||||
Create Date: 2025-12-09 17:35:00.000000
|
||||
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = 'c3d4e5f6g7h8'
|
||||
down_revision: Union[str, Sequence[str], None] = 'b2c3d4e5f6g7'
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Upgrade schema for Phase 4."""
|
||||
|
||||
# Create playbooks table
|
||||
op.create_table(
|
||||
'playbooks',
|
||||
sa.Column('id', sa.Integer(), nullable=False),
|
||||
sa.Column('tenant_id', sa.Integer(), nullable=False),
|
||||
sa.Column('name', sa.String(), nullable=False),
|
||||
sa.Column('description', sa.Text(), nullable=True),
|
||||
sa.Column('trigger_type', sa.String(), nullable=False),
|
||||
sa.Column('trigger_config', sa.JSON(), nullable=True),
|
||||
sa.Column('actions', sa.JSON(), nullable=False),
|
||||
sa.Column('is_enabled', sa.Boolean(), nullable=False),
|
||||
sa.Column('created_by', sa.Integer(), nullable=False),
|
||||
sa.Column('created_at', sa.DateTime(), nullable=True),
|
||||
sa.Column('updated_at', sa.DateTime(), nullable=True),
|
||||
sa.ForeignKeyConstraint(['tenant_id'], ['tenants.id'], ),
|
||||
sa.ForeignKeyConstraint(['created_by'], ['users.id'], ),
|
||||
sa.PrimaryKeyConstraint('id')
|
||||
)
|
||||
op.create_index(op.f('ix_playbooks_id'), 'playbooks', ['id'], unique=False)
|
||||
op.create_index(op.f('ix_playbooks_name'), 'playbooks', ['name'], unique=False)
|
||||
|
||||
# Create playbook_executions table
|
||||
op.create_table(
|
||||
'playbook_executions',
|
||||
sa.Column('id', sa.Integer(), nullable=False),
|
||||
sa.Column('playbook_id', sa.Integer(), nullable=False),
|
||||
sa.Column('tenant_id', sa.Integer(), nullable=False),
|
||||
sa.Column('status', sa.String(), nullable=False),
|
||||
sa.Column('started_at', sa.DateTime(), nullable=True),
|
||||
sa.Column('completed_at', sa.DateTime(), nullable=True),
|
||||
sa.Column('result', sa.JSON(), nullable=True),
|
||||
sa.Column('error_message', sa.Text(), nullable=True),
|
||||
sa.Column('triggered_by', sa.Integer(), nullable=True),
|
||||
sa.ForeignKeyConstraint(['playbook_id'], ['playbooks.id'], ),
|
||||
sa.ForeignKeyConstraint(['tenant_id'], ['tenants.id'], ),
|
||||
sa.ForeignKeyConstraint(['triggered_by'], ['users.id'], ),
|
||||
sa.PrimaryKeyConstraint('id')
|
||||
)
|
||||
op.create_index(op.f('ix_playbook_executions_id'), 'playbook_executions', ['id'], unique=False)
|
||||
|
||||
# Create threat_scores table
|
||||
op.create_table(
|
||||
'threat_scores',
|
||||
sa.Column('id', sa.Integer(), nullable=False),
|
||||
sa.Column('tenant_id', sa.Integer(), nullable=False),
|
||||
sa.Column('host_id', sa.Integer(), nullable=True),
|
||||
sa.Column('artifact_id', sa.Integer(), nullable=True),
|
||||
sa.Column('score', sa.Float(), nullable=False),
|
||||
sa.Column('confidence', sa.Float(), nullable=False),
|
||||
sa.Column('threat_type', sa.String(), nullable=False),
|
||||
sa.Column('description', sa.Text(), nullable=True),
|
||||
sa.Column('indicators', sa.JSON(), nullable=True),
|
||||
sa.Column('ml_model_version', sa.String(), nullable=True),
|
||||
sa.Column('created_at', sa.DateTime(), nullable=True),
|
||||
sa.ForeignKeyConstraint(['tenant_id'], ['tenants.id'], ),
|
||||
sa.ForeignKeyConstraint(['host_id'], ['hosts.id'], ),
|
||||
sa.ForeignKeyConstraint(['artifact_id'], ['artifacts.id'], ),
|
||||
sa.PrimaryKeyConstraint('id')
|
||||
)
|
||||
op.create_index(op.f('ix_threat_scores_id'), 'threat_scores', ['id'], unique=False)
|
||||
op.create_index(op.f('ix_threat_scores_score'), 'threat_scores', ['score'], unique=False)
|
||||
op.create_index(op.f('ix_threat_scores_created_at'), 'threat_scores', ['created_at'], unique=False)
|
||||
|
||||
# Create report_templates table
|
||||
op.create_table(
|
||||
'report_templates',
|
||||
sa.Column('id', sa.Integer(), nullable=False),
|
||||
sa.Column('tenant_id', sa.Integer(), nullable=False),
|
||||
sa.Column('name', sa.String(), nullable=False),
|
||||
sa.Column('description', sa.Text(), nullable=True),
|
||||
sa.Column('template_type', sa.String(), nullable=False),
|
||||
sa.Column('template_config', sa.JSON(), nullable=False),
|
||||
sa.Column('is_default', sa.Boolean(), nullable=False),
|
||||
sa.Column('created_by', sa.Integer(), nullable=False),
|
||||
sa.Column('created_at', sa.DateTime(), nullable=True),
|
||||
sa.ForeignKeyConstraint(['tenant_id'], ['tenants.id'], ),
|
||||
sa.ForeignKeyConstraint(['created_by'], ['users.id'], ),
|
||||
sa.PrimaryKeyConstraint('id')
|
||||
)
|
||||
op.create_index(op.f('ix_report_templates_id'), 'report_templates', ['id'], unique=False)
|
||||
op.create_index(op.f('ix_report_templates_name'), 'report_templates', ['name'], unique=False)
|
||||
|
||||
# Create reports table
|
||||
op.create_table(
|
||||
'reports',
|
||||
sa.Column('id', sa.Integer(), nullable=False),
|
||||
sa.Column('tenant_id', sa.Integer(), nullable=False),
|
||||
sa.Column('template_id', sa.Integer(), nullable=True),
|
||||
sa.Column('title', sa.String(), nullable=False),
|
||||
sa.Column('report_type', sa.String(), nullable=False),
|
||||
sa.Column('format', sa.String(), nullable=False),
|
||||
sa.Column('file_path', sa.String(), nullable=True),
|
||||
sa.Column('status', sa.String(), nullable=False),
|
||||
sa.Column('generated_by', sa.Integer(), nullable=False),
|
||||
sa.Column('generated_at', sa.DateTime(), nullable=True),
|
||||
sa.ForeignKeyConstraint(['tenant_id'], ['tenants.id'], ),
|
||||
sa.ForeignKeyConstraint(['template_id'], ['report_templates.id'], ),
|
||||
sa.ForeignKeyConstraint(['generated_by'], ['users.id'], ),
|
||||
sa.PrimaryKeyConstraint('id')
|
||||
)
|
||||
op.create_index(op.f('ix_reports_id'), 'reports', ['id'], unique=False)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Downgrade schema for Phase 4."""
|
||||
|
||||
# Drop reports table
|
||||
op.drop_index(op.f('ix_reports_id'), table_name='reports')
|
||||
op.drop_table('reports')
|
||||
|
||||
# Drop report_templates table
|
||||
op.drop_index(op.f('ix_report_templates_name'), table_name='report_templates')
|
||||
op.drop_index(op.f('ix_report_templates_id'), table_name='report_templates')
|
||||
op.drop_table('report_templates')
|
||||
|
||||
# Drop threat_scores table
|
||||
op.drop_index(op.f('ix_threat_scores_created_at'), table_name='threat_scores')
|
||||
op.drop_index(op.f('ix_threat_scores_score'), table_name='threat_scores')
|
||||
op.drop_index(op.f('ix_threat_scores_id'), table_name='threat_scores')
|
||||
op.drop_table('threat_scores')
|
||||
|
||||
# Drop playbook_executions table
|
||||
op.drop_index(op.f('ix_playbook_executions_id'), table_name='playbook_executions')
|
||||
op.drop_table('playbook_executions')
|
||||
|
||||
# Drop playbooks table
|
||||
op.drop_index(op.f('ix_playbooks_name'), table_name='playbooks')
|
||||
op.drop_index(op.f('ix_playbooks_id'), table_name='playbooks')
|
||||
op.drop_table('playbooks')
|
||||
@@ -1,114 +0,0 @@
|
||||
"""Initial migration
|
||||
|
||||
Revision ID: f82b3092d056
|
||||
Revises:
|
||||
Create Date: 2025-12-09 14:25:47.222289
|
||||
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = 'f82b3092d056'
|
||||
down_revision: Union[str, Sequence[str], None] = None
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Upgrade schema."""
|
||||
# Create tenants table
|
||||
op.create_table(
|
||||
'tenants',
|
||||
sa.Column('id', sa.Integer(), nullable=False),
|
||||
sa.Column('name', sa.String(), nullable=False),
|
||||
sa.Column('description', sa.String(), nullable=True),
|
||||
sa.Column('created_at', sa.DateTime(), nullable=True),
|
||||
sa.PrimaryKeyConstraint('id')
|
||||
)
|
||||
op.create_index(op.f('ix_tenants_id'), 'tenants', ['id'], unique=False)
|
||||
op.create_index(op.f('ix_tenants_name'), 'tenants', ['name'], unique=True)
|
||||
|
||||
# Create users table
|
||||
op.create_table(
|
||||
'users',
|
||||
sa.Column('id', sa.Integer(), nullable=False),
|
||||
sa.Column('username', sa.String(), nullable=False),
|
||||
sa.Column('password_hash', sa.String(), nullable=False),
|
||||
sa.Column('role', sa.String(), nullable=False),
|
||||
sa.Column('tenant_id', sa.Integer(), nullable=False),
|
||||
sa.Column('is_active', sa.Boolean(), nullable=False),
|
||||
sa.Column('created_at', sa.DateTime(), nullable=True),
|
||||
sa.ForeignKeyConstraint(['tenant_id'], ['tenants.id'], ),
|
||||
sa.PrimaryKeyConstraint('id')
|
||||
)
|
||||
op.create_index(op.f('ix_users_id'), 'users', ['id'], unique=False)
|
||||
op.create_index(op.f('ix_users_username'), 'users', ['username'], unique=True)
|
||||
|
||||
# Create hosts table
|
||||
op.create_table(
|
||||
'hosts',
|
||||
sa.Column('id', sa.Integer(), nullable=False),
|
||||
sa.Column('hostname', sa.String(), nullable=False),
|
||||
sa.Column('ip_address', sa.String(), nullable=True),
|
||||
sa.Column('os', sa.String(), nullable=True),
|
||||
sa.Column('tenant_id', sa.Integer(), nullable=False),
|
||||
sa.Column('host_metadata', sa.JSON(), nullable=True),
|
||||
sa.Column('created_at', sa.DateTime(), nullable=True),
|
||||
sa.Column('last_seen', sa.DateTime(), nullable=True),
|
||||
sa.ForeignKeyConstraint(['tenant_id'], ['tenants.id'], ),
|
||||
sa.PrimaryKeyConstraint('id')
|
||||
)
|
||||
op.create_index(op.f('ix_hosts_id'), 'hosts', ['id'], unique=False)
|
||||
op.create_index(op.f('ix_hosts_hostname'), 'hosts', ['hostname'], unique=False)
|
||||
|
||||
# Create cases table
|
||||
op.create_table(
|
||||
'cases',
|
||||
sa.Column('id', sa.Integer(), nullable=False),
|
||||
sa.Column('title', sa.String(), nullable=False),
|
||||
sa.Column('description', sa.Text(), nullable=True),
|
||||
sa.Column('status', sa.String(), nullable=False),
|
||||
sa.Column('severity', sa.String(), nullable=True),
|
||||
sa.Column('tenant_id', sa.Integer(), nullable=False),
|
||||
sa.Column('created_at', sa.DateTime(), nullable=True),
|
||||
sa.Column('updated_at', sa.DateTime(), nullable=True),
|
||||
sa.ForeignKeyConstraint(['tenant_id'], ['tenants.id'], ),
|
||||
sa.PrimaryKeyConstraint('id')
|
||||
)
|
||||
op.create_index(op.f('ix_cases_id'), 'cases', ['id'], unique=False)
|
||||
|
||||
# Create artifacts table
|
||||
op.create_table(
|
||||
'artifacts',
|
||||
sa.Column('id', sa.Integer(), nullable=False),
|
||||
sa.Column('artifact_type', sa.String(), nullable=False),
|
||||
sa.Column('value', sa.String(), nullable=False),
|
||||
sa.Column('description', sa.Text(), nullable=True),
|
||||
sa.Column('case_id', sa.Integer(), nullable=True),
|
||||
sa.Column('artifact_metadata', sa.JSON(), nullable=True),
|
||||
sa.Column('created_at', sa.DateTime(), nullable=True),
|
||||
sa.ForeignKeyConstraint(['case_id'], ['cases.id'], ),
|
||||
sa.PrimaryKeyConstraint('id')
|
||||
)
|
||||
op.create_index(op.f('ix_artifacts_id'), 'artifacts', ['id'], unique=False)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Downgrade schema."""
|
||||
op.drop_index(op.f('ix_artifacts_id'), table_name='artifacts')
|
||||
op.drop_table('artifacts')
|
||||
op.drop_index(op.f('ix_cases_id'), table_name='cases')
|
||||
op.drop_table('cases')
|
||||
op.drop_index(op.f('ix_hosts_hostname'), table_name='hosts')
|
||||
op.drop_index(op.f('ix_hosts_id'), table_name='hosts')
|
||||
op.drop_table('hosts')
|
||||
op.drop_index(op.f('ix_users_username'), table_name='users')
|
||||
op.drop_index(op.f('ix_users_id'), table_name='users')
|
||||
op.drop_table('users')
|
||||
op.drop_index(op.f('ix_tenants_name'), table_name='tenants')
|
||||
op.drop_index(op.f('ix_tenants_id'), table_name='tenants')
|
||||
op.drop_table('tenants')
|
||||
323
backend/app.py
Normal file
323
backend/app.py
Normal file
@@ -0,0 +1,323 @@
|
||||
import os
|
||||
import logging
|
||||
from datetime import datetime, timedelta
|
||||
from flask import Flask, request, jsonify, send_from_directory
|
||||
from flask_sqlalchemy import SQLAlchemy
|
||||
from flask_jwt_extended import JWTManager, jwt_required, create_access_token, get_jwt_identity
|
||||
from werkzeug.utils import secure_filename
|
||||
import bcrypt
|
||||
|
||||
# Try to import flask-cors
|
||||
try:
|
||||
from flask_cors import CORS
|
||||
CORS_AVAILABLE = True
|
||||
except ImportError:
|
||||
CORS_AVAILABLE = False
|
||||
print("Warning: flask-cors not available")
|
||||
|
||||
# Configure logging
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
app = Flask(__name__, static_folder="../frontend/dist")
|
||||
|
||||
# Configuration
|
||||
app.config['SQLALCHEMY_DATABASE_URI'] = os.getenv('DATABASE_URL', 'postgresql://admin:secure_password_123@localhost:5432/threat_hunter')
|
||||
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
|
||||
app.config['JWT_SECRET_KEY'] = os.getenv('SECRET_KEY', 'change-this-in-production')
|
||||
app.config['JWT_ACCESS_TOKEN_EXPIRES'] = timedelta(hours=24)
|
||||
app.config['MAX_CONTENT_LENGTH'] = 100 * 1024 * 1024
|
||||
app.config['UPLOAD_FOLDER'] = 'uploaded'
|
||||
app.config['OUTPUT_FOLDER'] = 'output'
|
||||
app.config['ALLOWED_EXTENSIONS'] = {'csv', 'json', 'txt', 'log'}
|
||||
|
||||
# Ensure upload directories exist
|
||||
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
|
||||
os.makedirs(app.config['OUTPUT_FOLDER'], exist_ok=True)
|
||||
|
||||
# Initialize extensions
|
||||
db = SQLAlchemy(app)
|
||||
jwt = JWTManager(app)
|
||||
|
||||
# Enable CORS
|
||||
if CORS_AVAILABLE:
|
||||
CORS(app)
|
||||
else:
|
||||
@app.after_request
|
||||
def after_request(response):
|
||||
response.headers.add('Access-Control-Allow-Origin', '*')
|
||||
response.headers.add('Access-Control-Allow-Headers', 'Content-Type,Authorization')
|
||||
response.headers.add('Access-Control-Allow-Methods', 'GET,PUT,POST,DELETE')
|
||||
return response
|
||||
|
||||
def allowed_file(filename):
|
||||
return '.' in filename and \
|
||||
filename.rsplit('.', 1)[1].lower() in app.config['ALLOWED_EXTENSIONS']
|
||||
|
||||
@app.errorhandler(RequestEntityTooLarge)
|
||||
def handle_file_too_large(e):
|
||||
return jsonify({'error': 'File too large. Maximum size is 100MB.'}), 413
|
||||
|
||||
@app.errorhandler(Exception)
|
||||
def handle_exception(e):
|
||||
logger.error(f"Unhandled exception: {e}")
|
||||
return jsonify({'error': 'Internal server error'}), 500
|
||||
|
||||
# Database Models
|
||||
class User(db.Model):
|
||||
__tablename__ = 'users'
|
||||
id = db.Column(db.String, primary_key=True, default=lambda: str(uuid.uuid4()))
|
||||
username = db.Column(db.String(50), unique=True, nullable=False)
|
||||
email = db.Column(db.String(100), unique=True, nullable=False)
|
||||
password_hash = db.Column(db.String(255), nullable=False)
|
||||
created_at = db.Column(db.DateTime, default=datetime.utcnow)
|
||||
last_login = db.Column(db.DateTime)
|
||||
|
||||
class Hunt(db.Model):
|
||||
__tablename__ = 'hunts'
|
||||
id = db.Column(db.String, primary_key=True, default=lambda: str(uuid.uuid4()))
|
||||
name = db.Column(db.String(100), nullable=False)
|
||||
description = db.Column(db.Text)
|
||||
created_by = db.Column(db.String, db.ForeignKey('users.id'))
|
||||
created_at = db.Column(db.DateTime, default=datetime.utcnow)
|
||||
updated_at = db.Column(db.DateTime, default=datetime.utcnow)
|
||||
status = db.Column(db.String(20), default='active')
|
||||
|
||||
# Authentication Routes
|
||||
@app.route('/api/auth/login', methods=['POST'])
|
||||
def login():
|
||||
try:
|
||||
data = request.get_json()
|
||||
username = data.get('username')
|
||||
password = data.get('password')
|
||||
|
||||
user = User.query.filter_by(username=username).first()
|
||||
|
||||
if user and bcrypt.checkpw(password.encode('utf-8'), user.password_hash.encode('utf-8')):
|
||||
access_token = create_access_token(identity=user.id)
|
||||
user.last_login = datetime.utcnow()
|
||||
db.session.commit()
|
||||
|
||||
return jsonify({
|
||||
'access_token': access_token,
|
||||
'user': {
|
||||
'id': user.id,
|
||||
'username': user.username,
|
||||
'email': user.email
|
||||
}
|
||||
})
|
||||
else:
|
||||
return jsonify({'error': 'Invalid credentials'}), 401
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Login error: {e}")
|
||||
return jsonify({'error': 'Login failed'}), 500
|
||||
|
||||
@app.route('/api/auth/register', methods=['POST'])
|
||||
def register():
|
||||
try:
|
||||
data = request.get_json()
|
||||
username = data.get('username')
|
||||
email = data.get('email')
|
||||
password = data.get('password')
|
||||
|
||||
# Check if user exists
|
||||
if User.query.filter_by(username=username).first():
|
||||
return jsonify({'error': 'Username already exists'}), 400
|
||||
|
||||
if User.query.filter_by(email=email).first():
|
||||
return jsonify({'error': 'Email already exists'}), 400
|
||||
|
||||
# Hash password
|
||||
password_hash = bcrypt.hashpw(password.encode('utf-8'), bcrypt.gensalt()).decode('utf-8')
|
||||
|
||||
# Create user
|
||||
user = User(username=username, email=email, password_hash=password_hash)
|
||||
db.session.add(user)
|
||||
db.session.commit()
|
||||
|
||||
access_token = create_access_token(identity=user.id)
|
||||
|
||||
return jsonify({
|
||||
'access_token': access_token,
|
||||
'user': {
|
||||
'id': user.id,
|
||||
'username': user.username,
|
||||
'email': user.email
|
||||
}
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Registration error: {e}")
|
||||
return jsonify({'error': 'Registration failed'}), 500
|
||||
|
||||
# Hunt Management Routes
|
||||
@app.route('/api/hunts', methods=['GET'])
|
||||
@jwt_required()
|
||||
def get_hunts():
|
||||
try:
|
||||
user_id = get_jwt_identity()
|
||||
hunts = Hunt.query.filter_by(created_by=user_id).all()
|
||||
|
||||
return jsonify({
|
||||
'hunts': [{
|
||||
'id': hunt.id,
|
||||
'name': hunt.name,
|
||||
'description': hunt.description,
|
||||
'created_at': hunt.created_at.isoformat(),
|
||||
'status': hunt.status
|
||||
} for hunt in hunts]
|
||||
})
|
||||
except Exception as e:
|
||||
logger.error(f"Get hunts error: {e}")
|
||||
return jsonify({'error': 'Failed to fetch hunts'}), 500
|
||||
|
||||
@app.route('/api/hunts', methods=['POST'])
|
||||
@jwt_required()
|
||||
def create_hunt():
|
||||
try:
|
||||
user_id = get_jwt_identity()
|
||||
data = request.get_json()
|
||||
|
||||
hunt = Hunt(
|
||||
name=data.get('name'),
|
||||
description=data.get('description'),
|
||||
created_by=user_id
|
||||
)
|
||||
|
||||
db.session.add(hunt)
|
||||
db.session.commit()
|
||||
|
||||
return jsonify({
|
||||
'hunt': {
|
||||
'id': hunt.id,
|
||||
'name': hunt.name,
|
||||
'description': hunt.description,
|
||||
'created_at': hunt.created_at.isoformat(),
|
||||
'status': hunt.status
|
||||
}
|
||||
})
|
||||
except Exception as e:
|
||||
logger.error(f"Create hunt error: {e}")
|
||||
return jsonify({'error': 'Failed to create hunt'}), 500
|
||||
|
||||
# API Routes
|
||||
@app.route('/api/health')
|
||||
def health_check():
|
||||
return jsonify({
|
||||
'status': 'healthy',
|
||||
'timestamp': datetime.utcnow().isoformat(),
|
||||
'version': '1.0.0',
|
||||
'service': 'Cyber Threat Hunter API'
|
||||
})
|
||||
|
||||
@app.route('/api/upload', methods=['POST'])
|
||||
def upload_file():
|
||||
try:
|
||||
if 'file' not in request.files:
|
||||
return jsonify({'error': 'No file provided'}), 400
|
||||
|
||||
file = request.files['file']
|
||||
if file.filename == '':
|
||||
return jsonify({'error': 'No file selected'}), 400
|
||||
|
||||
if not allowed_file(file.filename):
|
||||
return jsonify({'error': 'File type not allowed'}), 400
|
||||
|
||||
filename = secure_filename(file.filename)
|
||||
timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
|
||||
filename = f"{timestamp}_{filename}"
|
||||
filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
|
||||
|
||||
file.save(filepath)
|
||||
|
||||
logger.info(f"File uploaded successfully: {filename}")
|
||||
return jsonify({
|
||||
'message': 'File uploaded successfully',
|
||||
'filename': filename,
|
||||
'size': os.path.getsize(filepath)
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Upload error: {e}")
|
||||
return jsonify({'error': 'Upload failed'}), 500
|
||||
|
||||
@app.route('/api/files')
|
||||
def list_files():
|
||||
try:
|
||||
files = []
|
||||
upload_dir = app.config['UPLOAD_FOLDER']
|
||||
|
||||
if os.path.exists(upload_dir):
|
||||
for filename in os.listdir(upload_dir):
|
||||
filepath = os.path.join(upload_dir, filename)
|
||||
if os.path.isfile(filepath):
|
||||
stat = os.stat(filepath)
|
||||
files.append({
|
||||
'name': filename,
|
||||
'size': stat.st_size,
|
||||
'modified': datetime.fromtimestamp(stat.st_mtime).isoformat()
|
||||
})
|
||||
|
||||
return jsonify({'files': files})
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"List files error: {e}")
|
||||
return jsonify({'error': 'Failed to list files'}), 500
|
||||
|
||||
@app.route('/api/stats')
|
||||
def get_stats():
|
||||
try:
|
||||
upload_dir = app.config['UPLOAD_FOLDER']
|
||||
files_count = 0
|
||||
if os.path.exists(upload_dir):
|
||||
files_count = len([f for f in os.listdir(upload_dir) if os.path.isfile(os.path.join(upload_dir, f))])
|
||||
|
||||
return jsonify({
|
||||
'filesUploaded': files_count,
|
||||
'analysesCompleted': files_count,
|
||||
'threatsDetected': 0
|
||||
})
|
||||
except Exception as e:
|
||||
logger.error(f"Stats error: {e}")
|
||||
return jsonify({'error': 'Failed to get stats'}), 500
|
||||
|
||||
# Static file serving for React app
|
||||
@app.route("/assets/<path:path>")
|
||||
def send_assets(path):
|
||||
return send_from_directory(os.path.join(app.static_folder, "assets"), path)
|
||||
|
||||
@app.route("/")
|
||||
def index():
|
||||
if os.path.exists(os.path.join(app.static_folder, "index.html")):
|
||||
return send_from_directory(app.static_folder, "index.html")
|
||||
else:
|
||||
return jsonify({
|
||||
'message': 'Cyber Threat Hunter API',
|
||||
'status': 'running',
|
||||
'endpoints': [
|
||||
'GET /api/health',
|
||||
'POST /api/upload',
|
||||
'GET /api/files',
|
||||
'GET /api/stats'
|
||||
]
|
||||
})
|
||||
|
||||
# Catch-all route for React Router
|
||||
@app.route("/<path:path>")
|
||||
def catch_all(path):
|
||||
if os.path.exists(os.path.join(app.static_folder, "index.html")):
|
||||
return send_from_directory(app.static_folder, "index.html")
|
||||
else:
|
||||
return jsonify({'error': 'Frontend not built yet'})
|
||||
|
||||
if __name__ == "__main__":
|
||||
with app.app_context():
|
||||
db.create_all()
|
||||
|
||||
print("=" * 50)
|
||||
print("Starting Cyber Threat Hunter Backend...")
|
||||
print("API available at: http://localhost:5000")
|
||||
print("Database: Connected to PostgreSQL")
|
||||
print("=" * 50)
|
||||
app.run(host="0.0.0.0", port=5000, debug=True)
|
||||
@@ -1,75 +0,0 @@
|
||||
from typing import List, Optional
|
||||
from fastapi import APIRouter, Depends, Query
|
||||
from sqlalchemy.orm import Session
|
||||
from datetime import datetime
|
||||
|
||||
from app.core.database import get_db
|
||||
from app.core.deps import get_current_active_user, require_role
|
||||
from app.models.user import User
|
||||
from app.models.audit_log import AuditLog
|
||||
from app.schemas.audit import AuditLogRead
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.get("/", response_model=List[AuditLogRead])
|
||||
async def list_audit_logs(
|
||||
skip: int = 0,
|
||||
limit: int = 100,
|
||||
action: Optional[str] = Query(None, description="Filter by action type"),
|
||||
resource_type: Optional[str] = Query(None, description="Filter by resource type"),
|
||||
start_date: Optional[datetime] = Query(None, description="Filter from date"),
|
||||
end_date: Optional[datetime] = Query(None, description="Filter to date"),
|
||||
current_user: User = Depends(require_role(["admin"])),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
List audit logs (admin only, scoped to tenant)
|
||||
|
||||
Provides a complete audit trail of actions within the tenant.
|
||||
"""
|
||||
# Base query scoped to tenant
|
||||
query = db.query(AuditLog).filter(AuditLog.tenant_id == current_user.tenant_id)
|
||||
|
||||
# Apply filters
|
||||
if action:
|
||||
query = query.filter(AuditLog.action == action)
|
||||
if resource_type:
|
||||
query = query.filter(AuditLog.resource_type == resource_type)
|
||||
if start_date:
|
||||
query = query.filter(AuditLog.created_at >= start_date)
|
||||
if end_date:
|
||||
query = query.filter(AuditLog.created_at <= end_date)
|
||||
|
||||
# Order by most recent first
|
||||
query = query.order_by(AuditLog.created_at.desc())
|
||||
|
||||
# Paginate
|
||||
logs = query.offset(skip).limit(limit).all()
|
||||
|
||||
return logs
|
||||
|
||||
|
||||
@router.get("/{log_id}", response_model=AuditLogRead)
|
||||
async def get_audit_log(
|
||||
log_id: int,
|
||||
current_user: User = Depends(require_role(["admin"])),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Get a specific audit log entry (admin only)
|
||||
"""
|
||||
from fastapi import HTTPException, status
|
||||
|
||||
log = db.query(AuditLog).filter(
|
||||
AuditLog.id == log_id,
|
||||
AuditLog.tenant_id == current_user.tenant_id
|
||||
).first()
|
||||
|
||||
if not log:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail="Audit log not found"
|
||||
)
|
||||
|
||||
return log
|
||||
@@ -1,432 +0,0 @@
|
||||
from fastapi import APIRouter, Depends, HTTPException, status
|
||||
from fastapi.security import OAuth2PasswordRequestForm
|
||||
from sqlalchemy.orm import Session
|
||||
from datetime import datetime, timezone, timedelta
|
||||
import io
|
||||
import qrcode
|
||||
|
||||
from app.core.database import get_db
|
||||
from app.core.security import (
|
||||
verify_password, get_password_hash, create_access_token,
|
||||
create_refresh_token, create_reset_token, generate_totp_secret,
|
||||
verify_totp, get_totp_uri
|
||||
)
|
||||
from app.core.deps import get_current_active_user
|
||||
from app.models.user import User
|
||||
from app.models.tenant import Tenant
|
||||
from app.models.refresh_token import RefreshToken
|
||||
from app.models.password_reset_token import PasswordResetToken
|
||||
from app.schemas.auth import (
|
||||
Token, UserLogin, UserRegister, RefreshTokenRequest,
|
||||
PasswordResetRequest, PasswordResetConfirm,
|
||||
TwoFactorSetup, TwoFactorVerify
|
||||
)
|
||||
from app.schemas.user import UserRead, UserUpdate
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.post("/register", response_model=UserRead, status_code=status.HTTP_201_CREATED)
|
||||
async def register(
|
||||
user_data: UserRegister,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Register a new user
|
||||
|
||||
Creates a new user with hashed password. If tenant_id is not provided,
|
||||
a default tenant is created or used.
|
||||
"""
|
||||
# Check if username already exists
|
||||
existing_user = db.query(User).filter(User.username == user_data.username).first()
|
||||
if existing_user:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Username already registered"
|
||||
)
|
||||
|
||||
# Handle tenant_id
|
||||
tenant_id = user_data.tenant_id
|
||||
if tenant_id is None:
|
||||
# Create or get default tenant
|
||||
default_tenant = db.query(Tenant).filter(Tenant.name == "default").first()
|
||||
if not default_tenant:
|
||||
default_tenant = Tenant(name="default", description="Default tenant")
|
||||
db.add(default_tenant)
|
||||
db.commit()
|
||||
db.refresh(default_tenant)
|
||||
tenant_id = default_tenant.id
|
||||
else:
|
||||
# Verify tenant exists
|
||||
tenant = db.query(Tenant).filter(Tenant.id == tenant_id).first()
|
||||
if not tenant:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail="Tenant not found"
|
||||
)
|
||||
|
||||
# Create new user with hashed password
|
||||
hashed_password = get_password_hash(user_data.password)
|
||||
new_user = User(
|
||||
username=user_data.username,
|
||||
password_hash=hashed_password,
|
||||
role=user_data.role,
|
||||
tenant_id=tenant_id
|
||||
)
|
||||
|
||||
db.add(new_user)
|
||||
db.commit()
|
||||
db.refresh(new_user)
|
||||
|
||||
return new_user
|
||||
|
||||
|
||||
@router.post("/login", response_model=Token)
|
||||
async def login(
|
||||
form_data: OAuth2PasswordRequestForm = Depends(),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Authenticate user and return JWT token
|
||||
|
||||
Uses OAuth2 password flow for compatibility with OpenAPI docs.
|
||||
"""
|
||||
# Find user by username
|
||||
user = db.query(User).filter(User.username == form_data.username).first()
|
||||
if not user:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="Incorrect username or password",
|
||||
headers={"WWW-Authenticate": "Bearer"},
|
||||
)
|
||||
|
||||
# Verify password
|
||||
if not verify_password(form_data.password, user.password_hash):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="Incorrect username or password",
|
||||
headers={"WWW-Authenticate": "Bearer"},
|
||||
)
|
||||
|
||||
# Check if user is active
|
||||
if not user.is_active:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Inactive user"
|
||||
)
|
||||
|
||||
# Check 2FA if enabled (TOTP code should be in scopes for OAuth2)
|
||||
if user.totp_enabled:
|
||||
# For OAuth2 password flow, we'll check totp in scopes
|
||||
totp_code = form_data.scopes[0] if form_data.scopes else None
|
||||
if not totp_code or not verify_totp(user.totp_secret, totp_code):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="Invalid 2FA code",
|
||||
headers={"WWW-Authenticate": "Bearer"},
|
||||
)
|
||||
|
||||
# Create access token
|
||||
access_token = create_access_token(
|
||||
data={
|
||||
"sub": user.id,
|
||||
"tenant_id": user.tenant_id,
|
||||
"role": user.role
|
||||
}
|
||||
)
|
||||
|
||||
# Create refresh token
|
||||
refresh_token_str = create_refresh_token()
|
||||
refresh_token_obj = RefreshToken(
|
||||
token=refresh_token_str,
|
||||
user_id=user.id,
|
||||
expires_at=datetime.now(timezone.utc) + timedelta(days=30)
|
||||
)
|
||||
db.add(refresh_token_obj)
|
||||
db.commit()
|
||||
|
||||
return {
|
||||
"access_token": access_token,
|
||||
"refresh_token": refresh_token_str,
|
||||
"token_type": "bearer"
|
||||
}
|
||||
|
||||
|
||||
@router.get("/me", response_model=UserRead)
|
||||
async def get_current_user_profile(
|
||||
current_user: User = Depends(get_current_active_user)
|
||||
):
|
||||
"""
|
||||
Get current user profile
|
||||
|
||||
Returns the profile of the authenticated user.
|
||||
"""
|
||||
return current_user
|
||||
|
||||
|
||||
@router.put("/me", response_model=UserRead)
|
||||
async def update_current_user_profile(
|
||||
user_update: UserUpdate,
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Update current user profile
|
||||
|
||||
Allows users to update their own profile information.
|
||||
"""
|
||||
# Update username if provided
|
||||
if user_update.username is not None:
|
||||
# Check if new username is already taken
|
||||
existing_user = db.query(User).filter(
|
||||
User.username == user_update.username,
|
||||
User.id != current_user.id
|
||||
).first()
|
||||
if existing_user:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Username already taken"
|
||||
)
|
||||
current_user.username = user_update.username
|
||||
|
||||
# Update password if provided
|
||||
if user_update.password is not None:
|
||||
current_user.password_hash = get_password_hash(user_update.password)
|
||||
|
||||
# Users cannot change their own role through this endpoint
|
||||
# (admin users should use the admin endpoints in /users)
|
||||
|
||||
db.commit()
|
||||
db.refresh(current_user)
|
||||
|
||||
return current_user
|
||||
|
||||
|
||||
@router.post("/refresh", response_model=Token)
|
||||
async def refresh_access_token(
|
||||
refresh_request: RefreshTokenRequest,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Refresh access token using refresh token
|
||||
|
||||
Provides a new access token without requiring login.
|
||||
"""
|
||||
# Find refresh token
|
||||
refresh_token = db.query(RefreshToken).filter(
|
||||
RefreshToken.token == refresh_request.refresh_token,
|
||||
RefreshToken.is_revoked == False
|
||||
).first()
|
||||
|
||||
if not refresh_token:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="Invalid refresh token"
|
||||
)
|
||||
|
||||
# Check if expired
|
||||
if refresh_token.expires_at < datetime.now(timezone.utc):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="Refresh token expired"
|
||||
)
|
||||
|
||||
# Get user
|
||||
user = db.query(User).filter(User.id == refresh_token.user_id).first()
|
||||
if not user or not user.is_active:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="User not found or inactive"
|
||||
)
|
||||
|
||||
# Create new access token
|
||||
access_token = create_access_token(
|
||||
data={
|
||||
"sub": user.id,
|
||||
"tenant_id": user.tenant_id,
|
||||
"role": user.role
|
||||
}
|
||||
)
|
||||
|
||||
return {
|
||||
"access_token": access_token,
|
||||
"refresh_token": refresh_request.refresh_token,
|
||||
"token_type": "bearer"
|
||||
}
|
||||
|
||||
|
||||
@router.post("/2fa/setup", response_model=TwoFactorSetup)
|
||||
async def setup_two_factor(
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Set up two-factor authentication
|
||||
|
||||
Generates a TOTP secret and QR code URI for the user.
|
||||
"""
|
||||
if current_user.totp_enabled:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="2FA is already enabled"
|
||||
)
|
||||
|
||||
# Generate secret
|
||||
secret = generate_totp_secret()
|
||||
current_user.totp_secret = secret
|
||||
db.commit()
|
||||
|
||||
# Get QR code URI
|
||||
qr_uri = get_totp_uri(secret, current_user.username)
|
||||
|
||||
return {
|
||||
"secret": secret,
|
||||
"qr_code_uri": qr_uri
|
||||
}
|
||||
|
||||
|
||||
@router.post("/2fa/verify")
|
||||
async def verify_two_factor(
|
||||
verify_request: TwoFactorVerify,
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Verify and enable two-factor authentication
|
||||
|
||||
User must provide a valid TOTP code to enable 2FA.
|
||||
"""
|
||||
if current_user.totp_enabled:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="2FA is already enabled"
|
||||
)
|
||||
|
||||
if not current_user.totp_secret:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="2FA setup not initiated. Call /2fa/setup first."
|
||||
)
|
||||
|
||||
# Verify code
|
||||
if not verify_totp(current_user.totp_secret, verify_request.code):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Invalid 2FA code"
|
||||
)
|
||||
|
||||
# Enable 2FA
|
||||
current_user.totp_enabled = True
|
||||
db.commit()
|
||||
|
||||
return {"message": "2FA enabled successfully"}
|
||||
|
||||
|
||||
@router.post("/2fa/disable")
|
||||
async def disable_two_factor(
|
||||
verify_request: TwoFactorVerify,
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Disable two-factor authentication
|
||||
|
||||
User must provide a valid TOTP code to disable 2FA.
|
||||
"""
|
||||
if not current_user.totp_enabled:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="2FA is not enabled"
|
||||
)
|
||||
|
||||
# Verify code
|
||||
if not verify_totp(current_user.totp_secret, verify_request.code):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Invalid 2FA code"
|
||||
)
|
||||
|
||||
# Disable 2FA
|
||||
current_user.totp_enabled = False
|
||||
current_user.totp_secret = None
|
||||
db.commit()
|
||||
|
||||
return {"message": "2FA disabled successfully"}
|
||||
|
||||
|
||||
@router.post("/password-reset/request")
|
||||
async def request_password_reset(
|
||||
reset_request: PasswordResetRequest,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Request a password reset
|
||||
|
||||
Sends a password reset email to the user (mock implementation).
|
||||
"""
|
||||
# Find user by email
|
||||
user = db.query(User).filter(User.email == reset_request.email).first()
|
||||
|
||||
# Don't reveal if email exists or not (security best practice)
|
||||
# Always return success even if email doesn't exist
|
||||
if user:
|
||||
# Create reset token
|
||||
reset_token = create_reset_token()
|
||||
reset_token_obj = PasswordResetToken(
|
||||
token=reset_token,
|
||||
user_id=user.id,
|
||||
expires_at=datetime.now(timezone.utc) + timedelta(hours=1)
|
||||
)
|
||||
db.add(reset_token_obj)
|
||||
db.commit()
|
||||
|
||||
# TODO: Send email with reset link
|
||||
# For now, we'll just log it (in production, use an email service)
|
||||
print(f"Password reset token for {user.email}: {reset_token}")
|
||||
|
||||
return {"message": "If the email exists, a password reset link has been sent"}
|
||||
|
||||
|
||||
@router.post("/password-reset/confirm")
|
||||
async def confirm_password_reset(
|
||||
reset_confirm: PasswordResetConfirm,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Confirm password reset with token
|
||||
|
||||
Sets a new password for the user.
|
||||
"""
|
||||
# Find reset token
|
||||
reset_token = db.query(PasswordResetToken).filter(
|
||||
PasswordResetToken.token == reset_confirm.token,
|
||||
PasswordResetToken.is_used == False
|
||||
).first()
|
||||
|
||||
if not reset_token:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Invalid or expired reset token"
|
||||
)
|
||||
|
||||
# Check if expired
|
||||
if reset_token.expires_at < datetime.now(timezone.utc):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Reset token expired"
|
||||
)
|
||||
|
||||
# Get user
|
||||
user = db.query(User).filter(User.id == reset_token.user_id).first()
|
||||
if not user:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail="User not found"
|
||||
)
|
||||
|
||||
# Update password
|
||||
user.password_hash = get_password_hash(reset_confirm.new_password)
|
||||
reset_token.is_used = True
|
||||
db.commit()
|
||||
|
||||
return {"message": "Password reset successful"}
|
||||
@@ -1,126 +0,0 @@
|
||||
from typing import List, Optional
|
||||
from fastapi import APIRouter, Depends, HTTPException, status
|
||||
from sqlalchemy.orm import Session
|
||||
from pydantic import BaseModel
|
||||
from datetime import datetime
|
||||
|
||||
from app.core.database import get_db
|
||||
from app.core.deps import get_current_active_user, get_tenant_id
|
||||
from app.models.user import User
|
||||
from app.models.host import Host
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
class HostCreate(BaseModel):
|
||||
hostname: str
|
||||
ip_address: Optional[str] = None
|
||||
os: Optional[str] = None
|
||||
host_metadata: Optional[dict] = None
|
||||
|
||||
|
||||
class HostRead(BaseModel):
|
||||
id: int
|
||||
hostname: str
|
||||
ip_address: Optional[str] = None
|
||||
os: Optional[str] = None
|
||||
tenant_id: int
|
||||
host_metadata: Optional[dict] = None
|
||||
created_at: datetime
|
||||
last_seen: datetime
|
||||
|
||||
class Config:
|
||||
from_attributes = True
|
||||
|
||||
|
||||
@router.get("/", response_model=List[HostRead])
|
||||
async def list_hosts(
|
||||
skip: int = 0,
|
||||
limit: int = 100,
|
||||
hostname: Optional[str] = None,
|
||||
ip_address: Optional[str] = None,
|
||||
os: Optional[str] = None,
|
||||
sort_by: Optional[str] = None,
|
||||
sort_desc: bool = False,
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
tenant_id: int = Depends(get_tenant_id),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
List hosts scoped to user's tenant with advanced filtering
|
||||
|
||||
Supports:
|
||||
- Filtering by hostname, IP address, OS
|
||||
- Sorting by any field
|
||||
- Pagination
|
||||
"""
|
||||
query = db.query(Host).filter(Host.tenant_id == tenant_id)
|
||||
|
||||
# Apply filters
|
||||
if hostname:
|
||||
query = query.filter(Host.hostname.ilike(f"%{hostname}%"))
|
||||
if ip_address:
|
||||
query = query.filter(Host.ip_address.ilike(f"%{ip_address}%"))
|
||||
if os:
|
||||
query = query.filter(Host.os.ilike(f"%{os}%"))
|
||||
|
||||
# Apply sorting
|
||||
if sort_by:
|
||||
sort_column = getattr(Host, sort_by, None)
|
||||
if sort_column:
|
||||
if sort_desc:
|
||||
query = query.order_by(sort_column.desc())
|
||||
else:
|
||||
query = query.order_by(sort_column)
|
||||
|
||||
hosts = query.offset(skip).limit(limit).all()
|
||||
return hosts
|
||||
|
||||
|
||||
@router.post("/", response_model=HostRead, status_code=status.HTTP_201_CREATED)
|
||||
async def create_host(
|
||||
host_data: HostCreate,
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
tenant_id: int = Depends(get_tenant_id),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Create a new host
|
||||
"""
|
||||
new_host = Host(
|
||||
hostname=host_data.hostname,
|
||||
ip_address=host_data.ip_address,
|
||||
os=host_data.os,
|
||||
tenant_id=tenant_id,
|
||||
host_metadata=host_data.host_metadata
|
||||
)
|
||||
|
||||
db.add(new_host)
|
||||
db.commit()
|
||||
db.refresh(new_host)
|
||||
|
||||
return new_host
|
||||
|
||||
|
||||
@router.get("/{host_id}", response_model=HostRead)
|
||||
async def get_host(
|
||||
host_id: int,
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
tenant_id: int = Depends(get_tenant_id),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Get host by ID (scoped to tenant)
|
||||
"""
|
||||
host = db.query(Host).filter(
|
||||
Host.id == host_id,
|
||||
Host.tenant_id == tenant_id
|
||||
).first()
|
||||
|
||||
if not host:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail="Host not found"
|
||||
)
|
||||
|
||||
return host
|
||||
@@ -1,60 +0,0 @@
|
||||
from fastapi import APIRouter, Depends, HTTPException, status
|
||||
from sqlalchemy.orm import Session
|
||||
from pydantic import BaseModel
|
||||
from typing import Optional, Dict, Any
|
||||
|
||||
from app.core.database import get_db
|
||||
from app.core.deps import get_current_active_user, get_tenant_id
|
||||
from app.models.user import User
|
||||
from app.models.host import Host
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
class IngestionData(BaseModel):
|
||||
hostname: str
|
||||
data: Dict[str, Any]
|
||||
|
||||
|
||||
class IngestionResponse(BaseModel):
|
||||
message: str
|
||||
host_id: int
|
||||
|
||||
|
||||
@router.post("/ingest", response_model=IngestionResponse)
|
||||
async def ingest_data(
|
||||
ingestion: IngestionData,
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
tenant_id: int = Depends(get_tenant_id),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Ingest data from Velociraptor
|
||||
|
||||
Creates or updates host information scoped to the user's tenant.
|
||||
"""
|
||||
# Find or create host
|
||||
host = db.query(Host).filter(
|
||||
Host.hostname == ingestion.hostname,
|
||||
Host.tenant_id == tenant_id
|
||||
).first()
|
||||
|
||||
if host:
|
||||
# Update existing host
|
||||
host.host_metadata = ingestion.data
|
||||
else:
|
||||
# Create new host
|
||||
host = Host(
|
||||
hostname=ingestion.hostname,
|
||||
tenant_id=tenant_id,
|
||||
host_metadata=ingestion.data
|
||||
)
|
||||
db.add(host)
|
||||
|
||||
db.commit()
|
||||
db.refresh(host)
|
||||
|
||||
return {
|
||||
"message": "Data ingested successfully",
|
||||
"host_id": host.id
|
||||
}
|
||||
@@ -1,250 +0,0 @@
|
||||
from typing import Dict, Any, List, Optional
|
||||
from fastapi import APIRouter, Depends, HTTPException, status
|
||||
from sqlalchemy.orm import Session
|
||||
from pydantic import BaseModel
|
||||
|
||||
from app.core.database import get_db
|
||||
from app.core.deps import get_current_active_user, require_role
|
||||
from app.core.llm_router import get_llm_router, TaskType
|
||||
from app.core.job_scheduler import get_job_scheduler, Job, NodeStatus
|
||||
from app.core.llm_pool import get_llm_pool
|
||||
from app.core.merger_agent import get_merger_agent, MergeStrategy
|
||||
from app.models.user import User
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
class LLMRequest(BaseModel):
|
||||
"""Request for LLM processing"""
|
||||
prompt: str
|
||||
task_hints: Optional[List[str]] = []
|
||||
requires_parallel: bool = False
|
||||
requires_chaining: bool = False
|
||||
batch_size: int = 1
|
||||
operations: Optional[List[str]] = []
|
||||
parameters: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
class LLMResponse(BaseModel):
|
||||
"""Response from LLM processing"""
|
||||
job_id: str
|
||||
result: Any
|
||||
execution_mode: str
|
||||
models_used: List[str]
|
||||
strategy: Optional[str] = None
|
||||
|
||||
|
||||
class NodeStatusUpdate(BaseModel):
|
||||
"""Update node status"""
|
||||
node_id: str
|
||||
vram_used_gb: Optional[int] = None
|
||||
compute_utilization: Optional[float] = None
|
||||
status: Optional[str] = None
|
||||
|
||||
|
||||
@router.post("/process", response_model=Dict[str, Any])
|
||||
async def process_llm_request(
|
||||
request: LLMRequest,
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Process an LLM request through the distributed routing system
|
||||
|
||||
The request flows through:
|
||||
1. Router Agent - classifies and routes to appropriate model
|
||||
2. Job Scheduler - determines execution strategy
|
||||
3. LLM Pool - executes on appropriate endpoints
|
||||
4. Merger Agent - combines results if multiple models used
|
||||
"""
|
||||
# Step 1: Route the request
|
||||
router_agent = get_llm_router()
|
||||
routing_decision = router_agent.route_request(request.dict())
|
||||
|
||||
# Step 2: Schedule the job
|
||||
scheduler = get_job_scheduler()
|
||||
job = Job(
|
||||
job_id=f"job_{current_user.id}_{hash(request.prompt) % 10000}",
|
||||
model=routing_decision["model"],
|
||||
priority=routing_decision["priority"],
|
||||
estimated_vram_gb=10, # Estimate based on model
|
||||
requires_parallel=request.requires_parallel,
|
||||
requires_chaining=request.requires_chaining,
|
||||
payload=request.dict()
|
||||
)
|
||||
|
||||
scheduling_decision = await scheduler.schedule_job(job)
|
||||
|
||||
# Step 3: Execute on LLM pool
|
||||
pool = get_llm_pool()
|
||||
|
||||
if scheduling_decision["execution_mode"] == "parallel":
|
||||
# Execute on multiple nodes
|
||||
model_names = [routing_decision["model"]] * len(scheduling_decision["nodes"])
|
||||
results = await pool.call_multiple_models(
|
||||
model_names,
|
||||
request.prompt,
|
||||
request.parameters
|
||||
)
|
||||
|
||||
# Step 4: Merge results
|
||||
merger = get_merger_agent()
|
||||
final_result = merger.merge_results(
|
||||
results["results"],
|
||||
strategy=MergeStrategy.CONSENSUS
|
||||
)
|
||||
|
||||
return {
|
||||
"job_id": job.job_id,
|
||||
"status": "completed",
|
||||
"routing": routing_decision,
|
||||
"scheduling": scheduling_decision,
|
||||
"result": final_result,
|
||||
"execution_mode": "parallel"
|
||||
}
|
||||
|
||||
elif scheduling_decision["execution_mode"] == "queued":
|
||||
return {
|
||||
"job_id": job.job_id,
|
||||
"status": "queued",
|
||||
"queue_position": scheduling_decision["queue_position"],
|
||||
"message": "Job queued - no nodes available"
|
||||
}
|
||||
|
||||
else:
|
||||
# Single node execution
|
||||
result = await pool.call_model(
|
||||
routing_decision["model"],
|
||||
request.prompt,
|
||||
request.parameters
|
||||
)
|
||||
|
||||
return {
|
||||
"job_id": job.job_id,
|
||||
"status": "completed",
|
||||
"routing": routing_decision,
|
||||
"scheduling": scheduling_decision,
|
||||
"result": result,
|
||||
"execution_mode": scheduling_decision["execution_mode"]
|
||||
}
|
||||
|
||||
|
||||
@router.get("/models")
|
||||
async def list_available_models(
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
List all available LLM models in the pool
|
||||
"""
|
||||
pool = get_llm_pool()
|
||||
models = pool.list_available_models()
|
||||
|
||||
return {
|
||||
"models": models,
|
||||
"total": len(models)
|
||||
}
|
||||
|
||||
|
||||
@router.get("/nodes")
|
||||
async def list_gpu_nodes(
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
List all GPU nodes and their status
|
||||
"""
|
||||
scheduler = get_job_scheduler()
|
||||
nodes = scheduler.get_available_nodes()
|
||||
|
||||
return {
|
||||
"nodes": [
|
||||
{
|
||||
"node_id": node.node_id,
|
||||
"hostname": node.hostname,
|
||||
"vram_total_gb": node.vram_total_gb,
|
||||
"vram_used_gb": node.vram_used_gb,
|
||||
"vram_available_gb": node.vram_available_gb,
|
||||
"compute_utilization": node.compute_utilization,
|
||||
"status": node.status.value,
|
||||
"models_loaded": node.models_loaded
|
||||
}
|
||||
for node in scheduler.nodes.values()
|
||||
],
|
||||
"available_count": len(nodes)
|
||||
}
|
||||
|
||||
|
||||
@router.post("/nodes/status")
|
||||
async def update_node_status(
|
||||
update: NodeStatusUpdate,
|
||||
current_user: User = Depends(require_role(["admin"])),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Update GPU node status (admin only)
|
||||
"""
|
||||
scheduler = get_job_scheduler()
|
||||
|
||||
status_enum = None
|
||||
if update.status:
|
||||
try:
|
||||
status_enum = NodeStatus[update.status.upper()]
|
||||
except KeyError:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=f"Invalid status: {update.status}"
|
||||
)
|
||||
|
||||
scheduler.update_node_status(
|
||||
update.node_id,
|
||||
vram_used_gb=update.vram_used_gb,
|
||||
compute_utilization=update.compute_utilization,
|
||||
status=status_enum
|
||||
)
|
||||
|
||||
return {"message": "Node status updated"}
|
||||
|
||||
|
||||
@router.get("/routing/rules")
|
||||
async def get_routing_rules(
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Get current routing rules for task classification
|
||||
"""
|
||||
router_agent = get_llm_router()
|
||||
|
||||
return {
|
||||
"routing_rules": {
|
||||
task_type.value: {
|
||||
"model": rule["model"],
|
||||
"endpoint": rule["endpoint"],
|
||||
"priority": rule["priority"],
|
||||
"description": rule["description"]
|
||||
}
|
||||
for task_type, rule in router_agent.routing_rules.items()
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@router.post("/test-classification")
|
||||
async def test_classification(
|
||||
request: LLMRequest,
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Test task classification without executing
|
||||
"""
|
||||
router_agent = get_llm_router()
|
||||
task_type = router_agent.classify_request(request.dict())
|
||||
routing_decision = router_agent.route_request(request.dict())
|
||||
|
||||
return {
|
||||
"task_type": task_type.value,
|
||||
"routing_decision": routing_decision,
|
||||
"should_parallelize": router_agent.should_parallelize(request.dict()),
|
||||
"requires_chaining": router_agent.requires_serial_chaining(request.dict())
|
||||
}
|
||||
@@ -1,164 +0,0 @@
|
||||
from typing import List
|
||||
from fastapi import APIRouter, Depends, HTTPException, status, WebSocket, WebSocketDisconnect
|
||||
from sqlalchemy.orm import Session
|
||||
import json
|
||||
|
||||
from app.core.database import get_db
|
||||
from app.core.deps import get_current_active_user
|
||||
from app.models.user import User
|
||||
from app.models.notification import Notification
|
||||
from app.schemas.notification import NotificationRead, NotificationUpdate
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
# Store active WebSocket connections
|
||||
active_connections: dict[int, List[WebSocket]] = {}
|
||||
|
||||
|
||||
@router.get("/", response_model=List[NotificationRead])
|
||||
async def list_notifications(
|
||||
skip: int = 0,
|
||||
limit: int = 50,
|
||||
unread_only: bool = False,
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
List notifications for current user
|
||||
|
||||
Supports filtering by read/unread status.
|
||||
"""
|
||||
query = db.query(Notification).filter(
|
||||
Notification.user_id == current_user.id,
|
||||
Notification.tenant_id == current_user.tenant_id
|
||||
)
|
||||
|
||||
if unread_only:
|
||||
query = query.filter(Notification.is_read == False)
|
||||
|
||||
notifications = query.order_by(Notification.created_at.desc()).offset(skip).limit(limit).all()
|
||||
return notifications
|
||||
|
||||
|
||||
@router.put("/{notification_id}", response_model=NotificationRead)
|
||||
async def update_notification(
|
||||
notification_id: int,
|
||||
notification_update: NotificationUpdate,
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Update notification (mark as read/unread)
|
||||
"""
|
||||
notification = db.query(Notification).filter(
|
||||
Notification.id == notification_id,
|
||||
Notification.user_id == current_user.id
|
||||
).first()
|
||||
|
||||
if not notification:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail="Notification not found"
|
||||
)
|
||||
|
||||
notification.is_read = notification_update.is_read
|
||||
db.commit()
|
||||
db.refresh(notification)
|
||||
|
||||
return notification
|
||||
|
||||
|
||||
@router.post("/mark-all-read")
|
||||
async def mark_all_read(
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Mark all notifications as read for current user
|
||||
"""
|
||||
db.query(Notification).filter(
|
||||
Notification.user_id == current_user.id,
|
||||
Notification.is_read == False
|
||||
).update({"is_read": True})
|
||||
db.commit()
|
||||
|
||||
return {"message": "All notifications marked as read"}
|
||||
|
||||
|
||||
@router.websocket("/ws")
|
||||
async def websocket_endpoint(
|
||||
websocket: WebSocket,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
WebSocket endpoint for real-time notifications
|
||||
|
||||
Clients should send their token on connect, then receive notifications in real-time.
|
||||
"""
|
||||
await websocket.accept()
|
||||
user_id = None
|
||||
|
||||
try:
|
||||
# Wait for authentication message
|
||||
auth_data = await websocket.receive_text()
|
||||
auth_json = json.loads(auth_data)
|
||||
token = auth_json.get("token")
|
||||
|
||||
# Validate token and get user (simplified - in production use proper auth)
|
||||
from app.core.security import verify_token
|
||||
payload = verify_token(token)
|
||||
if payload:
|
||||
user_id = payload.get("sub")
|
||||
|
||||
# Register connection
|
||||
if user_id not in active_connections:
|
||||
active_connections[user_id] = []
|
||||
active_connections[user_id].append(websocket)
|
||||
|
||||
# Send confirmation
|
||||
await websocket.send_json({"type": "connected", "user_id": user_id})
|
||||
|
||||
# Keep connection alive
|
||||
while True:
|
||||
data = await websocket.receive_text()
|
||||
# Echo back for keepalive
|
||||
await websocket.send_json({"type": "pong"})
|
||||
else:
|
||||
await websocket.close(code=1008) # Policy violation
|
||||
|
||||
except WebSocketDisconnect:
|
||||
# Remove connection
|
||||
if user_id and user_id in active_connections:
|
||||
active_connections[user_id].remove(websocket)
|
||||
if not active_connections[user_id]:
|
||||
del active_connections[user_id]
|
||||
except Exception as e:
|
||||
print(f"WebSocket error: {e}")
|
||||
await websocket.close(code=1011) # Internal error
|
||||
|
||||
|
||||
async def send_notification_to_user(user_id: int, notification: dict):
|
||||
"""
|
||||
Send notification to all connected clients for a user
|
||||
|
||||
Args:
|
||||
user_id: User ID to send notification to
|
||||
notification: Notification data to send
|
||||
"""
|
||||
if user_id in active_connections:
|
||||
disconnected = []
|
||||
for websocket in active_connections[user_id]:
|
||||
try:
|
||||
await websocket.send_json({
|
||||
"type": "notification",
|
||||
"data": notification
|
||||
})
|
||||
except:
|
||||
disconnected.append(websocket)
|
||||
|
||||
# Clean up disconnected websockets
|
||||
for ws in disconnected:
|
||||
active_connections[user_id].remove(ws)
|
||||
|
||||
if not active_connections[user_id]:
|
||||
del active_connections[user_id]
|
||||
@@ -1,144 +0,0 @@
|
||||
from typing import List
|
||||
from fastapi import APIRouter, Depends, HTTPException, status
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from app.core.database import get_db
|
||||
from app.core.deps import get_current_active_user, require_role, get_tenant_id
|
||||
from app.core.playbook_engine import get_playbook_engine
|
||||
from app.models.user import User
|
||||
from app.models.playbook import Playbook, PlaybookExecution
|
||||
from app.schemas.playbook import PlaybookCreate, PlaybookRead, PlaybookUpdate, PlaybookExecutionRead
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.get("/", response_model=List[PlaybookRead])
|
||||
async def list_playbooks(
|
||||
skip: int = 0,
|
||||
limit: int = 100,
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
tenant_id: int = Depends(get_tenant_id),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""List playbooks scoped to user's tenant"""
|
||||
playbooks = db.query(Playbook).filter(
|
||||
Playbook.tenant_id == tenant_id
|
||||
).offset(skip).limit(limit).all()
|
||||
return playbooks
|
||||
|
||||
|
||||
@router.post("/", response_model=PlaybookRead, status_code=status.HTTP_201_CREATED)
|
||||
async def create_playbook(
|
||||
playbook_data: PlaybookCreate,
|
||||
current_user: User = Depends(require_role(["admin"])),
|
||||
tenant_id: int = Depends(get_tenant_id),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Create a new playbook (admin only)"""
|
||||
playbook = Playbook(
|
||||
tenant_id=tenant_id,
|
||||
created_by=current_user.id,
|
||||
**playbook_data.dict()
|
||||
)
|
||||
db.add(playbook)
|
||||
db.commit()
|
||||
db.refresh(playbook)
|
||||
return playbook
|
||||
|
||||
|
||||
@router.get("/{playbook_id}", response_model=PlaybookRead)
|
||||
async def get_playbook(
|
||||
playbook_id: int,
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
tenant_id: int = Depends(get_tenant_id),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Get playbook by ID"""
|
||||
playbook = db.query(Playbook).filter(
|
||||
Playbook.id == playbook_id,
|
||||
Playbook.tenant_id == tenant_id
|
||||
).first()
|
||||
|
||||
if not playbook:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail="Playbook not found"
|
||||
)
|
||||
|
||||
return playbook
|
||||
|
||||
|
||||
@router.post("/{playbook_id}/execute", response_model=PlaybookExecutionRead)
|
||||
async def execute_playbook(
|
||||
playbook_id: int,
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
tenant_id: int = Depends(get_tenant_id),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Execute a playbook"""
|
||||
playbook = db.query(Playbook).filter(
|
||||
Playbook.id == playbook_id,
|
||||
Playbook.tenant_id == tenant_id
|
||||
).first()
|
||||
|
||||
if not playbook:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail="Playbook not found"
|
||||
)
|
||||
|
||||
if not playbook.is_enabled:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Playbook is disabled"
|
||||
)
|
||||
|
||||
# Create execution record
|
||||
execution = PlaybookExecution(
|
||||
playbook_id=playbook_id,
|
||||
tenant_id=tenant_id,
|
||||
status="running",
|
||||
triggered_by=current_user.id
|
||||
)
|
||||
db.add(execution)
|
||||
db.commit()
|
||||
db.refresh(execution)
|
||||
|
||||
# Execute playbook asynchronously
|
||||
try:
|
||||
engine = get_playbook_engine()
|
||||
result = await engine.execute_playbook(
|
||||
{"actions": playbook.actions},
|
||||
{"tenant_id": tenant_id, "user_id": current_user.id}
|
||||
)
|
||||
|
||||
execution.status = result["status"]
|
||||
execution.result = result
|
||||
from datetime import datetime, timezone
|
||||
execution.completed_at = datetime.now(timezone.utc)
|
||||
except Exception as e:
|
||||
execution.status = "failed"
|
||||
execution.error_message = str(e)
|
||||
|
||||
db.commit()
|
||||
db.refresh(execution)
|
||||
|
||||
return execution
|
||||
|
||||
|
||||
@router.get("/{playbook_id}/executions", response_model=List[PlaybookExecutionRead])
|
||||
async def list_playbook_executions(
|
||||
playbook_id: int,
|
||||
skip: int = 0,
|
||||
limit: int = 100,
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
tenant_id: int = Depends(get_tenant_id),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""List executions for a playbook"""
|
||||
executions = db.query(PlaybookExecution).filter(
|
||||
PlaybookExecution.playbook_id == playbook_id,
|
||||
PlaybookExecution.tenant_id == tenant_id
|
||||
).order_by(PlaybookExecution.started_at.desc()).offset(skip).limit(limit).all()
|
||||
|
||||
return executions
|
||||
@@ -1,120 +0,0 @@
|
||||
from typing import List
|
||||
from fastapi import APIRouter, Depends, HTTPException, status
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from app.core.database import get_db
|
||||
from app.core.deps import get_current_active_user, get_tenant_id
|
||||
from app.models.user import User
|
||||
from app.models.report_template import ReportTemplate, Report
|
||||
from app.schemas.report import ReportTemplateCreate, ReportTemplateRead, ReportCreate, ReportRead
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.get("/templates", response_model=List[ReportTemplateRead])
|
||||
async def list_report_templates(
|
||||
skip: int = 0,
|
||||
limit: int = 100,
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
tenant_id: int = Depends(get_tenant_id),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""List report templates scoped to tenant"""
|
||||
templates = db.query(ReportTemplate).filter(
|
||||
ReportTemplate.tenant_id == tenant_id
|
||||
).offset(skip).limit(limit).all()
|
||||
return templates
|
||||
|
||||
|
||||
@router.post("/templates", response_model=ReportTemplateRead, status_code=status.HTTP_201_CREATED)
|
||||
async def create_report_template(
|
||||
template_data: ReportTemplateCreate,
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
tenant_id: int = Depends(get_tenant_id),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Create a new report template"""
|
||||
template = ReportTemplate(
|
||||
tenant_id=tenant_id,
|
||||
created_by=current_user.id,
|
||||
**template_data.dict()
|
||||
)
|
||||
db.add(template)
|
||||
db.commit()
|
||||
db.refresh(template)
|
||||
return template
|
||||
|
||||
|
||||
@router.post("/generate", response_model=ReportRead, status_code=status.HTTP_201_CREATED)
|
||||
async def generate_report(
|
||||
report_data: ReportCreate,
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
tenant_id: int = Depends(get_tenant_id),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Generate a new report
|
||||
|
||||
This is a simplified implementation. In production, this would:
|
||||
1. Fetch relevant data based on report type
|
||||
2. Apply template formatting
|
||||
3. Generate PDF/HTML output
|
||||
4. Store file and return path
|
||||
"""
|
||||
report = Report(
|
||||
tenant_id=tenant_id,
|
||||
template_id=report_data.template_id,
|
||||
title=report_data.title,
|
||||
report_type=report_data.report_type,
|
||||
format=report_data.format,
|
||||
status="generating",
|
||||
generated_by=current_user.id
|
||||
)
|
||||
db.add(report)
|
||||
db.commit()
|
||||
|
||||
# Simulate report generation
|
||||
# In production, this would be an async task
|
||||
report.status = "completed"
|
||||
report.file_path = f"/reports/{report.id}.{report_data.format}"
|
||||
db.commit()
|
||||
db.refresh(report)
|
||||
|
||||
return report
|
||||
|
||||
|
||||
@router.get("/", response_model=List[ReportRead])
|
||||
async def list_reports(
|
||||
skip: int = 0,
|
||||
limit: int = 100,
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
tenant_id: int = Depends(get_tenant_id),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""List generated reports"""
|
||||
reports = db.query(Report).filter(
|
||||
Report.tenant_id == tenant_id
|
||||
).order_by(Report.generated_at.desc()).offset(skip).limit(limit).all()
|
||||
return reports
|
||||
|
||||
|
||||
@router.get("/{report_id}", response_model=ReportRead)
|
||||
async def get_report(
|
||||
report_id: int,
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
tenant_id: int = Depends(get_tenant_id),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""Get a specific report"""
|
||||
report = db.query(Report).filter(
|
||||
Report.id == report_id,
|
||||
Report.tenant_id == tenant_id
|
||||
).first()
|
||||
|
||||
if not report:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail="Report not found"
|
||||
)
|
||||
|
||||
return report
|
||||
@@ -1,103 +0,0 @@
|
||||
from typing import List
|
||||
from fastapi import APIRouter, Depends, HTTPException, status
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from app.core.database import get_db
|
||||
from app.core.deps import get_current_active_user, require_role
|
||||
from app.models.user import User
|
||||
from app.models.tenant import Tenant
|
||||
from pydantic import BaseModel
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
class TenantCreate(BaseModel):
|
||||
name: str
|
||||
description: str = None
|
||||
|
||||
|
||||
class TenantRead(BaseModel):
|
||||
id: int
|
||||
name: str
|
||||
description: str = None
|
||||
|
||||
class Config:
|
||||
from_attributes = True
|
||||
|
||||
|
||||
@router.get("/", response_model=List[TenantRead])
|
||||
async def list_tenants(
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
List tenants
|
||||
|
||||
Regular users can only see their own tenant.
|
||||
Admins can see all tenants (cross-tenant access).
|
||||
"""
|
||||
if current_user.role == "admin":
|
||||
# Admins can see all tenants
|
||||
tenants = db.query(Tenant).all()
|
||||
else:
|
||||
# Regular users only see their tenant
|
||||
tenants = db.query(Tenant).filter(Tenant.id == current_user.tenant_id).all()
|
||||
|
||||
return tenants
|
||||
|
||||
|
||||
@router.post("/", response_model=TenantRead, status_code=status.HTTP_201_CREATED)
|
||||
async def create_tenant(
|
||||
tenant_data: TenantCreate,
|
||||
current_user: User = Depends(require_role(["admin"])),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Create a new tenant (admin only)
|
||||
"""
|
||||
# Check if tenant name already exists
|
||||
existing_tenant = db.query(Tenant).filter(Tenant.name == tenant_data.name).first()
|
||||
if existing_tenant:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Tenant name already exists"
|
||||
)
|
||||
|
||||
new_tenant = Tenant(
|
||||
name=tenant_data.name,
|
||||
description=tenant_data.description
|
||||
)
|
||||
|
||||
db.add(new_tenant)
|
||||
db.commit()
|
||||
db.refresh(new_tenant)
|
||||
|
||||
return new_tenant
|
||||
|
||||
|
||||
@router.get("/{tenant_id}", response_model=TenantRead)
|
||||
async def get_tenant(
|
||||
tenant_id: int,
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Get tenant by ID
|
||||
|
||||
Users can only view their own tenant unless they are admin.
|
||||
"""
|
||||
tenant = db.query(Tenant).filter(Tenant.id == tenant_id).first()
|
||||
if not tenant:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail="Tenant not found"
|
||||
)
|
||||
|
||||
# Check permissions
|
||||
if current_user.role != "admin" and tenant.id != current_user.tenant_id:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail="Not authorized to view this tenant"
|
||||
)
|
||||
|
||||
return tenant
|
||||
@@ -1,127 +0,0 @@
|
||||
from typing import List
|
||||
from fastapi import APIRouter, Depends, HTTPException, status
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from app.core.database import get_db
|
||||
from app.core.deps import get_current_active_user, get_tenant_id
|
||||
from app.core.threat_intel import get_threat_analyzer
|
||||
from app.models.user import User
|
||||
from app.models.threat_score import ThreatScore
|
||||
from app.models.host import Host
|
||||
from app.models.artifact import Artifact
|
||||
from app.schemas.threat_score import ThreatScoreRead, ThreatScoreCreate
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.post("/analyze/host/{host_id}", response_model=ThreatScoreRead)
|
||||
async def analyze_host(
|
||||
host_id: int,
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
tenant_id: int = Depends(get_tenant_id),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Analyze a host for threats using ML
|
||||
"""
|
||||
host = db.query(Host).filter(
|
||||
Host.id == host_id,
|
||||
Host.tenant_id == tenant_id
|
||||
).first()
|
||||
|
||||
if not host:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail="Host not found"
|
||||
)
|
||||
|
||||
# Analyze host
|
||||
analyzer = get_threat_analyzer()
|
||||
analysis = analyzer.analyze_host({
|
||||
"hostname": host.hostname,
|
||||
"ip_address": host.ip_address,
|
||||
"os": host.os,
|
||||
"host_metadata": host.host_metadata
|
||||
})
|
||||
|
||||
# Store threat score
|
||||
threat_score = ThreatScore(
|
||||
tenant_id=tenant_id,
|
||||
host_id=host_id,
|
||||
score=analysis["score"],
|
||||
confidence=analysis["confidence"],
|
||||
threat_type=analysis["threat_type"],
|
||||
indicators=analysis["indicators"],
|
||||
ml_model_version=analysis["ml_model_version"]
|
||||
)
|
||||
db.add(threat_score)
|
||||
db.commit()
|
||||
db.refresh(threat_score)
|
||||
|
||||
return threat_score
|
||||
|
||||
|
||||
@router.post("/analyze/artifact/{artifact_id}", response_model=ThreatScoreRead)
|
||||
async def analyze_artifact(
|
||||
artifact_id: int,
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
tenant_id: int = Depends(get_tenant_id),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Analyze an artifact for threats
|
||||
"""
|
||||
artifact = db.query(Artifact).filter(Artifact.id == artifact_id).first()
|
||||
|
||||
if not artifact:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail="Artifact not found"
|
||||
)
|
||||
|
||||
# Analyze artifact
|
||||
analyzer = get_threat_analyzer()
|
||||
analysis = analyzer.analyze_artifact({
|
||||
"artifact_type": artifact.artifact_type,
|
||||
"value": artifact.value
|
||||
})
|
||||
|
||||
# Store threat score
|
||||
threat_score = ThreatScore(
|
||||
tenant_id=tenant_id,
|
||||
artifact_id=artifact_id,
|
||||
score=analysis["score"],
|
||||
confidence=analysis["confidence"],
|
||||
threat_type=analysis["threat_type"],
|
||||
indicators=analysis["indicators"],
|
||||
ml_model_version=analysis["ml_model_version"]
|
||||
)
|
||||
db.add(threat_score)
|
||||
db.commit()
|
||||
db.refresh(threat_score)
|
||||
|
||||
return threat_score
|
||||
|
||||
|
||||
@router.get("/scores", response_model=List[ThreatScoreRead])
|
||||
async def list_threat_scores(
|
||||
skip: int = 0,
|
||||
limit: int = 100,
|
||||
min_score: float = 0.0,
|
||||
threat_type: str = None,
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
tenant_id: int = Depends(get_tenant_id),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
List threat scores with filtering
|
||||
"""
|
||||
query = db.query(ThreatScore).filter(ThreatScore.tenant_id == tenant_id)
|
||||
|
||||
if min_score:
|
||||
query = query.filter(ThreatScore.score >= min_score)
|
||||
if threat_type:
|
||||
query = query.filter(ThreatScore.threat_type == threat_type)
|
||||
|
||||
scores = query.order_by(ThreatScore.score.desc()).offset(skip).limit(limit).all()
|
||||
return scores
|
||||
@@ -1,154 +0,0 @@
|
||||
from typing import List
|
||||
from fastapi import APIRouter, Depends, HTTPException, status
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from app.core.database import get_db
|
||||
from app.core.deps import get_current_active_user, require_role
|
||||
from app.core.security import get_password_hash
|
||||
from app.models.user import User
|
||||
from app.schemas.user import UserRead, UserUpdate, UserCreate
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.get("/", response_model=List[UserRead])
|
||||
async def list_users(
|
||||
skip: int = 0,
|
||||
limit: int = 100,
|
||||
current_user: User = Depends(require_role(["admin"])),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
List all users (admin only, scoped to tenant)
|
||||
|
||||
Admins can only see users within their own tenant unless they have
|
||||
cross-tenant access.
|
||||
"""
|
||||
# Scope to tenant
|
||||
query = db.query(User).filter(User.tenant_id == current_user.tenant_id)
|
||||
users = query.offset(skip).limit(limit).all()
|
||||
return users
|
||||
|
||||
|
||||
@router.get("/{user_id}", response_model=UserRead)
|
||||
async def get_user(
|
||||
user_id: int,
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Get user by ID
|
||||
|
||||
Users can view their own profile or admins can view users in their tenant.
|
||||
"""
|
||||
user = db.query(User).filter(User.id == user_id).first()
|
||||
if not user:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail="User not found"
|
||||
)
|
||||
|
||||
# Check permissions: user can view themselves or admin can view users in their tenant
|
||||
if user.id != current_user.id and (
|
||||
current_user.role != "admin" or user.tenant_id != current_user.tenant_id
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail="Not authorized to view this user"
|
||||
)
|
||||
|
||||
return user
|
||||
|
||||
|
||||
@router.put("/{user_id}", response_model=UserRead)
|
||||
async def update_user(
|
||||
user_id: int,
|
||||
user_update: UserUpdate,
|
||||
current_user: User = Depends(require_role(["admin"])),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Update user (admin only)
|
||||
|
||||
Admins can update users within their tenant.
|
||||
"""
|
||||
user = db.query(User).filter(User.id == user_id).first()
|
||||
if not user:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail="User not found"
|
||||
)
|
||||
|
||||
# Check tenant access
|
||||
if user.tenant_id != current_user.tenant_id:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail="Not authorized to update this user"
|
||||
)
|
||||
|
||||
# Update fields
|
||||
if user_update.username is not None:
|
||||
# Check if new username is already taken
|
||||
existing_user = db.query(User).filter(
|
||||
User.username == user_update.username,
|
||||
User.id != user_id
|
||||
).first()
|
||||
if existing_user:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Username already taken"
|
||||
)
|
||||
user.username = user_update.username
|
||||
|
||||
if user_update.password is not None:
|
||||
user.password_hash = get_password_hash(user_update.password)
|
||||
|
||||
if user_update.role is not None:
|
||||
user.role = user_update.role
|
||||
|
||||
if user_update.is_active is not None:
|
||||
user.is_active = user_update.is_active
|
||||
|
||||
db.commit()
|
||||
db.refresh(user)
|
||||
|
||||
return user
|
||||
|
||||
|
||||
@router.delete("/{user_id}", status_code=status.HTTP_204_NO_CONTENT)
|
||||
async def delete_user(
|
||||
user_id: int,
|
||||
current_user: User = Depends(require_role(["admin"])),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Deactivate user (admin only)
|
||||
|
||||
Soft delete by setting is_active to False.
|
||||
"""
|
||||
user = db.query(User).filter(User.id == user_id).first()
|
||||
if not user:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail="User not found"
|
||||
)
|
||||
|
||||
# Check tenant access
|
||||
if user.tenant_id != current_user.tenant_id:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail="Not authorized to delete this user"
|
||||
)
|
||||
|
||||
# Prevent self-deletion
|
||||
if user.id == current_user.id:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Cannot delete your own account"
|
||||
)
|
||||
|
||||
# Soft delete
|
||||
user.is_active = False
|
||||
db.commit()
|
||||
|
||||
return None
|
||||
@@ -1,197 +0,0 @@
|
||||
from typing import List, Dict, Any, Optional
|
||||
from fastapi import APIRouter, Depends, HTTPException, status
|
||||
from sqlalchemy.orm import Session
|
||||
from pydantic import BaseModel
|
||||
|
||||
from app.core.database import get_db
|
||||
from app.core.deps import get_current_active_user, require_role
|
||||
from app.core.velociraptor import get_velociraptor_client
|
||||
from app.models.user import User
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
class VelociraptorConfig(BaseModel):
|
||||
"""Velociraptor server configuration"""
|
||||
base_url: str
|
||||
api_key: str
|
||||
|
||||
|
||||
class ArtifactCollectionRequest(BaseModel):
|
||||
"""Request to collect an artifact"""
|
||||
client_id: str
|
||||
artifact_name: str
|
||||
parameters: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
class HuntCreateRequest(BaseModel):
|
||||
"""Request to create a hunt"""
|
||||
hunt_name: str
|
||||
artifact_name: str
|
||||
description: str
|
||||
parameters: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
# In a real implementation, this would be stored per-tenant in database
|
||||
# For now, using a simple in-memory store
|
||||
velociraptor_configs: Dict[int, VelociraptorConfig] = {}
|
||||
|
||||
|
||||
@router.post("/config")
|
||||
async def set_velociraptor_config(
|
||||
config: VelociraptorConfig,
|
||||
current_user: User = Depends(require_role(["admin"])),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Configure Velociraptor integration (admin only)
|
||||
|
||||
Stores Velociraptor server URL and API key for the tenant.
|
||||
"""
|
||||
velociraptor_configs[current_user.tenant_id] = config
|
||||
return {"message": "Velociraptor configuration saved"}
|
||||
|
||||
|
||||
@router.get("/clients")
|
||||
async def list_velociraptor_clients(
|
||||
limit: int = 100,
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
List clients from Velociraptor server
|
||||
"""
|
||||
config = velociraptor_configs.get(current_user.tenant_id)
|
||||
if not config:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Velociraptor not configured for this tenant"
|
||||
)
|
||||
|
||||
client = get_velociraptor_client(config.base_url, config.api_key)
|
||||
try:
|
||||
clients = await client.list_clients(limit=limit)
|
||||
return {"clients": clients}
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail=f"Failed to fetch clients: {str(e)}"
|
||||
)
|
||||
|
||||
|
||||
@router.get("/clients/{client_id}")
|
||||
async def get_velociraptor_client_info(
|
||||
client_id: str,
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Get information about a specific Velociraptor client
|
||||
"""
|
||||
config = velociraptor_configs.get(current_user.tenant_id)
|
||||
if not config:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Velociraptor not configured for this tenant"
|
||||
)
|
||||
|
||||
client = get_velociraptor_client(config.base_url, config.api_key)
|
||||
try:
|
||||
client_info = await client.get_client(client_id)
|
||||
return client_info
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail=f"Failed to fetch client: {str(e)}"
|
||||
)
|
||||
|
||||
|
||||
@router.post("/collect")
|
||||
async def collect_artifact(
|
||||
request: ArtifactCollectionRequest,
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Collect an artifact from a Velociraptor client
|
||||
"""
|
||||
config = velociraptor_configs.get(current_user.tenant_id)
|
||||
if not config:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Velociraptor not configured for this tenant"
|
||||
)
|
||||
|
||||
client = get_velociraptor_client(config.base_url, config.api_key)
|
||||
try:
|
||||
result = await client.collect_artifact(
|
||||
request.client_id,
|
||||
request.artifact_name,
|
||||
request.parameters
|
||||
)
|
||||
return result
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail=f"Failed to collect artifact: {str(e)}"
|
||||
)
|
||||
|
||||
|
||||
@router.post("/hunts")
|
||||
async def create_hunt(
|
||||
request: HuntCreateRequest,
|
||||
current_user: User = Depends(require_role(["admin"])),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Create a new hunt (admin only)
|
||||
"""
|
||||
config = velociraptor_configs.get(current_user.tenant_id)
|
||||
if not config:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Velociraptor not configured for this tenant"
|
||||
)
|
||||
|
||||
client = get_velociraptor_client(config.base_url, config.api_key)
|
||||
try:
|
||||
result = await client.create_hunt(
|
||||
request.hunt_name,
|
||||
request.artifact_name,
|
||||
request.description,
|
||||
request.parameters
|
||||
)
|
||||
return result
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail=f"Failed to create hunt: {str(e)}"
|
||||
)
|
||||
|
||||
|
||||
@router.get("/hunts/{hunt_id}/results")
|
||||
async def get_hunt_results(
|
||||
hunt_id: str,
|
||||
limit: int = 1000,
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Get results from a hunt
|
||||
"""
|
||||
config = velociraptor_configs.get(current_user.tenant_id)
|
||||
if not config:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Velociraptor not configured for this tenant"
|
||||
)
|
||||
|
||||
client = get_velociraptor_client(config.base_url, config.api_key)
|
||||
try:
|
||||
results = await client.get_hunt_results(hunt_id, limit=limit)
|
||||
return {"results": results}
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail=f"Failed to fetch hunt results: {str(e)}"
|
||||
)
|
||||
@@ -1,40 +0,0 @@
|
||||
from fastapi import APIRouter, Depends, HTTPException, status
|
||||
from sqlalchemy.orm import Session
|
||||
from pydantic import BaseModel
|
||||
from typing import Optional
|
||||
|
||||
from app.core.database import get_db
|
||||
from app.core.deps import get_current_active_user
|
||||
from app.models.user import User
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
class VTLookupRequest(BaseModel):
|
||||
hash: str
|
||||
|
||||
|
||||
class VTLookupResponse(BaseModel):
|
||||
hash: str
|
||||
malicious: Optional[bool] = None
|
||||
message: str
|
||||
|
||||
|
||||
@router.post("/lookup", response_model=VTLookupResponse)
|
||||
async def virustotal_lookup(
|
||||
request: VTLookupRequest,
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Lookup hash in VirusTotal
|
||||
|
||||
Requires authentication. In a real implementation, this would call
|
||||
the VirusTotal API.
|
||||
"""
|
||||
# Placeholder implementation
|
||||
return {
|
||||
"hash": request.hash,
|
||||
"malicious": None,
|
||||
"message": "VirusTotal integration not yet implemented"
|
||||
}
|
||||
@@ -1,52 +0,0 @@
|
||||
from typing import Optional, Dict, Any
|
||||
from sqlalchemy.orm import Session
|
||||
from fastapi import Request
|
||||
|
||||
from app.models.audit_log import AuditLog
|
||||
|
||||
|
||||
def log_action(
|
||||
db: Session,
|
||||
user_id: Optional[int],
|
||||
tenant_id: int,
|
||||
action: str,
|
||||
resource_type: str,
|
||||
resource_id: Optional[int] = None,
|
||||
details: Optional[Dict[str, Any]] = None,
|
||||
request: Optional[Request] = None
|
||||
):
|
||||
"""
|
||||
Log an action to the audit log
|
||||
|
||||
Args:
|
||||
db: Database session
|
||||
user_id: ID of user performing action (None for system actions)
|
||||
tenant_id: Tenant ID
|
||||
action: Action type (CREATE, READ, UPDATE, DELETE, LOGIN, etc.)
|
||||
resource_type: Type of resource (user, host, case, etc.)
|
||||
resource_id: ID of the resource (if applicable)
|
||||
details: Additional details as JSON
|
||||
request: FastAPI request object (for IP and user agent)
|
||||
"""
|
||||
ip_address = None
|
||||
user_agent = None
|
||||
|
||||
if request:
|
||||
ip_address = request.client.host if request.client else None
|
||||
user_agent = request.headers.get("user-agent")
|
||||
|
||||
audit_log = AuditLog(
|
||||
user_id=user_id,
|
||||
tenant_id=tenant_id,
|
||||
action=action,
|
||||
resource_type=resource_type,
|
||||
resource_id=resource_id,
|
||||
details=details,
|
||||
ip_address=ip_address,
|
||||
user_agent=user_agent
|
||||
)
|
||||
|
||||
db.add(audit_log)
|
||||
db.commit()
|
||||
|
||||
return audit_log
|
||||
@@ -1,27 +0,0 @@
|
||||
from pydantic_settings import BaseSettings
|
||||
|
||||
|
||||
class Settings(BaseSettings):
|
||||
"""Application settings"""
|
||||
app_name: str = "VelociCompanion"
|
||||
database_url: str = "postgresql://postgres:postgres@db:5432/velocicompanion"
|
||||
secret_key: str = "your-secret-key-change-in-production-min-32-chars-long"
|
||||
access_token_expire_minutes: int = 30
|
||||
refresh_token_expire_days: int = 30
|
||||
algorithm: str = "HS256"
|
||||
|
||||
# Email settings (for password reset)
|
||||
smtp_host: str = "localhost"
|
||||
smtp_port: int = 587
|
||||
smtp_user: str = ""
|
||||
smtp_password: str = ""
|
||||
from_email: str = "noreply@velocicompanion.com"
|
||||
|
||||
# WebSocket settings
|
||||
ws_enabled: bool = True
|
||||
|
||||
class Config:
|
||||
env_file = ".env"
|
||||
|
||||
|
||||
settings = Settings()
|
||||
@@ -1,19 +0,0 @@
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.ext.declarative import declarative_base
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
|
||||
from app.core.config import settings
|
||||
|
||||
engine = create_engine(settings.database_url)
|
||||
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
|
||||
|
||||
Base = declarative_base()
|
||||
|
||||
|
||||
def get_db():
|
||||
"""Dependency for getting database session"""
|
||||
db = SessionLocal()
|
||||
try:
|
||||
yield db
|
||||
finally:
|
||||
db.close()
|
||||
@@ -1,108 +0,0 @@
|
||||
from typing import List, Optional
|
||||
from fastapi import Depends, HTTPException, status
|
||||
from fastapi.security import OAuth2PasswordBearer
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from app.core.database import get_db
|
||||
from app.core.security import verify_token
|
||||
from app.models.user import User
|
||||
|
||||
# OAuth2 scheme for token extraction
|
||||
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/api/auth/login")
|
||||
|
||||
|
||||
async def get_current_user(
|
||||
token: str = Depends(oauth2_scheme),
|
||||
db: Session = Depends(get_db)
|
||||
) -> User:
|
||||
"""
|
||||
Extract and validate JWT from Authorization header
|
||||
|
||||
Args:
|
||||
token: JWT token from Authorization header
|
||||
db: Database session
|
||||
|
||||
Returns:
|
||||
Current user object
|
||||
|
||||
Raises:
|
||||
HTTPException: If token is invalid or user not found
|
||||
"""
|
||||
credentials_exception = HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="Could not validate credentials",
|
||||
headers={"WWW-Authenticate": "Bearer"},
|
||||
)
|
||||
|
||||
# Verify and decode token
|
||||
payload = verify_token(token)
|
||||
if payload is None:
|
||||
raise credentials_exception
|
||||
|
||||
user_id: Optional[int] = payload.get("sub")
|
||||
if user_id is None:
|
||||
raise credentials_exception
|
||||
|
||||
# Get user from database
|
||||
user = db.query(User).filter(User.id == user_id).first()
|
||||
if user is None:
|
||||
raise credentials_exception
|
||||
|
||||
return user
|
||||
|
||||
|
||||
async def get_current_active_user(
|
||||
current_user: User = Depends(get_current_user)
|
||||
) -> User:
|
||||
"""
|
||||
Ensure user is active
|
||||
|
||||
Args:
|
||||
current_user: Current user from get_current_user
|
||||
|
||||
Returns:
|
||||
Active user object
|
||||
|
||||
Raises:
|
||||
HTTPException: If user is inactive
|
||||
"""
|
||||
if not current_user.is_active:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Inactive user"
|
||||
)
|
||||
return current_user
|
||||
|
||||
|
||||
def require_role(allowed_roles: List[str]):
|
||||
"""
|
||||
Role-based access control dependency factory
|
||||
|
||||
Args:
|
||||
allowed_roles: List of roles that are allowed to access the endpoint
|
||||
|
||||
Returns:
|
||||
Dependency function that checks user role
|
||||
"""
|
||||
async def role_checker(current_user: User = Depends(get_current_active_user)) -> User:
|
||||
if current_user.role not in allowed_roles:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail="Insufficient permissions"
|
||||
)
|
||||
return current_user
|
||||
|
||||
return role_checker
|
||||
|
||||
|
||||
async def get_tenant_id(current_user: User = Depends(get_current_active_user)) -> int:
|
||||
"""
|
||||
Extract tenant_id from current user for scoping queries
|
||||
|
||||
Args:
|
||||
current_user: Current active user
|
||||
|
||||
Returns:
|
||||
Tenant ID
|
||||
"""
|
||||
return current_user.tenant_id
|
||||
@@ -1,263 +0,0 @@
|
||||
"""
|
||||
Job Scheduler
|
||||
|
||||
Manages job distribution across GPU nodes based on availability and load.
|
||||
"""
|
||||
|
||||
from typing import Dict, Any, List, Optional
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
import asyncio
|
||||
|
||||
|
||||
class NodeStatus(Enum):
|
||||
"""Status of GPU node"""
|
||||
AVAILABLE = "available"
|
||||
BUSY = "busy"
|
||||
OFFLINE = "offline"
|
||||
|
||||
|
||||
@dataclass
|
||||
class GPUNode:
|
||||
"""Represents a GPU compute node"""
|
||||
node_id: str
|
||||
hostname: str
|
||||
port: int
|
||||
vram_total_gb: int
|
||||
vram_used_gb: int
|
||||
compute_utilization: float # 0.0 to 1.0
|
||||
status: NodeStatus
|
||||
models_loaded: List[str]
|
||||
|
||||
@property
|
||||
def vram_available_gb(self) -> int:
|
||||
"""Calculate available VRAM"""
|
||||
return self.vram_total_gb - self.vram_used_gb
|
||||
|
||||
@property
|
||||
def is_available(self) -> bool:
|
||||
"""Check if node is available for work"""
|
||||
return self.status == NodeStatus.AVAILABLE and self.compute_utilization < 0.9
|
||||
|
||||
|
||||
@dataclass
|
||||
class Job:
|
||||
"""Represents an LLM job"""
|
||||
job_id: str
|
||||
model: str
|
||||
priority: int
|
||||
estimated_vram_gb: int
|
||||
requires_parallel: bool
|
||||
requires_chaining: bool
|
||||
payload: Dict[str, Any]
|
||||
|
||||
|
||||
class JobScheduler:
|
||||
"""
|
||||
Job Scheduler - Manages distribution of LLM jobs across GPU nodes
|
||||
|
||||
Decides:
|
||||
- Which GB10 device is available
|
||||
- GPU load (VRAM, compute utilization)
|
||||
- Whether to parallelize across both nodes
|
||||
- Whether job requires serial reasoning (chained)
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize job scheduler"""
|
||||
self.nodes: Dict[str, GPUNode] = {}
|
||||
self.job_queue: List[Job] = []
|
||||
self._initialize_nodes()
|
||||
|
||||
def _initialize_nodes(self):
|
||||
"""Initialize GPU node configuration"""
|
||||
# GB10 Node 1
|
||||
self.nodes["gb10-node-1"] = GPUNode(
|
||||
node_id="gb10-node-1",
|
||||
hostname="gb10-node-1",
|
||||
port=8001,
|
||||
vram_total_gb=80,
|
||||
vram_used_gb=0,
|
||||
compute_utilization=0.0,
|
||||
status=NodeStatus.AVAILABLE,
|
||||
models_loaded=["deepseek", "qwen72"]
|
||||
)
|
||||
|
||||
# GB10 Node 2
|
||||
self.nodes["gb10-node-2"] = GPUNode(
|
||||
node_id="gb10-node-2",
|
||||
hostname="gb10-node-2",
|
||||
port=8001,
|
||||
vram_total_gb=80,
|
||||
vram_used_gb=0,
|
||||
compute_utilization=0.0,
|
||||
status=NodeStatus.AVAILABLE,
|
||||
models_loaded=["phi4", "qwen-coder", "llama31", "granite-guardian"]
|
||||
)
|
||||
|
||||
def get_available_nodes(self) -> List[GPUNode]:
|
||||
"""Get list of available nodes"""
|
||||
return [node for node in self.nodes.values() if node.is_available]
|
||||
|
||||
def find_best_node(self, job: Job) -> Optional[GPUNode]:
|
||||
"""
|
||||
Find best node for a job based on availability and requirements
|
||||
|
||||
Args:
|
||||
job: Job to schedule
|
||||
|
||||
Returns:
|
||||
Best GPU node or None if unavailable
|
||||
"""
|
||||
available_nodes = self.get_available_nodes()
|
||||
|
||||
# Filter nodes that have required model loaded
|
||||
suitable_nodes = [
|
||||
node for node in available_nodes
|
||||
if job.model in node.models_loaded
|
||||
and node.vram_available_gb >= job.estimated_vram_gb
|
||||
]
|
||||
|
||||
if not suitable_nodes:
|
||||
return None
|
||||
|
||||
# Sort by compute utilization (prefer less loaded nodes)
|
||||
suitable_nodes.sort(key=lambda n: n.compute_utilization)
|
||||
|
||||
return suitable_nodes[0]
|
||||
|
||||
def should_parallelize(self, job: Job) -> bool:
|
||||
"""
|
||||
Determine if job should be parallelized across multiple nodes
|
||||
|
||||
Args:
|
||||
job: Job to evaluate
|
||||
|
||||
Returns:
|
||||
True if should parallelize
|
||||
"""
|
||||
available_nodes = self.get_available_nodes()
|
||||
|
||||
# Need at least 2 nodes for parallelization
|
||||
if len(available_nodes) < 2:
|
||||
return False
|
||||
|
||||
# Job explicitly requires parallel execution
|
||||
if job.requires_parallel:
|
||||
return True
|
||||
|
||||
# High priority jobs with multiple available nodes
|
||||
if job.priority >= 1 and len(available_nodes) >= 2:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def get_parallel_nodes(self, job: Job) -> List[GPUNode]:
|
||||
"""
|
||||
Get nodes for parallel execution
|
||||
|
||||
Args:
|
||||
job: Job to parallelize
|
||||
|
||||
Returns:
|
||||
List of nodes to use
|
||||
"""
|
||||
available_nodes = self.get_available_nodes()
|
||||
|
||||
# Filter nodes with required model and sufficient VRAM
|
||||
suitable_nodes = [
|
||||
node for node in available_nodes
|
||||
if job.model in node.models_loaded
|
||||
and node.vram_available_gb >= job.estimated_vram_gb
|
||||
]
|
||||
|
||||
# Return up to 2 nodes for parallel execution
|
||||
return suitable_nodes[:2]
|
||||
|
||||
async def schedule_job(self, job: Job) -> Dict[str, Any]:
|
||||
"""
|
||||
Schedule a job for execution
|
||||
|
||||
Args:
|
||||
job: Job to schedule
|
||||
|
||||
Returns:
|
||||
Scheduling decision with node assignments
|
||||
"""
|
||||
# Check if job should be parallelized
|
||||
if self.should_parallelize(job):
|
||||
nodes = self.get_parallel_nodes(job)
|
||||
if len(nodes) >= 2:
|
||||
return {
|
||||
"job_id": job.job_id,
|
||||
"execution_mode": "parallel",
|
||||
"nodes": [
|
||||
{"node_id": node.node_id, "endpoint": f"http://{node.hostname}:{node.port}/{job.model}"}
|
||||
for node in nodes
|
||||
],
|
||||
"estimated_time": "distributed"
|
||||
}
|
||||
|
||||
# Serial execution on single node
|
||||
node = self.find_best_node(job)
|
||||
|
||||
if not node:
|
||||
# Add to queue if no nodes available
|
||||
self.job_queue.append(job)
|
||||
return {
|
||||
"job_id": job.job_id,
|
||||
"execution_mode": "queued",
|
||||
"status": "waiting_for_resources",
|
||||
"queue_position": len(self.job_queue)
|
||||
}
|
||||
|
||||
return {
|
||||
"job_id": job.job_id,
|
||||
"execution_mode": "serial" if job.requires_chaining else "single",
|
||||
"node": {
|
||||
"node_id": node.node_id,
|
||||
"endpoint": f"http://{node.hostname}:{node.port}/{job.model}"
|
||||
},
|
||||
"vram_allocated_gb": job.estimated_vram_gb,
|
||||
"estimated_time": "standard"
|
||||
}
|
||||
|
||||
def update_node_status(
|
||||
self,
|
||||
node_id: str,
|
||||
vram_used_gb: Optional[int] = None,
|
||||
compute_utilization: Optional[float] = None,
|
||||
status: Optional[NodeStatus] = None
|
||||
):
|
||||
"""
|
||||
Update node status metrics
|
||||
|
||||
Args:
|
||||
node_id: Node to update
|
||||
vram_used_gb: Current VRAM usage
|
||||
compute_utilization: Current compute utilization (0.0-1.0)
|
||||
status: Node status
|
||||
"""
|
||||
if node_id not in self.nodes:
|
||||
return
|
||||
|
||||
node = self.nodes[node_id]
|
||||
|
||||
if vram_used_gb is not None:
|
||||
node.vram_used_gb = vram_used_gb
|
||||
|
||||
if compute_utilization is not None:
|
||||
node.compute_utilization = compute_utilization
|
||||
|
||||
if status is not None:
|
||||
node.status = status
|
||||
|
||||
|
||||
def get_job_scheduler() -> JobScheduler:
|
||||
"""
|
||||
Factory function to create job scheduler
|
||||
|
||||
Returns:
|
||||
Configured JobScheduler instance
|
||||
"""
|
||||
return JobScheduler()
|
||||
@@ -1,211 +0,0 @@
|
||||
"""
|
||||
LLM Pool Manager
|
||||
|
||||
Manages pool of LLM endpoints with OpenAI-compatible interface.
|
||||
"""
|
||||
|
||||
from typing import Dict, Any, List, Optional
|
||||
import httpx
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass
|
||||
class LLMEndpoint:
|
||||
"""Represents an LLM endpoint"""
|
||||
model_name: str
|
||||
node_id: str
|
||||
base_url: str
|
||||
is_available: bool = True
|
||||
|
||||
@property
|
||||
def endpoint_url(self) -> str:
|
||||
"""Get full endpoint URL"""
|
||||
return f"{self.base_url}/{self.model_name}"
|
||||
|
||||
|
||||
class LLMPoolManager:
|
||||
"""
|
||||
Pool of LLM Endpoints
|
||||
|
||||
Each model is exposed via an OpenAI-compatible endpoint:
|
||||
- http://gb10-node-1:8001/deepseek
|
||||
- http://gb10-node-1:8001/qwen72
|
||||
- http://gb10-node-2:8001/phi4
|
||||
- http://gb10-node-2:8001/qwen-coder
|
||||
- http://gb10-node-2:8001/llama31
|
||||
- http://gb10-node-2:8001/granite-guardian
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize LLM pool"""
|
||||
self.endpoints: Dict[str, LLMEndpoint] = {}
|
||||
self._initialize_endpoints()
|
||||
|
||||
def _initialize_endpoints(self):
|
||||
"""Initialize all LLM endpoints"""
|
||||
# GB10 Node 1 endpoints
|
||||
self.endpoints["deepseek"] = LLMEndpoint(
|
||||
model_name="deepseek",
|
||||
node_id="gb10-node-1",
|
||||
base_url="http://gb10-node-1:8001"
|
||||
)
|
||||
|
||||
self.endpoints["qwen72"] = LLMEndpoint(
|
||||
model_name="qwen72",
|
||||
node_id="gb10-node-1",
|
||||
base_url="http://gb10-node-1:8001"
|
||||
)
|
||||
|
||||
# GB10 Node 2 endpoints
|
||||
self.endpoints["phi4"] = LLMEndpoint(
|
||||
model_name="phi4",
|
||||
node_id="gb10-node-2",
|
||||
base_url="http://gb10-node-2:8001"
|
||||
)
|
||||
|
||||
self.endpoints["qwen-coder"] = LLMEndpoint(
|
||||
model_name="qwen-coder",
|
||||
node_id="gb10-node-2",
|
||||
base_url="http://gb10-node-2:8001"
|
||||
)
|
||||
|
||||
self.endpoints["llama31"] = LLMEndpoint(
|
||||
model_name="llama31",
|
||||
node_id="gb10-node-2",
|
||||
base_url="http://gb10-node-2:8001"
|
||||
)
|
||||
|
||||
self.endpoints["granite-guardian"] = LLMEndpoint(
|
||||
model_name="granite-guardian",
|
||||
node_id="gb10-node-2",
|
||||
base_url="http://gb10-node-2:8001"
|
||||
)
|
||||
|
||||
def get_endpoint(self, model_name: str) -> Optional[LLMEndpoint]:
|
||||
"""
|
||||
Get endpoint for a specific model
|
||||
|
||||
Args:
|
||||
model_name: Name of the model
|
||||
|
||||
Returns:
|
||||
LLMEndpoint or None if not found
|
||||
"""
|
||||
return self.endpoints.get(model_name)
|
||||
|
||||
async def call_model(
|
||||
self,
|
||||
model_name: str,
|
||||
prompt: str,
|
||||
parameters: Optional[Dict[str, Any]] = None
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Call an LLM model via its endpoint
|
||||
|
||||
Args:
|
||||
model_name: Name of the model
|
||||
prompt: Input prompt
|
||||
parameters: Optional model parameters
|
||||
|
||||
Returns:
|
||||
Model response
|
||||
"""
|
||||
endpoint = self.get_endpoint(model_name)
|
||||
|
||||
if not endpoint:
|
||||
return {
|
||||
"error": f"Model {model_name} not found",
|
||||
"available_models": list(self.endpoints.keys())
|
||||
}
|
||||
|
||||
if not endpoint.is_available:
|
||||
return {
|
||||
"error": f"Endpoint {model_name} is currently unavailable",
|
||||
"status": "offline"
|
||||
}
|
||||
|
||||
# Prepare OpenAI-compatible request
|
||||
payload = {
|
||||
"model": model_name,
|
||||
"messages": [
|
||||
{"role": "user", "content": prompt}
|
||||
],
|
||||
"temperature": parameters.get("temperature", 0.7) if parameters else 0.7,
|
||||
"max_tokens": parameters.get("max_tokens", 2048) if parameters else 2048
|
||||
}
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=60.0) as client:
|
||||
response = await client.post(
|
||||
f"{endpoint.endpoint_url}/v1/chat/completions",
|
||||
json=payload
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except httpx.HTTPError as e:
|
||||
return {
|
||||
"error": f"Failed to call {model_name}",
|
||||
"details": str(e),
|
||||
"endpoint": endpoint.endpoint_url
|
||||
}
|
||||
|
||||
async def call_multiple_models(
|
||||
self,
|
||||
model_names: List[str],
|
||||
prompt: str,
|
||||
parameters: Optional[Dict[str, Any]] = None
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Call multiple models in parallel
|
||||
|
||||
Args:
|
||||
model_names: List of model names
|
||||
prompt: Input prompt
|
||||
parameters: Optional model parameters
|
||||
|
||||
Returns:
|
||||
Combined results from all models
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
tasks = [
|
||||
self.call_model(model, prompt, parameters)
|
||||
for model in model_names
|
||||
]
|
||||
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
return {
|
||||
"models": model_names,
|
||||
"results": [
|
||||
{"model": model, "response": result}
|
||||
for model, result in zip(model_names, results)
|
||||
]
|
||||
}
|
||||
|
||||
def list_available_models(self) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
List all available models
|
||||
|
||||
Returns:
|
||||
List of model information
|
||||
"""
|
||||
return [
|
||||
{
|
||||
"model_name": endpoint.model_name,
|
||||
"node_id": endpoint.node_id,
|
||||
"endpoint_url": endpoint.endpoint_url,
|
||||
"is_available": endpoint.is_available
|
||||
}
|
||||
for endpoint in self.endpoints.values()
|
||||
]
|
||||
|
||||
|
||||
def get_llm_pool() -> LLMPoolManager:
|
||||
"""
|
||||
Factory function to create LLM pool manager
|
||||
|
||||
Returns:
|
||||
Configured LLMPoolManager instance
|
||||
"""
|
||||
return LLMPoolManager()
|
||||
@@ -1,187 +0,0 @@
|
||||
"""
|
||||
LLM Router Agent
|
||||
|
||||
Routes requests to appropriate LLM models based on task classification.
|
||||
"""
|
||||
|
||||
from typing import Dict, Any, Optional, List
|
||||
from enum import Enum
|
||||
import httpx
|
||||
|
||||
|
||||
class TaskType(Enum):
|
||||
"""Types of tasks for LLM routing"""
|
||||
GENERAL_REASONING = "general_reasoning" # DeepSeek
|
||||
MULTILINGUAL = "multilingual" # Qwen / Aya
|
||||
STRUCTURED_PARSING = "structured_parsing" # Phi-4
|
||||
RULE_GENERATION = "rule_generation" # Qwen-Coder
|
||||
ADVERSARIAL_REASONING = "adversarial_reasoning" # LLaMA 3.1
|
||||
CLASSIFICATION = "classification" # Granite Guardian
|
||||
|
||||
|
||||
class LLMRouterAgent:
|
||||
"""
|
||||
Router Agent - Interprets incoming requests and routes to appropriate LLM
|
||||
|
||||
This agent classifies the incoming request and determines which specialized
|
||||
LLM should handle it based on the task type.
|
||||
"""
|
||||
|
||||
def __init__(self, policy_config: Optional[Dict[str, Any]] = None):
|
||||
"""
|
||||
Initialize router agent
|
||||
|
||||
Args:
|
||||
policy_config: Optional routing policy configuration
|
||||
"""
|
||||
self.policy_config = policy_config or {}
|
||||
self.routing_rules = self._initialize_routing_rules()
|
||||
|
||||
def _initialize_routing_rules(self) -> Dict[TaskType, Dict[str, Any]]:
|
||||
"""Initialize routing rules for each task type"""
|
||||
return {
|
||||
TaskType.GENERAL_REASONING: {
|
||||
"model": "deepseek",
|
||||
"endpoint": "deepseek",
|
||||
"priority": 1,
|
||||
"description": "General reasoning and complex analysis"
|
||||
},
|
||||
TaskType.MULTILINGUAL: {
|
||||
"model": "qwen72",
|
||||
"endpoint": "qwen72",
|
||||
"priority": 2,
|
||||
"description": "Multilingual translation and analysis"
|
||||
},
|
||||
TaskType.STRUCTURED_PARSING: {
|
||||
"model": "phi4",
|
||||
"endpoint": "phi4",
|
||||
"priority": 3,
|
||||
"description": "Structured data parsing and extraction"
|
||||
},
|
||||
TaskType.RULE_GENERATION: {
|
||||
"model": "qwen-coder",
|
||||
"endpoint": "qwen-coder",
|
||||
"priority": 2,
|
||||
"description": "Code and rule generation"
|
||||
},
|
||||
TaskType.ADVERSARIAL_REASONING: {
|
||||
"model": "llama31",
|
||||
"endpoint": "llama31",
|
||||
"priority": 1,
|
||||
"description": "Adversarial threat analysis"
|
||||
},
|
||||
TaskType.CLASSIFICATION: {
|
||||
"model": "granite-guardian",
|
||||
"endpoint": "granite-guardian",
|
||||
"priority": 4,
|
||||
"description": "Pure classification tasks"
|
||||
}
|
||||
}
|
||||
|
||||
def classify_request(self, request: Dict[str, Any]) -> TaskType:
|
||||
"""
|
||||
Classify incoming request to determine task type
|
||||
|
||||
Args:
|
||||
request: Request containing prompt and metadata
|
||||
|
||||
Returns:
|
||||
Classified task type
|
||||
"""
|
||||
prompt = request.get("prompt", "").lower()
|
||||
task_hints = request.get("task_hints", [])
|
||||
|
||||
# Classification logic based on keywords and hints
|
||||
if any(hint in task_hints for hint in ["translate", "multilingual", "language"]):
|
||||
return TaskType.MULTILINGUAL
|
||||
|
||||
if any(hint in task_hints for hint in ["parse", "extract", "structure"]):
|
||||
return TaskType.STRUCTURED_PARSING
|
||||
|
||||
if any(hint in task_hints for hint in ["code", "rule", "generate", "script"]):
|
||||
return TaskType.RULE_GENERATION
|
||||
|
||||
if any(hint in task_hints for hint in ["threat", "adversary", "attack", "malicious"]):
|
||||
return TaskType.ADVERSARIAL_REASONING
|
||||
|
||||
if any(hint in task_hints for hint in ["classify", "categorize", "label"]):
|
||||
return TaskType.CLASSIFICATION
|
||||
|
||||
# Default to general reasoning
|
||||
return TaskType.GENERAL_REASONING
|
||||
|
||||
def route_request(self, request: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""
|
||||
Route request to appropriate LLM endpoint
|
||||
|
||||
Args:
|
||||
request: Request to route
|
||||
|
||||
Returns:
|
||||
Routing decision with endpoint and model info
|
||||
"""
|
||||
task_type = self.classify_request(request)
|
||||
routing_rule = self.routing_rules[task_type]
|
||||
|
||||
return {
|
||||
"task_type": task_type.value,
|
||||
"model": routing_rule["model"],
|
||||
"endpoint": routing_rule["endpoint"],
|
||||
"priority": routing_rule["priority"],
|
||||
"description": routing_rule["description"],
|
||||
"requires_parallel": request.get("requires_parallel", False),
|
||||
"requires_chaining": request.get("requires_chaining", False)
|
||||
}
|
||||
|
||||
def should_parallelize(self, request: Dict[str, Any]) -> bool:
|
||||
"""
|
||||
Determine if request should be parallelized across multiple nodes
|
||||
|
||||
Args:
|
||||
request: Request to evaluate
|
||||
|
||||
Returns:
|
||||
True if should be parallelized
|
||||
"""
|
||||
# Large batch requests
|
||||
if request.get("batch_size", 1) > 10:
|
||||
return True
|
||||
|
||||
# Explicit parallel flag
|
||||
if request.get("requires_parallel", False):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def requires_serial_chaining(self, request: Dict[str, Any]) -> bool:
|
||||
"""
|
||||
Determine if request requires serial reasoning (chained operations)
|
||||
|
||||
Args:
|
||||
request: Request to evaluate
|
||||
|
||||
Returns:
|
||||
True if requires chaining
|
||||
"""
|
||||
# Complex multi-step reasoning
|
||||
if request.get("requires_chaining", False):
|
||||
return True
|
||||
|
||||
# Multiple dependent operations
|
||||
if len(request.get("operations", [])) > 1:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def get_llm_router(policy_config: Optional[Dict[str, Any]] = None) -> LLMRouterAgent:
|
||||
"""
|
||||
Factory function to create LLM router agent
|
||||
|
||||
Args:
|
||||
policy_config: Optional routing policy configuration
|
||||
|
||||
Returns:
|
||||
Configured LLMRouterAgent instance
|
||||
"""
|
||||
return LLMRouterAgent(policy_config)
|
||||
@@ -1,259 +0,0 @@
|
||||
"""
|
||||
Merger Agent
|
||||
|
||||
Combines and synthesizes results from multiple LLM models.
|
||||
"""
|
||||
|
||||
from typing import Dict, Any, List, Optional
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class MergeStrategy(Enum):
|
||||
"""Strategies for merging LLM results"""
|
||||
CONSENSUS = "consensus" # Take majority vote
|
||||
WEIGHTED = "weighted" # Weight by model confidence
|
||||
CONCATENATE = "concatenate" # Combine all outputs
|
||||
BEST_QUALITY = "best_quality" # Select highest quality response
|
||||
ENSEMBLE = "ensemble" # Ensemble multiple results
|
||||
|
||||
|
||||
class MergerAgent:
|
||||
"""
|
||||
Merger Agent - Combines results from multiple LLM executions
|
||||
|
||||
When multiple models process the same or related requests,
|
||||
this agent intelligently merges their outputs into a coherent response.
|
||||
"""
|
||||
|
||||
def __init__(self, default_strategy: MergeStrategy = MergeStrategy.CONSENSUS):
|
||||
"""
|
||||
Initialize merger agent
|
||||
|
||||
Args:
|
||||
default_strategy: Default merging strategy
|
||||
"""
|
||||
self.default_strategy = default_strategy
|
||||
|
||||
def merge_results(
|
||||
self,
|
||||
results: List[Dict[str, Any]],
|
||||
strategy: Optional[MergeStrategy] = None
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Merge multiple LLM results using specified strategy
|
||||
|
||||
Args:
|
||||
results: List of results from different models
|
||||
strategy: Merging strategy (uses default if not specified)
|
||||
|
||||
Returns:
|
||||
Merged result
|
||||
"""
|
||||
if not results:
|
||||
return {"error": "No results to merge"}
|
||||
|
||||
if len(results) == 1:
|
||||
return results[0]
|
||||
|
||||
merge_strategy = strategy or self.default_strategy
|
||||
|
||||
if merge_strategy == MergeStrategy.CONSENSUS:
|
||||
return self._merge_consensus(results)
|
||||
elif merge_strategy == MergeStrategy.WEIGHTED:
|
||||
return self._merge_weighted(results)
|
||||
elif merge_strategy == MergeStrategy.CONCATENATE:
|
||||
return self._merge_concatenate(results)
|
||||
elif merge_strategy == MergeStrategy.BEST_QUALITY:
|
||||
return self._merge_best_quality(results)
|
||||
elif merge_strategy == MergeStrategy.ENSEMBLE:
|
||||
return self._merge_ensemble(results)
|
||||
else:
|
||||
return self._merge_consensus(results)
|
||||
|
||||
def _merge_consensus(self, results: List[Dict[str, Any]]) -> Dict[str, Any]:
|
||||
"""
|
||||
Merge by consensus - take majority vote or most common response
|
||||
|
||||
Args:
|
||||
results: List of results
|
||||
|
||||
Returns:
|
||||
Consensus result
|
||||
"""
|
||||
# For classification tasks, take majority vote
|
||||
if all("classification" in r for r in results):
|
||||
classifications = [r["classification"] for r in results]
|
||||
most_common = max(set(classifications), key=classifications.count)
|
||||
|
||||
return {
|
||||
"strategy": "consensus",
|
||||
"result": most_common,
|
||||
"confidence": classifications.count(most_common) / len(classifications),
|
||||
"votes": dict((k, classifications.count(k)) for k in set(classifications))
|
||||
}
|
||||
|
||||
# For text generation, use first high-quality result
|
||||
valid_results = [r for r in results if "response" in r and r["response"]]
|
||||
if valid_results:
|
||||
return {
|
||||
"strategy": "consensus",
|
||||
"result": valid_results[0]["response"],
|
||||
"num_models": len(results)
|
||||
}
|
||||
|
||||
return {"strategy": "consensus", "result": None}
|
||||
|
||||
def _merge_weighted(self, results: List[Dict[str, Any]]) -> Dict[str, Any]:
|
||||
"""
|
||||
Merge by weighted average based on model confidence
|
||||
|
||||
Args:
|
||||
results: List of results with confidence scores
|
||||
|
||||
Returns:
|
||||
Weighted result
|
||||
"""
|
||||
# Extract results with confidence scores
|
||||
weighted_results = [
|
||||
(r.get("response", ""), r.get("confidence", 0.5))
|
||||
for r in results
|
||||
]
|
||||
|
||||
# Sort by confidence
|
||||
weighted_results.sort(key=lambda x: x[1], reverse=True)
|
||||
|
||||
return {
|
||||
"strategy": "weighted",
|
||||
"result": weighted_results[0][0],
|
||||
"confidence": weighted_results[0][1],
|
||||
"all_results": [
|
||||
{"response": resp, "confidence": conf}
|
||||
for resp, conf in weighted_results
|
||||
]
|
||||
}
|
||||
|
||||
def _merge_concatenate(self, results: List[Dict[str, Any]]) -> Dict[str, Any]:
|
||||
"""
|
||||
Concatenate all results
|
||||
|
||||
Args:
|
||||
results: List of results
|
||||
|
||||
Returns:
|
||||
Concatenated result
|
||||
"""
|
||||
responses = [
|
||||
r.get("response", "") for r in results if r.get("response")
|
||||
]
|
||||
|
||||
return {
|
||||
"strategy": "concatenate",
|
||||
"result": "\n\n---\n\n".join(responses),
|
||||
"num_responses": len(responses)
|
||||
}
|
||||
|
||||
def _merge_best_quality(self, results: List[Dict[str, Any]]) -> Dict[str, Any]:
|
||||
"""
|
||||
Select the highest quality response
|
||||
|
||||
Args:
|
||||
results: List of results
|
||||
|
||||
Returns:
|
||||
Best quality result
|
||||
"""
|
||||
# Score responses by length and presence of key indicators
|
||||
scored_results = []
|
||||
|
||||
for r in results:
|
||||
response = r.get("response", "")
|
||||
if not response:
|
||||
continue
|
||||
|
||||
score = 0
|
||||
score += len(response) / 100 # Longer responses get higher score
|
||||
score += response.count(".") * 0.5 # More complete sentences
|
||||
score += response.count("\n") * 0.3 # Better formatting
|
||||
|
||||
scored_results.append((response, score, r))
|
||||
|
||||
if not scored_results:
|
||||
return {"strategy": "best_quality", "result": None}
|
||||
|
||||
scored_results.sort(key=lambda x: x[1], reverse=True)
|
||||
best_response, best_score, best_result = scored_results[0]
|
||||
|
||||
return {
|
||||
"strategy": "best_quality",
|
||||
"result": best_response,
|
||||
"quality_score": best_score,
|
||||
"model": best_result.get("model", "unknown")
|
||||
}
|
||||
|
||||
def _merge_ensemble(self, results: List[Dict[str, Any]]) -> Dict[str, Any]:
|
||||
"""
|
||||
Ensemble multiple results by combining their insights
|
||||
|
||||
Args:
|
||||
results: List of results
|
||||
|
||||
Returns:
|
||||
Ensemble result
|
||||
"""
|
||||
# Collect unique insights from all models
|
||||
all_responses = [r.get("response", "") for r in results if r.get("response")]
|
||||
|
||||
# Create ensemble summary
|
||||
ensemble_summary = {
|
||||
"strategy": "ensemble",
|
||||
"num_models": len(results),
|
||||
"individual_results": [
|
||||
{
|
||||
"model": r.get("model", "unknown"),
|
||||
"response": r.get("response", ""),
|
||||
"confidence": r.get("confidence", 0.5)
|
||||
}
|
||||
for r in results
|
||||
],
|
||||
"synthesized_result": self._synthesize_insights(all_responses)
|
||||
}
|
||||
|
||||
return ensemble_summary
|
||||
|
||||
def _synthesize_insights(self, responses: List[str]) -> str:
|
||||
"""
|
||||
Synthesize insights from multiple responses
|
||||
|
||||
Args:
|
||||
responses: List of response strings
|
||||
|
||||
Returns:
|
||||
Synthesized summary
|
||||
"""
|
||||
if not responses:
|
||||
return ""
|
||||
|
||||
# Simple synthesis - in production, use another LLM to synthesize
|
||||
unique_points = []
|
||||
for response in responses:
|
||||
sentences = response.split(". ")
|
||||
for sentence in sentences:
|
||||
if sentence and sentence not in unique_points:
|
||||
unique_points.append(sentence)
|
||||
|
||||
return ". ".join(unique_points[:10]) # Top 10 insights
|
||||
|
||||
|
||||
def get_merger_agent(
|
||||
default_strategy: MergeStrategy = MergeStrategy.CONSENSUS
|
||||
) -> MergerAgent:
|
||||
"""
|
||||
Factory function to create merger agent
|
||||
|
||||
Args:
|
||||
default_strategy: Default merging strategy
|
||||
|
||||
Returns:
|
||||
Configured MergerAgent instance
|
||||
"""
|
||||
return MergerAgent(default_strategy)
|
||||
@@ -1,165 +0,0 @@
|
||||
"""
|
||||
Playbook Execution Engine
|
||||
|
||||
Executes automated response playbooks based on triggers.
|
||||
"""
|
||||
|
||||
from typing import Dict, Any, List
|
||||
from datetime import datetime, timezone
|
||||
import asyncio
|
||||
|
||||
|
||||
class PlaybookEngine:
|
||||
"""Engine for executing playbooks"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize playbook engine"""
|
||||
self.actions_registry = {
|
||||
"send_notification": self._action_send_notification,
|
||||
"create_case": self._action_create_case,
|
||||
"isolate_host": self._action_isolate_host,
|
||||
"collect_artifact": self._action_collect_artifact,
|
||||
"block_ip": self._action_block_ip,
|
||||
"send_email": self._action_send_email,
|
||||
}
|
||||
|
||||
async def execute_playbook(
|
||||
self,
|
||||
playbook: Dict[str, Any],
|
||||
context: Dict[str, Any]
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Execute a playbook
|
||||
|
||||
Args:
|
||||
playbook: Playbook definition
|
||||
context: Execution context with relevant data
|
||||
|
||||
Returns:
|
||||
Execution result
|
||||
"""
|
||||
results = []
|
||||
errors = []
|
||||
|
||||
actions = playbook.get("actions", [])
|
||||
|
||||
for action in actions:
|
||||
action_type = action.get("type")
|
||||
action_params = action.get("params", {})
|
||||
|
||||
try:
|
||||
# Get action handler
|
||||
handler = self.actions_registry.get(action_type)
|
||||
if not handler:
|
||||
errors.append(f"Unknown action type: {action_type}")
|
||||
continue
|
||||
|
||||
# Execute action
|
||||
result = await handler(action_params, context)
|
||||
results.append({
|
||||
"action": action_type,
|
||||
"status": "success",
|
||||
"result": result
|
||||
})
|
||||
except Exception as e:
|
||||
errors.append(f"Error in action {action_type}: {str(e)}")
|
||||
results.append({
|
||||
"action": action_type,
|
||||
"status": "failed",
|
||||
"error": str(e)
|
||||
})
|
||||
|
||||
return {
|
||||
"status": "completed" if not errors else "completed_with_errors",
|
||||
"results": results,
|
||||
"errors": errors
|
||||
}
|
||||
|
||||
async def _action_send_notification(
|
||||
self,
|
||||
params: Dict[str, Any],
|
||||
context: Dict[str, Any]
|
||||
) -> Dict[str, Any]:
|
||||
"""Send a notification"""
|
||||
# In production, this would create a notification in the database
|
||||
# and push it via WebSocket
|
||||
return {
|
||||
"notification_sent": True,
|
||||
"message": params.get("message", "Playbook notification")
|
||||
}
|
||||
|
||||
async def _action_create_case(
|
||||
self,
|
||||
params: Dict[str, Any],
|
||||
context: Dict[str, Any]
|
||||
) -> Dict[str, Any]:
|
||||
"""Create a new case"""
|
||||
# In production, this would create a case in the database
|
||||
return {
|
||||
"case_created": True,
|
||||
"case_title": params.get("title", "Automated Case"),
|
||||
"case_id": "placeholder_id"
|
||||
}
|
||||
|
||||
async def _action_isolate_host(
|
||||
self,
|
||||
params: Dict[str, Any],
|
||||
context: Dict[str, Any]
|
||||
) -> Dict[str, Any]:
|
||||
"""Isolate a host"""
|
||||
# In production, this would call Velociraptor or other tools
|
||||
# to isolate the host from the network
|
||||
host_id = params.get("host_id")
|
||||
return {
|
||||
"host_isolated": True,
|
||||
"host_id": host_id
|
||||
}
|
||||
|
||||
async def _action_collect_artifact(
|
||||
self,
|
||||
params: Dict[str, Any],
|
||||
context: Dict[str, Any]
|
||||
) -> Dict[str, Any]:
|
||||
"""Collect an artifact from a host"""
|
||||
# In production, this would trigger Velociraptor collection
|
||||
return {
|
||||
"collection_started": True,
|
||||
"artifact": params.get("artifact_name"),
|
||||
"client_id": params.get("client_id")
|
||||
}
|
||||
|
||||
async def _action_block_ip(
|
||||
self,
|
||||
params: Dict[str, Any],
|
||||
context: Dict[str, Any]
|
||||
) -> Dict[str, Any]:
|
||||
"""Block an IP address"""
|
||||
# In production, this would update firewall rules
|
||||
ip_address = params.get("ip_address")
|
||||
return {
|
||||
"ip_blocked": True,
|
||||
"ip_address": ip_address
|
||||
}
|
||||
|
||||
async def _action_send_email(
|
||||
self,
|
||||
params: Dict[str, Any],
|
||||
context: Dict[str, Any]
|
||||
) -> Dict[str, Any]:
|
||||
"""Send an email"""
|
||||
# In production, this would send an actual email
|
||||
return {
|
||||
"email_sent": True,
|
||||
"to": params.get("to"),
|
||||
"subject": params.get("subject")
|
||||
}
|
||||
|
||||
|
||||
def get_playbook_engine() -> PlaybookEngine:
|
||||
"""
|
||||
Factory function to create playbook engine
|
||||
|
||||
Returns:
|
||||
Configured PlaybookEngine instance
|
||||
"""
|
||||
return PlaybookEngine()
|
||||
@@ -1,120 +0,0 @@
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Optional
|
||||
import secrets
|
||||
import pyotp
|
||||
from jose import JWTError, jwt
|
||||
from passlib.context import CryptContext
|
||||
|
||||
from app.core.config import settings
|
||||
|
||||
# Password hashing context
|
||||
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
|
||||
|
||||
|
||||
def verify_password(plain_password: str, hashed_password: str) -> bool:
|
||||
"""Verify a plain password against a hashed password"""
|
||||
return pwd_context.verify(plain_password, hashed_password)
|
||||
|
||||
|
||||
def get_password_hash(password: str) -> str:
|
||||
"""Hash a password using bcrypt"""
|
||||
return pwd_context.hash(password)
|
||||
|
||||
|
||||
def create_access_token(data: dict, expires_delta: Optional[timedelta] = None) -> str:
|
||||
"""
|
||||
Create a JWT access token
|
||||
|
||||
Args:
|
||||
data: Dictionary containing user_id, tenant_id, role
|
||||
expires_delta: Optional expiration time delta
|
||||
|
||||
Returns:
|
||||
Encoded JWT token
|
||||
"""
|
||||
to_encode = data.copy()
|
||||
if expires_delta:
|
||||
expire = datetime.now(timezone.utc) + expires_delta
|
||||
else:
|
||||
expire = datetime.now(timezone.utc) + timedelta(minutes=settings.access_token_expire_minutes)
|
||||
|
||||
to_encode.update({"exp": expire})
|
||||
encoded_jwt = jwt.encode(to_encode, settings.secret_key, algorithm=settings.algorithm)
|
||||
return encoded_jwt
|
||||
|
||||
|
||||
def verify_token(token: str) -> Optional[dict]:
|
||||
"""
|
||||
Verify and decode a JWT token
|
||||
|
||||
Args:
|
||||
token: JWT token string
|
||||
|
||||
Returns:
|
||||
Decoded token payload or None if invalid
|
||||
"""
|
||||
try:
|
||||
payload = jwt.decode(token, settings.secret_key, algorithms=[settings.algorithm])
|
||||
return payload
|
||||
except JWTError:
|
||||
return None
|
||||
|
||||
|
||||
def create_refresh_token() -> str:
|
||||
"""
|
||||
Create a secure random refresh token
|
||||
|
||||
Returns:
|
||||
Random token string
|
||||
"""
|
||||
return secrets.token_urlsafe(32)
|
||||
|
||||
|
||||
def create_reset_token() -> str:
|
||||
"""
|
||||
Create a secure random password reset token
|
||||
|
||||
Returns:
|
||||
Random token string
|
||||
"""
|
||||
return secrets.token_urlsafe(32)
|
||||
|
||||
|
||||
def generate_totp_secret() -> str:
|
||||
"""
|
||||
Generate a TOTP secret for 2FA
|
||||
|
||||
Returns:
|
||||
Base32 encoded secret
|
||||
"""
|
||||
return pyotp.random_base32()
|
||||
|
||||
|
||||
def verify_totp(secret: str, code: str) -> bool:
|
||||
"""
|
||||
Verify a TOTP code
|
||||
|
||||
Args:
|
||||
secret: TOTP secret
|
||||
code: 6-digit code from authenticator app
|
||||
|
||||
Returns:
|
||||
True if code is valid
|
||||
"""
|
||||
totp = pyotp.TOTP(secret)
|
||||
return totp.verify(code, valid_window=1)
|
||||
|
||||
|
||||
def get_totp_uri(secret: str, username: str) -> str:
|
||||
"""
|
||||
Get TOTP provisioning URI for QR code
|
||||
|
||||
Args:
|
||||
secret: TOTP secret
|
||||
username: User's username
|
||||
|
||||
Returns:
|
||||
otpauth:// URI
|
||||
"""
|
||||
totp = pyotp.TOTP(secret)
|
||||
return totp.provisioning_uri(name=username, issuer_name="VelociCompanion")
|
||||
@@ -1,198 +0,0 @@
|
||||
"""
|
||||
Threat Intelligence and Machine Learning Module
|
||||
|
||||
This module provides threat scoring, anomaly detection, and predictive analytics.
|
||||
"""
|
||||
|
||||
from typing import Dict, Any, List, Optional
|
||||
import random # For demo purposes - would use actual ML models in production
|
||||
|
||||
|
||||
class ThreatAnalyzer:
|
||||
"""Analyzes threats using ML models and heuristics"""
|
||||
|
||||
def __init__(self, model_version: str = "1.0"):
|
||||
"""
|
||||
Initialize threat analyzer
|
||||
|
||||
Args:
|
||||
model_version: Version of ML models to use
|
||||
"""
|
||||
self.model_version = model_version
|
||||
|
||||
def analyze_host(self, host_data: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""
|
||||
Analyze a host for threats
|
||||
|
||||
Args:
|
||||
host_data: Host information and telemetry
|
||||
|
||||
Returns:
|
||||
Dictionary with threat score and indicators
|
||||
"""
|
||||
# In production, this would use ML models
|
||||
# For demo, using simple heuristics
|
||||
|
||||
score = 0.0
|
||||
confidence = 0.8
|
||||
indicators = []
|
||||
|
||||
# Check for suspicious patterns
|
||||
hostname = host_data.get("hostname", "")
|
||||
if "temp" in hostname.lower() or "test" in hostname.lower():
|
||||
score += 0.2
|
||||
indicators.append({
|
||||
"type": "suspicious_hostname",
|
||||
"description": "Hostname contains suspicious keywords",
|
||||
"severity": "low"
|
||||
})
|
||||
|
||||
# Check metadata for anomalies
|
||||
metadata = host_data.get("host_metadata", {})
|
||||
if metadata:
|
||||
# Check for unusual processes, connections, etc.
|
||||
if "suspicious_process" in str(metadata):
|
||||
score += 0.5
|
||||
indicators.append({
|
||||
"type": "suspicious_process",
|
||||
"description": "Unusual process detected",
|
||||
"severity": "high"
|
||||
})
|
||||
|
||||
# Normalize score
|
||||
score = min(score, 1.0)
|
||||
|
||||
return {
|
||||
"score": score,
|
||||
"confidence": confidence,
|
||||
"threat_type": self._classify_threat(score),
|
||||
"indicators": indicators,
|
||||
"ml_model_version": self.model_version
|
||||
}
|
||||
|
||||
def analyze_artifact(self, artifact_data: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""
|
||||
Analyze an artifact for threats
|
||||
|
||||
Args:
|
||||
artifact_data: Artifact information
|
||||
|
||||
Returns:
|
||||
Dictionary with threat score and indicators
|
||||
"""
|
||||
score = 0.0
|
||||
confidence = 0.7
|
||||
indicators = []
|
||||
|
||||
artifact_type = artifact_data.get("artifact_type", "")
|
||||
value = artifact_data.get("value", "")
|
||||
|
||||
# Hash analysis
|
||||
if artifact_type == "hash":
|
||||
# In production, check against threat intelligence feeds
|
||||
if len(value) == 32: # MD5
|
||||
score += 0.3
|
||||
indicators.append({
|
||||
"type": "weak_hash",
|
||||
"description": "MD5 hashes are considered weak",
|
||||
"severity": "low"
|
||||
})
|
||||
|
||||
# IP analysis
|
||||
elif artifact_type == "ip":
|
||||
# Check if IP is in known malicious ranges
|
||||
if value.startswith("10.") or value.startswith("192.168."):
|
||||
score += 0.1 # Private IP, lower risk
|
||||
else:
|
||||
score += 0.4 # Public IP, higher scrutiny
|
||||
indicators.append({
|
||||
"type": "public_ip",
|
||||
"description": "Communication with public IP",
|
||||
"severity": "medium"
|
||||
})
|
||||
|
||||
# Domain analysis
|
||||
elif artifact_type == "domain":
|
||||
# Check for suspicious TLDs or patterns
|
||||
suspicious_tlds = [".ru", ".cn", ".tk", ".xyz"]
|
||||
if any(value.endswith(tld) for tld in suspicious_tlds):
|
||||
score += 0.6
|
||||
indicators.append({
|
||||
"type": "suspicious_tld",
|
||||
"description": f"Domain uses potentially suspicious TLD",
|
||||
"severity": "high"
|
||||
})
|
||||
|
||||
score = min(score, 1.0)
|
||||
|
||||
return {
|
||||
"score": score,
|
||||
"confidence": confidence,
|
||||
"threat_type": self._classify_threat(score),
|
||||
"indicators": indicators,
|
||||
"ml_model_version": self.model_version
|
||||
}
|
||||
|
||||
def detect_anomalies(
|
||||
self,
|
||||
historical_data: List[Dict[str, Any]],
|
||||
current_data: Dict[str, Any]
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Detect anomalies in current data compared to historical baseline
|
||||
|
||||
Args:
|
||||
historical_data: Historical baseline data
|
||||
current_data: Current data to analyze
|
||||
|
||||
Returns:
|
||||
Anomaly detection results
|
||||
"""
|
||||
# Simple anomaly detection based on statistical deviation
|
||||
# In production, use more sophisticated methods
|
||||
|
||||
anomalies = []
|
||||
score = 0.0
|
||||
|
||||
# Compare metrics
|
||||
if historical_data and len(historical_data) >= 3:
|
||||
# Calculate baseline
|
||||
# This is a simplified example
|
||||
anomalies.append({
|
||||
"type": "behavioral_anomaly",
|
||||
"description": "Behavior deviates from baseline",
|
||||
"severity": "medium"
|
||||
})
|
||||
score = 0.5
|
||||
|
||||
return {
|
||||
"is_anomaly": score > 0.4,
|
||||
"anomaly_score": score,
|
||||
"anomalies": anomalies
|
||||
}
|
||||
|
||||
def _classify_threat(self, score: float) -> str:
|
||||
"""Classify threat based on score"""
|
||||
if score >= 0.8:
|
||||
return "critical"
|
||||
elif score >= 0.6:
|
||||
return "high"
|
||||
elif score >= 0.4:
|
||||
return "medium"
|
||||
elif score >= 0.2:
|
||||
return "low"
|
||||
else:
|
||||
return "benign"
|
||||
|
||||
|
||||
def get_threat_analyzer(model_version: str = "1.0") -> ThreatAnalyzer:
|
||||
"""
|
||||
Factory function to create threat analyzer
|
||||
|
||||
Args:
|
||||
model_version: Version of ML models
|
||||
|
||||
Returns:
|
||||
Configured ThreatAnalyzer instance
|
||||
"""
|
||||
return ThreatAnalyzer(model_version)
|
||||
@@ -1,194 +0,0 @@
|
||||
"""
|
||||
Velociraptor API Client
|
||||
|
||||
This module provides integration with Velociraptor servers for artifact
|
||||
collection, hunt management, and client operations.
|
||||
"""
|
||||
|
||||
from typing import List, Dict, Any, Optional
|
||||
import httpx
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
class VelociraptorClient:
|
||||
"""Client for interacting with Velociraptor API"""
|
||||
|
||||
def __init__(self, base_url: str, api_key: str):
|
||||
"""
|
||||
Initialize Velociraptor client
|
||||
|
||||
Args:
|
||||
base_url: Base URL of Velociraptor server
|
||||
api_key: API key for authentication
|
||||
"""
|
||||
self.base_url = base_url.rstrip('/')
|
||||
self.api_key = api_key
|
||||
self.headers = {
|
||||
"Authorization": f"******",
|
||||
"Content-Type": "application/json"
|
||||
}
|
||||
|
||||
async def list_clients(self, limit: int = 100) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
List all clients
|
||||
|
||||
Args:
|
||||
limit: Maximum number of clients to return
|
||||
|
||||
Returns:
|
||||
List of client information dictionaries
|
||||
"""
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.get(
|
||||
f"{self.base_url}/api/v1/clients",
|
||||
headers=self.headers,
|
||||
params={"count": limit}
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
async def get_client(self, client_id: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Get information about a specific client
|
||||
|
||||
Args:
|
||||
client_id: Client ID
|
||||
|
||||
Returns:
|
||||
Client information dictionary
|
||||
"""
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.get(
|
||||
f"{self.base_url}/api/v1/clients/{client_id}",
|
||||
headers=self.headers
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
async def collect_artifact(
|
||||
self,
|
||||
client_id: str,
|
||||
artifact_name: str,
|
||||
parameters: Optional[Dict[str, Any]] = None
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Collect an artifact from a client
|
||||
|
||||
Args:
|
||||
client_id: Client ID
|
||||
artifact_name: Name of artifact to collect
|
||||
parameters: Optional parameters for artifact collection
|
||||
|
||||
Returns:
|
||||
Collection flow information
|
||||
"""
|
||||
payload = {
|
||||
"client_id": client_id,
|
||||
"artifacts": [artifact_name],
|
||||
"parameters": parameters or {}
|
||||
}
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.post(
|
||||
f"{self.base_url}/api/v1/flows",
|
||||
headers=self.headers,
|
||||
json=payload
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
async def create_hunt(
|
||||
self,
|
||||
hunt_name: str,
|
||||
artifact_name: str,
|
||||
description: str,
|
||||
parameters: Optional[Dict[str, Any]] = None
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Create a new hunt
|
||||
|
||||
Args:
|
||||
hunt_name: Name of the hunt
|
||||
artifact_name: Artifact to collect
|
||||
description: Hunt description
|
||||
parameters: Optional parameters for artifact
|
||||
|
||||
Returns:
|
||||
Hunt information
|
||||
"""
|
||||
payload = {
|
||||
"hunt_name": hunt_name,
|
||||
"description": description,
|
||||
"artifacts": [artifact_name],
|
||||
"parameters": parameters or {}
|
||||
}
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.post(
|
||||
f"{self.base_url}/api/v1/hunts",
|
||||
headers=self.headers,
|
||||
json=payload
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
async def get_hunt_results(
|
||||
self,
|
||||
hunt_id: str,
|
||||
limit: int = 1000
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Get results from a hunt
|
||||
|
||||
Args:
|
||||
hunt_id: Hunt ID
|
||||
limit: Maximum number of results to return
|
||||
|
||||
Returns:
|
||||
List of hunt results
|
||||
"""
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.get(
|
||||
f"{self.base_url}/api/v1/hunts/{hunt_id}/results",
|
||||
headers=self.headers,
|
||||
params={"count": limit}
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
async def get_flow_results(
|
||||
self,
|
||||
client_id: str,
|
||||
flow_id: str
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Get results from a flow
|
||||
|
||||
Args:
|
||||
client_id: Client ID
|
||||
flow_id: Flow ID
|
||||
|
||||
Returns:
|
||||
List of flow results
|
||||
"""
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.get(
|
||||
f"{self.base_url}/api/v1/clients/{client_id}/flows/{flow_id}/results",
|
||||
headers=self.headers
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
|
||||
def get_velociraptor_client(base_url: str, api_key: str) -> VelociraptorClient:
|
||||
"""
|
||||
Factory function to create Velociraptor client
|
||||
|
||||
Args:
|
||||
base_url: Base URL of Velociraptor server
|
||||
api_key: API key for authentication
|
||||
|
||||
Returns:
|
||||
Configured VelociraptorClient instance
|
||||
"""
|
||||
return VelociraptorClient(base_url, api_key)
|
||||
@@ -1,65 +0,0 @@
|
||||
from fastapi import FastAPI
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
|
||||
from app.api.routes import (
|
||||
auth, users, tenants, hosts, ingestion, vt, audit,
|
||||
notifications, velociraptor, playbooks, threat_intel, reports, llm
|
||||
)
|
||||
from app.core.config import settings
|
||||
|
||||
app = FastAPI(
|
||||
title=settings.app_name,
|
||||
description="Multi-tenant threat hunting companion for Velociraptor with ML-powered threat detection and distributed LLM routing",
|
||||
version="1.1.0"
|
||||
)
|
||||
|
||||
# Configure CORS
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["http://localhost:3000", "http://localhost:5173"], # Frontend URLs
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
# Include routers
|
||||
app.include_router(auth.router, prefix="/api/auth", tags=["Authentication"])
|
||||
app.include_router(users.router, prefix="/api/users", tags=["Users"])
|
||||
app.include_router(tenants.router, prefix="/api/tenants", tags=["Tenants"])
|
||||
app.include_router(hosts.router, prefix="/api/hosts", tags=["Hosts"])
|
||||
app.include_router(ingestion.router, prefix="/api/ingestion", tags=["Ingestion"])
|
||||
app.include_router(vt.router, prefix="/api/vt", tags=["VirusTotal"])
|
||||
app.include_router(audit.router, prefix="/api/audit", tags=["Audit Logs"])
|
||||
app.include_router(notifications.router, prefix="/api/notifications", tags=["Notifications"])
|
||||
app.include_router(velociraptor.router, prefix="/api/velociraptor", tags=["Velociraptor"])
|
||||
app.include_router(playbooks.router, prefix="/api/playbooks", tags=["Playbooks"])
|
||||
app.include_router(threat_intel.router, prefix="/api/threat-intel", tags=["Threat Intelligence"])
|
||||
app.include_router(reports.router, prefix="/api/reports", tags=["Reports"])
|
||||
app.include_router(llm.router, prefix="/api/llm", tags=["Distributed LLM"])
|
||||
|
||||
|
||||
@app.get("/")
|
||||
async def root():
|
||||
"""Root endpoint"""
|
||||
return {
|
||||
"message": f"Welcome to {settings.app_name}",
|
||||
"version": "1.1.0",
|
||||
"docs": "/docs",
|
||||
"features": [
|
||||
"JWT Authentication with 2FA",
|
||||
"Multi-tenant isolation",
|
||||
"Audit logging",
|
||||
"Real-time notifications",
|
||||
"Velociraptor integration",
|
||||
"ML-powered threat detection",
|
||||
"Automated playbooks",
|
||||
"Advanced reporting",
|
||||
"Distributed LLM routing (Phase 5)"
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
async def health_check():
|
||||
"""Health check endpoint"""
|
||||
return {"status": "healthy"}
|
||||
@@ -1,20 +0,0 @@
|
||||
from sqlalchemy import Column, Integer, String, DateTime, ForeignKey, Text, JSON
|
||||
from sqlalchemy.orm import relationship
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from app.core.database import Base
|
||||
|
||||
|
||||
class Artifact(Base):
|
||||
__tablename__ = "artifacts"
|
||||
|
||||
id = Column(Integer, primary_key=True, index=True)
|
||||
artifact_type = Column(String, nullable=False) # hash, ip, domain, email, etc.
|
||||
value = Column(String, nullable=False)
|
||||
description = Column(Text, nullable=True)
|
||||
case_id = Column(Integer, ForeignKey("cases.id"), nullable=True)
|
||||
artifact_metadata = Column(JSON, nullable=True)
|
||||
created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc))
|
||||
|
||||
# Relationships
|
||||
case = relationship("Case", back_populates="artifacts")
|
||||
@@ -1,24 +0,0 @@
|
||||
from sqlalchemy import Column, Integer, String, DateTime, ForeignKey, Text, JSON
|
||||
from sqlalchemy.orm import relationship
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from app.core.database import Base
|
||||
|
||||
|
||||
class AuditLog(Base):
|
||||
__tablename__ = "audit_logs"
|
||||
|
||||
id = Column(Integer, primary_key=True, index=True)
|
||||
user_id = Column(Integer, ForeignKey("users.id"), nullable=True)
|
||||
tenant_id = Column(Integer, ForeignKey("tenants.id"), nullable=False)
|
||||
action = Column(String, nullable=False) # CREATE, READ, UPDATE, DELETE
|
||||
resource_type = Column(String, nullable=False) # user, host, case, etc.
|
||||
resource_id = Column(Integer, nullable=True)
|
||||
details = Column(JSON, nullable=True)
|
||||
ip_address = Column(String, nullable=True)
|
||||
user_agent = Column(String, nullable=True)
|
||||
created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc), index=True)
|
||||
|
||||
# Relationships
|
||||
user = relationship("User")
|
||||
tenant = relationship("Tenant")
|
||||
@@ -1,22 +0,0 @@
|
||||
from sqlalchemy import Column, Integer, String, DateTime, ForeignKey, Text
|
||||
from sqlalchemy.orm import relationship
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from app.core.database import Base
|
||||
|
||||
|
||||
class Case(Base):
|
||||
__tablename__ = "cases"
|
||||
|
||||
id = Column(Integer, primary_key=True, index=True)
|
||||
title = Column(String, nullable=False)
|
||||
description = Column(Text, nullable=True)
|
||||
status = Column(String, default="open", nullable=False) # open, closed, investigating
|
||||
severity = Column(String, nullable=True) # low, medium, high, critical
|
||||
tenant_id = Column(Integer, ForeignKey("tenants.id"), nullable=False)
|
||||
created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc))
|
||||
updated_at = Column(DateTime, default=lambda: datetime.now(timezone.utc), onupdate=lambda: datetime.now(timezone.utc))
|
||||
|
||||
# Relationships
|
||||
tenant = relationship("Tenant", back_populates="cases")
|
||||
artifacts = relationship("Artifact", back_populates="case")
|
||||
@@ -1,21 +0,0 @@
|
||||
from sqlalchemy import Column, Integer, String, DateTime, ForeignKey, JSON
|
||||
from sqlalchemy.orm import relationship
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from app.core.database import Base
|
||||
|
||||
|
||||
class Host(Base):
|
||||
__tablename__ = "hosts"
|
||||
|
||||
id = Column(Integer, primary_key=True, index=True)
|
||||
hostname = Column(String, index=True, nullable=False)
|
||||
ip_address = Column(String, nullable=True)
|
||||
os = Column(String, nullable=True)
|
||||
tenant_id = Column(Integer, ForeignKey("tenants.id"), nullable=False)
|
||||
host_metadata = Column(JSON, nullable=True)
|
||||
created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc))
|
||||
last_seen = Column(DateTime, default=lambda: datetime.now(timezone.utc))
|
||||
|
||||
# Relationships
|
||||
tenant = relationship("Tenant", back_populates="hosts")
|
||||
@@ -1,23 +0,0 @@
|
||||
from sqlalchemy import Column, Integer, String, DateTime, ForeignKey, Boolean, Text
|
||||
from sqlalchemy.orm import relationship
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from app.core.database import Base
|
||||
|
||||
|
||||
class Notification(Base):
|
||||
__tablename__ = "notifications"
|
||||
|
||||
id = Column(Integer, primary_key=True, index=True)
|
||||
user_id = Column(Integer, ForeignKey("users.id"), nullable=False)
|
||||
tenant_id = Column(Integer, ForeignKey("tenants.id"), nullable=False)
|
||||
title = Column(String, nullable=False)
|
||||
message = Column(Text, nullable=False)
|
||||
notification_type = Column(String, nullable=False) # info, warning, error, success
|
||||
is_read = Column(Boolean, default=False, nullable=False)
|
||||
link = Column(String, nullable=True)
|
||||
created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc), index=True)
|
||||
|
||||
# Relationships
|
||||
user = relationship("User")
|
||||
tenant = relationship("Tenant")
|
||||
@@ -1,19 +0,0 @@
|
||||
from sqlalchemy import Column, Integer, String, DateTime, ForeignKey, Boolean
|
||||
from sqlalchemy.orm import relationship
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from app.core.database import Base
|
||||
|
||||
|
||||
class PasswordResetToken(Base):
|
||||
__tablename__ = "password_reset_tokens"
|
||||
|
||||
id = Column(Integer, primary_key=True, index=True)
|
||||
token = Column(String, unique=True, index=True, nullable=False)
|
||||
user_id = Column(Integer, ForeignKey("users.id"), nullable=False)
|
||||
expires_at = Column(DateTime, nullable=False)
|
||||
is_used = Column(Boolean, default=False, nullable=False)
|
||||
created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc))
|
||||
|
||||
# Relationships
|
||||
user = relationship("User")
|
||||
@@ -1,45 +0,0 @@
|
||||
from sqlalchemy import Column, Integer, String, DateTime, ForeignKey, Boolean, Text, JSON
|
||||
from sqlalchemy.orm import relationship
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from app.core.database import Base
|
||||
|
||||
|
||||
class Playbook(Base):
|
||||
__tablename__ = "playbooks"
|
||||
|
||||
id = Column(Integer, primary_key=True, index=True)
|
||||
tenant_id = Column(Integer, ForeignKey("tenants.id"), nullable=False)
|
||||
name = Column(String, nullable=False, index=True)
|
||||
description = Column(Text, nullable=True)
|
||||
trigger_type = Column(String, nullable=False) # manual, scheduled, event
|
||||
trigger_config = Column(JSON, nullable=True)
|
||||
actions = Column(JSON, nullable=False) # List of action definitions
|
||||
is_enabled = Column(Boolean, default=True, nullable=False)
|
||||
created_by = Column(Integer, ForeignKey("users.id"), nullable=False)
|
||||
created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc))
|
||||
updated_at = Column(DateTime, default=lambda: datetime.now(timezone.utc), onupdate=lambda: datetime.now(timezone.utc))
|
||||
|
||||
# Relationships
|
||||
tenant = relationship("Tenant")
|
||||
creator = relationship("User", foreign_keys=[created_by])
|
||||
executions = relationship("PlaybookExecution", back_populates="playbook")
|
||||
|
||||
|
||||
class PlaybookExecution(Base):
|
||||
__tablename__ = "playbook_executions"
|
||||
|
||||
id = Column(Integer, primary_key=True, index=True)
|
||||
playbook_id = Column(Integer, ForeignKey("playbooks.id"), nullable=False)
|
||||
tenant_id = Column(Integer, ForeignKey("tenants.id"), nullable=False)
|
||||
status = Column(String, nullable=False) # pending, running, completed, failed
|
||||
started_at = Column(DateTime, default=lambda: datetime.now(timezone.utc))
|
||||
completed_at = Column(DateTime, nullable=True)
|
||||
result = Column(JSON, nullable=True)
|
||||
error_message = Column(Text, nullable=True)
|
||||
triggered_by = Column(Integer, ForeignKey("users.id"), nullable=True)
|
||||
|
||||
# Relationships
|
||||
playbook = relationship("Playbook", back_populates="executions")
|
||||
tenant = relationship("Tenant")
|
||||
trigger_user = relationship("User")
|
||||
@@ -1,19 +0,0 @@
|
||||
from sqlalchemy import Column, Integer, String, DateTime, ForeignKey, Boolean
|
||||
from sqlalchemy.orm import relationship
|
||||
from datetime import datetime, timezone, timedelta
|
||||
|
||||
from app.core.database import Base
|
||||
|
||||
|
||||
class RefreshToken(Base):
|
||||
__tablename__ = "refresh_tokens"
|
||||
|
||||
id = Column(Integer, primary_key=True, index=True)
|
||||
token = Column(String, unique=True, index=True, nullable=False)
|
||||
user_id = Column(Integer, ForeignKey("users.id"), nullable=False)
|
||||
expires_at = Column(DateTime, nullable=False)
|
||||
is_revoked = Column(Boolean, default=False, nullable=False)
|
||||
created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc))
|
||||
|
||||
# Relationships
|
||||
user = relationship("User", back_populates="refresh_tokens")
|
||||
@@ -1,43 +0,0 @@
|
||||
from sqlalchemy import Column, Integer, String, DateTime, ForeignKey, Boolean, Text, JSON
|
||||
from sqlalchemy.orm import relationship
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from app.core.database import Base
|
||||
|
||||
|
||||
class ReportTemplate(Base):
|
||||
__tablename__ = "report_templates"
|
||||
|
||||
id = Column(Integer, primary_key=True, index=True)
|
||||
tenant_id = Column(Integer, ForeignKey("tenants.id"), nullable=False)
|
||||
name = Column(String, nullable=False, index=True)
|
||||
description = Column(Text, nullable=True)
|
||||
template_type = Column(String, nullable=False) # case_summary, host_analysis, threat_report
|
||||
template_config = Column(JSON, nullable=False) # Configuration for report generation
|
||||
is_default = Column(Boolean, default=False, nullable=False)
|
||||
created_by = Column(Integer, ForeignKey("users.id"), nullable=False)
|
||||
created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc))
|
||||
|
||||
# Relationships
|
||||
tenant = relationship("Tenant")
|
||||
creator = relationship("User")
|
||||
|
||||
|
||||
class Report(Base):
|
||||
__tablename__ = "reports"
|
||||
|
||||
id = Column(Integer, primary_key=True, index=True)
|
||||
tenant_id = Column(Integer, ForeignKey("tenants.id"), nullable=False)
|
||||
template_id = Column(Integer, ForeignKey("report_templates.id"), nullable=True)
|
||||
title = Column(String, nullable=False)
|
||||
report_type = Column(String, nullable=False)
|
||||
format = Column(String, nullable=False) # pdf, html, json
|
||||
file_path = Column(String, nullable=True)
|
||||
status = Column(String, nullable=False) # generating, completed, failed
|
||||
generated_by = Column(Integer, ForeignKey("users.id"), nullable=False)
|
||||
generated_at = Column(DateTime, default=lambda: datetime.now(timezone.utc))
|
||||
|
||||
# Relationships
|
||||
tenant = relationship("Tenant")
|
||||
template = relationship("ReportTemplate")
|
||||
generator = relationship("User")
|
||||
@@ -1,19 +0,0 @@
|
||||
from sqlalchemy import Column, Integer, String, DateTime
|
||||
from sqlalchemy.orm import relationship
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from app.core.database import Base
|
||||
|
||||
|
||||
class Tenant(Base):
|
||||
__tablename__ = "tenants"
|
||||
|
||||
id = Column(Integer, primary_key=True, index=True)
|
||||
name = Column(String, unique=True, index=True, nullable=False)
|
||||
description = Column(String, nullable=True)
|
||||
created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc))
|
||||
|
||||
# Relationships
|
||||
users = relationship("User", back_populates="tenant")
|
||||
hosts = relationship("Host", back_populates="tenant")
|
||||
cases = relationship("Case", back_populates="tenant")
|
||||
@@ -1,26 +0,0 @@
|
||||
from sqlalchemy import Column, Integer, String, DateTime, ForeignKey, Float, Text, JSON
|
||||
from sqlalchemy.orm import relationship
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from app.core.database import Base
|
||||
|
||||
|
||||
class ThreatScore(Base):
|
||||
__tablename__ = "threat_scores"
|
||||
|
||||
id = Column(Integer, primary_key=True, index=True)
|
||||
tenant_id = Column(Integer, ForeignKey("tenants.id"), nullable=False)
|
||||
host_id = Column(Integer, ForeignKey("hosts.id"), nullable=True)
|
||||
artifact_id = Column(Integer, ForeignKey("artifacts.id"), nullable=True)
|
||||
score = Column(Float, nullable=False, index=True) # 0.0 to 1.0
|
||||
confidence = Column(Float, nullable=False) # 0.0 to 1.0
|
||||
threat_type = Column(String, nullable=False) # malware, suspicious, anomaly, etc.
|
||||
description = Column(Text, nullable=True)
|
||||
indicators = Column(JSON, nullable=True) # List of indicators that contributed to score
|
||||
ml_model_version = Column(String, nullable=True)
|
||||
created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc), index=True)
|
||||
|
||||
# Relationships
|
||||
tenant = relationship("Tenant")
|
||||
host = relationship("Host")
|
||||
artifact = relationship("Artifact")
|
||||
@@ -1,25 +0,0 @@
|
||||
from sqlalchemy import Column, Integer, String, DateTime, ForeignKey, Boolean
|
||||
from sqlalchemy.orm import relationship
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from app.core.database import Base
|
||||
|
||||
|
||||
class User(Base):
|
||||
__tablename__ = "users"
|
||||
|
||||
id = Column(Integer, primary_key=True, index=True)
|
||||
username = Column(String, unique=True, index=True, nullable=False)
|
||||
password_hash = Column(String, nullable=False)
|
||||
role = Column(String, default="user", nullable=False) # user, admin
|
||||
tenant_id = Column(Integer, ForeignKey("tenants.id"), nullable=False)
|
||||
is_active = Column(Boolean, default=True, nullable=False)
|
||||
email = Column(String, unique=True, nullable=True, index=True)
|
||||
email_verified = Column(Boolean, default=False, nullable=False)
|
||||
totp_secret = Column(String, nullable=True)
|
||||
totp_enabled = Column(Boolean, default=False, nullable=False)
|
||||
created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc))
|
||||
|
||||
# Relationships
|
||||
tenant = relationship("Tenant", back_populates="users")
|
||||
refresh_tokens = relationship("RefreshToken", back_populates="user")
|
||||
@@ -1,29 +0,0 @@
|
||||
from pydantic import BaseModel
|
||||
from typing import Optional, Dict, Any
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
class AuditLogBase(BaseModel):
|
||||
"""Base audit log schema"""
|
||||
action: str
|
||||
resource_type: str
|
||||
resource_id: Optional[int] = None
|
||||
details: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
class AuditLogCreate(AuditLogBase):
|
||||
"""Schema for creating an audit log entry"""
|
||||
pass
|
||||
|
||||
|
||||
class AuditLogRead(AuditLogBase):
|
||||
"""Schema for reading audit log data"""
|
||||
id: int
|
||||
user_id: Optional[int]
|
||||
tenant_id: int
|
||||
ip_address: Optional[str]
|
||||
user_agent: Optional[str]
|
||||
created_at: datetime
|
||||
|
||||
class Config:
|
||||
from_attributes = True
|
||||
@@ -1,59 +0,0 @@
|
||||
from pydantic import BaseModel, EmailStr
|
||||
from typing import Optional
|
||||
|
||||
|
||||
class Token(BaseModel):
|
||||
"""Token response schema"""
|
||||
access_token: str
|
||||
refresh_token: Optional[str] = None
|
||||
token_type: str = "bearer"
|
||||
|
||||
|
||||
class TokenData(BaseModel):
|
||||
"""Token payload data"""
|
||||
user_id: Optional[int] = None
|
||||
tenant_id: Optional[int] = None
|
||||
role: Optional[str] = None
|
||||
|
||||
|
||||
class UserLogin(BaseModel):
|
||||
"""User login request schema"""
|
||||
username: str
|
||||
password: str
|
||||
totp_code: Optional[str] = None
|
||||
|
||||
|
||||
class UserRegister(BaseModel):
|
||||
"""User registration request schema"""
|
||||
username: str
|
||||
password: str
|
||||
email: Optional[EmailStr] = None
|
||||
tenant_id: Optional[int] = None
|
||||
role: str = "user"
|
||||
|
||||
|
||||
class RefreshTokenRequest(BaseModel):
|
||||
"""Refresh token request schema"""
|
||||
refresh_token: str
|
||||
|
||||
|
||||
class PasswordResetRequest(BaseModel):
|
||||
"""Password reset request schema"""
|
||||
email: EmailStr
|
||||
|
||||
|
||||
class PasswordResetConfirm(BaseModel):
|
||||
"""Password reset confirmation schema"""
|
||||
token: str
|
||||
new_password: str
|
||||
|
||||
|
||||
class TwoFactorSetup(BaseModel):
|
||||
"""2FA setup response schema"""
|
||||
secret: str
|
||||
qr_code_uri: str
|
||||
|
||||
|
||||
class TwoFactorVerify(BaseModel):
|
||||
"""2FA verification schema"""
|
||||
code: str
|
||||
@@ -1,74 +0,0 @@
|
||||
from pydantic import BaseModel
|
||||
from typing import Optional, List, Dict, Any
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
class LLMRequestSchema(BaseModel):
|
||||
"""Schema for LLM processing request"""
|
||||
prompt: str
|
||||
task_hints: Optional[List[str]] = []
|
||||
requires_parallel: bool = False
|
||||
requires_chaining: bool = False
|
||||
batch_size: int = 1
|
||||
operations: Optional[List[str]] = []
|
||||
parameters: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
class RoutingDecision(BaseModel):
|
||||
"""Schema for routing decision"""
|
||||
task_type: str
|
||||
model: str
|
||||
endpoint: str
|
||||
priority: int
|
||||
description: str
|
||||
requires_parallel: bool
|
||||
requires_chaining: bool
|
||||
|
||||
|
||||
class NodeInfo(BaseModel):
|
||||
"""Schema for GPU node information"""
|
||||
node_id: str
|
||||
hostname: str
|
||||
vram_total_gb: int
|
||||
vram_used_gb: int
|
||||
vram_available_gb: int
|
||||
compute_utilization: float
|
||||
status: str
|
||||
models_loaded: List[str]
|
||||
|
||||
|
||||
class SchedulingDecision(BaseModel):
|
||||
"""Schema for job scheduling decision"""
|
||||
job_id: str
|
||||
execution_mode: str
|
||||
nodes: Optional[List[Dict[str, str]]] = None
|
||||
node: Optional[Dict[str, str]] = None
|
||||
status: Optional[str] = None
|
||||
queue_position: Optional[int] = None
|
||||
|
||||
|
||||
class LLMResponseSchema(BaseModel):
|
||||
"""Schema for LLM response"""
|
||||
job_id: str
|
||||
status: str
|
||||
routing: Optional[RoutingDecision] = None
|
||||
scheduling: Optional[SchedulingDecision] = None
|
||||
result: Any
|
||||
execution_mode: str
|
||||
|
||||
|
||||
class ModelInfo(BaseModel):
|
||||
"""Schema for model information"""
|
||||
model_name: str
|
||||
node_id: str
|
||||
endpoint_url: str
|
||||
is_available: bool
|
||||
|
||||
|
||||
class MergedResult(BaseModel):
|
||||
"""Schema for merged result"""
|
||||
strategy: str
|
||||
result: Any
|
||||
confidence: Optional[float] = None
|
||||
num_models: Optional[int] = None
|
||||
all_results: Optional[List[Dict[str, Any]]] = None
|
||||
@@ -1,34 +0,0 @@
|
||||
from pydantic import BaseModel
|
||||
from typing import Optional
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
class NotificationBase(BaseModel):
|
||||
"""Base notification schema"""
|
||||
title: str
|
||||
message: str
|
||||
notification_type: str
|
||||
link: Optional[str] = None
|
||||
|
||||
|
||||
class NotificationCreate(NotificationBase):
|
||||
"""Schema for creating a notification"""
|
||||
user_id: int
|
||||
tenant_id: int
|
||||
|
||||
|
||||
class NotificationRead(NotificationBase):
|
||||
"""Schema for reading notification data"""
|
||||
id: int
|
||||
user_id: int
|
||||
tenant_id: int
|
||||
is_read: bool
|
||||
created_at: datetime
|
||||
|
||||
class Config:
|
||||
from_attributes = True
|
||||
|
||||
|
||||
class NotificationUpdate(BaseModel):
|
||||
"""Schema for updating a notification"""
|
||||
is_read: bool
|
||||
@@ -1,55 +0,0 @@
|
||||
from pydantic import BaseModel
|
||||
from typing import Optional, List, Dict, Any
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
class PlaybookBase(BaseModel):
|
||||
"""Base playbook schema"""
|
||||
name: str
|
||||
description: Optional[str] = None
|
||||
trigger_type: str
|
||||
trigger_config: Optional[Dict[str, Any]] = None
|
||||
actions: List[Dict[str, Any]]
|
||||
is_enabled: bool = True
|
||||
|
||||
|
||||
class PlaybookCreate(PlaybookBase):
|
||||
"""Schema for creating a playbook"""
|
||||
pass
|
||||
|
||||
|
||||
class PlaybookUpdate(BaseModel):
|
||||
"""Schema for updating a playbook"""
|
||||
name: Optional[str] = None
|
||||
description: Optional[str] = None
|
||||
trigger_type: Optional[str] = None
|
||||
trigger_config: Optional[Dict[str, Any]] = None
|
||||
actions: Optional[List[Dict[str, Any]]] = None
|
||||
is_enabled: Optional[bool] = None
|
||||
|
||||
|
||||
class PlaybookRead(PlaybookBase):
|
||||
"""Schema for reading playbook data"""
|
||||
id: int
|
||||
tenant_id: int
|
||||
created_by: int
|
||||
created_at: datetime
|
||||
updated_at: datetime
|
||||
|
||||
class Config:
|
||||
from_attributes = True
|
||||
|
||||
|
||||
class PlaybookExecutionRead(BaseModel):
|
||||
"""Schema for playbook execution"""
|
||||
id: int
|
||||
playbook_id: int
|
||||
tenant_id: int
|
||||
status: str
|
||||
started_at: datetime
|
||||
completed_at: Optional[datetime]
|
||||
result: Optional[Dict[str, Any]]
|
||||
error_message: Optional[str]
|
||||
|
||||
class Config:
|
||||
from_attributes = True
|
||||
@@ -1,54 +0,0 @@
|
||||
from pydantic import BaseModel
|
||||
from typing import Optional, Dict, Any
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
class ReportTemplateBase(BaseModel):
|
||||
"""Base report template schema"""
|
||||
name: str
|
||||
description: Optional[str] = None
|
||||
template_type: str
|
||||
template_config: Dict[str, Any]
|
||||
is_default: bool = False
|
||||
|
||||
|
||||
class ReportTemplateCreate(ReportTemplateBase):
|
||||
"""Schema for creating a report template"""
|
||||
pass
|
||||
|
||||
|
||||
class ReportTemplateRead(ReportTemplateBase):
|
||||
"""Schema for reading report template data"""
|
||||
id: int
|
||||
tenant_id: int
|
||||
created_by: int
|
||||
created_at: datetime
|
||||
|
||||
class Config:
|
||||
from_attributes = True
|
||||
|
||||
|
||||
class ReportBase(BaseModel):
|
||||
"""Base report schema"""
|
||||
title: str
|
||||
report_type: str
|
||||
format: str
|
||||
|
||||
|
||||
class ReportCreate(ReportBase):
|
||||
"""Schema for creating a report"""
|
||||
template_id: Optional[int] = None
|
||||
|
||||
|
||||
class ReportRead(ReportBase):
|
||||
"""Schema for reading report data"""
|
||||
id: int
|
||||
tenant_id: int
|
||||
template_id: Optional[int]
|
||||
file_path: Optional[str]
|
||||
status: str
|
||||
generated_by: int
|
||||
generated_at: datetime
|
||||
|
||||
class Config:
|
||||
from_attributes = True
|
||||
@@ -1,32 +0,0 @@
|
||||
from pydantic import BaseModel
|
||||
from typing import Optional, List, Dict, Any
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
class ThreatScoreBase(BaseModel):
|
||||
"""Base threat score schema"""
|
||||
score: float
|
||||
confidence: float
|
||||
threat_type: str
|
||||
description: Optional[str] = None
|
||||
indicators: Optional[List[Dict[str, Any]]] = None
|
||||
|
||||
|
||||
class ThreatScoreCreate(ThreatScoreBase):
|
||||
"""Schema for creating a threat score"""
|
||||
host_id: Optional[int] = None
|
||||
artifact_id: Optional[int] = None
|
||||
ml_model_version: Optional[str] = None
|
||||
|
||||
|
||||
class ThreatScoreRead(ThreatScoreBase):
|
||||
"""Schema for reading threat score data"""
|
||||
id: int
|
||||
tenant_id: int
|
||||
host_id: Optional[int]
|
||||
artifact_id: Optional[int]
|
||||
ml_model_version: Optional[str]
|
||||
created_at: datetime
|
||||
|
||||
class Config:
|
||||
from_attributes = True
|
||||
@@ -1,37 +0,0 @@
|
||||
from pydantic import BaseModel, EmailStr
|
||||
from typing import Optional
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
class UserBase(BaseModel):
|
||||
"""Base user schema"""
|
||||
username: str
|
||||
email: Optional[EmailStr] = None
|
||||
role: str = "user"
|
||||
tenant_id: int
|
||||
|
||||
|
||||
class UserCreate(UserBase):
|
||||
"""Schema for creating a user"""
|
||||
password: str
|
||||
|
||||
|
||||
class UserUpdate(BaseModel):
|
||||
"""Schema for updating a user"""
|
||||
username: Optional[str] = None
|
||||
email: Optional[EmailStr] = None
|
||||
password: Optional[str] = None
|
||||
role: Optional[str] = None
|
||||
is_active: Optional[bool] = None
|
||||
|
||||
|
||||
class UserRead(UserBase):
|
||||
"""Schema for reading user data (excludes password_hash and secrets)"""
|
||||
id: int
|
||||
is_active: bool
|
||||
email_verified: bool
|
||||
totp_enabled: bool
|
||||
created_at: datetime
|
||||
|
||||
class Config:
|
||||
from_attributes = True
|
||||
@@ -1,15 +1,9 @@
|
||||
fastapi==0.109.0
|
||||
uvicorn[standard]==0.27.0
|
||||
sqlalchemy==2.0.25
|
||||
flask==3.0.0
|
||||
flask-cors==4.0.0
|
||||
flask-sqlalchemy==3.1.1
|
||||
flask-jwt-extended==4.6.0
|
||||
psycopg2-binary==2.9.9
|
||||
python-jose[cryptography]==3.3.0
|
||||
passlib[bcrypt]==1.7.4
|
||||
python-multipart==0.0.6
|
||||
alembic==1.13.1
|
||||
pydantic==2.5.3
|
||||
pydantic-settings==2.1.0
|
||||
pyotp==2.9.0
|
||||
qrcode[pil]==7.4.2
|
||||
websockets==12.0
|
||||
httpx==0.26.0
|
||||
email-validator==2.1.2
|
||||
python-dotenv==1.0.0
|
||||
requests==2.31.0
|
||||
werkzeug==3.0.1
|
||||
bcrypt==4.1.2
|
||||
|
||||
52
database/init.sql
Normal file
52
database/init.sql
Normal file
@@ -0,0 +1,52 @@
|
||||
CREATE EXTENSION IF NOT EXISTS "uuid-ossp";
|
||||
|
||||
-- Users table
|
||||
CREATE TABLE users (
|
||||
id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
|
||||
username VARCHAR(50) UNIQUE NOT NULL,
|
||||
email VARCHAR(100) UNIQUE NOT NULL,
|
||||
password_hash VARCHAR(255) NOT NULL,
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
last_login TIMESTAMP
|
||||
);
|
||||
|
||||
-- Hunts table
|
||||
CREATE TABLE hunts (
|
||||
id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
|
||||
name VARCHAR(100) NOT NULL,
|
||||
description TEXT,
|
||||
created_by UUID REFERENCES users(id),
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
status VARCHAR(20) DEFAULT 'active' CHECK (status IN ('active', 'completed', 'archived'))
|
||||
);
|
||||
|
||||
-- Files table
|
||||
CREATE TABLE files (
|
||||
id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
|
||||
hunt_id UUID REFERENCES hunts(id) ON DELETE CASCADE,
|
||||
filename VARCHAR(255) NOT NULL,
|
||||
original_filename VARCHAR(255) NOT NULL,
|
||||
file_size BIGINT,
|
||||
file_type VARCHAR(50),
|
||||
uploaded_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
analysis_status VARCHAR(20) DEFAULT 'pending' CHECK (analysis_status IN ('pending', 'processing', 'completed', 'failed'))
|
||||
);
|
||||
|
||||
-- Analysis results table
|
||||
CREATE TABLE analysis_results (
|
||||
id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
|
||||
file_id UUID REFERENCES files(id) ON DELETE CASCADE,
|
||||
analysis_type VARCHAR(50) NOT NULL,
|
||||
results JSONB,
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||||
);
|
||||
|
||||
-- Create indexes
|
||||
CREATE INDEX idx_hunts_created_by ON hunts(created_by);
|
||||
CREATE INDEX idx_files_hunt_id ON files(hunt_id);
|
||||
CREATE INDEX idx_analysis_file_id ON analysis_results(file_id);
|
||||
|
||||
-- Insert default admin user (password: admin123)
|
||||
INSERT INTO users (username, email, password_hash) VALUES
|
||||
('admin', 'admin@threathunter.local', '$2b$12$LQv3c1yqBWVHxkd0LHAkCOYz6TtxMQJqhN8/LewKyIqhFrMpGLgGi');
|
||||
31
deploy/aws-deploy.yml
Normal file
31
deploy/aws-deploy.yml
Normal file
@@ -0,0 +1,31 @@
|
||||
# AWS ECS Deployment Configuration
|
||||
version: '3.8'
|
||||
|
||||
services:
|
||||
database:
|
||||
image: postgres:15
|
||||
environment:
|
||||
POSTGRES_DB: threat_hunter
|
||||
POSTGRES_USER: ${DB_USER}
|
||||
POSTGRES_PASSWORD: ${DB_PASSWORD}
|
||||
volumes:
|
||||
- postgres_data:/var/lib/postgresql/data
|
||||
|
||||
backend:
|
||||
image: your-registry/threat-hunter-backend:latest
|
||||
environment:
|
||||
DATABASE_URL: postgresql://${DB_USER}:${DB_PASSWORD}@database:5432/threat_hunter
|
||||
SECRET_KEY: ${SECRET_KEY}
|
||||
FLASK_ENV: production
|
||||
depends_on:
|
||||
- database
|
||||
|
||||
frontend:
|
||||
image: your-registry/threat-hunter-frontend:latest
|
||||
ports:
|
||||
- "80:3000"
|
||||
depends_on:
|
||||
- backend
|
||||
|
||||
volumes:
|
||||
postgres_data:
|
||||
18
deploy/backup/backup.sh
Normal file
18
deploy/backup/backup.sh
Normal file
@@ -0,0 +1,18 @@
|
||||
#!/bin/bash
|
||||
|
||||
# Database backup
|
||||
BACKUP_DIR="/backups"
|
||||
DATE=$(date +%Y%m%d_%H%M%S)
|
||||
|
||||
echo "Creating database backup..."
|
||||
docker exec threat-hunter-db pg_dump -U admin threat_hunter > "$BACKUP_DIR/db_backup_$DATE.sql"
|
||||
|
||||
# File uploads backup
|
||||
echo "Backing up uploads..."
|
||||
tar -czf "$BACKUP_DIR/uploads_backup_$DATE.tar.gz" ./uploads
|
||||
|
||||
# Keep only last 7 days of backups
|
||||
find $BACKUP_DIR -name "*.sql" -mtime +7 -delete
|
||||
find $BACKUP_DIR -name "*.tar.gz" -mtime +7 -delete
|
||||
|
||||
echo "Backup completed: $DATE"
|
||||
28
deploy/deploy.sh
Normal file
28
deploy/deploy.sh
Normal file
@@ -0,0 +1,28 @@
|
||||
#!/bin/bash
|
||||
|
||||
set -e
|
||||
|
||||
echo "🚀 Deploying Cyber Threat Hunter..."
|
||||
|
||||
# Build and push images
|
||||
echo "📦 Building Docker images..."
|
||||
docker build -t your-registry/threat-hunter-backend:latest ./backend
|
||||
docker build -t your-registry/threat-hunter-frontend:latest ./frontend
|
||||
|
||||
echo "🔄 Pushing to registry..."
|
||||
docker push your-registry/threat-hunter-backend:latest
|
||||
docker push your-registry/threat-hunter-frontend:latest
|
||||
|
||||
# Deploy based on environment
|
||||
if [ "$1" = "kubernetes" ]; then
|
||||
echo "☸️ Deploying to Kubernetes..."
|
||||
kubectl apply -f deploy/kubernetes/
|
||||
elif [ "$1" = "swarm" ]; then
|
||||
echo "🐳 Deploying to Docker Swarm..."
|
||||
docker stack deploy -c deploy/docker-stack.yml threat-hunter
|
||||
else
|
||||
echo "🐙 Deploying with Docker Compose..."
|
||||
docker-compose -f docker-compose.prod.yml up -d
|
||||
fi
|
||||
|
||||
echo "✅ Deployment complete!"
|
||||
55
deploy/docker-stack.yml
Normal file
55
deploy/docker-stack.yml
Normal file
@@ -0,0 +1,55 @@
|
||||
version: '3.8'
|
||||
|
||||
services:
|
||||
database:
|
||||
image: postgres:15
|
||||
environment:
|
||||
POSTGRES_DB: threat_hunter
|
||||
POSTGRES_USER_FILE: /run/secrets/db_user
|
||||
POSTGRES_PASSWORD_FILE: /run/secrets/db_password
|
||||
volumes:
|
||||
- postgres_data:/var/lib/postgresql/data
|
||||
secrets:
|
||||
- db_user
|
||||
- db_password
|
||||
deploy:
|
||||
replicas: 1
|
||||
placement:
|
||||
constraints:
|
||||
- node.role == manager
|
||||
|
||||
backend:
|
||||
image: your-registry/threat-hunter-backend:latest
|
||||
environment:
|
||||
DATABASE_URL: postgresql://admin:secure_password_123@database:5432/threat_hunter
|
||||
SECRET_KEY_FILE: /run/secrets/secret_key
|
||||
secrets:
|
||||
- secret_key
|
||||
deploy:
|
||||
replicas: 3
|
||||
update_config:
|
||||
parallelism: 1
|
||||
delay: 10s
|
||||
restart_policy:
|
||||
condition: on-failure
|
||||
|
||||
frontend:
|
||||
image: your-registry/threat-hunter-frontend:latest
|
||||
ports:
|
||||
- "80:3000"
|
||||
deploy:
|
||||
replicas: 2
|
||||
update_config:
|
||||
parallelism: 1
|
||||
delay: 10s
|
||||
|
||||
volumes:
|
||||
postgres_data:
|
||||
|
||||
secrets:
|
||||
db_user:
|
||||
external: true
|
||||
db_password:
|
||||
external: true
|
||||
secret_key:
|
||||
external: true
|
||||
37
deploy/kubernetes/deployment.yaml
Normal file
37
deploy/kubernetes/deployment.yaml
Normal file
@@ -0,0 +1,37 @@
|
||||
apiVersion: apps/v1
|
||||
kind: Deployment
|
||||
metadata:
|
||||
name: threat-hunter-backend
|
||||
spec:
|
||||
replicas: 3
|
||||
selector:
|
||||
matchLabels:
|
||||
app: threat-hunter-backend
|
||||
template:
|
||||
metadata:
|
||||
labels:
|
||||
app: threat-hunter-backend
|
||||
spec:
|
||||
containers:
|
||||
- name: backend
|
||||
image: your-registry/threat-hunter-backend:latest
|
||||
ports:
|
||||
- containerPort: 5000
|
||||
env:
|
||||
- name: DATABASE_URL
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: threat-hunter-secrets
|
||||
key: database-url
|
||||
---
|
||||
apiVersion: v1
|
||||
kind: Service
|
||||
metadata:
|
||||
name: threat-hunter-backend-service
|
||||
spec:
|
||||
selector:
|
||||
app: threat-hunter-backend
|
||||
ports:
|
||||
- port: 5000
|
||||
targetPort: 5000
|
||||
type: LoadBalancer
|
||||
21
deploy/monitoring/docker-compose.monitoring.yml
Normal file
21
deploy/monitoring/docker-compose.monitoring.yml
Normal file
@@ -0,0 +1,21 @@
|
||||
version: '3.8'
|
||||
|
||||
services:
|
||||
prometheus:
|
||||
image: prom/prometheus
|
||||
ports:
|
||||
- "9090:9090"
|
||||
volumes:
|
||||
- ./prometheus.yml:/etc/prometheus/prometheus.yml
|
||||
|
||||
grafana:
|
||||
image: grafana/grafana
|
||||
ports:
|
||||
- "3001:3000"
|
||||
environment:
|
||||
GF_SECURITY_ADMIN_PASSWORD: admin
|
||||
volumes:
|
||||
- grafana_data:/var/lib/grafana
|
||||
|
||||
volumes:
|
||||
grafana_data:
|
||||
26
deploy/security/security-checklist.md
Normal file
26
deploy/security/security-checklist.md
Normal file
@@ -0,0 +1,26 @@
|
||||
# Security Deployment Checklist
|
||||
|
||||
## Pre-Deployment
|
||||
- [ ] Change all default passwords
|
||||
- [ ] Generate strong SECRET_KEY
|
||||
- [ ] Setup SSL/TLS certificates
|
||||
- [ ] Configure firewall rules
|
||||
- [ ] Set up backup strategy
|
||||
|
||||
## Database Security
|
||||
- [ ] Use strong database passwords
|
||||
- [ ] Enable database encryption
|
||||
- [ ] Configure database firewall
|
||||
- [ ] Set up regular backups
|
||||
|
||||
## Application Security
|
||||
- [ ] Update all dependencies
|
||||
- [ ] Configure CORS properly
|
||||
- [ ] Enable rate limiting
|
||||
- [ ] Set up monitoring/logging
|
||||
|
||||
## Infrastructure Security
|
||||
- [ ] Use private networks
|
||||
- [ ] Configure load balancer
|
||||
- [ ] Set up intrusion detection
|
||||
- [ ] Regular security updates
|
||||
29
deploy/setup-prod.bat
Normal file
29
deploy/setup-prod.bat
Normal file
@@ -0,0 +1,29 @@
|
||||
@echo off
|
||||
echo Setting up production environment...
|
||||
|
||||
REM Create environment file
|
||||
echo Creating .env.prod file...
|
||||
(
|
||||
echo DB_USER=threat_hunter_user
|
||||
echo DB_PASSWORD=%RANDOM%%RANDOM%
|
||||
echo SECRET_KEY=%RANDOM%%RANDOM%%RANDOM%
|
||||
echo FLASK_ENV=production
|
||||
) > .env.prod
|
||||
|
||||
REM Setup SSL certificates
|
||||
echo Setting up SSL certificates...
|
||||
mkdir ssl
|
||||
REM Add your SSL certificate generation here
|
||||
|
||||
REM Create backup directory
|
||||
mkdir backups
|
||||
mkdir logs
|
||||
|
||||
REM Setup firewall rules
|
||||
echo Configuring firewall...
|
||||
netsh advfirewall firewall add rule name="HTTP" dir=in action=allow protocol=TCP localport=80
|
||||
netsh advfirewall firewall add rule name="HTTPS" dir=in action=allow protocol=TCP localport=443
|
||||
|
||||
echo Production setup complete!
|
||||
echo Please update .env.prod with your actual values
|
||||
pause
|
||||
72
docker-compose.prod.yml
Normal file
72
docker-compose.prod.yml
Normal file
@@ -0,0 +1,72 @@
|
||||
version: '3.8'
|
||||
|
||||
services:
|
||||
database:
|
||||
image: postgres:15
|
||||
environment:
|
||||
POSTGRES_DB: threat_hunter
|
||||
POSTGRES_USER: ${DB_USER}
|
||||
POSTGRES_PASSWORD: ${DB_PASSWORD}
|
||||
volumes:
|
||||
- postgres_data:/var/lib/postgresql/data
|
||||
- ./backups:/backups
|
||||
restart: unless-stopped
|
||||
networks:
|
||||
- internal
|
||||
|
||||
backend:
|
||||
build:
|
||||
context: ./backend
|
||||
dockerfile: Dockerfile.prod
|
||||
environment:
|
||||
DATABASE_URL: postgresql://${DB_USER}:${DB_PASSWORD}@database:5432/threat_hunter
|
||||
SECRET_KEY: ${SECRET_KEY}
|
||||
FLASK_ENV: production
|
||||
depends_on:
|
||||
- database
|
||||
volumes:
|
||||
- ./uploads:/app/uploads
|
||||
- ./logs:/app/logs
|
||||
restart: unless-stopped
|
||||
networks:
|
||||
- internal
|
||||
- web
|
||||
|
||||
frontend:
|
||||
build:
|
||||
context: ./frontend
|
||||
dockerfile: Dockerfile.prod
|
||||
ports:
|
||||
- "80:80"
|
||||
- "443:443"
|
||||
volumes:
|
||||
- ./ssl:/etc/ssl/certs
|
||||
depends_on:
|
||||
- backend
|
||||
restart: unless-stopped
|
||||
networks:
|
||||
- web
|
||||
|
||||
nginx:
|
||||
image: nginx:alpine
|
||||
ports:
|
||||
- "80:80"
|
||||
- "443:443"
|
||||
volumes:
|
||||
- ./nginx/nginx.conf:/etc/nginx/nginx.conf
|
||||
- ./ssl:/etc/ssl/certs
|
||||
depends_on:
|
||||
- frontend
|
||||
- backend
|
||||
restart: unless-stopped
|
||||
networks:
|
||||
- web
|
||||
|
||||
volumes:
|
||||
postgres_data:
|
||||
|
||||
networks:
|
||||
web:
|
||||
external: true
|
||||
internal:
|
||||
internal: true
|
||||
@@ -1,51 +1,55 @@
|
||||
version: '3.8'
|
||||
version: "3.8"
|
||||
|
||||
services:
|
||||
db:
|
||||
database:
|
||||
image: postgres:15
|
||||
container_name: threat-hunter-db
|
||||
environment:
|
||||
POSTGRES_USER: postgres
|
||||
POSTGRES_PASSWORD: postgres
|
||||
POSTGRES_DB: velocicompanion
|
||||
ports:
|
||||
- "5432:5432"
|
||||
POSTGRES_DB: threat_hunter
|
||||
POSTGRES_USER: admin
|
||||
POSTGRES_PASSWORD: secure_password_123
|
||||
volumes:
|
||||
- postgres_data:/var/lib/postgresql/data
|
||||
healthcheck:
|
||||
test: ["CMD-SHELL", "pg_isready -U postgres"]
|
||||
interval: 5s
|
||||
timeout: 5s
|
||||
retries: 5
|
||||
- ./database/init.sql:/docker-entrypoint-initdb.d/init.sql
|
||||
ports:
|
||||
- "5432:5432"
|
||||
networks:
|
||||
- threat-hunter-network
|
||||
|
||||
backend:
|
||||
build:
|
||||
context: ./backend
|
||||
dockerfile: Dockerfile
|
||||
ports:
|
||||
- "8000:8000"
|
||||
container_name: threat-hunter-backend
|
||||
environment:
|
||||
- DATABASE_URL=postgresql://postgres:postgres@db:5432/velocicompanion
|
||||
- SECRET_KEY=your-secret-key-change-in-production-min-32-chars-long
|
||||
- ACCESS_TOKEN_EXPIRE_MINUTES=30
|
||||
DATABASE_URL: postgresql://admin:secure_password_123@database:5432/threat_hunter
|
||||
FLASK_ENV: production
|
||||
SECRET_KEY: your-secret-key-change-in-production
|
||||
depends_on:
|
||||
db:
|
||||
condition: service_healthy
|
||||
- database
|
||||
ports:
|
||||
- "5000:5000"
|
||||
volumes:
|
||||
- ./backend:/app
|
||||
command: sh -c "alembic upgrade head && uvicorn app.main:app --host 0.0.0.0 --port 8000 --reload"
|
||||
- ./uploads:/app/uploads
|
||||
- ./output:/app/output
|
||||
networks:
|
||||
- threat-hunter-network
|
||||
|
||||
frontend:
|
||||
build:
|
||||
context: ./frontend
|
||||
dockerfile: Dockerfile
|
||||
container_name: threat-hunter-frontend
|
||||
ports:
|
||||
- "3000:3000"
|
||||
environment:
|
||||
- REACT_APP_API_URL=http://localhost:8000
|
||||
volumes:
|
||||
- ./frontend:/app
|
||||
- /app/node_modules
|
||||
command: npm start
|
||||
depends_on:
|
||||
- backend
|
||||
networks:
|
||||
- threat-hunter-network
|
||||
|
||||
volumes:
|
||||
postgres_data:
|
||||
|
||||
networks:
|
||||
threat-hunter-network:
|
||||
driver: bridge
|
||||
|
||||
20
docker-start.bat
Normal file
20
docker-start.bat
Normal file
@@ -0,0 +1,20 @@
|
||||
@echo off
|
||||
echo Starting Cyber Threat Hunter with Docker...
|
||||
|
||||
docker-compose down
|
||||
docker-compose build
|
||||
docker-compose up -d
|
||||
|
||||
echo.
|
||||
echo ========================================
|
||||
echo Cyber Threat Hunter is starting...
|
||||
echo Frontend: http://localhost:3000
|
||||
echo Backend API: http://localhost:5000
|
||||
echo Database: PostgreSQL on port 5432
|
||||
echo ========================================
|
||||
echo.
|
||||
echo Default credentials:
|
||||
echo Username: admin
|
||||
echo Password: admin123
|
||||
echo.
|
||||
pause
|
||||
@@ -1,8 +0,0 @@
|
||||
node_modules
|
||||
npm-debug.log
|
||||
build
|
||||
.git
|
||||
.gitignore
|
||||
README.md
|
||||
.env
|
||||
.env.local
|
||||
@@ -3,14 +3,18 @@ FROM node:18-alpine
|
||||
WORKDIR /app
|
||||
|
||||
# Copy package files
|
||||
COPY package.json ./
|
||||
COPY package-lock.json* ./
|
||||
|
||||
# Install dependencies
|
||||
COPY package*.json ./
|
||||
RUN npm install
|
||||
|
||||
# Copy application code
|
||||
# Copy source code
|
||||
COPY . .
|
||||
|
||||
# Start development server
|
||||
CMD ["npm", "start"]
|
||||
# Build the application
|
||||
RUN npm run build
|
||||
|
||||
# Serve with a simple server
|
||||
RUN npm install -g serve
|
||||
|
||||
EXPOSE 3000
|
||||
|
||||
CMD ["serve", "-s", "dist", "-l", "3000"]
|
||||
|
||||
19
frontend/Dockerfile.prod
Normal file
19
frontend/Dockerfile.prod
Normal file
@@ -0,0 +1,19 @@
|
||||
# Build stage
|
||||
FROM node:18-alpine as build
|
||||
|
||||
WORKDIR /app
|
||||
COPY package*.json ./
|
||||
RUN npm ci --only=production
|
||||
|
||||
COPY . .
|
||||
RUN npm run build
|
||||
|
||||
# Production stage
|
||||
FROM nginx:alpine
|
||||
|
||||
COPY --from=build /app/dist /usr/share/nginx/html
|
||||
COPY nginx.conf /etc/nginx/nginx.conf
|
||||
|
||||
EXPOSE 80 443
|
||||
|
||||
CMD ["nginx", "-g", "daemon off;"]
|
||||
12
frontend/index.html
Normal file
12
frontend/index.html
Normal file
@@ -0,0 +1,12 @@
|
||||
<!doctype html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8" />
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0" />
|
||||
<title>Cyber Threat Hunter</title>
|
||||
</head>
|
||||
<body>
|
||||
<div id="root"></div>
|
||||
<script type="module" src="/src/main.jsx"></script>
|
||||
</body>
|
||||
</html>
|
||||
16
frontend/node_modules/.bin/acorn
generated
vendored
Normal file
16
frontend/node_modules/.bin/acorn
generated
vendored
Normal file
@@ -0,0 +1,16 @@
|
||||
#!/bin/sh
|
||||
basedir=$(dirname "$(echo "$0" | sed -e 's,\\,/,g')")
|
||||
|
||||
case `uname` in
|
||||
*CYGWIN*|*MINGW*|*MSYS*)
|
||||
if command -v cygpath > /dev/null 2>&1; then
|
||||
basedir=`cygpath -w "$basedir"`
|
||||
fi
|
||||
;;
|
||||
esac
|
||||
|
||||
if [ -x "$basedir/node" ]; then
|
||||
exec "$basedir/node" "$basedir/../acorn/bin/acorn" "$@"
|
||||
else
|
||||
exec node "$basedir/../acorn/bin/acorn" "$@"
|
||||
fi
|
||||
17
frontend/node_modules/.bin/acorn.cmd
generated
vendored
Normal file
17
frontend/node_modules/.bin/acorn.cmd
generated
vendored
Normal file
@@ -0,0 +1,17 @@
|
||||
@ECHO off
|
||||
GOTO start
|
||||
:find_dp0
|
||||
SET dp0=%~dp0
|
||||
EXIT /b
|
||||
:start
|
||||
SETLOCAL
|
||||
CALL :find_dp0
|
||||
|
||||
IF EXIST "%dp0%\node.exe" (
|
||||
SET "_prog=%dp0%\node.exe"
|
||||
) ELSE (
|
||||
SET "_prog=node"
|
||||
SET PATHEXT=%PATHEXT:;.JS;=;%
|
||||
)
|
||||
|
||||
endLocal & goto #_undefined_# 2>NUL || title %COMSPEC% & "%_prog%" "%dp0%\..\acorn\bin\acorn" %*
|
||||
28
frontend/node_modules/.bin/acorn.ps1
generated
vendored
Normal file
28
frontend/node_modules/.bin/acorn.ps1
generated
vendored
Normal file
@@ -0,0 +1,28 @@
|
||||
#!/usr/bin/env pwsh
|
||||
$basedir=Split-Path $MyInvocation.MyCommand.Definition -Parent
|
||||
|
||||
$exe=""
|
||||
if ($PSVersionTable.PSVersion -lt "6.0" -or $IsWindows) {
|
||||
# Fix case when both the Windows and Linux builds of Node
|
||||
# are installed in the same directory
|
||||
$exe=".exe"
|
||||
}
|
||||
$ret=0
|
||||
if (Test-Path "$basedir/node$exe") {
|
||||
# Support pipeline input
|
||||
if ($MyInvocation.ExpectingInput) {
|
||||
$input | & "$basedir/node$exe" "$basedir/../acorn/bin/acorn" $args
|
||||
} else {
|
||||
& "$basedir/node$exe" "$basedir/../acorn/bin/acorn" $args
|
||||
}
|
||||
$ret=$LASTEXITCODE
|
||||
} else {
|
||||
# Support pipeline input
|
||||
if ($MyInvocation.ExpectingInput) {
|
||||
$input | & "node$exe" "$basedir/../acorn/bin/acorn" $args
|
||||
} else {
|
||||
& "node$exe" "$basedir/../acorn/bin/acorn" $args
|
||||
}
|
||||
$ret=$LASTEXITCODE
|
||||
}
|
||||
exit $ret
|
||||
16
frontend/node_modules/.bin/autoprefixer
generated
vendored
Normal file
16
frontend/node_modules/.bin/autoprefixer
generated
vendored
Normal file
@@ -0,0 +1,16 @@
|
||||
#!/bin/sh
|
||||
basedir=$(dirname "$(echo "$0" | sed -e 's,\\,/,g')")
|
||||
|
||||
case `uname` in
|
||||
*CYGWIN*|*MINGW*|*MSYS*)
|
||||
if command -v cygpath > /dev/null 2>&1; then
|
||||
basedir=`cygpath -w "$basedir"`
|
||||
fi
|
||||
;;
|
||||
esac
|
||||
|
||||
if [ -x "$basedir/node" ]; then
|
||||
exec "$basedir/node" "$basedir/../autoprefixer/bin/autoprefixer" "$@"
|
||||
else
|
||||
exec node "$basedir/../autoprefixer/bin/autoprefixer" "$@"
|
||||
fi
|
||||
17
frontend/node_modules/.bin/autoprefixer.cmd
generated
vendored
Normal file
17
frontend/node_modules/.bin/autoprefixer.cmd
generated
vendored
Normal file
@@ -0,0 +1,17 @@
|
||||
@ECHO off
|
||||
GOTO start
|
||||
:find_dp0
|
||||
SET dp0=%~dp0
|
||||
EXIT /b
|
||||
:start
|
||||
SETLOCAL
|
||||
CALL :find_dp0
|
||||
|
||||
IF EXIST "%dp0%\node.exe" (
|
||||
SET "_prog=%dp0%\node.exe"
|
||||
) ELSE (
|
||||
SET "_prog=node"
|
||||
SET PATHEXT=%PATHEXT:;.JS;=;%
|
||||
)
|
||||
|
||||
endLocal & goto #_undefined_# 2>NUL || title %COMSPEC% & "%_prog%" "%dp0%\..\autoprefixer\bin\autoprefixer" %*
|
||||
28
frontend/node_modules/.bin/autoprefixer.ps1
generated
vendored
Normal file
28
frontend/node_modules/.bin/autoprefixer.ps1
generated
vendored
Normal file
@@ -0,0 +1,28 @@
|
||||
#!/usr/bin/env pwsh
|
||||
$basedir=Split-Path $MyInvocation.MyCommand.Definition -Parent
|
||||
|
||||
$exe=""
|
||||
if ($PSVersionTable.PSVersion -lt "6.0" -or $IsWindows) {
|
||||
# Fix case when both the Windows and Linux builds of Node
|
||||
# are installed in the same directory
|
||||
$exe=".exe"
|
||||
}
|
||||
$ret=0
|
||||
if (Test-Path "$basedir/node$exe") {
|
||||
# Support pipeline input
|
||||
if ($MyInvocation.ExpectingInput) {
|
||||
$input | & "$basedir/node$exe" "$basedir/../autoprefixer/bin/autoprefixer" $args
|
||||
} else {
|
||||
& "$basedir/node$exe" "$basedir/../autoprefixer/bin/autoprefixer" $args
|
||||
}
|
||||
$ret=$LASTEXITCODE
|
||||
} else {
|
||||
# Support pipeline input
|
||||
if ($MyInvocation.ExpectingInput) {
|
||||
$input | & "node$exe" "$basedir/../autoprefixer/bin/autoprefixer" $args
|
||||
} else {
|
||||
& "node$exe" "$basedir/../autoprefixer/bin/autoprefixer" $args
|
||||
}
|
||||
$ret=$LASTEXITCODE
|
||||
}
|
||||
exit $ret
|
||||
16
frontend/node_modules/.bin/browserslist
generated
vendored
Normal file
16
frontend/node_modules/.bin/browserslist
generated
vendored
Normal file
@@ -0,0 +1,16 @@
|
||||
#!/bin/sh
|
||||
basedir=$(dirname "$(echo "$0" | sed -e 's,\\,/,g')")
|
||||
|
||||
case `uname` in
|
||||
*CYGWIN*|*MINGW*|*MSYS*)
|
||||
if command -v cygpath > /dev/null 2>&1; then
|
||||
basedir=`cygpath -w "$basedir"`
|
||||
fi
|
||||
;;
|
||||
esac
|
||||
|
||||
if [ -x "$basedir/node" ]; then
|
||||
exec "$basedir/node" "$basedir/../browserslist/cli.js" "$@"
|
||||
else
|
||||
exec node "$basedir/../browserslist/cli.js" "$@"
|
||||
fi
|
||||
17
frontend/node_modules/.bin/browserslist.cmd
generated
vendored
Normal file
17
frontend/node_modules/.bin/browserslist.cmd
generated
vendored
Normal file
@@ -0,0 +1,17 @@
|
||||
@ECHO off
|
||||
GOTO start
|
||||
:find_dp0
|
||||
SET dp0=%~dp0
|
||||
EXIT /b
|
||||
:start
|
||||
SETLOCAL
|
||||
CALL :find_dp0
|
||||
|
||||
IF EXIST "%dp0%\node.exe" (
|
||||
SET "_prog=%dp0%\node.exe"
|
||||
) ELSE (
|
||||
SET "_prog=node"
|
||||
SET PATHEXT=%PATHEXT:;.JS;=;%
|
||||
)
|
||||
|
||||
endLocal & goto #_undefined_# 2>NUL || title %COMSPEC% & "%_prog%" "%dp0%\..\browserslist\cli.js" %*
|
||||
28
frontend/node_modules/.bin/browserslist.ps1
generated
vendored
Normal file
28
frontend/node_modules/.bin/browserslist.ps1
generated
vendored
Normal file
@@ -0,0 +1,28 @@
|
||||
#!/usr/bin/env pwsh
|
||||
$basedir=Split-Path $MyInvocation.MyCommand.Definition -Parent
|
||||
|
||||
$exe=""
|
||||
if ($PSVersionTable.PSVersion -lt "6.0" -or $IsWindows) {
|
||||
# Fix case when both the Windows and Linux builds of Node
|
||||
# are installed in the same directory
|
||||
$exe=".exe"
|
||||
}
|
||||
$ret=0
|
||||
if (Test-Path "$basedir/node$exe") {
|
||||
# Support pipeline input
|
||||
if ($MyInvocation.ExpectingInput) {
|
||||
$input | & "$basedir/node$exe" "$basedir/../browserslist/cli.js" $args
|
||||
} else {
|
||||
& "$basedir/node$exe" "$basedir/../browserslist/cli.js" $args
|
||||
}
|
||||
$ret=$LASTEXITCODE
|
||||
} else {
|
||||
# Support pipeline input
|
||||
if ($MyInvocation.ExpectingInput) {
|
||||
$input | & "node$exe" "$basedir/../browserslist/cli.js" $args
|
||||
} else {
|
||||
& "node$exe" "$basedir/../browserslist/cli.js" $args
|
||||
}
|
||||
$ret=$LASTEXITCODE
|
||||
}
|
||||
exit $ret
|
||||
16
frontend/node_modules/.bin/cssesc
generated
vendored
Normal file
16
frontend/node_modules/.bin/cssesc
generated
vendored
Normal file
@@ -0,0 +1,16 @@
|
||||
#!/bin/sh
|
||||
basedir=$(dirname "$(echo "$0" | sed -e 's,\\,/,g')")
|
||||
|
||||
case `uname` in
|
||||
*CYGWIN*|*MINGW*|*MSYS*)
|
||||
if command -v cygpath > /dev/null 2>&1; then
|
||||
basedir=`cygpath -w "$basedir"`
|
||||
fi
|
||||
;;
|
||||
esac
|
||||
|
||||
if [ -x "$basedir/node" ]; then
|
||||
exec "$basedir/node" "$basedir/../cssesc/bin/cssesc" "$@"
|
||||
else
|
||||
exec node "$basedir/../cssesc/bin/cssesc" "$@"
|
||||
fi
|
||||
17
frontend/node_modules/.bin/cssesc.cmd
generated
vendored
Normal file
17
frontend/node_modules/.bin/cssesc.cmd
generated
vendored
Normal file
@@ -0,0 +1,17 @@
|
||||
@ECHO off
|
||||
GOTO start
|
||||
:find_dp0
|
||||
SET dp0=%~dp0
|
||||
EXIT /b
|
||||
:start
|
||||
SETLOCAL
|
||||
CALL :find_dp0
|
||||
|
||||
IF EXIST "%dp0%\node.exe" (
|
||||
SET "_prog=%dp0%\node.exe"
|
||||
) ELSE (
|
||||
SET "_prog=node"
|
||||
SET PATHEXT=%PATHEXT:;.JS;=;%
|
||||
)
|
||||
|
||||
endLocal & goto #_undefined_# 2>NUL || title %COMSPEC% & "%_prog%" "%dp0%\..\cssesc\bin\cssesc" %*
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user