Compare commits
18 Commits
main
..
e657693e2a
| Author | SHA1 | Date | |
|---|---|---|---|
 kitos
|
e657693e2a | ||
|
kitos
|
d49c4aa009 | ||
|
kitos
|
24062bd009 | ||
|
kitos
|
207973aab8 | ||
|
kitos
|
035b51b3d6 | ||
|
kitos
|
b64b06f7e9 | ||
|
kitos
|
1f136a846c | ||
|
kitos
|
1c5ece1593 | ||
|
kitos
|
a05f37a99b | ||
|
kitos
|
2a278a612a | ||
|
kitos
|
372183cd96 | ||
|
kitos
|
2f0ff7f8a3 | ||
|
kitos
|
a07d1d5f74 | ||
|
kitos
|
13055b24a2 | ||
|
kitos
|
a93c81674d | ||
|
kitos
|
03e06591c6 | ||
|
kitos
|
966aad689c | ||
|
kitos
|
d392d41bf8 |
+6
-26
@@ -1,44 +1,24 @@
|
||||
# =============================================================================
|
||||
# Aegis Environment Variables
|
||||
# =============================================================================
|
||||
# Copy this file to .env and fill in the values BEFORE deploying.
|
||||
#
|
||||
# Generate secure random values with:
|
||||
# openssl rand -hex 32 (for SECRET_KEY)
|
||||
# openssl rand -base64 18 (for passwords)
|
||||
# Copy this file to .env and fill in the values
|
||||
# =============================================================================
|
||||
|
||||
# ── Database ─────────────────────────────────────────────────────────────────
|
||||
DB_USER=postgres
|
||||
DB_PASSWORD= # REQUIRED — set a strong password
|
||||
DB_PASSWORD=change-me-in-production
|
||||
DB_NAME=attackdb
|
||||
|
||||
# ── Security ─────────────────────────────────────────────────────────────────
|
||||
# REQUIRED in production — the app will refuse to start without it.
|
||||
# Generate with: openssl rand -hex 32
|
||||
SECRET_KEY=
|
||||
|
||||
# IMPORTANT: Generate a strong random key for production
|
||||
# Example: openssl rand -hex 32
|
||||
SECRET_KEY=change-me-in-production-use-a-long-random-string
|
||||
TOKEN_EXPIRE_MINUTES=60
|
||||
|
||||
# ── Initial Admin Account ────────────────────────────────────────────────────
|
||||
# If ADMIN_PASSWORD is empty, a random password is auto-generated and
|
||||
# printed to the backend container logs on first startup.
|
||||
ADMIN_USERNAME=admin
|
||||
ADMIN_PASSWORD=
|
||||
|
||||
# ── MinIO Object Storage ─────────────────────────────────────────────────────
|
||||
MINIO_ACCESS_KEY=minioadmin
|
||||
MINIO_SECRET_KEY= # REQUIRED — set a strong password
|
||||
MINIO_SECRET_KEY=change-me-in-production
|
||||
MINIO_BUCKET=evidence
|
||||
MINIO_SECURE=false # Set to true if MinIO is behind TLS
|
||||
|
||||
# ── CORS ──────────────────────────────────────────────────────────────────────
|
||||
# Comma-separated list of allowed frontend origins
|
||||
CORS_ORIGINS=https://your-domain.com
|
||||
|
||||
# ── Frontend ─────────────────────────────────────────────────────────────────
|
||||
FRONTEND_PORT=80
|
||||
|
||||
# ── Environment flag ─────────────────────────────────────────────────────────
|
||||
# Set to "production" for production deployments (enforces SECRET_KEY, etc.)
|
||||
AEGIS_ENV=production
|
||||
|
||||
@@ -1,64 +0,0 @@
|
||||
name: Aegis CI
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [main, develop]
|
||||
pull_request:
|
||||
branches: [main]
|
||||
|
||||
jobs:
|
||||
lint-and-test:
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
services:
|
||||
postgres:
|
||||
image: postgres:15-alpine
|
||||
env:
|
||||
POSTGRES_DB: testdb
|
||||
POSTGRES_USER: test
|
||||
POSTGRES_PASSWORD: test
|
||||
ports:
|
||||
- 5432:5432
|
||||
options: >-
|
||||
--health-cmd pg_isready
|
||||
--health-interval 10s
|
||||
--health-timeout 5s
|
||||
--health-retries 5
|
||||
redis:
|
||||
image: redis:7-alpine
|
||||
ports:
|
||||
- 6379:6379
|
||||
options: >-
|
||||
--health-cmd "redis-cli ping"
|
||||
--health-interval 10s
|
||||
--health-timeout 5s
|
||||
--health-retries 5
|
||||
|
||||
defaults:
|
||||
run:
|
||||
working-directory: backend
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: "3.11"
|
||||
cache: pip
|
||||
cache-dependency-path: backend/requirements.txt
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
pip install --upgrade pip
|
||||
pip install -r requirements.txt
|
||||
pip install ruff
|
||||
|
||||
- name: Lint
|
||||
run: ruff check app/ tests/
|
||||
|
||||
- name: Test
|
||||
env:
|
||||
DATABASE_URL: postgresql://test:test@localhost:5432/testdb
|
||||
REDIS_URL: redis://localhost:6379/0
|
||||
SECRET_KEY: ci-test-secret-key-not-for-production
|
||||
run: pytest tests/ -v --tb=short
|
||||
@@ -1,71 +0,0 @@
|
||||
name: Snyk Security Scan
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [main, develop]
|
||||
pull_request:
|
||||
branches: [main]
|
||||
schedule:
|
||||
- cron: '0 6 * * 1' # Weekly on Monday 06:00 UTC
|
||||
|
||||
jobs:
|
||||
snyk-backend:
|
||||
name: Python vulnerabilities (backend)
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: "3.11"
|
||||
|
||||
- name: Install backend dependencies
|
||||
run: pip install -r backend/requirements-lock.txt
|
||||
|
||||
- name: Snyk — scan Python packages
|
||||
uses: snyk/actions/python@master
|
||||
env:
|
||||
SNYK_TOKEN: ${{ secrets.SNYK_TOKEN }}
|
||||
with:
|
||||
args: --file=backend/requirements-lock.txt --severity-threshold=high
|
||||
continue-on-error: true # report without blocking CI during initial cleanup
|
||||
|
||||
snyk-frontend:
|
||||
name: npm vulnerabilities (frontend)
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- uses: actions/setup-node@v4
|
||||
with:
|
||||
node-version: '20'
|
||||
|
||||
- name: Install frontend dependencies
|
||||
run: npm ci
|
||||
working-directory: frontend
|
||||
|
||||
- name: Snyk — scan npm packages
|
||||
uses: snyk/actions/node@master
|
||||
env:
|
||||
SNYK_TOKEN: ${{ secrets.SNYK_TOKEN }}
|
||||
with:
|
||||
args: --file=frontend/package.json --severity-threshold=high
|
||||
continue-on-error: true
|
||||
|
||||
snyk-docker-backend:
|
||||
name: Docker image vulnerabilities (backend)
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Build backend image for scanning
|
||||
run: docker build -t aegis-backend:scan backend/
|
||||
|
||||
- name: Snyk — scan Docker image
|
||||
uses: snyk/actions/docker@master
|
||||
env:
|
||||
SNYK_TOKEN: ${{ secrets.SNYK_TOKEN }}
|
||||
with:
|
||||
image: aegis-backend:scan
|
||||
args: --severity-threshold=high
|
||||
continue-on-error: true
|
||||
@@ -60,12 +60,3 @@ Thumbs.db
|
||||
|
||||
# Local development
|
||||
*.local
|
||||
|
||||
# Documentation drafts — never commit, delivered directly in chat
|
||||
docs/confluence/
|
||||
docs/drafts/
|
||||
|
||||
# Editor / AI assistant working files — never commit
|
||||
.claude/
|
||||
.cursor/
|
||||
CLAUDE.md
|
||||
|
||||
+1232
File diff suppressed because it is too large
Load Diff
@@ -1,97 +1,24 @@
|
||||
# Aegis — MITRE ATT&CK Coverage Platform
|
||||
# Aegis - MITRE ATT&CK Coverage Platform
|
||||
|
||||
Continuous integration (lint + tests against PostgreSQL and Redis) is defined in [`.github/workflows/ci.yml`](.github/workflows/ci.yml).
|
||||
|
||||
Aegis is a comprehensive platform for tracking and managing security coverage against the MITRE ATT&CK framework. It enables security teams to document, validate, and visualize their defensive capabilities against known adversary techniques through a structured Red Team / Blue Team validation workflow.
|
||||
Aegis is a comprehensive platform for tracking and managing security coverage against the MITRE ATT&CK framework. It enables security teams to document, validate, and visualize their defensive capabilities against known adversary techniques.
|
||||
|
||||
## Features
|
||||
|
||||
### Core (V1)
|
||||
- **MITRE ATT&CK Integration** — Automatic synchronization via TAXII 2.0 (with GitHub fallback), scheduled every 24h
|
||||
- **Red/Blue Validation Workflow** — Structured dual-validation lifecycle: draft → red_executing → blue_evaluating → in_review → validated/rejected
|
||||
- **Dual Validation** — Independent Red Lead / Blue Lead approval before finalization
|
||||
- **Coverage Tracking** — Per-technique status (validated, partial, not covered, in progress)
|
||||
- **Evidence Storage** — Secure evidence with SHA256 integrity, separated by team (red/blue)
|
||||
- **Role-Based Access Control** — Granular permissions for 6 roles (admin, red_tech, blue_tech, red_lead, blue_lead, viewer)
|
||||
|
||||
### Enhanced (V2)
|
||||
- **Test Template Catalog** — Import from Atomic Red Team, CALDERA, LOLBAS, GTFOBins; create custom templates; bulk activate/deactivate
|
||||
- **In-App Notifications** — Real-time notification bell with polling and automatic state-change alerts
|
||||
- **Reports & Export** — Coverage summary, test results, and remediation reports in JSON and CSV
|
||||
- **Remediation Tracking** — Step-by-step remediation assignments with status tracking
|
||||
- **Metrics Dashboard** — Pipeline funnel, team activity, validation rates
|
||||
|
||||
### Advanced (V3)
|
||||
- **Multi-Source Data Import** — Sigma, CALDERA, LOLBAS, GTFOBins, D3FEND, MITRE CTI threat actors, compliance mappings (NIST 800-53, CIS Controls v8)
|
||||
- **Detection Rule Tracking** — Import and evaluate Sigma/Elastic detection rules per test
|
||||
- **ATT&CK Heatmap** — Interactive Navigator-style heatmap with layers, filters, and export
|
||||
- **Threat Actor Intelligence** — Track intrusion sets and their technique coverage
|
||||
- **Campaign Management** — Group tests into campaigns with dependencies, scheduling, and recurring execution
|
||||
- **Compliance Mapping** — Map NIST 800-53 and CIS Controls v8 to ATT&CK techniques with gap analysis
|
||||
- **Granular Scoring** — 0–100 scoring for techniques, tactics, actors, and organization with configurable weights
|
||||
- **Operational Metrics** — MTTD, MTTR, detection efficacy, alert fidelity, coverage velocity
|
||||
- **Executive Dashboard** — High-level KPIs for leadership (leads + admin)
|
||||
- **Temporal Comparison** — Coverage snapshots with historical comparison and trend analysis
|
||||
- **Auto Re-testing** — Automatic retest creation after remediation completion (configurable limit)
|
||||
- **Performance Optimizations** — Score caching, lazy loading, React.memo, database index optimization
|
||||
|
||||
## Red Team / Blue Team Validation Flow
|
||||
|
||||
```
|
||||
┌──────┐ ┌──────────────┐ ┌─────────────────┐ ┌───────────┐
|
||||
│ DRAFT│───▶│RED_EXECUTING │───▶│ BLUE_EVALUATING │───▶│ IN_REVIEW │
|
||||
└──────┘ └──────────────┘ └─────────────────┘ └─────┬─────┘
|
||||
│
|
||||
┌───────────────────┤
|
||||
▼ ▼
|
||||
┌──────────┐ ┌──────────┐
|
||||
│ REJECTED │ │VALIDATED │
|
||||
└────┬─────┘ └──────────┘
|
||||
│ │
|
||||
└──▶ Back to DRAFT ├──▶ Remediation
|
||||
└──▶ Auto Re-test
|
||||
```
|
||||
|
||||
### States
|
||||
|
||||
| State | Description | Who acts |
|
||||
|-------|-------------|----------|
|
||||
| `draft` | Created, pending execution | Red Tech |
|
||||
| `red_executing` | Red Team documents attack & uploads evidence | Red Tech |
|
||||
| `blue_evaluating` | Blue Team documents detection & uploads evidence | Blue Tech |
|
||||
| `in_review` | Both managers review evidence | Red Lead, Blue Lead |
|
||||
| `validated` | Approved by both managers | — (terminal) |
|
||||
| `rejected` | Rejected — returns to draft for redo | Red/Blue Lead can reopen |
|
||||
|
||||
### Dual Validation
|
||||
|
||||
Both Red Lead and Blue Lead must independently vote:
|
||||
- **Both approve** → test moves to `validated`
|
||||
- **Either rejects** → test moves to `rejected`
|
||||
- **One votes, other pending** → stays in `in_review`
|
||||
|
||||
## User Roles
|
||||
|
||||
| Role | Description | Capabilities |
|
||||
|------|-------------|-------------|
|
||||
| `admin` | Full system access | Everything |
|
||||
| `red_tech` | Red team technician | Create tests, document attacks, upload red evidence |
|
||||
| `blue_tech` | Blue team technician | Document detection, upload blue evidence |
|
||||
| `red_lead` | Red team lead | Validate/reject the red side of tests |
|
||||
| `blue_lead` | Blue team lead | Validate/reject the blue side of tests |
|
||||
| `viewer` | Read-only access | View all data |
|
||||
- **MITRE ATT&CK Integration**: Automatic synchronization with the MITRE ATT&CK framework via TAXII (with GitHub fallback), scheduled every 24h
|
||||
- **Coverage Tracking**: Track validation status for each technique (validated, partial, not covered, in progress)
|
||||
- **Test Management**: Document and manage security tests with full audit trail
|
||||
- **Evidence Storage**: Secure evidence file storage with SHA256 integrity verification
|
||||
- **Role-Based Access Control**: Granular permissions for red team, blue team, and leadership roles
|
||||
- **Intel Monitoring**: Automated scanning for new threat intelligence related to techniques
|
||||
- **Metrics Dashboard**: Real-time coverage metrics and reporting by tactic
|
||||
|
||||
## Tech Stack
|
||||
|
||||
- **Backend**: FastAPI (Python 3.11) — Clean Modular Monolith with domain entities, services, and repository pattern
|
||||
- **Database**: PostgreSQL 16 with UUID primary keys and JSONB columns
|
||||
- **Backend**: FastAPI (Python 3.11)
|
||||
- **Database**: PostgreSQL 15
|
||||
- **Object Storage**: MinIO (S3-compatible)
|
||||
- **ORM**: SQLAlchemy 2.x with Alembic migrations
|
||||
- **Frontend**: React 19 + TypeScript + Vite 7 + Tailwind CSS v4 + TanStack Query + TanStack Virtual
|
||||
- **Cache / Token Store**: Redis (token blacklist, score caching)
|
||||
- **Scheduler**: APScheduler (MITRE sync, Intel scan, Notification cleanup, Snapshots, Recurring campaigns)
|
||||
- **Testing**: Pytest (367+ tests), Ruff (linting), GitHub Actions CI
|
||||
- **Charts**: Recharts
|
||||
- **ORM**: SQLAlchemy with Alembic migrations
|
||||
- **Frontend**: React 19 + TypeScript + Vite + Tailwind CSS v4 + TanStack Query
|
||||
|
||||
## Quick Start
|
||||
|
||||
@@ -99,341 +26,340 @@ Both Red Lead and Blue Lead must independently vote:
|
||||
|
||||
- Docker and Docker Compose
|
||||
- Git
|
||||
- Linux / macOS (or WSL on Windows)
|
||||
|
||||
### Production Deployment
|
||||
|
||||
The recommended way to deploy Aegis in production:
|
||||
### Installation
|
||||
|
||||
1. Clone the repository:
|
||||
```bash
|
||||
git clone <repository-url>
|
||||
cd Aegis
|
||||
chmod +x scripts/install.sh
|
||||
./scripts/install.sh
|
||||
```
|
||||
|
||||
The interactive install wizard will guide you through:
|
||||
1. **Domain configuration** — your domain or IP, protocol (HTTP/HTTPS), and port
|
||||
2. **Admin account** — custom username and password (or auto-generated secure password)
|
||||
3. **Database** — name, user, and password (or auto-generated)
|
||||
4. **Session duration** — JWT token expiry (default: 15 minutes)
|
||||
5. **MITRE ATT&CK sync** — optionally import ~700 techniques on first run
|
||||
|
||||
The script automatically generates cryptographically secure random secrets for `SECRET_KEY`, database password, and MinIO credentials. A summary with all credentials is displayed at the end of the installation.
|
||||
|
||||
Access the application at the URL shown in the installation summary.
|
||||
|
||||
### Development Setup
|
||||
|
||||
For local development with hot-reload:
|
||||
|
||||
2. Start all services:
|
||||
```bash
|
||||
git clone <repository-url>
|
||||
cd Aegis
|
||||
docker-compose up -d
|
||||
./scripts/init.sh
|
||||
```
|
||||
|
||||
Access at **http://localhost:5173** (frontend dev server) and **http://localhost:8000/docs** (API docs).
|
||||
3. Run database migrations:
|
||||
```bash
|
||||
docker exec -w /app aegis-backend-1 alembic upgrade head
|
||||
```
|
||||
|
||||
4. Seed the admin user:
|
||||
```bash
|
||||
docker exec -w /app aegis-backend-1 python -m app.seed
|
||||
```
|
||||
|
||||
5. Start the frontend (requires Node.js 20+ or Docker):
|
||||
```bash
|
||||
# Option A — with Node.js installed locally
|
||||
cd frontend && npm install && npm run dev
|
||||
|
||||
# Option B — via Docker
|
||||
docker run --rm -v ./frontend:/app -w /app -p 5173:5173 node:20-alpine sh -c "npm run dev"
|
||||
```
|
||||
|
||||
6. Verify the installation:
|
||||
```bash
|
||||
# Backend health
|
||||
curl http://localhost:8000/health
|
||||
# Expected: {"status":"ok"}
|
||||
|
||||
# Frontend
|
||||
# Open http://localhost:5173 — should show the Aegis login page
|
||||
```
|
||||
|
||||
### Authentication
|
||||
|
||||
JWT-based authentication with HttpOnly cookies. Admin credentials are configured during installation:
|
||||
|
||||
- If you set a custom password in the wizard, use that.
|
||||
- If you left it blank, a secure random password was auto-generated and displayed in the installation summary and backend logs.
|
||||
|
||||
To retrieve auto-generated credentials after installation:
|
||||
The platform uses JWT-based authentication. After seeding, log in with the default admin credentials:
|
||||
|
||||
```bash
|
||||
docker logs aegis-backend 2>&1 | grep -A 5 "Initial Admin User Created"
|
||||
# Obtain a token
|
||||
curl -X POST http://localhost:8000/api/v1/auth/login \
|
||||
-d "username=admin&password=admin123"
|
||||
|
||||
# Use the token to access protected endpoints
|
||||
curl http://localhost:8000/api/v1/auth/me \
|
||||
-H "Authorization: Bearer <your-token>"
|
||||
```
|
||||
|
||||
> **Note:** Passwords must meet complexity requirements: minimum 12 characters with at least one uppercase letter, one lowercase letter, one digit, and one special character.
|
||||
|
||||
### Importing Data Sources
|
||||
|
||||
On startup, the backend automatically seeds the initial data sources (Atomic Red Team, SigmaHQ, CALDERA, LOLBAS, GTFOBins, D3FEND). You can then sync each source from the UI:
|
||||
|
||||
1. Navigate to **Data Sources** in the sidebar
|
||||
2. Click **Sync** on each data source to import its content
|
||||
3. Trigger a **MITRE ATT&CK Sync** from the **System** page
|
||||
|
||||
Alternatively, use the API:
|
||||
|
||||
```bash
|
||||
# Sync MITRE ATT&CK techniques
|
||||
curl -X POST http://your-server/api/v1/system/sync-mitre -H "Authorization: Bearer $TOKEN"
|
||||
|
||||
# Sync all data sources at once
|
||||
curl -X POST http://your-server/api/v1/data-sources/sync-all -H "Authorization: Bearer $TOKEN"
|
||||
```
|
||||
|
||||
### Production Considerations
|
||||
|
||||
- **HTTPS/TLS:** For internet-facing deployments, place a reverse proxy with TLS in front (e.g., Traefik, Caddy, or Nginx with Let's Encrypt). Uncomment the HSTS header in `frontend/nginx.conf` once HTTPS is configured.
|
||||
- **Backups:** Set up regular PostgreSQL backups: `docker exec aegis-postgres pg_dump -U postgres attackdb > backup.sql`
|
||||
- **Updates:** To update, pull the latest code and run: `docker compose -f docker-compose.prod.yml up -d --build`
|
||||
- **Firewall:** Only expose port 80/443. All other services (DB, MinIO, backend) are internal only.
|
||||
- **Reconfigure:** Run `./scripts/install.sh` again to reconfigure the environment (domain, credentials, etc.).
|
||||
|
||||
### Configuring Scoring Weights
|
||||
|
||||
Default scoring weights can be adjusted via environment variables:
|
||||
|
||||
```env
|
||||
SCORING_WEIGHT_TESTS=40
|
||||
SCORING_WEIGHT_DETECTION_RULES=20
|
||||
SCORING_WEIGHT_D3FEND=15
|
||||
SCORING_WEIGHT_FRESHNESS=15
|
||||
SCORING_WEIGHT_PLATFORM_DIVERSITY=10
|
||||
```
|
||||
|
||||
Or at runtime via the admin API — see [docs/SCORING.md](docs/SCORING.md).
|
||||
|
||||
### Configuring Campaigns
|
||||
|
||||
1. Navigate to **Campaigns** in the sidebar
|
||||
2. Create a new campaign (custom or from a threat actor)
|
||||
3. Add tests with dependencies and phases
|
||||
4. Optionally enable **recurring scheduling** (weekly, monthly, quarterly)
|
||||
> **Important:** Change the default `admin123` password and `SECRET_KEY` in production.
|
||||
|
||||
## Services
|
||||
|
||||
| Service | Port | Description |
|
||||
|---------|------|-------------|
|
||||
| Service | Port | Description |
|
||||
|----------|------|-------------|
|
||||
| Frontend | 5173 | React dev server (Vite) |
|
||||
| Backend | 8000 | FastAPI REST API |
|
||||
| PostgreSQL | 5433 | Database |
|
||||
| Backend | 8000 | FastAPI REST API |
|
||||
| PostgreSQL | 5433 | Database (mapped to 5433 to avoid conflicts) |
|
||||
| MinIO API | 9000 | S3-compatible object storage |
|
||||
| MinIO Console | 9001 | MinIO web interface |
|
||||
|
||||
## Navigation
|
||||
|
||||
```
|
||||
📊 Dashboard
|
||||
📊 Executive Dashboard (leads + admin)
|
||||
🔲 ATT&CK Matrix (advanced heatmap)
|
||||
🧪 Tests
|
||||
├─ All Tests
|
||||
├─ My Pending Tasks
|
||||
└─ Test Catalog
|
||||
📋 Campaigns
|
||||
👤 Threat Actors
|
||||
📜 Compliance
|
||||
📈 Comparison (leads + admin)
|
||||
📄 Reports
|
||||
⚙️ System (admin only)
|
||||
├─ Data Sources (sync Atomic, Sigma, CALDERA, LOLBAS, GTFOBins, D3FEND)
|
||||
├─ MITRE Sync (ATT&CK sync, intel scan, template management)
|
||||
├─ Users
|
||||
└─ Audit Log
|
||||
```
|
||||
|
||||
## API Documentation
|
||||
|
||||
Interactive API documentation is available **in development only** (disabled in production for security):
|
||||
Once the backend is running, access the interactive API documentation at:
|
||||
|
||||
- **Swagger UI**: http://localhost:8000/docs
|
||||
- **ReDoc**: http://localhost:8000/redoc
|
||||
|
||||
> In production (`AEGIS_ENV=production`), these endpoints are disabled. Use the development environment or refer to [docs/API.md](docs/API.md).
|
||||
## API Endpoints
|
||||
|
||||
### API Endpoints
|
||||
### Auth
|
||||
| Method | Route | Auth | Description |
|
||||
|--------|-------|------|-------------|
|
||||
| POST | `/api/v1/auth/login` | Public | Obtain JWT token |
|
||||
| GET | `/api/v1/auth/me` | Authenticated | Current user profile |
|
||||
|
||||
| Group | Prefix | Key Operations |
|
||||
|-------|--------|---------------|
|
||||
| Auth | `/api/v1/auth` | Login, get current user |
|
||||
| Techniques | `/api/v1/techniques` | CRUD, list with filters, mark reviewed |
|
||||
| Tests | `/api/v1/tests` | Full Red/Blue workflow, remediation, retest chain |
|
||||
| Test Templates | `/api/v1/test-templates` | CRUD, stats, toggle active, bulk activate/deactivate |
|
||||
| Evidence | `/api/v1/tests/{id}/evidence` | Upload evidence, get presigned URLs |
|
||||
| Campaigns | `/api/v1/campaigns` | CRUD, scheduling, history |
|
||||
| Threat Actors | `/api/v1/threat-actors` | CRUD, technique mappings |
|
||||
| Detection Rules | `/api/v1/detection-rules` | List, filter by source/technique |
|
||||
| D3FEND | `/api/v1/d3fend` | Defensive techniques and mappings |
|
||||
| Compliance | `/api/v1/compliance` | Frameworks, controls, gaps |
|
||||
| Scores | `/api/v1/scores` | Technique/tactic/actor/org scores, config |
|
||||
| Operational Metrics | `/api/v1/metrics/operational` | MTTD, MTTR, trends, team breakdown |
|
||||
| Heatmap | `/api/v1/heatmap` | ATT&CK Navigator-style data |
|
||||
| Snapshots | `/api/v1/snapshots` | Create, compare, list snapshots |
|
||||
| Reports | `/api/v1/reports` | Coverage, results, remediation exports |
|
||||
| Notifications | `/api/v1/notifications` | List, read, mark all read |
|
||||
| Metrics | `/api/v1/metrics` | Summary, by-tactic, pipeline, team activity |
|
||||
| System | `/api/v1/system` | MITRE sync, import, scheduler status |
|
||||
| Users | `/api/v1/users` | User CRUD (admin) |
|
||||
| Audit Logs | `/api/v1/audit-logs` | Audit trail (admin) |
|
||||
| Data Sources | `/api/v1/data-sources` | Data source management (admin) |
|
||||
### Techniques
|
||||
| Method | Route | Auth | Description |
|
||||
|--------|-------|------|-------------|
|
||||
| GET | `/api/v1/techniques` | Authenticated | List all (filters: `?tactic=`, `?status=`, `?review_required=`) |
|
||||
| GET | `/api/v1/techniques/{mitre_id}` | Authenticated | Detail with associated tests |
|
||||
| POST | `/api/v1/techniques` | Admin | Create technique |
|
||||
| PATCH | `/api/v1/techniques/{mitre_id}` | Admin | Update technique fields |
|
||||
| PATCH | `/api/v1/techniques/{mitre_id}/review` | Lead, Admin | Mark as reviewed |
|
||||
|
||||
See [docs/API.md](docs/API.md) for the full endpoint reference.
|
||||
### Tests
|
||||
| Method | Route | Auth | Description |
|
||||
|--------|-------|------|-------------|
|
||||
| POST | `/api/v1/tests` | Red Tech, Admin | Create test (state=draft) |
|
||||
| GET | `/api/v1/tests/{id}` | Authenticated | Detail with evidences |
|
||||
| PATCH | `/api/v1/tests/{id}` | Creator, Admin | Update (only draft/rejected) |
|
||||
| POST | `/api/v1/tests/{id}/validate` | Lead, Admin | Validate + recalculate technique status |
|
||||
| POST | `/api/v1/tests/{id}/reject` | Lead, Admin | Reject test |
|
||||
|
||||
## Configuration
|
||||
### Evidence
|
||||
| Method | Route | Auth | Description |
|
||||
|--------|-------|------|-------------|
|
||||
| POST | `/api/v1/tests/{test_id}/evidence` | Authenticated | Upload evidence file (SHA-256 verified) |
|
||||
| GET | `/api/v1/evidence/{id}` | Authenticated | Get metadata + presigned download URL |
|
||||
|
||||
All variables are configured automatically by `scripts/install.sh`. For manual setup, copy `.env.example` to `.env` and fill in the values.
|
||||
### System
|
||||
| Method | Route | Auth | Description |
|
||||
|--------|-------|------|-------------|
|
||||
| POST | `/api/v1/system/sync-mitre` | Admin | Manually trigger MITRE ATT&CK sync |
|
||||
| POST | `/api/v1/system/run-intel-scan` | Admin | Manually trigger threat-intel RSS scan |
|
||||
| GET | `/api/v1/system/scheduler-status` | Admin | Background scheduler health & job list |
|
||||
|
||||
### Required (production)
|
||||
### Metrics
|
||||
| Method | Route | Auth | Description |
|
||||
|--------|-------|------|-------------|
|
||||
| GET | `/api/v1/metrics/summary` | Authenticated | Global coverage summary (counts + percentage) |
|
||||
| GET | `/api/v1/metrics/by-tactic` | Authenticated | Coverage breakdown per MITRE tactic |
|
||||
|
||||
| Variable | Description |
|
||||
|----------|-------------|
|
||||
| `SECRET_KEY` | JWT signing key — **required** in production (app refuses to start without it). Generate with `openssl rand -hex 32` |
|
||||
| `DB_PASSWORD` | PostgreSQL password |
|
||||
| `MINIO_SECRET_KEY` | MinIO secret key |
|
||||
### Users (Admin)
|
||||
| Method | Route | Auth | Description |
|
||||
|--------|-------|------|-------------|
|
||||
| GET | `/api/v1/users` | Admin | List all users |
|
||||
| POST | `/api/v1/users` | Admin | Create new user |
|
||||
| GET | `/api/v1/users/{id}` | Admin | Get user by ID |
|
||||
| PATCH | `/api/v1/users/{id}` | Admin | Update user (role, email, active status) |
|
||||
|
||||
### Security & Auth
|
||||
|
||||
| Variable | Default | Description |
|
||||
|----------|---------|-------------|
|
||||
| `AEGIS_ENV` | — | Set to `production` to enforce security settings |
|
||||
| `ADMIN_USERNAME` | `admin` | Initial admin account username |
|
||||
| `ADMIN_PASSWORD` | *(auto-generated)* | Initial admin password. If empty, a secure random password is generated and shown in logs |
|
||||
| `ACCESS_TOKEN_EXPIRE_MINUTES` | `15` | JWT token lifetime in minutes |
|
||||
| `CORS_ORIGINS` | `http://localhost:5173` | Comma-separated list of allowed frontend origins |
|
||||
|
||||
### Infrastructure
|
||||
|
||||
| Variable | Default | Description |
|
||||
|----------|---------|-------------|
|
||||
| `DB_USER` | `postgres` | PostgreSQL username |
|
||||
| `DB_NAME` | `attackdb` | PostgreSQL database name |
|
||||
| `MINIO_ENDPOINT` | `minio:9000` | MinIO server address |
|
||||
| `MINIO_ACCESS_KEY` | `minioadmin` | MinIO access key |
|
||||
| `MINIO_BUCKET` | `evidence` | MinIO bucket for evidence files |
|
||||
| `MINIO_SECURE` | `false` | Set to `true` if MinIO is behind TLS |
|
||||
| `FRONTEND_PORT` | `80` | Port exposed by the frontend container |
|
||||
|
||||
### Scoring Weights
|
||||
|
||||
| Variable | Default | Description |
|
||||
|----------|---------|-------------|
|
||||
| `SCORING_WEIGHT_TESTS` | `40` | Weight for test validation component |
|
||||
| `SCORING_WEIGHT_DETECTION_RULES` | `20` | Weight for detection rules component |
|
||||
| `SCORING_WEIGHT_D3FEND` | `15` | Weight for D3FEND coverage component |
|
||||
| `SCORING_WEIGHT_FRESHNESS` | `15` | Weight for freshness component |
|
||||
| `SCORING_WEIGHT_PLATFORM_DIVERSITY` | `10` | Weight for platform diversity component |
|
||||
| `MAX_RETEST_COUNT` | `3` | Max automatic retests per original test |
|
||||
|
||||
## Security
|
||||
|
||||
Aegis includes several security hardening measures:
|
||||
|
||||
- **Authentication:** JWT tokens stored in HttpOnly/Secure/SameSite cookies (immune to XSS theft). Token revocation via Redis-backed blacklist on logout.
|
||||
- **Rate limiting:** Login endpoint limited to 5 attempts per minute per IP (via slowapi).
|
||||
- **Password policy:** Minimum 12 characters with uppercase, lowercase, digit, and special character.
|
||||
- **CORS:** Configurable origins via `CORS_ORIGINS` environment variable. Restrictive method and header lists.
|
||||
- **Nginx security headers:** CSP, X-Frame-Options, X-Content-Type-Options, Referrer-Policy, Permissions-Policy.
|
||||
- **Non-root container:** Backend runs as `appuser` (UID 1001), not root.
|
||||
- **File uploads:** 50 MB size limit, extension whitelist, filename sanitization.
|
||||
- **ZIP imports:** Zip Slip (path traversal) and Zip Bomb (size/entry limit) protection.
|
||||
- **API surface:** Swagger UI, ReDoc, and OpenAPI schema disabled in production.
|
||||
- **Health endpoint:** Restricted to internal networks via Nginx ACL.
|
||||
- **Input sanitization:** LIKE wildcard escaping in all search queries; Pydantic schemas on all endpoints.
|
||||
- **XML parsing:** Uses `defusedxml` to prevent Billion Laughs / XXE attacks.
|
||||
- **Error handling:** Internal exception details are logged server-side only, never exposed to clients.
|
||||
### Audit Logs (Admin)
|
||||
| Method | Route | Auth | Description |
|
||||
|--------|-------|------|-------------|
|
||||
| GET | `/api/v1/audit-logs` | Admin | List audit logs (filters: `?action=`, `?entity_type=`, `?start_date=`, `?end_date=`) |
|
||||
| GET | `/api/v1/audit-logs/actions` | Admin | List distinct action types |
|
||||
| GET | `/api/v1/audit-logs/entity-types` | Admin | List distinct entity types |
|
||||
|
||||
## Project Structure
|
||||
|
||||
```
|
||||
Aegis/
|
||||
├── docker-compose.yml
|
||||
├── docker-compose.prod.yml
|
||||
├── .github/workflows/ci.yml # GitHub Actions: ruff + pytest on PostgreSQL + Redis
|
||||
├── docs/
|
||||
│ ├── API.md # Full API endpoint reference
|
||||
│ ├── ARCHITECTURE.md # System architecture, DB schema, service map
|
||||
│ ├── ADR.md # Architecture Decision Records
|
||||
│ ├── DATA_SOURCES.md # External data source documentation
|
||||
│ ├── SCORING.md # Scoring system and metrics
|
||||
│ ├── TECHNOLOGY_JUSTIFICATION.md
|
||||
│ ├── C4_CONTEXT_DIAGRAM.md # System context (C4 Level 1)
|
||||
│ └── C4_CONTAINER_DIAGRAM.md # Container architecture (C4 Level 2)
|
||||
├── docker-compose.yml # Docker services configuration
|
||||
├── backend/
|
||||
│ ├── Dockerfile
|
||||
│ ├── requirements.txt
|
||||
│ ├── alembic.ini
|
||||
│ ├── alembic/versions/ # Database migration files
|
||||
│ ├── pytest.ini
|
||||
│ ├── tests/ # 367+ pytest tests (domain, service, API)
|
||||
│ ├── Dockerfile # Backend container definition
|
||||
│ ├── requirements.txt # Python dependencies
|
||||
│ ├── alembic.ini # Alembic configuration
|
||||
│ ├── alembic/ # Database migrations
|
||||
│ │ ├── env.py
|
||||
│ │ ├── versions/ # Migration files
|
||||
│ │ └── ...
|
||||
│ └── app/
|
||||
│ ├── main.py # FastAPI app with all routers + lifespan
|
||||
│ ├── config.py # Pydantic Settings from environment
|
||||
│ ├── database.py # SQLAlchemy engine + session (lazy init)
|
||||
│ ├── storage.py # MinIO/S3 helpers
|
||||
│ ├── auth.py # Password hashing + JWT tokens
|
||||
│ ├── domain/ # Pure business logic (zero framework imports)
|
||||
│ │ ├── entities/ # Rich domain entities (Technique, Campaign, etc.)
|
||||
│ │ ├── ports/ # Protocol interfaces (repos, ImportService)
|
||||
│ │ ├── value_objects/ # Immutable types (MitreId, ScoringWeights)
|
||||
│ │ ├── errors.py # Domain exception hierarchy
|
||||
│ │ └── unit_of_work.py # Transaction management
|
||||
│ ├── infrastructure/ # SQLAlchemy repos, Redis, mappers
|
||||
│ ├── models/ # SQLAlchemy ORM models
|
||||
│ ├── schemas/ # Pydantic request/response schemas
|
||||
│ ├── routers/ # 27 thin HTTP adapter routers
|
||||
│ ├── services/ # 46 framework-agnostic business services
|
||||
│ ├── middleware/ # Error handler (domain exceptions → HTTP)
|
||||
│ ├── dependencies/ # FastAPI dependency injection (auth, repos)
|
||||
│ └── jobs/ # APScheduler background jobs
|
||||
└── frontend/src/
|
||||
├── App.tsx # Routes with lazy loading + role protection
|
||||
├── api/ # API client modules (Axios + TanStack Query)
|
||||
├── components/ # Reusable UI components
|
||||
├── hooks/ # Custom hooks (useDebounce, etc.)
|
||||
├── context/ # Auth state management
|
||||
└── pages/ # Page components
|
||||
│ ├── __init__.py
|
||||
│ ├── main.py # FastAPI application entry point
|
||||
│ ├── config.py # Application settings
|
||||
│ ├── database.py # SQLAlchemy configuration
|
||||
│ ├── auth.py # Password hashing & JWT utilities
|
||||
│ ├── seed.py # Admin seed script (python -m app.seed)
|
||||
│ ├── models/ # SQLAlchemy models
|
||||
│ │ ├── user.py # User authentication model
|
||||
│ │ ├── technique.py # MITRE ATT&CK techniques
|
||||
│ │ ├── test.py # Security tests
|
||||
│ │ ├── evidence.py # Test evidence files
|
||||
│ │ ├── intel.py # Threat intelligence items
|
||||
│ │ ├── audit.py # Audit logging
|
||||
│ │ └── enums.py # Shared enumerations
|
||||
│ ├── storage.py # MinIO/S3 client (upload, presigned URLs)
|
||||
│ ├── schemas/ # Pydantic request/response schemas
|
||||
│ │ ├── auth.py # LoginRequest, TokenResponse, UserOut
|
||||
│ │ ├── technique.py # TechniqueCreate/Update/Out/Summary
|
||||
│ │ ├── test.py # TestCreate/Update/Out/Validate
|
||||
│ │ └── evidence.py # EvidenceOut
|
||||
│ ├── routers/ # API endpoint routers
|
||||
│ │ ├── auth.py # POST /auth/login, GET /auth/me
|
||||
│ │ ├── techniques.py # CRUD techniques (list, detail, create, update, review)
|
||||
│ │ ├── tests.py # CRUD tests (create, detail, update, validate, reject)
|
||||
│ │ ├── evidence.py # Upload evidence, presigned download
|
||||
│ │ ├── system.py # MITRE sync trigger, scheduler status
|
||||
│ │ ├── metrics.py # Coverage summary & per-tactic breakdown
|
||||
│ │ ├── users.py # User management (admin only)
|
||||
│ │ └── audit.py # Audit log viewer (admin only)
|
||||
│ ├── dependencies/ # FastAPI dependencies (DI)
|
||||
│ │ └── auth.py # get_current_user, require_role, require_any_role
|
||||
│ ├── jobs/ # Background scheduled jobs
|
||||
│ │ └── mitre_sync_job.py # APScheduler: MITRE sync (24h) + Intel scan (7d)
|
||||
│ └── services/ # Business logic services
|
||||
│ ├── audit_service.py
|
||||
│ ├── status_service.py # Recalculate technique status from tests
|
||||
│ ├── mitre_sync_service.py # MITRE ATT&CK sync via TAXII / GitHub
|
||||
│ └── intel_service.py # Automated intel scan via RSS feeds
|
||||
└── frontend/ # React + TypeScript frontend
|
||||
├── index.html
|
||||
├── package.json
|
||||
├── tsconfig.json
|
||||
├── vite.config.ts
|
||||
└── src/
|
||||
├── main.tsx # App entry point
|
||||
├── App.tsx # Route definitions
|
||||
├── index.css # Tailwind CSS entry
|
||||
├── api/ # Axios clients
|
||||
│ ├── client.ts # Base axios instance with JWT interceptor
|
||||
│ ├── auth.ts # login(), getMe()
|
||||
│ ├── metrics.ts # getCoverageSummary(), getCoverageByTactic()
|
||||
│ ├── techniques.ts # getTechniques(), getTechniqueByMitreId()
|
||||
│ ├── tests.ts # createTest(), validateTest(), rejectTest()
|
||||
│ ├── evidence.ts # uploadEvidence(), getEvidence()
|
||||
│ ├── system.ts # triggerMitreSync(), triggerIntelScan()
|
||||
│ ├── users.ts # getUsers(), createUser(), updateUser()
|
||||
│ └── audit.ts # getAuditLogs(), getAuditActions()
|
||||
├── context/
|
||||
│ └── AuthContext.tsx # Auth state: user, login, logout, isLoading
|
||||
├── components/
|
||||
│ ├── Layout.tsx # Sidebar + header + <Outlet/>
|
||||
│ ├── Sidebar.tsx # Nav links (role-aware)
|
||||
│ ├── ProtectedRoute.tsx # Auth route guard with role support
|
||||
│ ├── CoverageSummaryCard.tsx # Metric card component
|
||||
│ ├── TacticCoverageChart.tsx # Coverage breakdown table
|
||||
│ ├── AttackMatrix.tsx # Interactive technique grid
|
||||
│ ├── TechniqueCell.tsx # Individual technique cell in matrix
|
||||
│ ├── TestForm.tsx # Reusable test creation/edit form
|
||||
│ ├── EvidenceUpload.tsx # Drag & drop file upload
|
||||
│ ├── EvidenceList.tsx # Evidence file listing
|
||||
│ ├── ErrorBoundary.tsx # Global error boundary
|
||||
│ ├── ErrorMessage.tsx # Reusable error display
|
||||
│ ├── LoadingSpinner.tsx # Reusable loading indicator
|
||||
│ └── Toast.tsx # Toast notification system
|
||||
├── pages/
|
||||
│ ├── LoginPage.tsx # User authentication form
|
||||
│ ├── DashboardPage.tsx # Coverage metrics dashboard with summary cards
|
||||
│ ├── TechniquesPage.tsx # Interactive ATT&CK matrix view with filters
|
||||
│ ├── TechniqueDetailPage.tsx # Individual technique detail with tests
|
||||
│ ├── TestsPage.tsx # Tests overview and navigation
|
||||
│ ├── TestCreatePage.tsx # Test creation form
|
||||
│ ├── TestDetailPage.tsx # Test details with evidence upload
|
||||
│ ├── SystemPage.tsx # Admin panel for MITRE sync & intel scan
|
||||
│ ├── UsersPage.tsx # User management (admin only)
|
||||
│ └── AuditLogPage.tsx # Audit log viewer (admin only)
|
||||
├── types/
|
||||
│ └── models.ts # TS interfaces matching backend schemas
|
||||
├── hooks/
|
||||
└── lib/
|
||||
```
|
||||
|
||||
## Database Schema
|
||||
|
||||
The platform uses the following data models:
|
||||
|
||||
| Table | Description |
|
||||
|-------|-------------|
|
||||
| `users` | User accounts with role-based access |
|
||||
| `techniques` | MITRE ATT&CK techniques with coverage status |
|
||||
| `tests` | Security tests validating technique coverage |
|
||||
| `evidences` | File evidence attached to tests (stored in MinIO) |
|
||||
| `intel_items` | Threat intelligence items linked to techniques |
|
||||
| `audit_logs` | System-wide audit trail for all actions |
|
||||
|
||||
## Configuration
|
||||
|
||||
The application can be configured via environment variables:
|
||||
|
||||
| Variable | Default | Description |
|
||||
|----------|---------|-------------|
|
||||
| `DATABASE_URL` | `postgresql://postgres:postgres@postgres:5432/attackdb` | PostgreSQL connection string |
|
||||
| `SECRET_KEY` | `change-me-in-production` | JWT signing key |
|
||||
| `ALGORITHM` | `HS256` | JWT signing algorithm |
|
||||
| `ACCESS_TOKEN_EXPIRE_MINUTES` | `60` | JWT token lifetime in minutes |
|
||||
| `MINIO_ENDPOINT` | `minio:9000` | MinIO server endpoint |
|
||||
| `MINIO_ACCESS_KEY` | `minioadmin` | MinIO access key |
|
||||
| `MINIO_SECRET_KEY` | `minioadmin` | MinIO secret key |
|
||||
| `MINIO_BUCKET` | `evidence` | Bucket for evidence files |
|
||||
|
||||
## Development
|
||||
|
||||
### Running Migrations
|
||||
|
||||
```bash
|
||||
docker exec aegis-backend alembic upgrade head
|
||||
docker exec aegis-backend alembic revision --autogenerate -m "description"
|
||||
docker exec aegis-backend alembic downgrade -1
|
||||
# Generate a new migration after model changes
|
||||
docker exec -w /app aegis-backend-1 alembic revision --autogenerate -m "description"
|
||||
|
||||
# Apply migrations
|
||||
docker exec -w /app aegis-backend-1 alembic upgrade head
|
||||
|
||||
# Rollback one migration
|
||||
docker exec -w /app aegis-backend-1 alembic downgrade -1
|
||||
|
||||
# Check current migration
|
||||
docker exec -w /app aegis-backend-1 alembic current
|
||||
```
|
||||
|
||||
### Accessing Services
|
||||
|
||||
- **MinIO Console**: http://localhost:9001 (login: `minioadmin` / `minioadmin`)
|
||||
- **PostgreSQL**: `psql -h localhost -p 5433 -U postgres -d attackdb`
|
||||
|
||||
### Running Tests
|
||||
|
||||
```bash
|
||||
# Run all V3 tests inside the container (recommended)
|
||||
docker exec aegis-backend python -m pytest tests/ -v --tb=short
|
||||
|
||||
# Run specific test suites
|
||||
docker exec aegis-backend python -m pytest tests/test_data_sources.py -v
|
||||
docker exec aegis-backend python -m pytest tests/test_scoring_and_compliance.py -v
|
||||
docker exec aegis-backend python -m pytest tests/test_campaigns_and_snapshots.py -v
|
||||
|
||||
# Skip integration tests (require network)
|
||||
docker exec aegis-backend python -m pytest tests/ -v -m "not integration"
|
||||
```
|
||||
|
||||
### Generating Reports
|
||||
The backend includes a test suite using pytest:
|
||||
|
||||
```bash
|
||||
# Coverage summary (JSON)
|
||||
GET /api/v1/reports/coverage-summary
|
||||
# Install test dependencies (if running locally)
|
||||
pip install pytest pytest-asyncio httpx
|
||||
|
||||
# Coverage CSV export
|
||||
GET /api/v1/reports/coverage-csv
|
||||
# Run all tests
|
||||
docker exec -w /app aegis-backend-1 pytest
|
||||
|
||||
# Compliance gap analysis
|
||||
GET /api/v1/compliance/{framework_id}/gaps
|
||||
# Run tests with verbose output
|
||||
docker exec -w /app aegis-backend-1 pytest -v
|
||||
|
||||
# Run specific test file
|
||||
docker exec -w /app aegis-backend-1 pytest tests/test_auth.py
|
||||
|
||||
# Run locally (requires SQLite)
|
||||
cd backend && pytest
|
||||
```
|
||||
|
||||
## Further Documentation
|
||||
Test files:
|
||||
- `test_health.py` - Health endpoint tests
|
||||
- `test_auth.py` - Authentication and authorization tests
|
||||
- `test_techniques.py` - Technique CRUD tests
|
||||
- `test_tests.py` - Security test CRUD and validation tests
|
||||
|
||||
- **[Architecture](docs/ARCHITECTURE.md)** — Database schema, backend layers, domain entities, service map
|
||||
- **[API Reference](docs/API.md)** — Full endpoint documentation
|
||||
- **[Scoring](docs/SCORING.md)** — Scoring system explained with examples and configuration
|
||||
- **[Data Sources](docs/DATA_SOURCES.md)** — All external data sources with import instructions
|
||||
- **[ADRs](docs/ADR.md)** — Architecture Decision Records
|
||||
- **[Technology Justification](docs/TECHNOLOGY_JUSTIFICATION.md)** — Technology choices and rationale
|
||||
- **[C4 Diagrams](docs/C4_CONTEXT_DIAGRAM.md)** — System context and container architecture
|
||||
## User Roles
|
||||
|
||||
| Role | Description |
|
||||
|------|-------------|
|
||||
| `admin` | Full system access |
|
||||
| `red_tech` | Red team technician - can create and edit tests |
|
||||
| `blue_tech` | Blue team technician - can create and edit tests |
|
||||
| `red_lead` | Red team lead - can validate tests |
|
||||
| `blue_lead` | Blue team lead - can validate tests |
|
||||
| `viewer` | Read-only access |
|
||||
|
||||
## License
|
||||
|
||||
This project is proprietary software. All rights reserved.
|
||||
|
||||
## Disclaimer
|
||||
## Contributing
|
||||
|
||||
This project has been developed with the assistance of [Cursor](https://cursor.com) and Claude Opus 4.6 (Anthropic).
|
||||
Please read the contribution guidelines before submitting pull requests.
|
||||
|
||||
@@ -1,2 +0,0 @@
|
||||
skips:
|
||||
- B311
|
||||
+3
-13
@@ -3,14 +3,10 @@ FROM python:3.11-slim
|
||||
WORKDIR /app
|
||||
|
||||
# Install system dependencies
|
||||
RUN apt-get update && apt-get upgrade -y && apt-get install -y \
|
||||
RUN apt-get update && apt-get install -y \
|
||||
gcc \
|
||||
libpq-dev \
|
||||
curl \
|
||||
pkg-config \
|
||||
libxml2-dev \
|
||||
libxmlsec1-dev \
|
||||
libxmlsec1-openssl \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Copy requirements first for better caching
|
||||
@@ -20,14 +16,8 @@ RUN pip install --no-cache-dir -r requirements.txt
|
||||
# Copy application code
|
||||
COPY . .
|
||||
|
||||
# Make entrypoints executable
|
||||
RUN chmod +x /app/entrypoint.sh /app/entrypoint.prod.sh
|
||||
|
||||
# Create a non-root user and give it ownership of /app
|
||||
RUN adduser --disabled-password --gecos '' --uid 1001 appuser \
|
||||
&& chown -R appuser:appuser /app
|
||||
|
||||
USER appuser
|
||||
# Make entrypoint executable
|
||||
RUN chmod +x /app/entrypoint.sh
|
||||
|
||||
# Expose port
|
||||
EXPOSE 8000
|
||||
|
||||
@@ -1,46 +0,0 @@
|
||||
"""add_notifications_table
|
||||
|
||||
Revision ID: b006notifications
|
||||
Revises: b005v2indexes
|
||||
Create Date: 2026-02-09 11:00:00.000000
|
||||
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.dialects.postgresql import UUID
|
||||
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = 'b006notifications'
|
||||
down_revision: Union[str, Sequence[str], None] = 'b005v2indexes'
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Create notifications table."""
|
||||
op.create_table(
|
||||
'notifications',
|
||||
sa.Column('id', UUID(as_uuid=True), primary_key=True, server_default=sa.text('gen_random_uuid()')),
|
||||
sa.Column('user_id', UUID(as_uuid=True), sa.ForeignKey('users.id'), nullable=False),
|
||||
sa.Column('type', sa.String(), nullable=False),
|
||||
sa.Column('title', sa.String(), nullable=False),
|
||||
sa.Column('message', sa.Text(), nullable=True),
|
||||
sa.Column('entity_type', sa.String(), nullable=True),
|
||||
sa.Column('entity_id', UUID(as_uuid=True), nullable=True),
|
||||
sa.Column('read', sa.Boolean(), server_default='false'),
|
||||
sa.Column('created_at', sa.DateTime(), server_default=sa.func.now()),
|
||||
)
|
||||
op.create_index('ix_notifications_user_id', 'notifications', ['user_id'])
|
||||
op.create_index('ix_notifications_read', 'notifications', ['read'])
|
||||
op.create_index('ix_notifications_created_at', 'notifications', ['created_at'])
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Drop notifications table."""
|
||||
op.drop_index('ix_notifications_created_at', table_name='notifications')
|
||||
op.drop_index('ix_notifications_read', table_name='notifications')
|
||||
op.drop_index('ix_notifications_user_id', table_name='notifications')
|
||||
op.drop_table('notifications')
|
||||
@@ -1,44 +0,0 @@
|
||||
"""add_remediation_fields
|
||||
|
||||
Revision ID: b007remediation
|
||||
Revises: b006notifications
|
||||
Create Date: 2026-02-09 11:30:00.000000
|
||||
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.dialects.postgresql import UUID
|
||||
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = 'b007remediation'
|
||||
down_revision: Union[str, Sequence[str], None] = 'b006notifications'
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Add remediation fields to tests and test_templates."""
|
||||
# Tests — remediation fields
|
||||
op.add_column('tests', sa.Column('remediation_steps', sa.Text(), nullable=True))
|
||||
op.add_column('tests', sa.Column('remediation_status', sa.String(), nullable=True))
|
||||
op.add_column('tests', sa.Column('remediation_assignee', UUID(as_uuid=True), nullable=True))
|
||||
op.create_foreign_key(
|
||||
'fk_tests_remediation_assignee',
|
||||
'tests', 'users',
|
||||
['remediation_assignee'], ['id'],
|
||||
)
|
||||
|
||||
# TestTemplates — suggested_remediation
|
||||
op.add_column('test_templates', sa.Column('suggested_remediation', sa.Text(), nullable=True))
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Remove remediation fields."""
|
||||
op.drop_column('test_templates', 'suggested_remediation')
|
||||
op.drop_constraint('fk_tests_remediation_assignee', 'tests', type_='foreignkey')
|
||||
op.drop_column('tests', 'remediation_assignee')
|
||||
op.drop_column('tests', 'remediation_status')
|
||||
op.drop_column('tests', 'remediation_steps')
|
||||
@@ -1,48 +0,0 @@
|
||||
"""add_data_sources_table
|
||||
|
||||
Revision ID: b008datasources
|
||||
Revises: b007remediation
|
||||
Create Date: 2026-02-09 14:00:00.000000
|
||||
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.dialects.postgresql import UUID, JSONB
|
||||
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = 'b008datasources'
|
||||
down_revision: Union[str, Sequence[str], None] = 'b007remediation'
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Create data_sources table."""
|
||||
op.create_table(
|
||||
'data_sources',
|
||||
sa.Column('id', UUID(as_uuid=True), primary_key=True),
|
||||
sa.Column('name', sa.String(), unique=True, nullable=False),
|
||||
sa.Column('display_name', sa.String(), nullable=False),
|
||||
sa.Column('type', sa.String(), nullable=False),
|
||||
sa.Column('url', sa.String(), nullable=True),
|
||||
sa.Column('description', sa.Text(), nullable=True),
|
||||
sa.Column('is_enabled', sa.Boolean(), server_default='true'),
|
||||
sa.Column('last_sync_at', sa.DateTime(), nullable=True),
|
||||
sa.Column('last_sync_status', sa.String(), nullable=True),
|
||||
sa.Column('last_sync_stats', JSONB(), nullable=True),
|
||||
sa.Column('sync_frequency', sa.String(), nullable=True),
|
||||
sa.Column('config', JSONB(), nullable=True),
|
||||
sa.Column('created_at', sa.DateTime(), server_default=sa.func.now()),
|
||||
)
|
||||
op.create_index('ix_data_sources_type', 'data_sources', ['type'])
|
||||
op.create_index('ix_data_sources_is_enabled', 'data_sources', ['is_enabled'])
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Drop data_sources table."""
|
||||
op.drop_index('ix_data_sources_is_enabled', table_name='data_sources')
|
||||
op.drop_index('ix_data_sources_type', table_name='data_sources')
|
||||
op.drop_table('data_sources')
|
||||
@@ -1,52 +0,0 @@
|
||||
"""add_detection_rules_table
|
||||
|
||||
Revision ID: b009detectionrules
|
||||
Revises: b008datasources
|
||||
Create Date: 2026-02-09 14:10:00.000000
|
||||
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.dialects.postgresql import UUID, JSONB
|
||||
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = 'b009detectionrules'
|
||||
down_revision: Union[str, Sequence[str], None] = 'b008datasources'
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Create detection_rules table."""
|
||||
op.create_table(
|
||||
'detection_rules',
|
||||
sa.Column('id', UUID(as_uuid=True), primary_key=True),
|
||||
sa.Column('mitre_technique_id', sa.String(), nullable=False),
|
||||
sa.Column('title', sa.String(), nullable=False),
|
||||
sa.Column('description', sa.Text(), nullable=True),
|
||||
sa.Column('source', sa.String(), nullable=False),
|
||||
sa.Column('source_id', sa.String(), nullable=True),
|
||||
sa.Column('source_url', sa.String(), nullable=True),
|
||||
sa.Column('rule_content', sa.Text(), nullable=False),
|
||||
sa.Column('rule_format', sa.String(), nullable=False),
|
||||
sa.Column('severity', sa.String(), nullable=True),
|
||||
sa.Column('platforms', JSONB(), nullable=True),
|
||||
sa.Column('log_sources', JSONB(), nullable=True),
|
||||
sa.Column('false_positive_rate', sa.String(), nullable=True),
|
||||
sa.Column('is_active', sa.Boolean(), server_default='true'),
|
||||
sa.Column('created_at', sa.DateTime(), server_default=sa.func.now()),
|
||||
)
|
||||
op.create_index('ix_detection_rules_mitre_technique_id', 'detection_rules', ['mitre_technique_id'])
|
||||
op.create_index('ix_detection_rules_source', 'detection_rules', ['source'])
|
||||
op.create_index('ix_detection_rules_severity', 'detection_rules', ['severity'])
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Drop detection_rules table."""
|
||||
op.drop_index('ix_detection_rules_severity', table_name='detection_rules')
|
||||
op.drop_index('ix_detection_rules_source', table_name='detection_rules')
|
||||
op.drop_index('ix_detection_rules_mitre_technique_id', table_name='detection_rules')
|
||||
op.drop_table('detection_rules')
|
||||
@@ -1,72 +0,0 @@
|
||||
"""add_threat_actors_tables
|
||||
|
||||
Revision ID: b010threatactors
|
||||
Revises: b009detectionrules
|
||||
Create Date: 2026-02-09 15:00:00.000000
|
||||
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.dialects.postgresql import UUID, JSONB
|
||||
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = 'b010threatactors'
|
||||
down_revision: Union[str, Sequence[str], None] = 'b009detectionrules'
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Create threat_actors and threat_actor_techniques tables."""
|
||||
# threat_actors
|
||||
op.create_table(
|
||||
'threat_actors',
|
||||
sa.Column('id', UUID(as_uuid=True), primary_key=True),
|
||||
sa.Column('mitre_id', sa.String(), unique=True, nullable=True),
|
||||
sa.Column('name', sa.String(), nullable=False),
|
||||
sa.Column('aliases', JSONB(), nullable=True),
|
||||
sa.Column('description', sa.Text(), nullable=True),
|
||||
sa.Column('country', sa.String(), nullable=True),
|
||||
sa.Column('target_sectors', JSONB(), nullable=True),
|
||||
sa.Column('target_regions', JSONB(), nullable=True),
|
||||
sa.Column('motivation', sa.String(), nullable=True),
|
||||
sa.Column('sophistication', sa.String(), nullable=True),
|
||||
sa.Column('first_seen', sa.String(), nullable=True),
|
||||
sa.Column('last_seen', sa.String(), nullable=True),
|
||||
sa.Column('references', JSONB(), nullable=True),
|
||||
sa.Column('mitre_url', sa.String(), nullable=True),
|
||||
sa.Column('is_active', sa.Boolean(), server_default='true'),
|
||||
sa.Column('created_at', sa.DateTime(), server_default=sa.func.now()),
|
||||
)
|
||||
op.create_index('ix_threat_actors_country', 'threat_actors', ['country'])
|
||||
op.create_index('ix_threat_actors_motivation', 'threat_actors', ['motivation'])
|
||||
|
||||
# threat_actor_techniques (junction table)
|
||||
op.create_table(
|
||||
'threat_actor_techniques',
|
||||
sa.Column('id', UUID(as_uuid=True), primary_key=True),
|
||||
sa.Column('threat_actor_id', UUID(as_uuid=True),
|
||||
sa.ForeignKey('threat_actors.id', ondelete='CASCADE'), nullable=False),
|
||||
sa.Column('technique_id', UUID(as_uuid=True),
|
||||
sa.ForeignKey('techniques.id', ondelete='CASCADE'), nullable=False),
|
||||
sa.Column('usage_description', sa.Text(), nullable=True),
|
||||
sa.Column('first_seen_using', sa.String(), nullable=True),
|
||||
)
|
||||
op.create_index('ix_threat_actor_techniques_actor', 'threat_actor_techniques', ['threat_actor_id'])
|
||||
op.create_index('ix_threat_actor_techniques_technique', 'threat_actor_techniques', ['technique_id'])
|
||||
op.create_unique_constraint('uq_actor_technique', 'threat_actor_techniques',
|
||||
['threat_actor_id', 'technique_id'])
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Drop threat_actor_techniques and threat_actors tables."""
|
||||
op.drop_constraint('uq_actor_technique', 'threat_actor_techniques', type_='unique')
|
||||
op.drop_index('ix_threat_actor_techniques_technique', table_name='threat_actor_techniques')
|
||||
op.drop_index('ix_threat_actor_techniques_actor', table_name='threat_actor_techniques')
|
||||
op.drop_table('threat_actor_techniques')
|
||||
op.drop_index('ix_threat_actors_motivation', table_name='threat_actors')
|
||||
op.drop_index('ix_threat_actors_country', table_name='threat_actors')
|
||||
op.drop_table('threat_actors')
|
||||
@@ -1,59 +0,0 @@
|
||||
"""add_defensive_techniques_tables
|
||||
|
||||
Revision ID: b011defensive
|
||||
Revises: b010threatactors
|
||||
Create Date: 2026-02-09 16:00:00.000000
|
||||
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.dialects.postgresql import UUID
|
||||
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = 'b011defensive'
|
||||
down_revision: Union[str, Sequence[str], None] = 'b010threatactors'
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Create defensive_techniques and defensive_technique_mappings tables."""
|
||||
# defensive_techniques
|
||||
op.create_table(
|
||||
'defensive_techniques',
|
||||
sa.Column('id', UUID(as_uuid=True), primary_key=True),
|
||||
sa.Column('d3fend_id', sa.String(), unique=True, nullable=False),
|
||||
sa.Column('name', sa.String(), nullable=False),
|
||||
sa.Column('description', sa.Text(), nullable=True),
|
||||
sa.Column('tactic', sa.String(), nullable=True),
|
||||
sa.Column('d3fend_url', sa.String(), nullable=True),
|
||||
sa.Column('created_at', sa.DateTime(), server_default=sa.func.now()),
|
||||
)
|
||||
op.create_index('ix_defensive_techniques_tactic', 'defensive_techniques', ['tactic'])
|
||||
|
||||
# defensive_technique_mappings (ATT&CK → D3FEND)
|
||||
op.create_table(
|
||||
'defensive_technique_mappings',
|
||||
sa.Column('id', UUID(as_uuid=True), primary_key=True),
|
||||
sa.Column('attack_technique_id', UUID(as_uuid=True),
|
||||
sa.ForeignKey('techniques.id', ondelete='CASCADE'), nullable=False),
|
||||
sa.Column('defensive_technique_id', UUID(as_uuid=True),
|
||||
sa.ForeignKey('defensive_techniques.id', ondelete='CASCADE'), nullable=False),
|
||||
)
|
||||
op.create_index('ix_dtm_attack_technique', 'defensive_technique_mappings', ['attack_technique_id'])
|
||||
op.create_index('ix_dtm_defensive_technique', 'defensive_technique_mappings', ['defensive_technique_id'])
|
||||
op.create_unique_constraint('uq_attack_defensive_technique', 'defensive_technique_mappings',
|
||||
['attack_technique_id', 'defensive_technique_id'])
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Drop defensive_technique_mappings and defensive_techniques tables."""
|
||||
op.drop_constraint('uq_attack_defensive_technique', 'defensive_technique_mappings', type_='unique')
|
||||
op.drop_index('ix_dtm_defensive_technique', table_name='defensive_technique_mappings')
|
||||
op.drop_index('ix_dtm_attack_technique', table_name='defensive_technique_mappings')
|
||||
op.drop_table('defensive_technique_mappings')
|
||||
op.drop_index('ix_defensive_techniques_tactic', table_name='defensive_techniques')
|
||||
op.drop_table('defensive_techniques')
|
||||
@@ -1,66 +0,0 @@
|
||||
"""add_detection_rule_associations
|
||||
|
||||
Revision ID: b012detectionassoc
|
||||
Revises: b011defensive
|
||||
Create Date: 2026-02-09 17:00:00.000000
|
||||
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.dialects.postgresql import UUID
|
||||
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = 'b012detectionassoc'
|
||||
down_revision: Union[str, Sequence[str], None] = 'b011defensive'
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Create test_template_detection_rules and test_detection_results tables."""
|
||||
|
||||
# test_template_detection_rules (template ↔ detection rule association)
|
||||
op.create_table(
|
||||
'test_template_detection_rules',
|
||||
sa.Column('id', UUID(as_uuid=True), primary_key=True),
|
||||
sa.Column('test_template_id', UUID(as_uuid=True),
|
||||
sa.ForeignKey('test_templates.id', ondelete='CASCADE'), nullable=True),
|
||||
sa.Column('detection_rule_id', UUID(as_uuid=True),
|
||||
sa.ForeignKey('detection_rules.id', ondelete='CASCADE'), nullable=False),
|
||||
sa.Column('is_primary', sa.Boolean(), server_default='false'),
|
||||
)
|
||||
op.create_index('ix_ttdr_template', 'test_template_detection_rules', ['test_template_id'])
|
||||
op.create_index('ix_ttdr_rule', 'test_template_detection_rules', ['detection_rule_id'])
|
||||
op.create_unique_constraint('uq_template_detection_rule', 'test_template_detection_rules',
|
||||
['test_template_id', 'detection_rule_id'])
|
||||
|
||||
# test_detection_results (per-test, per-rule evaluation results)
|
||||
op.create_table(
|
||||
'test_detection_results',
|
||||
sa.Column('id', UUID(as_uuid=True), primary_key=True),
|
||||
sa.Column('test_id', UUID(as_uuid=True),
|
||||
sa.ForeignKey('tests.id', ondelete='CASCADE'), nullable=False),
|
||||
sa.Column('detection_rule_id', UUID(as_uuid=True),
|
||||
sa.ForeignKey('detection_rules.id', ondelete='CASCADE'), nullable=False),
|
||||
sa.Column('triggered', sa.Boolean(), nullable=True),
|
||||
sa.Column('notes', sa.Text(), nullable=True),
|
||||
sa.Column('evaluated_by', UUID(as_uuid=True),
|
||||
sa.ForeignKey('users.id', ondelete='SET NULL'), nullable=True),
|
||||
sa.Column('evaluated_at', sa.DateTime(), nullable=True),
|
||||
)
|
||||
op.create_index('ix_tdr_test', 'test_detection_results', ['test_id'])
|
||||
op.create_index('ix_tdr_rule', 'test_detection_results', ['detection_rule_id'])
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Drop test_detection_results and test_template_detection_rules tables."""
|
||||
op.drop_index('ix_tdr_rule', table_name='test_detection_results')
|
||||
op.drop_index('ix_tdr_test', table_name='test_detection_results')
|
||||
op.drop_table('test_detection_results')
|
||||
op.drop_constraint('uq_template_detection_rule', 'test_template_detection_rules', type_='unique')
|
||||
op.drop_index('ix_ttdr_rule', table_name='test_template_detection_rules')
|
||||
op.drop_index('ix_ttdr_template', table_name='test_template_detection_rules')
|
||||
op.drop_table('test_template_detection_rules')
|
||||
@@ -1,74 +0,0 @@
|
||||
"""add_campaigns_tables
|
||||
|
||||
Revision ID: b013campaigns
|
||||
Revises: b012detectionassoc
|
||||
Create Date: 2026-02-09 18:00:00.000000
|
||||
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.dialects.postgresql import UUID, JSONB
|
||||
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = 'b013campaigns'
|
||||
down_revision: Union[str, Sequence[str], None] = 'b012detectionassoc'
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Create campaigns and campaign_tests tables."""
|
||||
|
||||
# campaigns
|
||||
op.create_table(
|
||||
'campaigns',
|
||||
sa.Column('id', UUID(as_uuid=True), primary_key=True),
|
||||
sa.Column('name', sa.String(), nullable=False),
|
||||
sa.Column('description', sa.Text(), nullable=True),
|
||||
sa.Column('type', sa.String(), nullable=False, server_default='custom'),
|
||||
sa.Column('threat_actor_id', UUID(as_uuid=True),
|
||||
sa.ForeignKey('threat_actors.id', ondelete='SET NULL'), nullable=True),
|
||||
sa.Column('status', sa.String(), nullable=False, server_default='draft'),
|
||||
sa.Column('created_by', UUID(as_uuid=True),
|
||||
sa.ForeignKey('users.id', ondelete='SET NULL'), nullable=True),
|
||||
sa.Column('scheduled_at', sa.DateTime(), nullable=True),
|
||||
sa.Column('completed_at', sa.DateTime(), nullable=True),
|
||||
sa.Column('target_platform', sa.String(), nullable=True),
|
||||
sa.Column('tags', JSONB(), nullable=True),
|
||||
sa.Column('created_at', sa.DateTime(), server_default=sa.func.now()),
|
||||
)
|
||||
op.create_index('ix_campaigns_status', 'campaigns', ['status'])
|
||||
op.create_index('ix_campaigns_type', 'campaigns', ['type'])
|
||||
op.create_index('ix_campaigns_threat_actor', 'campaigns', ['threat_actor_id'])
|
||||
op.create_index('ix_campaigns_created_by', 'campaigns', ['created_by'])
|
||||
|
||||
# campaign_tests
|
||||
op.create_table(
|
||||
'campaign_tests',
|
||||
sa.Column('id', UUID(as_uuid=True), primary_key=True),
|
||||
sa.Column('campaign_id', UUID(as_uuid=True),
|
||||
sa.ForeignKey('campaigns.id', ondelete='CASCADE'), nullable=False),
|
||||
sa.Column('test_id', UUID(as_uuid=True),
|
||||
sa.ForeignKey('tests.id', ondelete='CASCADE'), nullable=False),
|
||||
sa.Column('order_index', sa.Integer(), nullable=False, server_default='0'),
|
||||
sa.Column('depends_on', UUID(as_uuid=True),
|
||||
sa.ForeignKey('campaign_tests.id', ondelete='SET NULL'), nullable=True),
|
||||
sa.Column('phase', sa.String(), nullable=True),
|
||||
)
|
||||
op.create_index('ix_campaign_tests_campaign', 'campaign_tests', ['campaign_id'])
|
||||
op.create_index('ix_campaign_tests_test', 'campaign_tests', ['test_id'])
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Drop campaign_tests and campaigns tables."""
|
||||
op.drop_index('ix_campaign_tests_test', table_name='campaign_tests')
|
||||
op.drop_index('ix_campaign_tests_campaign', table_name='campaign_tests')
|
||||
op.drop_table('campaign_tests')
|
||||
op.drop_index('ix_campaigns_created_by', table_name='campaigns')
|
||||
op.drop_index('ix_campaigns_threat_actor', table_name='campaigns')
|
||||
op.drop_index('ix_campaigns_type', table_name='campaigns')
|
||||
op.drop_index('ix_campaigns_status', table_name='campaigns')
|
||||
op.drop_table('campaigns')
|
||||
@@ -1,92 +0,0 @@
|
||||
"""add_compliance_tables
|
||||
|
||||
Revision ID: b014compliance
|
||||
Revises: b013campaigns
|
||||
Create Date: 2026-02-09 20:00:00.000000
|
||||
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.dialects import postgresql
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "b014compliance"
|
||||
down_revision: Union[str, None] = "b013campaigns"
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# ── compliance_frameworks ─────────────────────────────────────
|
||||
op.create_table(
|
||||
"compliance_frameworks",
|
||||
sa.Column("id", postgresql.UUID(as_uuid=True), primary_key=True),
|
||||
sa.Column("name", sa.String, unique=True, nullable=False),
|
||||
sa.Column("version", sa.String, nullable=True),
|
||||
sa.Column("description", sa.Text, nullable=True),
|
||||
sa.Column("url", sa.String, nullable=True),
|
||||
sa.Column("is_active", sa.Boolean, server_default="true"),
|
||||
sa.Column("created_at", sa.DateTime, server_default=sa.func.now()),
|
||||
)
|
||||
|
||||
# ── compliance_controls ───────────────────────────────────────
|
||||
op.create_table(
|
||||
"compliance_controls",
|
||||
sa.Column("id", postgresql.UUID(as_uuid=True), primary_key=True),
|
||||
sa.Column(
|
||||
"framework_id",
|
||||
postgresql.UUID(as_uuid=True),
|
||||
sa.ForeignKey("compliance_frameworks.id", ondelete="CASCADE"),
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column("control_id", sa.String, nullable=False),
|
||||
sa.Column("title", sa.String, nullable=False),
|
||||
sa.Column("description", sa.Text, nullable=True),
|
||||
sa.Column("category", sa.String, nullable=True),
|
||||
)
|
||||
op.create_index(
|
||||
"ix_compliance_controls_framework",
|
||||
"compliance_controls",
|
||||
["framework_id"],
|
||||
)
|
||||
|
||||
# ── compliance_control_mappings ───────────────────────────────
|
||||
op.create_table(
|
||||
"compliance_control_mappings",
|
||||
sa.Column("id", postgresql.UUID(as_uuid=True), primary_key=True),
|
||||
sa.Column(
|
||||
"compliance_control_id",
|
||||
postgresql.UUID(as_uuid=True),
|
||||
sa.ForeignKey("compliance_controls.id", ondelete="CASCADE"),
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column(
|
||||
"technique_id",
|
||||
postgresql.UUID(as_uuid=True),
|
||||
sa.ForeignKey("techniques.id", ondelete="CASCADE"),
|
||||
nullable=False,
|
||||
),
|
||||
)
|
||||
op.create_index(
|
||||
"ix_compliance_mappings_control",
|
||||
"compliance_control_mappings",
|
||||
["compliance_control_id"],
|
||||
)
|
||||
op.create_index(
|
||||
"ix_compliance_mappings_technique",
|
||||
"compliance_control_mappings",
|
||||
["technique_id"],
|
||||
)
|
||||
op.create_unique_constraint(
|
||||
"uq_control_technique",
|
||||
"compliance_control_mappings",
|
||||
["compliance_control_id", "technique_id"],
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_table("compliance_control_mappings")
|
||||
op.drop_table("compliance_controls")
|
||||
op.drop_table("compliance_frameworks")
|
||||
@@ -1,77 +0,0 @@
|
||||
"""add_coverage_snapshots
|
||||
|
||||
Revision ID: b015snapshots
|
||||
Revises: b014compliance
|
||||
Create Date: 2026-02-10 00:00:00.000000
|
||||
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.dialects import postgresql
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "b015snapshots"
|
||||
down_revision: Union[str, None] = "b014compliance"
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# ── coverage_snapshots ────────────────────────────────────────
|
||||
op.create_table(
|
||||
"coverage_snapshots",
|
||||
sa.Column("id", postgresql.UUID(as_uuid=True), primary_key=True),
|
||||
sa.Column("name", sa.String, nullable=True),
|
||||
sa.Column("organization_score", sa.Float, nullable=False),
|
||||
sa.Column("total_techniques", sa.Integer, nullable=False),
|
||||
sa.Column("validated_count", sa.Integer, nullable=False),
|
||||
sa.Column("partial_count", sa.Integer, nullable=False),
|
||||
sa.Column("not_covered_count", sa.Integer, nullable=False),
|
||||
sa.Column("in_progress_count", sa.Integer, nullable=False),
|
||||
sa.Column("not_evaluated_count", sa.Integer, nullable=False),
|
||||
sa.Column(
|
||||
"created_by",
|
||||
postgresql.UUID(as_uuid=True),
|
||||
sa.ForeignKey("users.id", ondelete="SET NULL"),
|
||||
nullable=True,
|
||||
),
|
||||
sa.Column("created_at", sa.DateTime, server_default=sa.func.now()),
|
||||
)
|
||||
|
||||
# ── snapshot_technique_states ─────────────────────────────────
|
||||
op.create_table(
|
||||
"snapshot_technique_states",
|
||||
sa.Column("id", postgresql.UUID(as_uuid=True), primary_key=True),
|
||||
sa.Column(
|
||||
"snapshot_id",
|
||||
postgresql.UUID(as_uuid=True),
|
||||
sa.ForeignKey("coverage_snapshots.id", ondelete="CASCADE"),
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column(
|
||||
"technique_id",
|
||||
postgresql.UUID(as_uuid=True),
|
||||
sa.ForeignKey("techniques.id", ondelete="CASCADE"),
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column("mitre_id", sa.String, nullable=False),
|
||||
sa.Column("status", sa.String, nullable=False),
|
||||
sa.Column("score", sa.Float, nullable=True),
|
||||
)
|
||||
op.create_index(
|
||||
"ix_snapshot_technique_states_snapshot",
|
||||
"snapshot_technique_states",
|
||||
["snapshot_id"],
|
||||
)
|
||||
op.create_index(
|
||||
"ix_snapshot_technique_states_technique",
|
||||
"snapshot_technique_states",
|
||||
["technique_id"],
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_table("snapshot_technique_states")
|
||||
op.drop_table("coverage_snapshots")
|
||||
@@ -1,41 +0,0 @@
|
||||
"""add_retest_fields
|
||||
|
||||
Revision ID: b016retests
|
||||
Revises: b015snapshots
|
||||
Create Date: 2026-02-10 01:00:00.000000
|
||||
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.dialects import postgresql
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "b016retests"
|
||||
down_revision: Union[str, None] = "b015snapshots"
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.add_column(
|
||||
"tests",
|
||||
sa.Column(
|
||||
"retest_of",
|
||||
postgresql.UUID(as_uuid=True),
|
||||
sa.ForeignKey("tests.id"),
|
||||
nullable=True,
|
||||
),
|
||||
)
|
||||
op.add_column(
|
||||
"tests",
|
||||
sa.Column("retest_count", sa.Integer, server_default="0", nullable=False),
|
||||
)
|
||||
op.create_index("ix_tests_retest_of", "tests", ["retest_of"])
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_index("ix_tests_retest_of", table_name="tests")
|
||||
op.drop_column("tests", "retest_count")
|
||||
op.drop_column("tests", "retest_of")
|
||||
@@ -1,58 +0,0 @@
|
||||
"""add_campaign_scheduling
|
||||
|
||||
Revision ID: b017scheduling
|
||||
Revises: b016retests
|
||||
Create Date: 2026-02-10 02:00:00.000000
|
||||
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.dialects import postgresql
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "b017scheduling"
|
||||
down_revision: Union[str, None] = "b016retests"
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.add_column(
|
||||
"campaigns",
|
||||
sa.Column("is_recurring", sa.Boolean, server_default="false", nullable=False),
|
||||
)
|
||||
op.add_column(
|
||||
"campaigns",
|
||||
sa.Column("recurrence_pattern", sa.String, nullable=True),
|
||||
)
|
||||
op.add_column(
|
||||
"campaigns",
|
||||
sa.Column("next_run_at", sa.DateTime, nullable=True),
|
||||
)
|
||||
op.add_column(
|
||||
"campaigns",
|
||||
sa.Column("last_run_at", sa.DateTime, nullable=True),
|
||||
)
|
||||
op.add_column(
|
||||
"campaigns",
|
||||
sa.Column(
|
||||
"parent_campaign_id",
|
||||
postgresql.UUID(as_uuid=True),
|
||||
sa.ForeignKey("campaigns.id", ondelete="SET NULL"),
|
||||
nullable=True,
|
||||
),
|
||||
)
|
||||
op.create_index("ix_campaigns_next_run", "campaigns", ["next_run_at"])
|
||||
op.create_index("ix_campaigns_parent", "campaigns", ["parent_campaign_id"])
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_index("ix_campaigns_parent", table_name="campaigns")
|
||||
op.drop_index("ix_campaigns_next_run", table_name="campaigns")
|
||||
op.drop_column("campaigns", "parent_campaign_id")
|
||||
op.drop_column("campaigns", "last_run_at")
|
||||
op.drop_column("campaigns", "next_run_at")
|
||||
op.drop_column("campaigns", "recurrence_pattern")
|
||||
op.drop_column("campaigns", "is_recurring")
|
||||
@@ -1,78 +0,0 @@
|
||||
"""add_performance_indexes
|
||||
|
||||
Revision ID: b018perfidx
|
||||
Revises: b017scheduling
|
||||
Create Date: 2026-02-10 06:00:00.000000
|
||||
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "b018perfidx"
|
||||
down_revision: Union[str, None] = "b017scheduling"
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# Composite index for detection rules filtered by technique + source
|
||||
op.create_index(
|
||||
"ix_detection_rules_technique_source",
|
||||
"detection_rules",
|
||||
["mitre_technique_id", "source"],
|
||||
)
|
||||
|
||||
# Composite index for snapshot technique states
|
||||
op.create_index(
|
||||
"ix_snapshot_technique_states_snap_tech",
|
||||
"snapshot_technique_states",
|
||||
["snapshot_id", "technique_id"],
|
||||
unique=True,
|
||||
)
|
||||
|
||||
# Covering index for tests frequently filtered by technique + state
|
||||
op.create_index(
|
||||
"ix_tests_technique_state",
|
||||
"tests",
|
||||
["technique_id", "state"],
|
||||
)
|
||||
|
||||
# Audit logs — timestamp-based lookups
|
||||
op.create_index(
|
||||
"ix_audit_logs_timestamp",
|
||||
"audit_logs",
|
||||
["timestamp"],
|
||||
)
|
||||
|
||||
# Audit logs — entity lookups
|
||||
op.create_index(
|
||||
"ix_audit_logs_entity",
|
||||
"audit_logs",
|
||||
["entity_type", "entity_id"],
|
||||
)
|
||||
|
||||
# Test detection results — triggered flag for maturity queries
|
||||
op.create_index(
|
||||
"ix_test_detection_results_triggered",
|
||||
"test_detection_results",
|
||||
["triggered"],
|
||||
)
|
||||
|
||||
# Compliance control mappings — composite for joins
|
||||
op.create_index(
|
||||
"ix_compliance_mappings_control_technique",
|
||||
"compliance_control_mappings",
|
||||
["compliance_control_id", "technique_id"],
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_index("ix_compliance_mappings_control_technique", table_name="compliance_control_mappings")
|
||||
op.drop_index("ix_test_detection_results_triggered", table_name="test_detection_results")
|
||||
op.drop_index("ix_audit_logs_entity", table_name="audit_logs")
|
||||
op.drop_index("ix_audit_logs_timestamp", table_name="audit_logs")
|
||||
op.drop_index("ix_tests_technique_state", table_name="tests")
|
||||
op.drop_index("ix_snapshot_technique_states_snap_tech", table_name="snapshot_technique_states")
|
||||
op.drop_index("ix_detection_rules_technique_source", table_name="detection_rules")
|
||||
@@ -1,67 +0,0 @@
|
||||
"""add_composite_indexes
|
||||
|
||||
Additional composite indexes for scoring, heatmap, metrics, reports,
|
||||
MTTD/MTTR calculations, and notification queries.
|
||||
|
||||
Revision ID: b019composite
|
||||
Revises: b018perfidx
|
||||
Create Date: 2026-02-17 14:00:00.000000
|
||||
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "b019composite"
|
||||
down_revision: Union[str, None] = "b018perfidx"
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# ── Tests ────────────────────────────────────────────────────────
|
||||
# Used by scoring queries that filter by state + validation date
|
||||
op.create_index(
|
||||
"ix_tests_state_red_validated_at",
|
||||
"tests",
|
||||
["state", "red_validated_at"],
|
||||
)
|
||||
|
||||
# Used by remediation dashboard and metrics
|
||||
op.create_index(
|
||||
"ix_tests_remediation_status",
|
||||
"tests",
|
||||
["remediation_status"],
|
||||
)
|
||||
|
||||
# ── Audit logs ───────────────────────────────────────────────────
|
||||
# Three-column index for MTTD/MTTR queries that filter by entity + action
|
||||
op.create_index(
|
||||
"ix_audit_logs_entity_type_entity_id_action",
|
||||
"audit_logs",
|
||||
["entity_type", "entity_id", "action"],
|
||||
)
|
||||
|
||||
# Used for per-user audit trail queries
|
||||
op.create_index(
|
||||
"ix_audit_logs_user_id",
|
||||
"audit_logs",
|
||||
["user_id"],
|
||||
)
|
||||
|
||||
# ── Notifications ────────────────────────────────────────────────
|
||||
# Used by "unread notifications" badge and inbox queries
|
||||
op.create_index(
|
||||
"ix_notifications_user_id_read",
|
||||
"notifications",
|
||||
["user_id", "read"],
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_index("ix_notifications_user_id_read", table_name="notifications")
|
||||
op.drop_index("ix_audit_logs_user_id", table_name="audit_logs")
|
||||
op.drop_index("ix_audit_logs_entity_type_entity_id_action", table_name="audit_logs")
|
||||
op.drop_index("ix_tests_remediation_status", table_name="tests")
|
||||
op.drop_index("ix_tests_state_red_validated_at", table_name="tests")
|
||||
@@ -1,95 +0,0 @@
|
||||
"""add_jira_links_and_worklogs
|
||||
|
||||
Revision ID: b020jiraworklogs
|
||||
Revises: b019composite
|
||||
Create Date: 2026-02-17 16:00:00.000000
|
||||
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.dialects import postgresql
|
||||
|
||||
from alembic import op
|
||||
|
||||
revision: str = "b020jiraworklogs"
|
||||
down_revision: Union[str, None] = "b019composite"
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# ── jira_links: 100 % raw SQL to avoid all SQLAlchemy enum hooks ──
|
||||
op.execute("""
|
||||
DO $$ BEGIN
|
||||
IF NOT EXISTS (SELECT 1 FROM pg_type WHERE typname = 'jiralinkentitytype') THEN
|
||||
CREATE TYPE jiralinkentitytype AS ENUM ('test', 'technique', 'campaign', 'evidence');
|
||||
END IF;
|
||||
END $$;
|
||||
|
||||
DO $$ BEGIN
|
||||
IF NOT EXISTS (SELECT 1 FROM pg_type WHERE typname = 'jirasyncdirection') THEN
|
||||
CREATE TYPE jirasyncdirection AS ENUM ('aegis_to_jira', 'jira_to_aegis', 'bidirectional');
|
||||
END IF;
|
||||
END $$;
|
||||
|
||||
CREATE TABLE IF NOT EXISTS jira_links (
|
||||
id UUID PRIMARY KEY,
|
||||
entity_type jiralinkentitytype NOT NULL,
|
||||
entity_id UUID NOT NULL,
|
||||
jira_issue_key VARCHAR(50) NOT NULL,
|
||||
jira_issue_id VARCHAR(50),
|
||||
jira_project_key VARCHAR(20),
|
||||
jira_status VARCHAR(100),
|
||||
jira_priority VARCHAR(50),
|
||||
jira_assignee VARCHAR(255),
|
||||
jira_story_points VARCHAR(10),
|
||||
sync_direction jirasyncdirection DEFAULT 'bidirectional',
|
||||
last_synced_at TIMESTAMP,
|
||||
sync_metadata JSONB DEFAULT '{}',
|
||||
created_by UUID REFERENCES users(id),
|
||||
created_at TIMESTAMP DEFAULT now(),
|
||||
updated_at TIMESTAMP DEFAULT now()
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS ix_jira_links_entity_id
|
||||
ON jira_links (entity_id);
|
||||
CREATE INDEX IF NOT EXISTS ix_jira_links_issue_key
|
||||
ON jira_links (jira_issue_key);
|
||||
CREATE INDEX IF NOT EXISTS ix_jira_links_entity_type_entity_id
|
||||
ON jira_links (entity_type, entity_id);
|
||||
""")
|
||||
|
||||
# ── worklogs table (no enums, straightforward) ───────────────────
|
||||
op.execute("""
|
||||
CREATE TABLE IF NOT EXISTS worklogs (
|
||||
id UUID PRIMARY KEY,
|
||||
entity_type VARCHAR(50) NOT NULL,
|
||||
entity_id UUID NOT NULL,
|
||||
user_id UUID NOT NULL REFERENCES users(id),
|
||||
activity_type VARCHAR(100) NOT NULL,
|
||||
started_at TIMESTAMP NOT NULL,
|
||||
ended_at TIMESTAMP,
|
||||
duration_seconds INTEGER NOT NULL,
|
||||
description TEXT,
|
||||
tempo_synced TIMESTAMP,
|
||||
tempo_worklog_id VARCHAR(100),
|
||||
integrity_hash VARCHAR(64),
|
||||
created_at TIMESTAMP DEFAULT now(),
|
||||
metadata JSONB DEFAULT '{}'
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS ix_worklogs_entity_id
|
||||
ON worklogs (entity_id);
|
||||
CREATE INDEX IF NOT EXISTS ix_worklogs_user_id
|
||||
ON worklogs (user_id);
|
||||
CREATE INDEX IF NOT EXISTS ix_worklogs_entity_type_entity_id
|
||||
ON worklogs (entity_type, entity_id);
|
||||
""")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.execute("DROP TABLE IF EXISTS worklogs")
|
||||
op.execute("DROP TABLE IF EXISTS jira_links")
|
||||
op.execute("DROP TYPE IF EXISTS jirasyncdirection")
|
||||
op.execute("DROP TYPE IF EXISTS jiralinkentitytype")
|
||||
@@ -1,38 +0,0 @@
|
||||
"""add_phase_timing_fields
|
||||
|
||||
Revision ID: b021phasetiming
|
||||
Revises: b020jiraworklogs
|
||||
Create Date: 2026-02-17 18:00:00.000000
|
||||
|
||||
Add red_started_at and blue_started_at columns to the tests table
|
||||
so that automatic worklogs can record real elapsed time per phase.
|
||||
"""
|
||||
|
||||
from alembic import op
|
||||
|
||||
revision = "b021phasetiming"
|
||||
down_revision = "b020jiraworklogs"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.execute("""
|
||||
ALTER TABLE tests
|
||||
ADD COLUMN IF NOT EXISTS red_started_at TIMESTAMP,
|
||||
ADD COLUMN IF NOT EXISTS blue_started_at TIMESTAMP,
|
||||
ADD COLUMN IF NOT EXISTS paused_at TIMESTAMP,
|
||||
ADD COLUMN IF NOT EXISTS red_paused_seconds INTEGER DEFAULT 0,
|
||||
ADD COLUMN IF NOT EXISTS blue_paused_seconds INTEGER DEFAULT 0;
|
||||
""")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.execute("""
|
||||
ALTER TABLE tests
|
||||
DROP COLUMN IF EXISTS red_started_at,
|
||||
DROP COLUMN IF EXISTS blue_started_at,
|
||||
DROP COLUMN IF EXISTS paused_at,
|
||||
DROP COLUMN IF EXISTS red_paused_seconds,
|
||||
DROP COLUMN IF EXISTS blue_paused_seconds;
|
||||
""")
|
||||
@@ -1,47 +0,0 @@
|
||||
"""add_osint_items
|
||||
|
||||
Revision ID: b022osintitems
|
||||
Revises: b021phasetiming
|
||||
Create Date: 2026-02-17 22:00:00.000000
|
||||
|
||||
Add osint_items table for OSINT enrichment data linked to techniques.
|
||||
"""
|
||||
|
||||
from alembic import op
|
||||
|
||||
revision = "b022osintitems"
|
||||
down_revision = "b021phasetiming"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.execute("""
|
||||
CREATE TABLE IF NOT EXISTS osint_items (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
technique_id UUID NOT NULL REFERENCES techniques(id),
|
||||
source_type VARCHAR(50) NOT NULL,
|
||||
source_url TEXT NOT NULL,
|
||||
title VARCHAR(500) NOT NULL,
|
||||
description TEXT,
|
||||
severity VARCHAR(20),
|
||||
discovered_at TIMESTAMP NOT NULL DEFAULT now(),
|
||||
reviewed BOOLEAN NOT NULL DEFAULT false,
|
||||
metadata JSONB DEFAULT '{}'::jsonb
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS ix_osint_items_technique_id
|
||||
ON osint_items (technique_id);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS ix_osint_items_source_type
|
||||
ON osint_items (source_type);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS ix_osint_items_reviewed
|
||||
ON osint_items (reviewed) WHERE NOT reviewed;
|
||||
""")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.execute("""
|
||||
DROP TABLE IF EXISTS osint_items CASCADE;
|
||||
""")
|
||||
@@ -1,30 +0,0 @@
|
||||
"""add_must_change_password
|
||||
|
||||
Revision ID: b023mustchgpwd
|
||||
Revises: b022osintitems
|
||||
Create Date: 2026-02-17 23:00:00.000000
|
||||
|
||||
Add must_change_password column to users table to force password
|
||||
change on first login.
|
||||
"""
|
||||
|
||||
from alembic import op
|
||||
|
||||
revision = "b023mustchgpwd"
|
||||
down_revision = "b022osintitems"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.execute("""
|
||||
ALTER TABLE users
|
||||
ADD COLUMN IF NOT EXISTS must_change_password BOOLEAN DEFAULT true;
|
||||
""")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.execute("""
|
||||
ALTER TABLE users
|
||||
DROP COLUMN IF EXISTS must_change_password;
|
||||
""")
|
||||
@@ -1,37 +0,0 @@
|
||||
"""add_critical_test_audit_indexes
|
||||
|
||||
Add missing critical indexes for tests and audit_logs tables to match
|
||||
model __table_args__ declarations. Existing indexes (from b005, b018,
|
||||
b019) are left untouched; only the two genuinely new indexes are created.
|
||||
|
||||
Revision ID: b024critidx
|
||||
Revises: b023mustchgpwd
|
||||
Create Date: 2026-02-18 12:00:00.000000
|
||||
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
|
||||
revision: str = "b024critidx"
|
||||
down_revision: Union[str, None] = "b023mustchgpwd"
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.create_index(
|
||||
"ix_tests_created_at",
|
||||
"tests",
|
||||
["created_at"],
|
||||
)
|
||||
op.create_index(
|
||||
"ix_tests_state_created_at",
|
||||
"tests",
|
||||
["state", "created_at"],
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_index("ix_tests_state_created_at", table_name="tests")
|
||||
op.drop_index("ix_tests_created_at", table_name="tests")
|
||||
@@ -1,41 +0,0 @@
|
||||
"""add_unique_test_detection_result
|
||||
|
||||
Enforce one evaluation per (test, detection_rule) pair. Before creating
|
||||
the constraint, duplicate rows (if any) are collapsed so the migration
|
||||
never fails on an existing database.
|
||||
|
||||
Revision ID: b025uqtdr
|
||||
Revises: b024critidx
|
||||
Create Date: 2026-02-18 14:00:00.000000
|
||||
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
|
||||
revision: str = "b025uqtdr"
|
||||
down_revision: Union[str, None] = "b024critidx"
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# Remove duplicates keeping the most recently evaluated row
|
||||
op.execute("""
|
||||
DELETE FROM test_detection_results
|
||||
WHERE id NOT IN (
|
||||
SELECT DISTINCT ON (test_id, detection_rule_id) id
|
||||
FROM test_detection_results
|
||||
ORDER BY test_id, detection_rule_id, evaluated_at DESC NULLS LAST
|
||||
)
|
||||
""")
|
||||
|
||||
op.create_unique_constraint(
|
||||
"uq_tdr_test_rule",
|
||||
"test_detection_results",
|
||||
["test_id", "detection_rule_id"],
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_constraint("uq_tdr_test_rule", "test_detection_results", type_="unique")
|
||||
@@ -1,38 +0,0 @@
|
||||
"""add_technique_query_indexes
|
||||
|
||||
Add indexes on techniques table for common query patterns
|
||||
(filter by tactic, filter by status_global) used in heatmap, scoring,
|
||||
and list-all-techniques operations.
|
||||
|
||||
These may already exist if the ORM model auto-created them; the
|
||||
``if_not_exists`` flag makes this migration safe to run regardless.
|
||||
|
||||
Revision ID: b026techidx
|
||||
Revises: b025uqtdr
|
||||
Create Date: 2026-02-18 18:00:00.000000
|
||||
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
|
||||
revision: str = "b026techidx"
|
||||
down_revision: Union[str, None] = "b025uqtdr"
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.execute(
|
||||
"CREATE INDEX IF NOT EXISTS ix_techniques_tactic "
|
||||
"ON techniques (tactic)"
|
||||
)
|
||||
op.execute(
|
||||
"CREATE INDEX IF NOT EXISTS ix_techniques_status_global "
|
||||
"ON techniques (status_global)"
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_index("ix_techniques_status_global", table_name="techniques")
|
||||
op.drop_index("ix_techniques_tactic", table_name="techniques")
|
||||
@@ -1,37 +0,0 @@
|
||||
"""add_scoring_config
|
||||
|
||||
Single-row table to persist scoring weights in the database,
|
||||
replacing the mutable in-process Settings approach.
|
||||
|
||||
Revision ID: b027scorecfg
|
||||
Revises: b026techidx
|
||||
Create Date: 2026-02-19 10:00:00.000000
|
||||
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.dialects import postgresql
|
||||
|
||||
revision: str = "b027scorecfg"
|
||||
down_revision: Union[str, None] = "b026techidx"
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.create_table(
|
||||
"scoring_config",
|
||||
sa.Column("id", postgresql.UUID(as_uuid=True), primary_key=True),
|
||||
sa.Column("weight_tests", sa.Float(), nullable=False, server_default="40.0"),
|
||||
sa.Column("weight_detection_rules", sa.Float(), nullable=False, server_default="20.0"),
|
||||
sa.Column("weight_d3fend", sa.Float(), nullable=False, server_default="15.0"),
|
||||
sa.Column("weight_freshness", sa.Float(), nullable=False, server_default="15.0"),
|
||||
sa.Column("weight_platform_diversity", sa.Float(), nullable=False, server_default="10.0"),
|
||||
sa.Column("updated_at", sa.DateTime(timezone=True), server_default=sa.func.now()),
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_table("scoring_config")
|
||||
@@ -1,35 +0,0 @@
|
||||
"""phase0 SR-006 — campaign_tests composite index
|
||||
|
||||
Most SR-006 indexes already ship in b005, b009, b018, b019, and b026.
|
||||
``tests`` has no ``campaign_id`` column (membership is ``campaign_tests``),
|
||||
so this revision adds a composite index to speed “tests in campaign” joins.
|
||||
|
||||
Revision ID: b028phase0
|
||||
Revises: b027scorecfg
|
||||
Create Date: 2026-05-18 12:00:00.000000
|
||||
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
|
||||
revision: str = "b028phase0"
|
||||
down_revision: Union[str, None] = "b027scorecfg"
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.create_index(
|
||||
"ix_campaign_tests_campaign_id_test_id",
|
||||
"campaign_tests",
|
||||
["campaign_id", "test_id"],
|
||||
unique=False,
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_index(
|
||||
"ix_campaign_tests_campaign_id_test_id",
|
||||
table_name="campaign_tests",
|
||||
)
|
||||
@@ -1,58 +0,0 @@
|
||||
"""Phase 3: audit trail columns and data classification fields.
|
||||
|
||||
Revision ID: b029phase3
|
||||
Revises: b028phase0
|
||||
"""
|
||||
|
||||
from typing import Sequence, Union
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
revision: str = "b029phase3"
|
||||
down_revision: Union[str, None] = "b028phase0"
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def _column_names(table: str) -> set[str]:
|
||||
bind = op.get_bind()
|
||||
insp = sa.inspect(bind)
|
||||
return {c["name"] for c in insp.get_columns(table)}
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
audit_cols = _column_names("audit_logs")
|
||||
if "ip_address" not in audit_cols:
|
||||
op.add_column("audit_logs", sa.Column("ip_address", sa.String(45), nullable=True))
|
||||
if "user_agent" not in audit_cols:
|
||||
op.add_column("audit_logs", sa.Column("user_agent", sa.String(500), nullable=True))
|
||||
if "integrity_hash" not in audit_cols:
|
||||
op.add_column("audit_logs", sa.Column("integrity_hash", sa.String(64), nullable=True))
|
||||
if "session_id" not in audit_cols:
|
||||
op.add_column("audit_logs", sa.Column("session_id", sa.String(100), nullable=True))
|
||||
|
||||
for table in ("tests", "evidences", "campaigns"):
|
||||
cols = _column_names(table)
|
||||
if "data_classification" not in cols:
|
||||
op.add_column(
|
||||
table,
|
||||
sa.Column(
|
||||
"data_classification",
|
||||
sa.String(20),
|
||||
nullable=False,
|
||||
server_default="internal",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
for table in ("campaigns", "evidences", "tests"):
|
||||
cols = _column_names(table)
|
||||
if "data_classification" in cols:
|
||||
op.drop_column(table, "data_classification")
|
||||
|
||||
audit_cols = _column_names("audit_logs")
|
||||
for col in ("session_id", "integrity_hash", "user_agent", "ip_address"):
|
||||
if col in audit_cols:
|
||||
op.drop_column("audit_logs", col)
|
||||
@@ -1,117 +0,0 @@
|
||||
"""Phase 5: scoring recency/severity columns and snapshot breakdown fields.
|
||||
|
||||
Revision ID: b030phase5
|
||||
Revises: b029phase3
|
||||
"""
|
||||
|
||||
from typing import Sequence, Union
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
from sqlalchemy.dialects import postgresql
|
||||
|
||||
revision: str = "b030phase5"
|
||||
down_revision: Union[str, None] = "b029phase3"
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def _column_names(table: str) -> set[str]:
|
||||
bind = op.get_bind()
|
||||
insp = sa.inspect(bind)
|
||||
return {c["name"] for c in insp.get_columns(table)}
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
snap_cols = _column_names("coverage_snapshots")
|
||||
if "by_tactic" not in snap_cols:
|
||||
op.add_column(
|
||||
"coverage_snapshots",
|
||||
sa.Column("by_tactic", postgresql.JSONB(), nullable=False, server_default="{}"),
|
||||
)
|
||||
if "by_status" not in snap_cols:
|
||||
op.add_column(
|
||||
"coverage_snapshots",
|
||||
sa.Column("by_status", postgresql.JSONB(), nullable=False, server_default="{}"),
|
||||
)
|
||||
if "stale_count" not in snap_cols:
|
||||
op.add_column(
|
||||
"coverage_snapshots",
|
||||
sa.Column("stale_count", sa.Integer(), nullable=False, server_default="0"),
|
||||
)
|
||||
if "never_tested_count" not in snap_cols:
|
||||
op.add_column(
|
||||
"coverage_snapshots",
|
||||
sa.Column("never_tested_count", sa.Integer(), nullable=False, server_default="0"),
|
||||
)
|
||||
if "coverage_percentage" not in snap_cols:
|
||||
op.add_column(
|
||||
"coverage_snapshots",
|
||||
sa.Column("coverage_percentage", sa.Float(), nullable=False, server_default="0"),
|
||||
)
|
||||
|
||||
cfg_cols = _column_names("scoring_config")
|
||||
if "weight_recency" not in cfg_cols and "weight_freshness" in cfg_cols:
|
||||
op.alter_column(
|
||||
"scoring_config",
|
||||
"weight_freshness",
|
||||
new_column_name="weight_recency",
|
||||
)
|
||||
cfg_cols.remove("weight_freshness")
|
||||
cfg_cols.add("weight_recency")
|
||||
elif "weight_recency" not in cfg_cols:
|
||||
op.add_column(
|
||||
"scoring_config",
|
||||
sa.Column("weight_recency", sa.Float(), nullable=False, server_default="10.0"),
|
||||
)
|
||||
|
||||
if "weight_severity" not in cfg_cols and "weight_platform_diversity" in cfg_cols:
|
||||
op.alter_column(
|
||||
"scoring_config",
|
||||
"weight_platform_diversity",
|
||||
new_column_name="weight_severity",
|
||||
)
|
||||
elif "weight_severity" not in cfg_cols:
|
||||
op.add_column(
|
||||
"scoring_config",
|
||||
sa.Column("weight_severity", sa.Float(), nullable=False, server_default="10.0"),
|
||||
)
|
||||
|
||||
if "updated_by" not in cfg_cols:
|
||||
op.add_column(
|
||||
"scoring_config",
|
||||
sa.Column(
|
||||
"updated_by",
|
||||
postgresql.UUID(as_uuid=True),
|
||||
sa.ForeignKey("users.id", ondelete="SET NULL"),
|
||||
nullable=True,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
cfg_cols = _column_names("scoring_config")
|
||||
if "updated_by" in cfg_cols:
|
||||
op.drop_column("scoring_config", "updated_by")
|
||||
if "weight_severity" in cfg_cols:
|
||||
op.alter_column(
|
||||
"scoring_config",
|
||||
"weight_severity",
|
||||
new_column_name="weight_platform_diversity",
|
||||
)
|
||||
if "weight_recency" in cfg_cols:
|
||||
op.alter_column(
|
||||
"scoring_config",
|
||||
"weight_recency",
|
||||
new_column_name="weight_freshness",
|
||||
)
|
||||
|
||||
for col in (
|
||||
"coverage_percentage",
|
||||
"never_tested_count",
|
||||
"stale_count",
|
||||
"by_status",
|
||||
"by_tactic",
|
||||
):
|
||||
if col in _column_names("coverage_snapshots"):
|
||||
op.drop_column("coverage_snapshots", col)
|
||||
@@ -1,32 +0,0 @@
|
||||
"""Phase 6.1: webhook_configs table.
|
||||
|
||||
Revision ID: b031phase6
|
||||
Revises: b030phase5
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.dialects import postgresql
|
||||
from alembic import op
|
||||
|
||||
revision: str = "b031phase6"
|
||||
down_revision: Union[str, None] = "b030phase5"
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
def upgrade() -> None:
|
||||
op.create_table(
|
||||
"webhook_configs",
|
||||
sa.Column("id", postgresql.UUID(as_uuid=True), primary_key=True),
|
||||
sa.Column("name", sa.String(200), nullable=False),
|
||||
sa.Column("url", sa.Text, nullable=False),
|
||||
sa.Column("secret", sa.String(256), nullable=True),
|
||||
sa.Column("events", postgresql.JSONB, nullable=False, server_default="[]"),
|
||||
sa.Column("is_active", sa.Boolean, nullable=False, server_default="true"),
|
||||
sa.Column("created_by", postgresql.UUID(as_uuid=True), sa.ForeignKey("users.id", ondelete="SET NULL"), nullable=True),
|
||||
sa.Column("last_triggered_at", sa.DateTime, nullable=True),
|
||||
sa.Column("failure_count", sa.Integer, nullable=False, server_default="0"),
|
||||
sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.func.now()),
|
||||
)
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_table("webhook_configs")
|
||||
@@ -1,41 +0,0 @@
|
||||
"""Phase 7.2: user notification_preferences and jira_account_id columns.
|
||||
|
||||
Revision ID: b032phase7
|
||||
Revises: b031phase6
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.dialects import postgresql
|
||||
from alembic import op
|
||||
|
||||
revision: str = "b032phase7"
|
||||
down_revision: Union[str, None] = "b031phase6"
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
_DEFAULT_PREFS = '{"email_on_test_validated": true, "email_on_campaign_completed": true, "email_on_new_mitre_techniques": false, "in_app_all": true}'
|
||||
|
||||
def _column_names(table: str) -> set[str]:
|
||||
bind = op.get_bind()
|
||||
insp = sa.inspect(bind)
|
||||
return {c["name"] for c in insp.get_columns(table)}
|
||||
|
||||
def upgrade() -> None:
|
||||
user_cols = _column_names("users")
|
||||
if "notification_preferences" not in user_cols:
|
||||
op.add_column(
|
||||
"users",
|
||||
sa.Column("notification_preferences", postgresql.JSONB, nullable=True, server_default=_DEFAULT_PREFS),
|
||||
)
|
||||
if "jira_account_id" not in user_cols:
|
||||
op.add_column(
|
||||
"users",
|
||||
sa.Column("jira_account_id", sa.String(100), nullable=True),
|
||||
)
|
||||
|
||||
def downgrade() -> None:
|
||||
user_cols = _column_names("users")
|
||||
if "jira_account_id" in user_cols:
|
||||
op.drop_column("users", "jira_account_id")
|
||||
if "notification_preferences" in user_cols:
|
||||
op.drop_column("users", "notification_preferences")
|
||||
@@ -1,43 +0,0 @@
|
||||
"""Phase 8: system_configs table for runtime configuration.
|
||||
|
||||
Revision ID: b033syscfg
|
||||
Revises: b032phase7
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.dialects import postgresql
|
||||
from alembic import op
|
||||
|
||||
revision: str = "b033syscfg"
|
||||
down_revision: Union[str, None] = "b032phase7"
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def _table_exists(name: str) -> bool:
|
||||
bind = op.get_bind()
|
||||
insp = sa.inspect(bind)
|
||||
return name in insp.get_table_names()
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
if not _table_exists("system_configs"):
|
||||
op.create_table(
|
||||
"system_configs",
|
||||
sa.Column("id", postgresql.UUID(as_uuid=True), primary_key=True),
|
||||
sa.Column("key", sa.String(200), unique=True, nullable=False),
|
||||
sa.Column("value", sa.Text, nullable=True),
|
||||
sa.Column("description", sa.String(500), nullable=True),
|
||||
sa.Column(
|
||||
"updated_at",
|
||||
sa.DateTime(timezone=True),
|
||||
server_default=sa.text("now()"),
|
||||
),
|
||||
)
|
||||
op.create_index("ix_system_configs_key", "system_configs", ["key"])
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
if _table_exists("system_configs"):
|
||||
op.drop_index("ix_system_configs_key", table_name="system_configs")
|
||||
op.drop_table("system_configs")
|
||||
@@ -1,174 +0,0 @@
|
||||
"""Phase 8: Detection Lifecycle Management tables.
|
||||
|
||||
Revision ID: b034dlm
|
||||
Revises: b033syscfg
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.dialects import postgresql
|
||||
from alembic import op
|
||||
|
||||
revision: str = "b034dlm"
|
||||
down_revision: Union[str, None] = "b033syscfg"
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def _table_exists(name: str) -> bool:
|
||||
bind = op.get_bind()
|
||||
insp = sa.inspect(bind)
|
||||
return name in insp.get_table_names()
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
if not _table_exists("detection_assets"):
|
||||
op.create_table(
|
||||
"detection_assets",
|
||||
sa.Column("id", postgresql.UUID(as_uuid=True), primary_key=True),
|
||||
sa.Column("name", sa.String(500), nullable=False),
|
||||
sa.Column("description", sa.Text),
|
||||
sa.Column("asset_type", sa.String(50), nullable=False),
|
||||
sa.Column("platform", sa.String(100)),
|
||||
sa.Column("rule_content", sa.Text),
|
||||
sa.Column("rule_language", sa.String(50)),
|
||||
sa.Column("rule_repository_url", sa.Text),
|
||||
sa.Column("rule_file_path", sa.String(500)),
|
||||
sa.Column("rule_version", sa.String(50)),
|
||||
sa.Column("rule_hash", sa.String(64)),
|
||||
sa.Column("last_rule_change_at", sa.DateTime),
|
||||
sa.Column("log_source_name", sa.String(200)),
|
||||
sa.Column("log_source_version", sa.String(50)),
|
||||
sa.Column("log_source_config", postgresql.JSONB, server_default="{}"),
|
||||
sa.Column("infrastructure_hash", sa.String(64)),
|
||||
sa.Column("infrastructure_details", postgresql.JSONB, server_default="{}"),
|
||||
sa.Column("health_status", sa.String(20), server_default="untested", nullable=False),
|
||||
sa.Column("last_alert_at", sa.DateTime),
|
||||
sa.Column("alert_count_30d", sa.Integer, server_default="0"),
|
||||
sa.Column("false_positive_rate", sa.Float),
|
||||
sa.Column("expected_alert_frequency", sa.String(50)),
|
||||
sa.Column("owner_id", postgresql.UUID(as_uuid=True), sa.ForeignKey("users.id", ondelete="SET NULL")),
|
||||
sa.Column("backup_owner_id", postgresql.UUID(as_uuid=True), sa.ForeignKey("users.id", ondelete="SET NULL")),
|
||||
sa.Column("team", sa.String(100)),
|
||||
sa.Column("is_active", sa.Boolean, server_default="true", nullable=False),
|
||||
sa.Column("tags", postgresql.JSONB, server_default="[]"),
|
||||
sa.Column("asset_metadata", postgresql.JSONB, server_default="{}"),
|
||||
sa.Column("created_by", postgresql.UUID(as_uuid=True), sa.ForeignKey("users.id", ondelete="SET NULL")),
|
||||
sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.text("now()")),
|
||||
sa.Column("updated_at", sa.DateTime(timezone=True), server_default=sa.text("now()")),
|
||||
)
|
||||
op.create_index("ix_detection_assets_platform", "detection_assets", ["platform"])
|
||||
op.create_index("ix_detection_assets_health_status", "detection_assets", ["health_status"])
|
||||
op.create_index("ix_detection_assets_owner_id", "detection_assets", ["owner_id"])
|
||||
|
||||
if not _table_exists("detection_technique_mappings"):
|
||||
op.create_table(
|
||||
"detection_technique_mappings",
|
||||
sa.Column("id", postgresql.UUID(as_uuid=True), primary_key=True),
|
||||
sa.Column("detection_asset_id", postgresql.UUID(as_uuid=True), sa.ForeignKey("detection_assets.id", ondelete="CASCADE"), nullable=False),
|
||||
sa.Column("technique_id", postgresql.UUID(as_uuid=True), sa.ForeignKey("techniques.id", ondelete="CASCADE"), nullable=False),
|
||||
sa.Column("coverage_type", sa.String(50), server_default="detect"),
|
||||
sa.Column("confidence_level", sa.String(20), server_default="medium"),
|
||||
sa.Column("notes", sa.Text),
|
||||
sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.text("now()")),
|
||||
)
|
||||
op.create_index("ix_detection_technique_mappings_technique_id", "detection_technique_mappings", ["technique_id"])
|
||||
op.create_index("ix_detection_technique_mappings_asset_id", "detection_technique_mappings", ["detection_asset_id"])
|
||||
|
||||
if not _table_exists("detection_validations"):
|
||||
op.create_table(
|
||||
"detection_validations",
|
||||
sa.Column("id", postgresql.UUID(as_uuid=True), primary_key=True),
|
||||
sa.Column("detection_asset_id", postgresql.UUID(as_uuid=True), sa.ForeignKey("detection_assets.id", ondelete="CASCADE"), nullable=False),
|
||||
sa.Column("technique_id", postgresql.UUID(as_uuid=True), sa.ForeignKey("techniques.id", ondelete="SET NULL")),
|
||||
sa.Column("test_id", postgresql.UUID(as_uuid=True), sa.ForeignKey("tests.id", ondelete="SET NULL")),
|
||||
sa.Column("validated_at", sa.DateTime),
|
||||
sa.Column("expires_at", sa.DateTime, nullable=False),
|
||||
sa.Column("is_valid", sa.Boolean, server_default="true", nullable=False),
|
||||
sa.Column("validation_result", sa.String(50)),
|
||||
sa.Column("validation_method", sa.String(100)),
|
||||
sa.Column("rule_hash_at_validation", sa.String(64)),
|
||||
sa.Column("log_source_version_at_validation", sa.String(50)),
|
||||
sa.Column("infrastructure_hash_at_validation", sa.String(64)),
|
||||
sa.Column("environment_snapshot", postgresql.JSONB, server_default="{}"),
|
||||
sa.Column("invalidated_at", sa.DateTime),
|
||||
sa.Column("invalidation_reason", sa.String(50)),
|
||||
sa.Column("invalidation_details", sa.Text),
|
||||
sa.Column("invalidated_by", postgresql.UUID(as_uuid=True), sa.ForeignKey("users.id", ondelete="SET NULL")),
|
||||
sa.Column("validated_by", postgresql.UUID(as_uuid=True), sa.ForeignKey("users.id", ondelete="SET NULL"), nullable=False),
|
||||
sa.Column("integrity_hash", sa.String(64)),
|
||||
sa.Column("notes", sa.Text),
|
||||
sa.Column("evidence_ids", postgresql.JSONB, server_default="[]"),
|
||||
)
|
||||
op.create_index("ix_detection_validations_asset_id_valid", "detection_validations", ["detection_asset_id", "is_valid"])
|
||||
op.create_index("ix_detection_validations_expires_at", "detection_validations", ["expires_at"])
|
||||
|
||||
if not _table_exists("technique_confidence_scores"):
|
||||
op.create_table(
|
||||
"technique_confidence_scores",
|
||||
sa.Column("id", postgresql.UUID(as_uuid=True), primary_key=True),
|
||||
sa.Column("technique_id", postgresql.UUID(as_uuid=True), sa.ForeignKey("techniques.id", ondelete="CASCADE"), nullable=False, unique=True),
|
||||
sa.Column("confidence_level", sa.String(20), server_default="unknown"),
|
||||
sa.Column("confidence_score", sa.Float, server_default="0.0"),
|
||||
sa.Column("detection_count", sa.Integer, server_default="0"),
|
||||
sa.Column("valid_detection_count", sa.Integer, server_default="0"),
|
||||
sa.Column("last_validated_at", sa.DateTime),
|
||||
sa.Column("next_validation_due", sa.DateTime),
|
||||
sa.Column("last_recalculated_at", sa.DateTime),
|
||||
sa.Column("recency_factor", sa.Float, server_default="0.0"),
|
||||
sa.Column("coverage_factor", sa.Float, server_default="0.0"),
|
||||
sa.Column("health_factor", sa.Float, server_default="0.0"),
|
||||
sa.Column("diversity_factor", sa.Float, server_default="0.0"),
|
||||
sa.Column("score_breakdown", postgresql.JSONB, server_default="{}"),
|
||||
sa.Column("risk_factors", postgresql.JSONB, server_default="[]"),
|
||||
sa.Column("updated_at", sa.DateTime),
|
||||
)
|
||||
op.create_index("ix_technique_confidence_scores_technique_id", "technique_confidence_scores", ["technique_id"])
|
||||
op.create_index("ix_technique_confidence_scores_confidence_level", "technique_confidence_scores", ["confidence_level"])
|
||||
|
||||
if not _table_exists("infrastructure_change_logs"):
|
||||
op.create_table(
|
||||
"infrastructure_change_logs",
|
||||
sa.Column("id", postgresql.UUID(as_uuid=True), primary_key=True),
|
||||
sa.Column("change_type", sa.String(100), nullable=False),
|
||||
sa.Column("description", sa.Text, nullable=False),
|
||||
sa.Column("affected_platforms", postgresql.JSONB, server_default="[]"),
|
||||
sa.Column("affected_log_sources", postgresql.JSONB, server_default="[]"),
|
||||
sa.Column("change_date", sa.DateTime),
|
||||
sa.Column("reported_by", postgresql.UUID(as_uuid=True), sa.ForeignKey("users.id", ondelete="SET NULL")),
|
||||
sa.Column("auto_invalidate", sa.Boolean, server_default="true"),
|
||||
sa.Column("invalidated_count", sa.Integer, server_default="0"),
|
||||
sa.Column("change_metadata", postgresql.JSONB, server_default="{}"),
|
||||
sa.Column("created_at", sa.DateTime),
|
||||
)
|
||||
op.create_index("ix_infrastructure_change_logs_change_date", "infrastructure_change_logs", ["change_date"])
|
||||
|
||||
if not _table_exists("decay_policies"):
|
||||
op.create_table(
|
||||
"decay_policies",
|
||||
sa.Column("id", postgresql.UUID(as_uuid=True), primary_key=True),
|
||||
sa.Column("name", sa.String(200), nullable=False),
|
||||
sa.Column("description", sa.Text),
|
||||
sa.Column("applies_to_platform", sa.String(100)),
|
||||
sa.Column("applies_to_asset_type", sa.String(50)),
|
||||
sa.Column("applies_to_tactic", sa.String(100)),
|
||||
sa.Column("fresh_days", sa.Integer, server_default="90"),
|
||||
sa.Column("aging_days", sa.Integer, server_default="180"),
|
||||
sa.Column("stale_days", sa.Integer, server_default="365"),
|
||||
sa.Column("default_validity_days", sa.Integer, server_default="180"),
|
||||
sa.Column("silent_threshold_days", sa.Integer, server_default="30"),
|
||||
sa.Column("noisy_threshold_daily", sa.Integer, server_default="100"),
|
||||
sa.Column("recency_weight", sa.Float, server_default="0.3"),
|
||||
sa.Column("coverage_weight", sa.Float, server_default="0.3"),
|
||||
sa.Column("health_weight", sa.Float, server_default="0.25"),
|
||||
sa.Column("diversity_weight", sa.Float, server_default="0.15"),
|
||||
sa.Column("is_default", sa.Boolean, server_default="false"),
|
||||
sa.Column("is_active", sa.Boolean, server_default="true"),
|
||||
sa.Column("created_at", sa.DateTime),
|
||||
sa.Column("updated_at", sa.DateTime),
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
for table in ["decay_policies", "infrastructure_change_logs", "technique_confidence_scores", "detection_validations", "detection_technique_mappings", "detection_assets"]:
|
||||
if _table_exists(table):
|
||||
op.drop_table(table)
|
||||
@@ -1,118 +0,0 @@
|
||||
"""Phase 9: Ownership & Revalidation Queue
|
||||
|
||||
Revision ID: b035ownerq
|
||||
Revises: b034dlm
|
||||
Create Date: 2026-05-19
|
||||
|
||||
Uses raw SQL for all DDL to avoid SQLAlchemy before_create hook issues
|
||||
with existing enum types.
|
||||
"""
|
||||
|
||||
from typing import Union
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
|
||||
revision: str = "b035ownerq"
|
||||
down_revision: Union[str, None] = "b034dlm"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
|
||||
# ── Enums (idempotent) ────────────────────────────────────────────────────
|
||||
conn.execute(sa.text("""
|
||||
DO $$ BEGIN
|
||||
CREATE TYPE queue_priority AS ENUM ('critical', 'high', 'medium', 'low');
|
||||
EXCEPTION WHEN duplicate_object THEN NULL;
|
||||
END $$
|
||||
"""))
|
||||
conn.execute(sa.text("""
|
||||
DO $$ BEGIN
|
||||
CREATE TYPE queue_status AS ENUM ('pending', 'in_progress', 'completed', 'dismissed');
|
||||
EXCEPTION WHEN duplicate_object THEN NULL;
|
||||
END $$
|
||||
"""))
|
||||
conn.execute(sa.text("""
|
||||
DO $$ BEGIN
|
||||
CREATE TYPE queue_reason AS ENUM (
|
||||
'validation_expired', 'infra_change', 'osint_alert',
|
||||
'mitre_update', 'rule_modified', 'low_confidence', 'manual');
|
||||
EXCEPTION WHEN duplicate_object THEN NULL;
|
||||
END $$
|
||||
"""))
|
||||
|
||||
# ── technique_ownerships ──────────────────────────────────────────────────
|
||||
conn.execute(sa.text("""
|
||||
CREATE TABLE IF NOT EXISTS technique_ownerships (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
technique_id UUID NOT NULL UNIQUE
|
||||
REFERENCES techniques(id) ON DELETE CASCADE,
|
||||
owner_id UUID
|
||||
REFERENCES users(id) ON DELETE SET NULL,
|
||||
backup_owner_id UUID
|
||||
REFERENCES users(id) ON DELETE SET NULL,
|
||||
team VARCHAR(200),
|
||||
notes TEXT,
|
||||
assigned_at TIMESTAMP,
|
||||
assigned_by UUID
|
||||
REFERENCES users(id) ON DELETE SET NULL,
|
||||
created_at TIMESTAMP DEFAULT now(),
|
||||
updated_at TIMESTAMP DEFAULT now()
|
||||
)
|
||||
"""))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_techown_owner_id ON technique_ownerships (owner_id)"
|
||||
))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_techown_technique_id ON technique_ownerships (technique_id)"
|
||||
))
|
||||
|
||||
# ── revalidation_queue_items ──────────────────────────────────────────────
|
||||
conn.execute(sa.text("""
|
||||
CREATE TABLE IF NOT EXISTS revalidation_queue_items (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
technique_id UUID
|
||||
REFERENCES techniques(id) ON DELETE CASCADE,
|
||||
detection_asset_id UUID
|
||||
REFERENCES detection_assets(id) ON DELETE CASCADE,
|
||||
priority queue_priority NOT NULL DEFAULT 'medium',
|
||||
reason queue_reason NOT NULL,
|
||||
reason_detail TEXT,
|
||||
status queue_status NOT NULL DEFAULT 'pending',
|
||||
assigned_to UUID
|
||||
REFERENCES users(id) ON DELETE SET NULL,
|
||||
due_date TIMESTAMP,
|
||||
created_at TIMESTAMP DEFAULT now(),
|
||||
completed_at TIMESTAMP,
|
||||
dismissed_at TIMESTAMP,
|
||||
completed_by UUID
|
||||
REFERENCES users(id) ON DELETE SET NULL,
|
||||
extra JSONB
|
||||
)
|
||||
"""))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_rqueue_status ON revalidation_queue_items (status)"
|
||||
))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_rqueue_priority ON revalidation_queue_items (priority)"
|
||||
))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_rqueue_assigned_to ON revalidation_queue_items (assigned_to)"
|
||||
))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_rqueue_technique_id ON revalidation_queue_items (technique_id)"
|
||||
))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_rqueue_asset_id ON revalidation_queue_items (detection_asset_id)"
|
||||
))
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
conn.execute(sa.text("DROP TABLE IF EXISTS revalidation_queue_items"))
|
||||
conn.execute(sa.text("DROP TABLE IF EXISTS technique_ownerships"))
|
||||
conn.execute(sa.text("DROP TYPE IF EXISTS queue_reason"))
|
||||
conn.execute(sa.text("DROP TYPE IF EXISTS queue_status"))
|
||||
conn.execute(sa.text("DROP TYPE IF EXISTS queue_priority"))
|
||||
@@ -1,184 +0,0 @@
|
||||
"""Phase 10: Attack Paths & Advanced Purple Team
|
||||
|
||||
Revision ID: b036atk
|
||||
Revises: b035ownerq
|
||||
Create Date: 2026-05-19
|
||||
|
||||
Uses raw SQL to avoid SQLAlchemy DDL hook issues with enum types.
|
||||
"""
|
||||
|
||||
from typing import Union
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
|
||||
revision: str = "b036atk"
|
||||
down_revision: Union[str, None] = "b035ownerq"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
|
||||
# ── Enums ─────────────────────────────────────────────────────────────────
|
||||
conn.execute(sa.text("""
|
||||
DO $$ BEGIN
|
||||
CREATE TYPE execution_status AS ENUM
|
||||
('planned','in_progress','completed','aborted');
|
||||
EXCEPTION WHEN duplicate_object THEN NULL; END $$
|
||||
"""))
|
||||
conn.execute(sa.text("""
|
||||
DO $$ BEGIN
|
||||
CREATE TYPE step_result_status AS ENUM
|
||||
('pending','executing','detected','not_detected','skipped');
|
||||
EXCEPTION WHEN duplicate_object THEN NULL; END $$
|
||||
"""))
|
||||
conn.execute(sa.text("""
|
||||
DO $$ BEGIN
|
||||
CREATE TYPE timeline_actor_side AS ENUM ('red','blue','system');
|
||||
EXCEPTION WHEN duplicate_object THEN NULL; END $$
|
||||
"""))
|
||||
conn.execute(sa.text("""
|
||||
DO $$ BEGIN
|
||||
CREATE TYPE timeline_entry_type AS ENUM
|
||||
('action','detection','note','phase_transition','flag');
|
||||
EXCEPTION WHEN duplicate_object THEN NULL; END $$
|
||||
"""))
|
||||
|
||||
# ── attack_paths ──────────────────────────────────────────────────────────
|
||||
conn.execute(sa.text("""
|
||||
CREATE TABLE IF NOT EXISTS attack_paths (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
name VARCHAR(300) NOT NULL,
|
||||
description TEXT,
|
||||
objective TEXT,
|
||||
is_template BOOLEAN DEFAULT FALSE,
|
||||
threat_actor_id UUID REFERENCES threat_actors(id) ON DELETE SET NULL,
|
||||
created_by UUID REFERENCES users(id) ON DELETE SET NULL,
|
||||
tags JSONB,
|
||||
is_active BOOLEAN DEFAULT TRUE,
|
||||
created_at TIMESTAMP DEFAULT now(),
|
||||
updated_at TIMESTAMP DEFAULT now()
|
||||
)
|
||||
"""))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_attack_paths_created_by ON attack_paths (created_by)"
|
||||
))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_attack_paths_is_template ON attack_paths (is_template)"
|
||||
))
|
||||
|
||||
# ── attack_path_steps ─────────────────────────────────────────────────────
|
||||
conn.execute(sa.text("""
|
||||
CREATE TABLE IF NOT EXISTS attack_path_steps (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
attack_path_id UUID NOT NULL REFERENCES attack_paths(id) ON DELETE CASCADE,
|
||||
order_index INTEGER NOT NULL DEFAULT 0,
|
||||
kill_chain_phase VARCHAR(60),
|
||||
technique_id UUID REFERENCES techniques(id) ON DELETE SET NULL,
|
||||
test_id UUID REFERENCES tests(id) ON DELETE SET NULL,
|
||||
name VARCHAR(300),
|
||||
description TEXT,
|
||||
expected_detection BOOLEAN DEFAULT TRUE,
|
||||
notes TEXT
|
||||
)
|
||||
"""))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_ap_steps_path_id ON attack_path_steps (attack_path_id)"
|
||||
))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_ap_steps_technique_id ON attack_path_steps (technique_id)"
|
||||
))
|
||||
|
||||
# ── attack_path_executions ────────────────────────────────────────────────
|
||||
conn.execute(sa.text("""
|
||||
CREATE TABLE IF NOT EXISTS attack_path_executions (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
attack_path_id UUID NOT NULL REFERENCES attack_paths(id) ON DELETE CASCADE,
|
||||
status execution_status NOT NULL DEFAULT 'planned',
|
||||
environment VARCHAR(100),
|
||||
red_team_lead UUID REFERENCES users(id) ON DELETE SET NULL,
|
||||
blue_team_lead UUID REFERENCES users(id) ON DELETE SET NULL,
|
||||
started_by UUID REFERENCES users(id) ON DELETE SET NULL,
|
||||
started_at TIMESTAMP,
|
||||
completed_at TIMESTAMP,
|
||||
notes TEXT,
|
||||
created_at TIMESTAMP DEFAULT now(),
|
||||
-- kill-chain metrics (populated on completion)
|
||||
total_steps INTEGER,
|
||||
detected_steps INTEGER,
|
||||
not_detected_steps INTEGER,
|
||||
skipped_steps INTEGER,
|
||||
detection_rate FLOAT,
|
||||
mttd_seconds FLOAT,
|
||||
furthest_undetected_step INTEGER
|
||||
)
|
||||
"""))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_ap_exec_path_id ON attack_path_executions (attack_path_id)"
|
||||
))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_ap_exec_status ON attack_path_executions (status)"
|
||||
))
|
||||
|
||||
# ── attack_path_step_results ──────────────────────────────────────────────
|
||||
conn.execute(sa.text("""
|
||||
CREATE TABLE IF NOT EXISTS attack_path_step_results (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
execution_id UUID NOT NULL
|
||||
REFERENCES attack_path_executions(id) ON DELETE CASCADE,
|
||||
step_id UUID NOT NULL
|
||||
REFERENCES attack_path_steps(id) ON DELETE CASCADE,
|
||||
step_order INTEGER NOT NULL DEFAULT 0,
|
||||
status step_result_status NOT NULL DEFAULT 'pending',
|
||||
executed_by UUID REFERENCES users(id) ON DELETE SET NULL,
|
||||
executed_at TIMESTAMP,
|
||||
detected_at TIMESTAMP,
|
||||
time_to_detect_seconds FLOAT,
|
||||
detection_asset_id UUID
|
||||
REFERENCES detection_assets(id) ON DELETE SET NULL,
|
||||
notes TEXT,
|
||||
evidence_ids JSONB
|
||||
)
|
||||
"""))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_ap_stepres_exec ON attack_path_step_results (execution_id)"
|
||||
))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_ap_stepres_step ON attack_path_step_results (step_id)"
|
||||
))
|
||||
|
||||
# ── attack_path_timeline ──────────────────────────────────────────────────
|
||||
conn.execute(sa.text("""
|
||||
CREATE TABLE IF NOT EXISTS attack_path_timeline (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
execution_id UUID NOT NULL
|
||||
REFERENCES attack_path_executions(id) ON DELETE CASCADE,
|
||||
step_id UUID REFERENCES attack_path_steps(id) ON DELETE SET NULL,
|
||||
timestamp TIMESTAMP NOT NULL DEFAULT now(),
|
||||
actor_side timeline_actor_side NOT NULL,
|
||||
actor_id UUID REFERENCES users(id) ON DELETE SET NULL,
|
||||
entry_type timeline_entry_type NOT NULL,
|
||||
content TEXT NOT NULL,
|
||||
extra JSONB
|
||||
)
|
||||
"""))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_timeline_execution_id ON attack_path_timeline (execution_id)"
|
||||
))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_timeline_timestamp ON attack_path_timeline (timestamp)"
|
||||
))
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
conn.execute(sa.text("DROP TABLE IF EXISTS attack_path_timeline"))
|
||||
conn.execute(sa.text("DROP TABLE IF EXISTS attack_path_step_results"))
|
||||
conn.execute(sa.text("DROP TABLE IF EXISTS attack_path_executions"))
|
||||
conn.execute(sa.text("DROP TABLE IF EXISTS attack_path_steps"))
|
||||
conn.execute(sa.text("DROP TABLE IF EXISTS attack_paths"))
|
||||
conn.execute(sa.text("DROP TYPE IF EXISTS timeline_entry_type"))
|
||||
conn.execute(sa.text("DROP TYPE IF EXISTS timeline_actor_side"))
|
||||
conn.execute(sa.text("DROP TYPE IF EXISTS step_result_status"))
|
||||
conn.execute(sa.text("DROP TYPE IF EXISTS execution_status"))
|
||||
@@ -1,106 +0,0 @@
|
||||
"""Phase 11: Knowledge Management — Playbooks + Lessons Learned
|
||||
|
||||
Revision ID: b037know
|
||||
Revises: b036atk
|
||||
Create Date: 2026-05-20
|
||||
|
||||
Uses raw SQL to bypass SQLAlchemy DDL hooks (no enum types — string columns
|
||||
with Pydantic-layer validation instead, so no PostgreSQL enums needed).
|
||||
"""
|
||||
|
||||
from typing import Union
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
|
||||
revision: str = "b037know"
|
||||
down_revision: Union[str, None] = "b036atk"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
|
||||
# ── playbooks ──────────────────────────────────────────────────────────────
|
||||
conn.execute(sa.text("""
|
||||
CREATE TABLE IF NOT EXISTS playbooks (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
technique_id UUID NOT NULL REFERENCES techniques(id) ON DELETE CASCADE,
|
||||
playbook_type VARCHAR(32) NOT NULL,
|
||||
title VARCHAR(255) NOT NULL,
|
||||
content TEXT NOT NULL DEFAULT '',
|
||||
version INTEGER NOT NULL DEFAULT 1,
|
||||
tools JSONB,
|
||||
prerequisites JSONB,
|
||||
created_by UUID REFERENCES users(id) ON DELETE SET NULL,
|
||||
updated_by UUID REFERENCES users(id) ON DELETE SET NULL,
|
||||
created_at TIMESTAMP DEFAULT now(),
|
||||
updated_at TIMESTAMP DEFAULT now(),
|
||||
is_active BOOLEAN DEFAULT TRUE,
|
||||
CONSTRAINT uq_playbook_technique_type UNIQUE (technique_id, playbook_type)
|
||||
)
|
||||
"""))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_playbooks_technique_id ON playbooks (technique_id)"
|
||||
))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_playbooks_type ON playbooks (playbook_type)"
|
||||
))
|
||||
|
||||
# ── playbook_versions ──────────────────────────────────────────────────────
|
||||
conn.execute(sa.text("""
|
||||
CREATE TABLE IF NOT EXISTS playbook_versions (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
playbook_id UUID NOT NULL REFERENCES playbooks(id) ON DELETE CASCADE,
|
||||
version INTEGER NOT NULL,
|
||||
title VARCHAR(255) NOT NULL,
|
||||
content TEXT NOT NULL DEFAULT '',
|
||||
tools JSONB,
|
||||
prerequisites JSONB,
|
||||
changed_by UUID REFERENCES users(id) ON DELETE SET NULL,
|
||||
change_note VARCHAR(500),
|
||||
created_at TIMESTAMP DEFAULT now()
|
||||
)
|
||||
"""))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_pb_versions_playbook_id ON playbook_versions (playbook_id)"
|
||||
))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_pb_versions_version ON playbook_versions (playbook_id, version)"
|
||||
))
|
||||
|
||||
# ── lessons_learned ────────────────────────────────────────────────────────
|
||||
conn.execute(sa.text("""
|
||||
CREATE TABLE IF NOT EXISTS lessons_learned (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
title VARCHAR(255) NOT NULL,
|
||||
what_happened TEXT NOT NULL DEFAULT '',
|
||||
root_cause TEXT NOT NULL DEFAULT '',
|
||||
fix_applied TEXT,
|
||||
severity VARCHAR(16) NOT NULL DEFAULT 'medium',
|
||||
entity_type VARCHAR(32) NOT NULL DEFAULT 'manual',
|
||||
entity_id UUID,
|
||||
technique_ids JSONB,
|
||||
tags JSONB,
|
||||
created_by UUID REFERENCES users(id) ON DELETE SET NULL,
|
||||
created_at TIMESTAMP DEFAULT now(),
|
||||
updated_at TIMESTAMP DEFAULT now(),
|
||||
is_active BOOLEAN DEFAULT TRUE
|
||||
)
|
||||
"""))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_ll_entity ON lessons_learned (entity_type, entity_id)"
|
||||
))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_ll_severity ON lessons_learned (severity)"
|
||||
))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_ll_created_by ON lessons_learned (created_by)"
|
||||
))
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
conn.execute(sa.text("DROP TABLE IF EXISTS lessons_learned"))
|
||||
conn.execute(sa.text("DROP TABLE IF EXISTS playbook_versions"))
|
||||
conn.execute(sa.text("DROP TABLE IF EXISTS playbooks"))
|
||||
@@ -1,62 +0,0 @@
|
||||
"""Phase 12: Risk Intelligence — technique_risk_profiles table
|
||||
|
||||
Revision ID: b038risk
|
||||
Revises: b037know
|
||||
Create Date: 2026-05-20
|
||||
|
||||
Uses raw SQL to bypass SQLAlchemy DDL hooks.
|
||||
"""
|
||||
|
||||
from typing import Union
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
|
||||
revision: str = "b038risk"
|
||||
down_revision: Union[str, None] = "b037know"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
|
||||
conn.execute(sa.text("""
|
||||
CREATE TABLE IF NOT EXISTS technique_risk_profiles (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
technique_id UUID NOT NULL REFERENCES techniques(id) ON DELETE CASCADE,
|
||||
risk_score FLOAT NOT NULL DEFAULT 0.0,
|
||||
likelihood FLOAT NOT NULL DEFAULT 0.0,
|
||||
impact FLOAT NOT NULL DEFAULT 0.0,
|
||||
risk_level VARCHAR(16) NOT NULL DEFAULT 'info',
|
||||
detection_gap FLOAT NOT NULL DEFAULT 1.0,
|
||||
threat_actor_count INTEGER NOT NULL DEFAULT 0,
|
||||
osint_signal_count INTEGER NOT NULL DEFAULT 0,
|
||||
test_fail_count INTEGER NOT NULL DEFAULT 0,
|
||||
test_total_count INTEGER NOT NULL DEFAULT 0,
|
||||
test_failure_rate FLOAT NOT NULL DEFAULT 0.0,
|
||||
confidence_level FLOAT NOT NULL DEFAULT 0.0,
|
||||
scoring_breakdown JSONB,
|
||||
recommendations JSONB,
|
||||
computed_at TIMESTAMP DEFAULT now(),
|
||||
is_stale BOOLEAN DEFAULT TRUE,
|
||||
CONSTRAINT uq_risk_profile_technique UNIQUE (technique_id)
|
||||
)
|
||||
"""))
|
||||
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_risk_profiles_risk_score "
|
||||
"ON technique_risk_profiles (risk_score)"
|
||||
))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_risk_profiles_risk_level "
|
||||
"ON technique_risk_profiles (risk_level)"
|
||||
))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_risk_profiles_stale "
|
||||
"ON technique_risk_profiles (is_stale)"
|
||||
))
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
conn.execute(sa.text("DROP TABLE IF EXISTS technique_risk_profiles"))
|
||||
@@ -1,77 +0,0 @@
|
||||
"""Phase 13: Executive Dashboard — posture_snapshots table.
|
||||
|
||||
Revision ID: b039exec
|
||||
Revises: b038risk
|
||||
Create Date: 2026-05-20
|
||||
"""
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
|
||||
revision = "b039exec"
|
||||
down_revision = "b038risk"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
conn.execute(sa.text("""
|
||||
CREATE TABLE IF NOT EXISTS posture_snapshots (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
snapshot_date DATE NOT NULL,
|
||||
|
||||
-- Coverage
|
||||
total_techniques INTEGER NOT NULL DEFAULT 0,
|
||||
validated_count INTEGER NOT NULL DEFAULT 0,
|
||||
partial_count INTEGER NOT NULL DEFAULT 0,
|
||||
not_covered_count INTEGER NOT NULL DEFAULT 0,
|
||||
coverage_pct FLOAT NOT NULL DEFAULT 0.0,
|
||||
|
||||
-- Risk
|
||||
avg_risk_score FLOAT NOT NULL DEFAULT 0.0,
|
||||
critical_count INTEGER NOT NULL DEFAULT 0,
|
||||
high_count INTEGER NOT NULL DEFAULT 0,
|
||||
medium_count INTEGER NOT NULL DEFAULT 0,
|
||||
low_count INTEGER NOT NULL DEFAULT 0,
|
||||
|
||||
-- Operations
|
||||
open_queue_items INTEGER NOT NULL DEFAULT 0,
|
||||
orphan_techniques INTEGER NOT NULL DEFAULT 0,
|
||||
|
||||
-- Knowledge
|
||||
playbook_count INTEGER NOT NULL DEFAULT 0,
|
||||
lesson_count INTEGER NOT NULL DEFAULT 0,
|
||||
|
||||
-- MTTD
|
||||
mttd_avg_seconds FLOAT,
|
||||
executions_30d INTEGER NOT NULL DEFAULT 0,
|
||||
detection_rate_30d FLOAT,
|
||||
|
||||
-- Meta
|
||||
created_by UUID REFERENCES users(id) ON DELETE SET NULL,
|
||||
created_at TIMESTAMP WITHOUT TIME ZONE DEFAULT now(),
|
||||
extra JSONB
|
||||
)
|
||||
"""))
|
||||
|
||||
# Unique constraint: one snapshot per calendar day
|
||||
conn.execute(sa.text("""
|
||||
DO $$ BEGIN
|
||||
ALTER TABLE posture_snapshots
|
||||
ADD CONSTRAINT uq_posture_snapshot_date UNIQUE (snapshot_date);
|
||||
EXCEPTION WHEN duplicate_table THEN NULL;
|
||||
WHEN duplicate_object THEN NULL;
|
||||
END $$
|
||||
"""))
|
||||
|
||||
# Index for date-range trend queries
|
||||
conn.execute(sa.text("""
|
||||
CREATE INDEX IF NOT EXISTS ix_posture_snapshots_date
|
||||
ON posture_snapshots (snapshot_date)
|
||||
"""))
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
conn.execute(sa.text("DROP TABLE IF EXISTS posture_snapshots CASCADE"))
|
||||
@@ -1,75 +0,0 @@
|
||||
"""Phase 14: Enterprise Readiness — api_keys and sso_configs tables.
|
||||
|
||||
Revision ID: b040ent
|
||||
Revises: b039exec
|
||||
Create Date: 2026-05-20
|
||||
"""
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
|
||||
revision = "b040ent"
|
||||
down_revision = "b039exec"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
|
||||
# ── api_keys ──────────────────────────────────────────────────────────────
|
||||
conn.execute(sa.text("""
|
||||
CREATE TABLE IF NOT EXISTS api_keys (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
name VARCHAR(200) NOT NULL,
|
||||
description TEXT,
|
||||
key_prefix VARCHAR(13) NOT NULL,
|
||||
key_hash VARCHAR(64) NOT NULL UNIQUE,
|
||||
user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE,
|
||||
scopes JSONB NOT NULL DEFAULT '["read"]',
|
||||
last_used_at TIMESTAMP WITHOUT TIME ZONE,
|
||||
expires_at TIMESTAMP WITHOUT TIME ZONE,
|
||||
is_active BOOLEAN NOT NULL DEFAULT TRUE,
|
||||
created_at TIMESTAMP WITHOUT TIME ZONE DEFAULT now()
|
||||
)
|
||||
"""))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_api_keys_user_id ON api_keys (user_id)"
|
||||
))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_api_keys_key_hash ON api_keys (key_hash)"
|
||||
))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_api_keys_active ON api_keys (is_active)"
|
||||
))
|
||||
|
||||
# ── sso_configs ───────────────────────────────────────────────────────────
|
||||
conn.execute(sa.text("""
|
||||
CREATE TABLE IF NOT EXISTS sso_configs (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
is_enabled BOOLEAN NOT NULL DEFAULT FALSE,
|
||||
provider_name VARCHAR(200),
|
||||
sp_entity_id VARCHAR(500),
|
||||
sp_acs_url VARCHAR(500),
|
||||
sp_slo_url VARCHAR(500),
|
||||
sp_certificate TEXT,
|
||||
sp_private_key TEXT,
|
||||
idp_entity_id VARCHAR(500),
|
||||
idp_sso_url VARCHAR(500),
|
||||
idp_slo_url VARCHAR(500),
|
||||
idp_certificate TEXT,
|
||||
attr_email VARCHAR(200) DEFAULT 'email',
|
||||
attr_username VARCHAR(200) DEFAULT 'username',
|
||||
attr_role VARCHAR(200) DEFAULT 'role',
|
||||
default_role VARCHAR(50) DEFAULT 'viewer',
|
||||
auto_provision BOOLEAN NOT NULL DEFAULT TRUE,
|
||||
created_at TIMESTAMP WITHOUT TIME ZONE DEFAULT now(),
|
||||
updated_at TIMESTAMP WITHOUT TIME ZONE DEFAULT now()
|
||||
)
|
||||
"""))
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
conn.execute(sa.text("DROP TABLE IF EXISTS api_keys CASCADE"))
|
||||
conn.execute(sa.text("DROP TABLE IF EXISTS sso_configs CASCADE"))
|
||||
@@ -1,82 +0,0 @@
|
||||
"""Phase 13: Operational Alerts — alert_rules and alert_instances tables.
|
||||
|
||||
Revision ID: b041alerts
|
||||
Revises: b040ent
|
||||
Create Date: 2026-05-21
|
||||
"""
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
|
||||
revision = "b041alerts"
|
||||
down_revision = "b040ent"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
|
||||
# ── alert_rules ───────────────────────────────────────────────────────────
|
||||
conn.execute(sa.text("""
|
||||
CREATE TABLE IF NOT EXISTS alert_rules (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
name VARCHAR(300) NOT NULL,
|
||||
description TEXT,
|
||||
rule_type VARCHAR(50) NOT NULL,
|
||||
severity VARCHAR(20) NOT NULL DEFAULT 'medium',
|
||||
is_enabled BOOLEAN NOT NULL DEFAULT TRUE,
|
||||
is_system BOOLEAN NOT NULL DEFAULT FALSE,
|
||||
config JSONB NOT NULL DEFAULT '{}',
|
||||
notify_in_app BOOLEAN NOT NULL DEFAULT TRUE,
|
||||
notify_webhook BOOLEAN NOT NULL DEFAULT FALSE,
|
||||
webhook_id UUID REFERENCES webhook_configs(id) ON DELETE SET NULL,
|
||||
cooldown_hours INTEGER NOT NULL DEFAULT 24,
|
||||
created_by UUID REFERENCES users(id) ON DELETE SET NULL,
|
||||
created_at TIMESTAMP WITHOUT TIME ZONE DEFAULT now(),
|
||||
last_fired_at TIMESTAMP WITHOUT TIME ZONE
|
||||
)
|
||||
"""))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_alert_rules_type ON alert_rules (rule_type)"
|
||||
))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_alert_rules_enabled ON alert_rules (is_enabled)"
|
||||
))
|
||||
|
||||
# ── alert_instances ───────────────────────────────────────────────────────
|
||||
conn.execute(sa.text("""
|
||||
CREATE TABLE IF NOT EXISTS alert_instances (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
rule_id UUID REFERENCES alert_rules(id) ON DELETE SET NULL,
|
||||
rule_name VARCHAR(300) NOT NULL,
|
||||
rule_type VARCHAR(50) NOT NULL,
|
||||
severity VARCHAR(20) NOT NULL,
|
||||
title VARCHAR(500) NOT NULL,
|
||||
message TEXT NOT NULL,
|
||||
details JSONB,
|
||||
status VARCHAR(20) NOT NULL DEFAULT 'open',
|
||||
acknowledged_by UUID REFERENCES users(id) ON DELETE SET NULL,
|
||||
acknowledged_at TIMESTAMP WITHOUT TIME ZONE,
|
||||
resolved_at TIMESTAMP WITHOUT TIME ZONE,
|
||||
created_at TIMESTAMP WITHOUT TIME ZONE DEFAULT now()
|
||||
)
|
||||
"""))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_alert_instances_rule_id ON alert_instances (rule_id)"
|
||||
))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_alert_instances_status ON alert_instances (status)"
|
||||
))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_alert_instances_severity ON alert_instances (severity)"
|
||||
))
|
||||
conn.execute(sa.text(
|
||||
"CREATE INDEX IF NOT EXISTS ix_alert_instances_created ON alert_instances (created_at)"
|
||||
))
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
conn.execute(sa.text("DROP TABLE IF EXISTS alert_instances CASCADE"))
|
||||
conn.execute(sa.text("DROP TABLE IF EXISTS alert_rules CASCADE"))
|
||||
@@ -1,25 +0,0 @@
|
||||
"""Add jira_api_token to users table.
|
||||
|
||||
Revision ID: b042
|
||||
Revises: b041_operational_alerts
|
||||
Create Date: 2026-05-26
|
||||
"""
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
|
||||
revision = "b042"
|
||||
down_revision = "b041alerts"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.add_column(
|
||||
"users",
|
||||
sa.Column("jira_api_token", sa.String(500), nullable=True),
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_column("users", "jira_api_token")
|
||||
@@ -1,28 +0,0 @@
|
||||
"""Add jira_email to users table.
|
||||
|
||||
Allows each user to specify a separate email for Jira authentication,
|
||||
independent of their Aegis account email.
|
||||
|
||||
Revision ID: b043
|
||||
Revises: b042
|
||||
Create Date: 2026-05-26
|
||||
"""
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
|
||||
revision = "b043"
|
||||
down_revision = "b042"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.add_column(
|
||||
"users",
|
||||
sa.Column("jira_email", sa.String(255), nullable=True),
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_column("users", "jira_email")
|
||||
@@ -1,25 +0,0 @@
|
||||
"""add tempo_api_token to users
|
||||
|
||||
Revision ID: b044
|
||||
Revises: b043
|
||||
Create Date: 2026-05-27
|
||||
|
||||
"""
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
|
||||
revision = "b044"
|
||||
down_revision = "b043"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade():
|
||||
op.add_column(
|
||||
"users",
|
||||
sa.Column("tempo_api_token", sa.String(500), nullable=True),
|
||||
)
|
||||
|
||||
|
||||
def downgrade():
|
||||
op.drop_column("users", "tempo_api_token")
|
||||
@@ -1,16 +0,0 @@
|
||||
"""Add blue_work_started_at to tests table."""
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
|
||||
revision = "b045"
|
||||
down_revision = "b044"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade():
|
||||
op.add_column("tests", sa.Column("blue_work_started_at", sa.DateTime(), nullable=True))
|
||||
|
||||
|
||||
def downgrade():
|
||||
op.drop_column("tests", "blue_work_started_at")
|
||||
@@ -1,22 +0,0 @@
|
||||
"""Add 'disputed' value to teststate enum.
|
||||
|
||||
Revision ID: b046
|
||||
Revises: b045
|
||||
Create Date: 2026-06-03
|
||||
"""
|
||||
|
||||
from alembic import op
|
||||
|
||||
revision = "b046"
|
||||
down_revision = "b045"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.execute("ALTER TYPE teststate ADD VALUE IF NOT EXISTS 'disputed'")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
# PostgreSQL does not support removing enum values; downgrade is a no-op.
|
||||
pass
|
||||
@@ -1,27 +0,0 @@
|
||||
"""Add start_date to campaigns.
|
||||
|
||||
Revision ID: b047
|
||||
Revises: b046
|
||||
Create Date: 2026-06-03
|
||||
"""
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
|
||||
revision = "b047"
|
||||
down_revision = "b046"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.add_column(
|
||||
"campaigns",
|
||||
sa.Column("start_date", sa.DateTime(), nullable=True),
|
||||
)
|
||||
op.create_index("ix_campaigns_start_date", "campaigns", ["start_date"])
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_index("ix_campaigns_start_date", table_name="campaigns")
|
||||
op.drop_column("campaigns", "start_date")
|
||||
@@ -1,39 +0,0 @@
|
||||
"""Add evaluation_imports table.
|
||||
|
||||
Revision ID: b048
|
||||
Revises: b047
|
||||
Create Date: 2026-06-05
|
||||
"""
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.dialects.postgresql import UUID
|
||||
|
||||
revision = "b048"
|
||||
down_revision = "b047"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.create_table(
|
||||
"evaluation_imports",
|
||||
sa.Column("id", UUID(as_uuid=True), primary_key=True),
|
||||
sa.Column("adversary_name", sa.String, nullable=False),
|
||||
sa.Column("adversary_display", sa.String, nullable=False),
|
||||
sa.Column("eval_round", sa.Integer, nullable=False),
|
||||
sa.Column("imported_at", sa.DateTime, nullable=False),
|
||||
sa.Column("imported_by", UUID(as_uuid=True), sa.ForeignKey("users.id"), nullable=True),
|
||||
sa.Column("tests_created", sa.Integer, default=0),
|
||||
sa.Column("techniques_covered", sa.Integer, default=0),
|
||||
sa.Column("status", sa.String, default="completed"),
|
||||
sa.Column("notes", sa.Text, nullable=True),
|
||||
)
|
||||
op.create_index("ix_evaluation_imports_adversary", "evaluation_imports", ["adversary_name"])
|
||||
op.create_index("ix_evaluation_imports_round", "evaluation_imports", ["eval_round"])
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_index("ix_evaluation_imports_round", table_name="evaluation_imports")
|
||||
op.drop_index("ix_evaluation_imports_adversary", table_name="evaluation_imports")
|
||||
op.drop_table("evaluation_imports")
|
||||
@@ -1 +0,0 @@
|
||||
"""Aegis — MITRE ATT&CK Coverage Platform application package."""
|
||||
|
||||
+7
-95
@@ -1,34 +1,20 @@
|
||||
"""Security utilities: password hashing and JWT token management.
|
||||
"""
|
||||
Security utilities: password hashing and JWT token management.
|
||||
|
||||
This module provides pure functions for:
|
||||
- Hashing and verifying passwords using bcrypt via passlib.
|
||||
- Creating JWT access tokens using PyJWT.
|
||||
- Managing a Redis-backed token blacklist for revocation.
|
||||
- Creating JWT access tokens using python-jose.
|
||||
|
||||
No endpoints are defined here.
|
||||
"""
|
||||
|
||||
# Import logging
|
||||
import logging
|
||||
|
||||
# Import uuid
|
||||
import uuid as _uuid
|
||||
|
||||
# Import datetime, timedelta, timezone from datetime
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
# Import jwt (PyJWT)
|
||||
import jwt
|
||||
|
||||
# Import CryptContext from passlib.context
|
||||
from jose import jwt
|
||||
from passlib.context import CryptContext
|
||||
|
||||
# Import settings from app.config
|
||||
from app.config import settings
|
||||
|
||||
# Assign logger = logging.getLogger(__name__)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Password hashing
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -36,17 +22,13 @@ logger = logging.getLogger(__name__)
|
||||
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
|
||||
|
||||
|
||||
# Define function hash_password
|
||||
def hash_password(password: str) -> str:
|
||||
"""Return a bcrypt hash of *password*."""
|
||||
# Return pwd_context.hash(password)
|
||||
return pwd_context.hash(password)
|
||||
|
||||
|
||||
# Define function verify_password
|
||||
def verify_password(plain: str, hashed: str) -> bool:
|
||||
"""Return ``True`` if *plain* matches the bcrypt *hashed* value."""
|
||||
# Return pwd_context.verify(plain, hashed)
|
||||
return pwd_context.verify(plain, hashed)
|
||||
|
||||
|
||||
@@ -56,83 +38,13 @@ def verify_password(plain: str, hashed: str) -> bool:
|
||||
|
||||
|
||||
def create_access_token(data: dict) -> str:
|
||||
"""Create a signed JWT containing *data* plus ``exp`` and ``jti`` claims.
|
||||
"""Create a signed JWT containing *data* plus an ``exp`` claim.
|
||||
|
||||
- ``jti`` (JWT ID): unique identifier that enables token revocation.
|
||||
- ``exp``: expiration timestamp based on ``ACCESS_TOKEN_EXPIRE_MINUTES``.
|
||||
The token expires after ``ACCESS_TOKEN_EXPIRE_MINUTES`` (from settings).
|
||||
"""
|
||||
# Assign to_encode = data.copy()
|
||||
to_encode = data.copy()
|
||||
# Assign expire = datetime.now(timezone.utc) + timedelta(
|
||||
expire = datetime.now(timezone.utc) + timedelta(
|
||||
# Keyword argument: minutes
|
||||
minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES,
|
||||
)
|
||||
# Call to_encode.update()
|
||||
to_encode.update({
|
||||
# Literal argument value
|
||||
"exp": expire,
|
||||
# Literal argument value
|
||||
"jti": str(_uuid.uuid4()),
|
||||
})
|
||||
# Return jwt.encode(to_encode, settings.SECRET_KEY, algorithm=settings.ALGOR...
|
||||
to_encode.update({"exp": expire})
|
||||
return jwt.encode(to_encode, settings.SECRET_KEY, algorithm=settings.ALGORITHM)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Token blacklist (Redis-backed)
|
||||
# ---------------------------------------------------------------------------
|
||||
# Each revoked token's ``jti`` is stored in Redis with a TTL equal to the
|
||||
# token's remaining lifetime. This means entries auto-expire exactly when
|
||||
# the token would have become invalid anyway — no manual cleanup needed.
|
||||
#
|
||||
# Redis survives backend restarts, so blacklisted tokens stay revoked
|
||||
# across deploys and multi-worker setups.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_BLACKLIST_PREFIX = "blacklist:"
|
||||
|
||||
|
||||
# Define function blacklist_token
|
||||
def blacklist_token(jti: str, exp: float) -> None:
|
||||
"""Add *jti* to the Redis blacklist with a TTL derived from *exp*.
|
||||
|
||||
*exp* is the token's ``exp`` claim (epoch timestamp). The TTL is set
|
||||
to ``exp - now`` so the key vanishes when the token would have expired
|
||||
naturally.
|
||||
"""
|
||||
# Import get_redis_blacklist from app.infrastructure.redis_client
|
||||
from app.infrastructure.redis_client import get_redis_blacklist
|
||||
|
||||
# Assign ttl = max(int(exp - datetime.now(timezone.utc).timestamp()), 1)
|
||||
ttl = max(int(exp - datetime.now(timezone.utc).timestamp()), 1)
|
||||
# Attempt the following; catch errors below
|
||||
try:
|
||||
# Assign r = get_redis_blacklist()
|
||||
r = get_redis_blacklist()
|
||||
# Call r.setex()
|
||||
r.setex(f"{_BLACKLIST_PREFIX}{jti}", ttl, "1")
|
||||
# Handle Exception
|
||||
except Exception:
|
||||
# Log warning: "Failed to blacklist token %s in Redis", jti, exc_
|
||||
logger.warning("Failed to blacklist token %s in Redis", jti, exc_info=True)
|
||||
|
||||
|
||||
# Define function is_token_blacklisted
|
||||
def is_token_blacklisted(jti: str) -> bool:
|
||||
"""Return ``True`` if *jti* has been revoked (exists in Redis)."""
|
||||
# Import get_redis_blacklist from app.infrastructure.redis_client
|
||||
from app.infrastructure.redis_client import get_redis_blacklist
|
||||
|
||||
# Attempt the following; catch errors below
|
||||
try:
|
||||
# Assign r = get_redis_blacklist()
|
||||
r = get_redis_blacklist()
|
||||
# Return r.exists(f"{_BLACKLIST_PREFIX}{jti}") > 0
|
||||
return r.exists(f"{_BLACKLIST_PREFIX}{jti}") > 0
|
||||
# Handle Exception
|
||||
except Exception:
|
||||
# Log warning: "Failed to check blacklist for %s in Redis", jti,
|
||||
logger.warning("Failed to check blacklist for %s in Redis", jti, exc_info=True)
|
||||
# Return False
|
||||
return False
|
||||
|
||||
+2
-197
@@ -1,213 +1,18 @@
|
||||
"""Application configuration for the Aegis MITRE ATT&CK Coverage Platform.
|
||||
|
||||
Loads settings from environment variables and ``.env`` files via
|
||||
``pydantic-settings``. Validates critical secrets at import time and raises
|
||||
``RuntimeError`` (production) or issues a ``UserWarning`` (development) when
|
||||
unsafe defaults are detected.
|
||||
"""
|
||||
|
||||
# Import os
|
||||
import os
|
||||
|
||||
# Import secrets
|
||||
import secrets
|
||||
|
||||
# Import warnings
|
||||
import warnings
|
||||
|
||||
# Import BaseSettings from pydantic_settings
|
||||
from pydantic_settings import BaseSettings
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Detect environment: "production" when AEGIS_ENV or common indicators are set
|
||||
# ---------------------------------------------------------------------------
|
||||
_is_production = os.environ.get("AEGIS_ENV", "").lower() == "production"
|
||||
|
||||
|
||||
# Define class Settings
|
||||
class Settings(BaseSettings):
|
||||
"""Application settings loaded from environment variables and .env file."""
|
||||
|
||||
# Assign DATABASE_URL = "postgresql://postgres:postgres@postgres:5432/attackdb"
|
||||
DATABASE_URL: str = "postgresql://postgres:postgres@postgres:5432/attackdb"
|
||||
|
||||
# ── Security ──────────────────────────────────────────────────────
|
||||
# SECRET_KEY has NO safe default. In development a random key is
|
||||
# generated at startup (tokens invalidate on restart — acceptable
|
||||
# for local dev). In production it MUST be supplied via env/.env
|
||||
# so tokens survive restarts.
|
||||
SECRET_KEY: str = ""
|
||||
# Assign ALGORITHM = "HS256"
|
||||
SECRET_KEY: str = "change-me-in-production"
|
||||
ALGORITHM: str = "HS256"
|
||||
ACCESS_TOKEN_EXPIRE_MINUTES: int = 480 # 8 hours — /auth/refresh extends active sessions
|
||||
|
||||
# ── Redis ─────────────────────────────────────────────────────────
|
||||
REDIS_URL: str = "redis://redis:6379/0"
|
||||
# Logical DB indices on the same Redis instance (PATH in URL is overridden).
|
||||
REDIS_TOKEN_BLACKLIST_DB: int = 1
|
||||
# Assign REDIS_CACHE_DB = 2
|
||||
REDIS_CACHE_DB: int = 2
|
||||
|
||||
# ── CORS ─────────────────────────────────────────────────────────
|
||||
# Comma-separated list of allowed origins, or a JSON array.
|
||||
# In dev this defaults to common local ports; in production set it
|
||||
# to the actual frontend domain(s).
|
||||
CORS_ORIGINS: str = "http://localhost:3000,http://localhost:5173"
|
||||
|
||||
# ── MinIO / S3 ───────────────────────────────────────────────────
|
||||
ACCESS_TOKEN_EXPIRE_MINUTES: int = 60
|
||||
MINIO_ENDPOINT: str = "minio:9000"
|
||||
# Public hostname used in presigned URLs returned to browsers.
|
||||
# In production set this to <server-ip>:9000 (or a public FQDN) so
|
||||
# the browser can reach MinIO directly. Defaults to MINIO_ENDPOINT.
|
||||
MINIO_PUBLIC_ENDPOINT: str = ""
|
||||
MINIO_ACCESS_KEY: str = "minioadmin"
|
||||
# Assign MINIO_SECRET_KEY = "minioadmin"
|
||||
MINIO_SECRET_KEY: str = "minioadmin"
|
||||
# Assign MINIO_BUCKET = "evidence"
|
||||
MINIO_BUCKET: str = "evidence"
|
||||
# Assign MINIO_SECURE = False # True → use HTTPS to connect to MinIO
|
||||
MINIO_SECURE: bool = False # True → use HTTPS to connect to MinIO
|
||||
|
||||
# ── Re-testing ───────────────────────────────────────────────────
|
||||
MAX_RETEST_COUNT: int = 3 # maximum automatic retests per original test
|
||||
|
||||
# ── Jira Integration ────────────────────────────────────────────
|
||||
JIRA_ENABLED: bool = False
|
||||
# Assign JIRA_URL = ""
|
||||
JIRA_URL: str = ""
|
||||
# Assign JIRA_USERNAME = ""
|
||||
JIRA_USERNAME: str = ""
|
||||
# Assign JIRA_API_TOKEN = ""
|
||||
JIRA_API_TOKEN: str = ""
|
||||
# Assign JIRA_IS_CLOUD = True
|
||||
JIRA_IS_CLOUD: bool = True
|
||||
# Assign JIRA_DEFAULT_PROJECT = ""
|
||||
JIRA_DEFAULT_PROJECT: str = ""
|
||||
JIRA_ISSUE_TYPE_TEST: str = "Task" # tests (campaign or standalone)
|
||||
JIRA_ISSUE_TYPE_CAMPAIGN: str = "Epic" # campaigns (under Initiative)
|
||||
# Jira custom field ID for "Start date" — Jira Cloud team-managed: customfield_10015
|
||||
# Override with the correct field ID for your Jira instance if different.
|
||||
JIRA_START_DATE_FIELD: str = "customfield_10015"
|
||||
|
||||
# ── Tempo Integration ─────────────────────────────────────────────
|
||||
TEMPO_ENABLED: bool = False
|
||||
# Assign TEMPO_API_TOKEN = ""
|
||||
TEMPO_API_TOKEN: str = ""
|
||||
# Assign TEMPO_API_VERSION = 4
|
||||
TEMPO_API_VERSION: int = 4
|
||||
# Assign TEMPO_DEFAULT_WORK_TYPE = "Red Team"
|
||||
TEMPO_DEFAULT_WORK_TYPE: str = "Red Team"
|
||||
# Tempo API base URL — use https://api.eu.tempo.io/4 for EU workspaces.
|
||||
# Can also be set via system_configs key "tempo.base_url" at runtime.
|
||||
TEMPO_BASE_URL: str = "" # empty → falls back to https://api.tempo.io/4
|
||||
|
||||
# ── OSINT / Intelligence ────────────────────────────────────────
|
||||
NVD_API_KEY: str = "" # optional; increases NVD rate limit from 5/30s to 50/30s
|
||||
# Assign STALE_THRESHOLD_DAYS = 365 # days before coverage is considered stale
|
||||
STALE_THRESHOLD_DAYS: int = 365 # days before coverage is considered stale
|
||||
|
||||
# ── Reporting ─────────────────────────────────────────────────────
|
||||
REPORT_TEMPLATES_DIR: str = "app/templates/reports"
|
||||
# Assign REPORT_OUTPUT_DIR = "/tmp/aegis_reports"
|
||||
REPORT_OUTPUT_DIR: str = "/app/reports"
|
||||
# Assign COMPANY_NAME = "Organization"
|
||||
COMPANY_NAME: str = "Organization"
|
||||
# Assign COMPANY_LOGO_PATH = "app/templates/reports/assets/logo.png"
|
||||
COMPANY_LOGO_PATH: str = "app/templates/reports/assets/logo.png"
|
||||
|
||||
# ── Email / SMTP ──────────────────────────────────────────────────
|
||||
SMTP_ENABLED: bool = False
|
||||
SMTP_HOST: str = ""
|
||||
SMTP_PORT: int = 587
|
||||
SMTP_USERNAME: str = ""
|
||||
SMTP_PASSWORD: str = ""
|
||||
SMTP_FROM_EMAIL: str = "aegis@company.com"
|
||||
SMTP_USE_TLS: bool = True
|
||||
PLATFORM_URL: str = "http://localhost:5173" # base URL for links in emails
|
||||
|
||||
# ── Scoring weights (must sum to 100) ────────────────────────────
|
||||
SCORING_WEIGHT_TESTS: int = 40
|
||||
# Assign SCORING_WEIGHT_DETECTION_RULES = 25
|
||||
SCORING_WEIGHT_DETECTION_RULES: int = 25
|
||||
# Assign SCORING_WEIGHT_D3FEND = 15
|
||||
SCORING_WEIGHT_D3FEND: int = 15
|
||||
# Assign SCORING_WEIGHT_RECENCY = 10
|
||||
SCORING_WEIGHT_RECENCY: int = 10
|
||||
# Assign SCORING_WEIGHT_SEVERITY = 10
|
||||
SCORING_WEIGHT_SEVERITY: int = 10
|
||||
# Legacy env names (mapped in scoring_config_service)
|
||||
SCORING_WEIGHT_FRESHNESS: int = 10
|
||||
# Assign SCORING_WEIGHT_PLATFORM_DIVERSITY = 10
|
||||
SCORING_WEIGHT_PLATFORM_DIVERSITY: int = 10
|
||||
|
||||
# Define class Config
|
||||
class Config:
|
||||
"""Pydantic BaseSettings configuration — load from .env file."""
|
||||
|
||||
# Assign env_file = ".env"
|
||||
env_file = ".env"
|
||||
|
||||
|
||||
# Assign settings = Settings()
|
||||
settings = Settings()
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Post-init validation for SECRET_KEY
|
||||
# ---------------------------------------------------------------------------
|
||||
_UNSAFE_SECRETS = {
|
||||
# Literal argument value
|
||||
"",
|
||||
# Literal argument value
|
||||
"change-me-in-production",
|
||||
# Literal argument value
|
||||
"change-me-in-production-use-a-long-random-string",
|
||||
}
|
||||
|
||||
# Check: settings.SECRET_KEY in _UNSAFE_SECRETS
|
||||
if settings.SECRET_KEY in _UNSAFE_SECRETS:
|
||||
# Check: _is_production
|
||||
if _is_production:
|
||||
# Raise RuntimeError
|
||||
raise RuntimeError(
|
||||
# Literal argument value
|
||||
"CRITICAL: SECRET_KEY is not configured. "
|
||||
# Literal argument value
|
||||
"Set a strong random value (>= 32 chars) via the SECRET_KEY "
|
||||
# Literal argument value
|
||||
"environment variable or in your .env file before running in "
|
||||
# Literal argument value
|
||||
"production. Example: openssl rand -hex 32"
|
||||
)
|
||||
# Development: auto-generate an ephemeral key and warn
|
||||
settings.SECRET_KEY = secrets.token_hex(32)
|
||||
# Call warnings.warn()
|
||||
warnings.warn(
|
||||
# Literal argument value
|
||||
"SECRET_KEY was not set — using an auto-generated ephemeral key. "
|
||||
# Literal argument value
|
||||
"JWT tokens will be invalidated on every restart. "
|
||||
# Literal argument value
|
||||
"Set SECRET_KEY in your environment for persistent sessions.",
|
||||
# Keyword argument: stacklevel
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# SEC-002: Reject default credentials in production
|
||||
# ---------------------------------------------------------------------------
|
||||
if _is_production:
|
||||
# Assign _DEFAULT_CREDS = {
|
||||
_DEFAULT_CREDS = {
|
||||
("MINIO_ACCESS_KEY", settings.MINIO_ACCESS_KEY, "minioadmin"),
|
||||
("MINIO_SECRET_KEY", settings.MINIO_SECRET_KEY, "minioadmin"),
|
||||
}
|
||||
# Iterate over _DEFAULT_CREDS
|
||||
for name, current, default in _DEFAULT_CREDS:
|
||||
# Check: current == default
|
||||
if current == default:
|
||||
# Raise RuntimeError
|
||||
raise RuntimeError(
|
||||
f"CRITICAL: {name} is using the default value '{default}'. "
|
||||
f"Set a strong value via the {name} environment variable "
|
||||
f"before running in production."
|
||||
)
|
||||
|
||||
+5
-153
@@ -1,164 +1,16 @@
|
||||
"""Database engine and session management for the Aegis platform.
|
||||
|
||||
The engine and session factory are created lazily so that tests can override
|
||||
``DATABASE_URL`` via environment variables before any import triggers real
|
||||
PostgreSQL engine creation (which requires psycopg2).
|
||||
"""
|
||||
|
||||
# Import Generator from collections.abc
|
||||
from collections.abc import Generator
|
||||
|
||||
# Import create_engine from sqlalchemy
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.orm import sessionmaker, declarative_base
|
||||
|
||||
# Import Engine from sqlalchemy.engine
|
||||
from sqlalchemy.engine import Engine
|
||||
from app.config import settings
|
||||
|
||||
# Import Session, declarative_base, sessionmaker from sqlalchemy.orm
|
||||
from sqlalchemy.orm import Session, declarative_base, sessionmaker
|
||||
|
||||
# Assign Base = declarative_base()
|
||||
engine = create_engine(settings.DATABASE_URL)
|
||||
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
|
||||
Base = declarative_base()
|
||||
|
||||
# Engine and session factory are created lazily so that tests can
|
||||
# override DATABASE_URL via environment *before* any import triggers
|
||||
# the real PostgreSQL engine creation (which requires psycopg2).
|
||||
_engine = None
|
||||
# Assign _SessionLocal = None
|
||||
_SessionLocal = None
|
||||
|
||||
|
||||
# Define function _get_engine
|
||||
def _get_engine() -> Engine:
|
||||
"""Return the shared SQLAlchemy engine, creating it on first call.
|
||||
|
||||
Returns:
|
||||
Engine: Configured SQLAlchemy engine for the application database.
|
||||
"""
|
||||
# Declare global variable
|
||||
global _engine
|
||||
# Check: _engine is None
|
||||
if _engine is None:
|
||||
# Import settings from app.config
|
||||
from app.config import settings
|
||||
|
||||
# Assign url = settings.DATABASE_URL
|
||||
url = settings.DATABASE_URL
|
||||
# Assign kwargs = {}
|
||||
kwargs: dict = {}
|
||||
# Check: url.startswith("postgresql")
|
||||
if url.startswith("postgresql"):
|
||||
# Call kwargs.update()
|
||||
kwargs.update(
|
||||
# Keyword argument: pool_size
|
||||
pool_size=20,
|
||||
# Keyword argument: max_overflow
|
||||
max_overflow=10,
|
||||
# Keyword argument: pool_recycle
|
||||
pool_recycle=3600,
|
||||
# Keyword argument: pool_pre_ping
|
||||
pool_pre_ping=True,
|
||||
)
|
||||
# Assign _engine = create_engine(url, **kwargs)
|
||||
_engine = create_engine(url, **kwargs)
|
||||
# Return _engine
|
||||
return _engine
|
||||
|
||||
|
||||
# Define function _get_session_factory
|
||||
def _get_session_factory() -> sessionmaker:
|
||||
"""Return the shared sessionmaker, creating it on first call.
|
||||
|
||||
Returns:
|
||||
sessionmaker: Configured sessionmaker bound to the application engine.
|
||||
"""
|
||||
# Declare global variable
|
||||
global _SessionLocal
|
||||
# Check: _SessionLocal is None
|
||||
if _SessionLocal is None:
|
||||
# Assign _SessionLocal = sessionmaker(
|
||||
_SessionLocal = sessionmaker(
|
||||
# Keyword argument: autocommit
|
||||
autocommit=False, autoflush=False, bind=_get_engine()
|
||||
)
|
||||
# Return _SessionLocal
|
||||
return _SessionLocal
|
||||
|
||||
|
||||
# Define class _LazySessionLocal
|
||||
class _LazySessionLocal:
|
||||
"""Proxy so ``SessionLocal()`` keeps working as before but the real sessionmaker is only created on first call."""
|
||||
|
||||
# Define function __call__
|
||||
def __call__(self, *args: object, **kwargs: object) -> Session:
|
||||
"""Create and return a new database session.
|
||||
|
||||
Args:
|
||||
*args (object): Positional arguments forwarded to the sessionmaker.
|
||||
**kwargs (object): Keyword arguments forwarded to the sessionmaker.
|
||||
|
||||
Returns:
|
||||
Session: A new SQLAlchemy database session.
|
||||
"""
|
||||
# Return _get_session_factory()(*args, **kwargs)
|
||||
return _get_session_factory()(*args, **kwargs)
|
||||
|
||||
# Define function __getattr__
|
||||
def __getattr__(self, name: str) -> object:
|
||||
"""Delegate attribute access to the underlying sessionmaker.
|
||||
|
||||
Args:
|
||||
name (str): Attribute name to look up on the sessionmaker.
|
||||
|
||||
Returns:
|
||||
object: The attribute value from the underlying sessionmaker.
|
||||
"""
|
||||
# Return getattr(_get_session_factory(), name)
|
||||
return getattr(_get_session_factory(), name)
|
||||
|
||||
|
||||
# Assign SessionLocal = _LazySessionLocal()
|
||||
SessionLocal = _LazySessionLocal()
|
||||
|
||||
|
||||
# Define class _EngineProxy
|
||||
class _EngineProxy:
|
||||
"""Thin proxy so ``from app.database import engine`` still works."""
|
||||
|
||||
# Define function __getattr__
|
||||
def __getattr__(self, name: str) -> object:
|
||||
"""Delegate attribute access to the lazily-created engine.
|
||||
|
||||
Args:
|
||||
name (str): Attribute name to look up on the real engine.
|
||||
|
||||
Returns:
|
||||
object: The attribute value from the underlying SQLAlchemy engine.
|
||||
"""
|
||||
# Return getattr(_get_engine(), name)
|
||||
return getattr(_get_engine(), name)
|
||||
|
||||
|
||||
# Assign engine = _EngineProxy() # type: ignore[assignment]
|
||||
engine = _EngineProxy() # type: ignore[assignment]
|
||||
|
||||
|
||||
# Define function get_db
|
||||
def get_db() -> Generator[Session, None, None]:
|
||||
"""Yield a database session and close it when the request is done.
|
||||
|
||||
Intended for use as a FastAPI dependency.
|
||||
|
||||
Yields:
|
||||
Session: An active SQLAlchemy session for the current request.
|
||||
"""
|
||||
# Assign db = SessionLocal()
|
||||
def get_db():
|
||||
db = SessionLocal()
|
||||
# Attempt the following; catch errors below
|
||||
try:
|
||||
# Yield db
|
||||
yield db
|
||||
# Always execute this cleanup block
|
||||
finally:
|
||||
# Close the database session
|
||||
db.close()
|
||||
|
||||
@@ -1 +0,0 @@
|
||||
"""FastAPI dependency injection helpers for auth, DB, and shared state."""
|
||||
|
||||
@@ -1,52 +1,26 @@
|
||||
"""Authentication and RBAC dependencies for FastAPI.
|
||||
"""
|
||||
Authentication and RBAC dependencies for FastAPI.
|
||||
|
||||
Provides:
|
||||
- ``get_current_user``: decodes JWT from HttpOnly cookie (preferred) or
|
||||
Authorization header (fallback), fetches user from DB, raises 401 on failure.
|
||||
Also accepts Aegis API keys (``aegis_…`` prefix) as Bearer tokens.
|
||||
- ``get_current_user``: decodes JWT, fetches user from DB, raises 401 on failure.
|
||||
- ``require_role``: factory that returns a dependency enforcing a specific role
|
||||
(admins always pass).
|
||||
"""
|
||||
|
||||
# Import Callable from collections.abc
|
||||
from collections.abc import Callable
|
||||
|
||||
# Import Optional from typing
|
||||
from typing import Optional
|
||||
|
||||
# Import Cookie, Depends, HTTPException, status from fastapi
|
||||
from fastapi import Cookie, Depends, HTTPException, status
|
||||
|
||||
# Import OAuth2PasswordBearer from fastapi.security
|
||||
from fastapi import Depends, HTTPException, status
|
||||
from fastapi.security import OAuth2PasswordBearer
|
||||
|
||||
# Import jwt (PyJWT)
|
||||
import jwt
|
||||
|
||||
# Import Session from sqlalchemy.orm
|
||||
from jose import JWTError, jwt
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
# Import auth as auth_lib from app
|
||||
from app import auth as auth_lib
|
||||
|
||||
# Import settings from app.config
|
||||
from app.config import settings
|
||||
|
||||
# Import get_db from app.database
|
||||
from app.database import get_db
|
||||
|
||||
# Import User from app.models.user
|
||||
from app.models.user import User
|
||||
from app.models.api_key import KEY_PREFIX
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# OAuth2 scheme (reads Authorization header — used as fallback / Swagger UI)
|
||||
# OAuth2 scheme
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/api/v1/auth/login", auto_error=False)
|
||||
|
||||
# Cookie name — must match the one set in the auth router
|
||||
_COOKIE_NAME = "aegis_token"
|
||||
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/api/v1/auth/login")
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Current-user dependency
|
||||
@@ -54,94 +28,38 @@ _COOKIE_NAME = "aegis_token"
|
||||
|
||||
|
||||
async def get_current_user(
|
||||
# Entry: aegis_token
|
||||
aegis_token: Optional[str] = Cookie(None),
|
||||
# Entry: bearer_token
|
||||
bearer_token: Optional[str] = Depends(oauth2_scheme),
|
||||
# Entry: db
|
||||
token: str = Depends(oauth2_scheme),
|
||||
db: Session = Depends(get_db),
|
||||
) -> User:
|
||||
"""Decode the JWT, look up the user in *db*, and return it.
|
||||
|
||||
Token resolution order:
|
||||
1. ``aegis_token`` **HttpOnly cookie** (preferred — immune to XSS).
|
||||
2. ``Authorization: Bearer <token>`` header (fallback for API clients
|
||||
and Swagger UI).
|
||||
"""Decode the JWT *token*, look up the user in *db*, and return it.
|
||||
|
||||
Raises :class:`~fastapi.HTTPException` **401** when:
|
||||
- no token is found in either location,
|
||||
- the token cannot be decoded,
|
||||
- the ``sub`` claim is missing, or
|
||||
- no matching active user exists in the database.
|
||||
"""
|
||||
# Assign credentials_exception = HTTPException(
|
||||
credentials_exception = HTTPException(
|
||||
# Keyword argument: status_code
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
# Keyword argument: detail
|
||||
detail="Could not validate credentials",
|
||||
# Keyword argument: headers
|
||||
headers={"WWW-Authenticate": "Bearer"},
|
||||
)
|
||||
# Assign revoked_exception = HTTPException(
|
||||
revoked_exception = HTTPException(
|
||||
# Keyword argument: status_code
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
# Keyword argument: detail
|
||||
detail="Token has been revoked",
|
||||
# Keyword argument: headers
|
||||
headers={"WWW-Authenticate": "Bearer"},
|
||||
)
|
||||
|
||||
# Prefer cookie, fall back to header
|
||||
token = aegis_token or bearer_token
|
||||
# Check: token is None
|
||||
if token is None:
|
||||
# Raise credentials_exception
|
||||
raise credentials_exception
|
||||
|
||||
# ── API Key path (Bearer token starts with "aegis_") ──────────────────
|
||||
if token.startswith(KEY_PREFIX):
|
||||
from app.services.api_key_service import authenticate_raw_key
|
||||
user = authenticate_raw_key(db, token)
|
||||
if user is None:
|
||||
raise credentials_exception
|
||||
return user
|
||||
|
||||
# ── JWT path ──────────────────────────────────────────────────────────
|
||||
try:
|
||||
# Assign payload = jwt.decode(
|
||||
payload = jwt.decode(
|
||||
token,
|
||||
settings.SECRET_KEY,
|
||||
# Keyword argument: algorithms
|
||||
algorithms=[settings.ALGORITHM],
|
||||
)
|
||||
# Assign username = payload.get("sub")
|
||||
username: str | None = payload.get("sub")
|
||||
# Check: username is None
|
||||
if username is None:
|
||||
# Raise credentials_exception
|
||||
raise credentials_exception
|
||||
# Check token blacklist (revoked tokens)
|
||||
jti: str | None = payload.get("jti")
|
||||
# Check: jti and auth_lib.is_token_blacklisted(jti)
|
||||
if jti and auth_lib.is_token_blacklisted(jti):
|
||||
# Raise revoked_exception
|
||||
raise revoked_exception
|
||||
# Handle any JWT validation error (expired, invalid signature, malformed)
|
||||
except jwt.exceptions.InvalidTokenError:
|
||||
# Raise credentials_exception
|
||||
except JWTError:
|
||||
raise credentials_exception
|
||||
|
||||
# Assign user = db.query(User).filter(User.username == username).first()
|
||||
user = db.query(User).filter(User.username == username).first()
|
||||
# Check: user is None or not user.is_active
|
||||
if user is None or not user.is_active:
|
||||
# Raise credentials_exception
|
||||
if user is None:
|
||||
raise credentials_exception
|
||||
|
||||
# Return user
|
||||
return user
|
||||
|
||||
|
||||
@@ -150,123 +68,43 @@ async def get_current_user(
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def require_password_changed(
|
||||
# Entry: current_user
|
||||
current_user: User = Depends(get_current_user),
|
||||
) -> User:
|
||||
"""Block all requests when the user still needs to change their password.
|
||||
|
||||
Only ``/auth/change-password`` and ``/auth/me`` are exempt — those
|
||||
endpoints do **not** depend on this function.
|
||||
"""
|
||||
# Check: getattr(current_user, "must_change_password", False)
|
||||
if getattr(current_user, "must_change_password", False):
|
||||
# Raise HTTPException
|
||||
raise HTTPException(
|
||||
# Keyword argument: status_code
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
# Keyword argument: detail
|
||||
detail="PASSWORD_CHANGE_REQUIRED",
|
||||
)
|
||||
# Return current_user
|
||||
return current_user
|
||||
|
||||
|
||||
def _check_api_key_scope(user: User, required_scope: str) -> None:
|
||||
"""Raise 403 if the request was authenticated via an API key that lacks *required_scope*.
|
||||
|
||||
When authenticated via JWT (browser session), ``_api_key_scopes`` is not set
|
||||
and the check is skipped — full access is granted based on role alone.
|
||||
"""
|
||||
key_scopes = getattr(user, "_api_key_scopes", None)
|
||||
if key_scopes is not None and required_scope not in key_scopes:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail=f"API key scope '{required_scope}' required for this operation",
|
||||
)
|
||||
|
||||
|
||||
def require_role(required_role: str):
|
||||
"""Return a FastAPI dependency that enforces *required_role*.
|
||||
|
||||
The dependency allows the request to proceed when
|
||||
``user.role == required_role`` **or** ``user.role == "admin"``.
|
||||
Also enforces API key scopes: admin-role endpoints require the ``admin``
|
||||
scope; all other role-restricted endpoints require ``write``.
|
||||
Otherwise it raises :class:`~fastapi.HTTPException` **403**.
|
||||
"""
|
||||
|
||||
# Define async function role_checker
|
||||
async def role_checker(
|
||||
# Entry: current_user
|
||||
current_user: User = Depends(get_current_user),
|
||||
) -> User:
|
||||
# Check: current_user.role != required_role and current_user.role != "admin"
|
||||
if current_user.role != required_role and current_user.role != "admin":
|
||||
# Raise HTTPException
|
||||
raise HTTPException(
|
||||
# Keyword argument: status_code
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
# Keyword argument: detail
|
||||
detail="Not enough permissions",
|
||||
)
|
||||
scope = "admin" if required_role == "admin" else "write"
|
||||
_check_api_key_scope(current_user, scope)
|
||||
return current_user
|
||||
|
||||
# Return role_checker
|
||||
return role_checker
|
||||
|
||||
|
||||
# Define function require_any_role
|
||||
def require_any_role(*roles: str) -> Callable[..., object]:
|
||||
def require_any_role(*roles: str):
|
||||
"""Return a FastAPI dependency that enforces **any** of the given *roles*.
|
||||
|
||||
Admins always pass. Also enforces API key scopes: if the only accepted
|
||||
role is ``admin``, the key must carry the ``admin`` scope; otherwise the
|
||||
``write`` scope is required.
|
||||
|
||||
Usage example::
|
||||
Admins always pass. Usage example::
|
||||
|
||||
@router.patch("/resource", dependencies=[Depends(require_any_role("red_lead", "blue_lead"))])
|
||||
"""
|
||||
|
||||
# Define async function role_checker
|
||||
async def role_checker(
|
||||
# Entry: current_user
|
||||
current_user: User = Depends(get_current_user),
|
||||
) -> User:
|
||||
# Check: current_user.role != "admin" and current_user.role not in roles
|
||||
if current_user.role != "admin" and current_user.role not in roles:
|
||||
# Raise HTTPException
|
||||
raise HTTPException(
|
||||
# Keyword argument: status_code
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
# Keyword argument: detail
|
||||
detail="Not enough permissions",
|
||||
)
|
||||
scope = "admin" if set(roles) == {"admin"} else "write"
|
||||
_check_api_key_scope(current_user, scope)
|
||||
return current_user
|
||||
|
||||
# Return role_checker
|
||||
return role_checker
|
||||
|
||||
|
||||
def require_scope(scope: str):
|
||||
"""Return a dependency that enforces the API key carries *scope*.
|
||||
|
||||
JWT-authenticated requests (browser sessions) bypass this check entirely.
|
||||
Use on mutation endpoints that don't already use ``require_role`` /
|
||||
``require_any_role``::
|
||||
|
||||
@router.post("/resource", dependencies=[Depends(require_scope("write"))])
|
||||
"""
|
||||
|
||||
async def scope_checker(
|
||||
current_user: User = Depends(get_current_user),
|
||||
) -> User:
|
||||
_check_api_key_scope(current_user, scope)
|
||||
return current_user
|
||||
|
||||
return scope_checker
|
||||
|
||||
@@ -1,44 +0,0 @@
|
||||
"""FastAPI dependency providers for repositories.
|
||||
|
||||
Wiring lives ONLY in the presentation layer — use cases and services
|
||||
never know which concrete repository implementation they receive.
|
||||
"""
|
||||
|
||||
# Import Depends from fastapi
|
||||
from fastapi import Depends
|
||||
|
||||
# Import Session from sqlalchemy.orm
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
# Import get_db from app.database
|
||||
from app.database import get_db
|
||||
|
||||
# Import from app.infrastructure.persistence.repositories.sa_technique_repository
|
||||
from app.infrastructure.persistence.repositories.sa_technique_repository import (
|
||||
SATechniqueRepository,
|
||||
)
|
||||
|
||||
# Import from app.infrastructure.persistence.repositories.sa_test_repository
|
||||
from app.infrastructure.persistence.repositories.sa_test_repository import (
|
||||
SATestRepository,
|
||||
)
|
||||
|
||||
|
||||
# Define function get_technique_repository
|
||||
def get_technique_repository(
|
||||
# Entry: db
|
||||
db: Session = Depends(get_db),
|
||||
) -> SATechniqueRepository:
|
||||
"""Provide a TechniqueRepository backed by the current DB session."""
|
||||
# Return SATechniqueRepository(db)
|
||||
return SATechniqueRepository(db)
|
||||
|
||||
|
||||
# Define function get_test_repository
|
||||
def get_test_repository(
|
||||
# Entry: db
|
||||
db: Session = Depends(get_db),
|
||||
) -> SATestRepository:
|
||||
"""Provide a TestRepository backed by the current DB session."""
|
||||
# Return SATestRepository(db)
|
||||
return SATestRepository(db)
|
||||
@@ -1 +0,0 @@
|
||||
"""Domain layer — entities, value objects, errors, and repository ports."""
|
||||
@@ -1,34 +0,0 @@
|
||||
"""Domain entity classes representing core business objects."""
|
||||
# Import CampaignEntity from app.domain.entities.campaign
|
||||
from app.domain.entities.campaign import CampaignEntity
|
||||
|
||||
# Import from app.domain.entities.compliance
|
||||
from app.domain.entities.compliance import (
|
||||
ComplianceControlEntity,
|
||||
ComplianceFrameworkEntity,
|
||||
ControlCoverageStatus,
|
||||
)
|
||||
|
||||
# Import TechniqueEntity from app.domain.entities.technique
|
||||
from app.domain.entities.technique import TechniqueEntity
|
||||
|
||||
# Import ThreatActorEntity, ThreatActorTechniqueRef from app.domain.entities.threat_actor
|
||||
from app.domain.entities.threat_actor import ThreatActorEntity, ThreatActorTechniqueRef
|
||||
|
||||
# Assign __all__ = [
|
||||
__all__ = [
|
||||
# Literal argument value
|
||||
"CampaignEntity",
|
||||
# Literal argument value
|
||||
"ComplianceControlEntity",
|
||||
# Literal argument value
|
||||
"ComplianceFrameworkEntity",
|
||||
# Literal argument value
|
||||
"ControlCoverageStatus",
|
||||
# Literal argument value
|
||||
"TechniqueEntity",
|
||||
# Literal argument value
|
||||
"ThreatActorEntity",
|
||||
# Literal argument value
|
||||
"ThreatActorTechniqueRef",
|
||||
]
|
||||
@@ -1,219 +0,0 @@
|
||||
"""Campaign domain entity with lifecycle validation.
|
||||
|
||||
Pure domain logic — no framework imports.
|
||||
"""
|
||||
|
||||
# Enable future language features for compatibility
|
||||
from __future__ import annotations
|
||||
|
||||
# Import enum
|
||||
import enum
|
||||
|
||||
# Import uuid
|
||||
import uuid
|
||||
|
||||
# Import dataclass, field from dataclasses
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
# Import TYPE_CHECKING from typing
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
# Import BusinessRuleViolation, InvalidStateTransition from app.domain.errors
|
||||
from app.domain.errors import BusinessRuleViolation, InvalidStateTransition
|
||||
|
||||
# Check: TYPE_CHECKING
|
||||
if TYPE_CHECKING:
|
||||
# Import Campaign as CampaignORM from app.models.campaign
|
||||
from app.models.campaign import Campaign as CampaignORM
|
||||
|
||||
|
||||
# Define class CampaignStatus
|
||||
class CampaignStatus(str, enum.Enum):
|
||||
"""Lifecycle states for a campaign."""
|
||||
|
||||
# Assign draft = "draft"
|
||||
draft = "draft"
|
||||
# Assign active = "active"
|
||||
active = "active"
|
||||
# Assign completed = "completed"
|
||||
completed = "completed"
|
||||
# Assign archived = "archived"
|
||||
archived = "archived"
|
||||
|
||||
|
||||
# Define class CampaignType
|
||||
class CampaignType(str, enum.Enum):
|
||||
"""Classification of the campaign's testing methodology."""
|
||||
|
||||
# Assign custom = "custom"
|
||||
custom = "custom"
|
||||
# Assign apt_emulation = "apt_emulation"
|
||||
apt_emulation = "apt_emulation"
|
||||
# Assign kill_chain = "kill_chain"
|
||||
kill_chain = "kill_chain"
|
||||
# Assign compliance = "compliance"
|
||||
compliance = "compliance"
|
||||
|
||||
|
||||
# Assign VALID_TRANSITIONS = {
|
||||
VALID_TRANSITIONS: dict[CampaignStatus, list[CampaignStatus]] = {
|
||||
CampaignStatus.draft: [CampaignStatus.active],
|
||||
CampaignStatus.active: [CampaignStatus.completed],
|
||||
CampaignStatus.completed: [CampaignStatus.archived],
|
||||
CampaignStatus.archived: [],
|
||||
}
|
||||
|
||||
|
||||
# Apply the @dataclass decorator
|
||||
@dataclass
|
||||
# Define class CampaignEntity
|
||||
class CampaignEntity:
|
||||
"""Pure domain representation of a security testing campaign.
|
||||
|
||||
Owns all lifecycle state-machine logic for campaign activation,
|
||||
completion, and archival.
|
||||
"""
|
||||
|
||||
# name: str
|
||||
name: str
|
||||
# Assign type = CampaignType.custom
|
||||
type: CampaignType = CampaignType.custom
|
||||
# Assign status = CampaignStatus.draft
|
||||
status: CampaignStatus = CampaignStatus.draft
|
||||
# Assign id = None
|
||||
id: uuid.UUID | None = None
|
||||
# Assign description = None
|
||||
description: str | None = None
|
||||
# Assign threat_actor_id = None
|
||||
threat_actor_id: uuid.UUID | None = None
|
||||
# Assign created_by = None
|
||||
created_by: uuid.UUID | None = None
|
||||
# Assign target_platform = None
|
||||
target_platform: str | None = None
|
||||
# Assign tags = field(default_factory=list)
|
||||
tags: list[str] = field(default_factory=list)
|
||||
# Assign test_count = 0
|
||||
test_count: int = 0
|
||||
|
||||
# Define function can_transition_to
|
||||
def can_transition_to(self, target: CampaignStatus) -> bool:
|
||||
"""Check whether transitioning from the current status to *target* is valid.
|
||||
|
||||
Args:
|
||||
target (CampaignStatus): The desired next status.
|
||||
|
||||
Returns:
|
||||
bool: True if the transition is allowed, False otherwise.
|
||||
"""
|
||||
# Return target in VALID_TRANSITIONS.get(self.status, [])
|
||||
return target in VALID_TRANSITIONS.get(self.status, [])
|
||||
|
||||
# Define function activate
|
||||
def activate(self) -> None:
|
||||
"""Transition the campaign from ``draft`` to ``active``.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# Check: not self.can_transition_to(CampaignStatus.active)
|
||||
if not self.can_transition_to(CampaignStatus.active):
|
||||
# Raise InvalidStateTransition
|
||||
raise InvalidStateTransition(
|
||||
self.status.value, CampaignStatus.active.value,
|
||||
[s.value for s in VALID_TRANSITIONS[self.status]],
|
||||
)
|
||||
# Check: self.test_count == 0
|
||||
if self.test_count == 0:
|
||||
# Raise BusinessRuleViolation
|
||||
raise BusinessRuleViolation(
|
||||
# Literal argument value
|
||||
"Campaign must have at least one test to activate"
|
||||
)
|
||||
# Assign self.status = CampaignStatus.active
|
||||
self.status = CampaignStatus.active
|
||||
|
||||
# Define function complete
|
||||
def complete(self) -> None:
|
||||
"""Transition the campaign from ``active`` to ``completed``.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# Check: not self.can_transition_to(CampaignStatus.completed)
|
||||
if not self.can_transition_to(CampaignStatus.completed):
|
||||
# Raise InvalidStateTransition
|
||||
raise InvalidStateTransition(
|
||||
self.status.value, CampaignStatus.completed.value,
|
||||
[s.value for s in VALID_TRANSITIONS[self.status]],
|
||||
)
|
||||
# Assign self.status = CampaignStatus.completed
|
||||
self.status = CampaignStatus.completed
|
||||
|
||||
# Define function archive
|
||||
def archive(self) -> None:
|
||||
"""Transition the campaign from ``completed`` to ``archived``.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# Check: not self.can_transition_to(CampaignStatus.archived)
|
||||
if not self.can_transition_to(CampaignStatus.archived):
|
||||
# Raise InvalidStateTransition
|
||||
raise InvalidStateTransition(
|
||||
self.status.value, CampaignStatus.archived.value,
|
||||
[s.value for s in VALID_TRANSITIONS[self.status]],
|
||||
)
|
||||
# Assign self.status = CampaignStatus.archived
|
||||
self.status = CampaignStatus.archived
|
||||
|
||||
# Define function ensure_modifiable
|
||||
def ensure_modifiable(self) -> None:
|
||||
"""Raise BusinessRuleViolation if the campaign is not in a modifiable state.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# Check: self.status not in (CampaignStatus.draft, CampaignStatus.active)
|
||||
if self.status not in (CampaignStatus.draft, CampaignStatus.active):
|
||||
# Raise BusinessRuleViolation
|
||||
raise BusinessRuleViolation(
|
||||
f"Cannot modify campaign in '{self.status.value}' state"
|
||||
)
|
||||
|
||||
# Apply the @classmethod decorator
|
||||
@classmethod
|
||||
# Define function from_orm
|
||||
def from_orm(cls, orm: CampaignORM) -> CampaignEntity:
|
||||
"""Build a CampaignEntity from a SQLAlchemy Campaign model.
|
||||
|
||||
Args:
|
||||
orm (CampaignORM): The SQLAlchemy Campaign ORM model instance.
|
||||
|
||||
Returns:
|
||||
CampaignEntity: A fully populated domain entity reflecting the ORM state.
|
||||
"""
|
||||
# Assign test_count = len(getattr(orm, "campaign_tests", None) or [])
|
||||
test_count = len(getattr(orm, "campaign_tests", None) or [])
|
||||
# Return cls(
|
||||
return cls(
|
||||
# Keyword argument: id
|
||||
id=orm.id,
|
||||
# Keyword argument: name
|
||||
name=orm.name,
|
||||
# Keyword argument: type
|
||||
type=CampaignType(orm.type) if orm.type else CampaignType.custom,
|
||||
# Keyword argument: status
|
||||
status=CampaignStatus(orm.status) if orm.status else CampaignStatus.draft,
|
||||
# Keyword argument: description
|
||||
description=orm.description,
|
||||
# Keyword argument: threat_actor_id
|
||||
threat_actor_id=orm.threat_actor_id,
|
||||
# Keyword argument: created_by
|
||||
created_by=orm.created_by,
|
||||
# Keyword argument: target_platform
|
||||
target_platform=orm.target_platform,
|
||||
# Keyword argument: tags
|
||||
tags=orm.tags or [],
|
||||
# Keyword argument: test_count
|
||||
test_count=test_count,
|
||||
)
|
||||
@@ -1,164 +0,0 @@
|
||||
"""Compliance domain entities with coverage calculation logic.
|
||||
|
||||
Pure domain logic — no framework imports.
|
||||
"""
|
||||
|
||||
# Enable future language features for compatibility
|
||||
from __future__ import annotations
|
||||
|
||||
# Import enum
|
||||
import enum
|
||||
|
||||
# Import uuid
|
||||
import uuid
|
||||
|
||||
# Import dataclass, field from dataclasses
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
|
||||
# Define class ControlCoverageStatus
|
||||
class ControlCoverageStatus(str, enum.Enum):
|
||||
"""Computed coverage level for a single compliance control."""
|
||||
|
||||
# Assign covered = "covered"
|
||||
covered = "covered"
|
||||
# Assign partially_covered = "partially_covered"
|
||||
partially_covered = "partially_covered"
|
||||
# Assign not_covered = "not_covered"
|
||||
not_covered = "not_covered"
|
||||
|
||||
|
||||
# Apply the @dataclass decorator
|
||||
@dataclass
|
||||
# Define class ComplianceControlEntity
|
||||
class ComplianceControlEntity:
|
||||
"""Pure domain representation of a single compliance framework control.
|
||||
|
||||
Derives its coverage status from the technique statuses associated
|
||||
with it via the ``technique_statuses`` list.
|
||||
"""
|
||||
|
||||
# control_id: str
|
||||
control_id: str
|
||||
# title: str
|
||||
title: str
|
||||
# Assign id = None
|
||||
id: uuid.UUID | None = None
|
||||
# Assign description = None
|
||||
description: str | None = None
|
||||
# Assign category = None
|
||||
category: str | None = None
|
||||
# Assign technique_statuses = field(default_factory=list)
|
||||
technique_statuses: list[str] = field(default_factory=list)
|
||||
|
||||
# Apply the @property decorator
|
||||
@property
|
||||
# Define function coverage_status
|
||||
def coverage_status(self) -> ControlCoverageStatus:
|
||||
"""Compute the coverage status for this control based on linked technique statuses.
|
||||
|
||||
Returns:
|
||||
ControlCoverageStatus: ``covered`` when all techniques are covered,
|
||||
``partially_covered`` when at least one is covered, and
|
||||
``not_covered`` when none are covered or the control has no techniques.
|
||||
"""
|
||||
# Check: not self.technique_statuses
|
||||
if not self.technique_statuses:
|
||||
# Return ControlCoverageStatus.not_covered
|
||||
return ControlCoverageStatus.not_covered
|
||||
# Assign covered_statuses = {"validated", "partial"}
|
||||
covered_statuses = {"validated", "partial"}
|
||||
# Assign covered = [s for s in self.technique_statuses if s in covered_statuses]
|
||||
covered = [s for s in self.technique_statuses if s in covered_statuses]
|
||||
# Check: len(covered) == len(self.technique_statuses)
|
||||
if len(covered) == len(self.technique_statuses):
|
||||
# Return ControlCoverageStatus.covered
|
||||
return ControlCoverageStatus.covered
|
||||
# Alternative: len(covered) > 0
|
||||
elif len(covered) > 0:
|
||||
# Return ControlCoverageStatus.partially_covered
|
||||
return ControlCoverageStatus.partially_covered
|
||||
# Return ControlCoverageStatus.not_covered
|
||||
return ControlCoverageStatus.not_covered
|
||||
|
||||
|
||||
# Apply the @dataclass decorator
|
||||
@dataclass
|
||||
# Define class ComplianceFrameworkEntity
|
||||
class ComplianceFrameworkEntity:
|
||||
"""Pure domain representation of a compliance framework (e.g. NIST 800-53, PCI-DSS).
|
||||
|
||||
Aggregates a collection of controls and provides aggregate coverage statistics.
|
||||
"""
|
||||
|
||||
# name: str
|
||||
name: str
|
||||
# Assign id = None
|
||||
id: uuid.UUID | None = None
|
||||
# Assign version = None
|
||||
version: str | None = None
|
||||
# Assign description = None
|
||||
description: str | None = None
|
||||
# Assign is_active = True
|
||||
is_active: bool = True
|
||||
# Assign controls = field(default_factory=list)
|
||||
controls: list[ComplianceControlEntity] = field(default_factory=list)
|
||||
|
||||
# Apply the @property decorator
|
||||
@property
|
||||
# Define function total_controls
|
||||
def total_controls(self) -> int:
|
||||
"""Return the total number of controls in this framework.
|
||||
|
||||
Returns:
|
||||
int: Count of all controls regardless of coverage status.
|
||||
"""
|
||||
# Return len(self.controls)
|
||||
return len(self.controls)
|
||||
|
||||
# Apply the @property decorator
|
||||
@property
|
||||
# Define function covered_controls
|
||||
def covered_controls(self) -> int:
|
||||
"""Return the number of fully covered controls in this framework.
|
||||
|
||||
Returns:
|
||||
int: Count of controls with ``ControlCoverageStatus.covered`` status.
|
||||
"""
|
||||
# Return sum(
|
||||
return sum(
|
||||
# Literal argument value
|
||||
1 for c in self.controls
|
||||
if c.coverage_status == ControlCoverageStatus.covered
|
||||
)
|
||||
|
||||
# Apply the @property decorator
|
||||
@property
|
||||
# Define function coverage_pct
|
||||
def coverage_pct(self) -> float:
|
||||
"""Return the percentage of controls that are fully covered.
|
||||
|
||||
Returns:
|
||||
float: A value from 0.0 to 100.0, rounded to one decimal place.
|
||||
Returns 0.0 when the framework has no controls.
|
||||
"""
|
||||
# Check: self.total_controls == 0
|
||||
if self.total_controls == 0:
|
||||
# Return 0.0
|
||||
return 0.0
|
||||
# Return round(self.covered_controls / self.total_controls * 100, 1)
|
||||
return round(self.covered_controls / self.total_controls * 100, 1)
|
||||
|
||||
# Define function get_gap_controls
|
||||
def get_gap_controls(self) -> list[ComplianceControlEntity]:
|
||||
"""Return controls that are not fully covered.
|
||||
|
||||
Returns:
|
||||
list[ComplianceControlEntity]: Controls with ``partially_covered`` or
|
||||
``not_covered`` status.
|
||||
"""
|
||||
# Return [
|
||||
return [
|
||||
c for c in self.controls
|
||||
if c.coverage_status != ControlCoverageStatus.covered
|
||||
]
|
||||
@@ -1,316 +0,0 @@
|
||||
"""TechniqueEntity — pure domain object for a MITRE ATT&CK technique.
|
||||
|
||||
Owns the status recalculation logic that was previously in
|
||||
``status_service.py``. Has **no** dependency on FastAPI, SQLAlchemy,
|
||||
or any infrastructure concern.
|
||||
|
||||
Usage::
|
||||
|
||||
entity = TechniqueEntity.from_orm(technique_orm_model)
|
||||
entity.recalculate_status(test_states_and_results)
|
||||
entity.mark_reviewed()
|
||||
entity.apply_to(technique_orm_model)
|
||||
"""
|
||||
|
||||
# Enable future language features for compatibility
|
||||
from __future__ import annotations
|
||||
|
||||
# Import uuid
|
||||
import uuid
|
||||
|
||||
# Import dataclass, field from dataclasses
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
# Import datetime from datetime
|
||||
from datetime import datetime
|
||||
|
||||
# Import TYPE_CHECKING from typing
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
# Import TechniqueStatus, TestResult, TestState from app.domain.enums
|
||||
from app.domain.enums import TechniqueStatus, TestResult, TestState
|
||||
|
||||
# Import MitreId from app.domain.value_objects.mitre_id
|
||||
from app.domain.value_objects.mitre_id import MitreId
|
||||
|
||||
# Check: TYPE_CHECKING
|
||||
if TYPE_CHECKING:
|
||||
# Import Technique as TechniqueORM from app.models.technique
|
||||
from app.models.technique import Technique as TechniqueORM
|
||||
|
||||
|
||||
# Apply the @dataclass decorator
|
||||
@dataclass(frozen=True)
|
||||
# Define class _TestSnapshot
|
||||
class _TestSnapshot:
|
||||
"""Minimal read-only view of a test for status calculation."""
|
||||
|
||||
# state: TestState
|
||||
state: TestState
|
||||
# detection_result: str | None
|
||||
detection_result: str | None
|
||||
|
||||
|
||||
# Apply the @dataclass decorator
|
||||
@dataclass
|
||||
# Define class TechniqueEntity
|
||||
class TechniqueEntity:
|
||||
"""Pure domain representation of a MITRE ATT&CK technique."""
|
||||
|
||||
# id: uuid.UUID
|
||||
id: uuid.UUID
|
||||
# mitre_id: str
|
||||
mitre_id: str
|
||||
# name: str
|
||||
name: str
|
||||
# Assign tactic = None
|
||||
tactic: str | None = None
|
||||
# Assign description = None
|
||||
description: str | None = None
|
||||
# Assign platforms = field(default_factory=list)
|
||||
platforms: list[str] = field(default_factory=list)
|
||||
# Assign is_subtechnique = False
|
||||
is_subtechnique: bool = False
|
||||
# Assign parent_mitre_id = None
|
||||
parent_mitre_id: str | None = None
|
||||
# Assign status_global = TechniqueStatus.not_evaluated
|
||||
status_global: TechniqueStatus = TechniqueStatus.not_evaluated
|
||||
# Assign review_required = False
|
||||
review_required: bool = False
|
||||
# Assign last_review_date = None
|
||||
last_review_date: datetime | None = None
|
||||
# Assign mitre_version = None
|
||||
mitre_version: str | None = None
|
||||
# Assign mitre_last_modified = None
|
||||
mitre_last_modified: datetime | None = None
|
||||
|
||||
# -- Factory -----------------------------------------------------------
|
||||
|
||||
@classmethod
|
||||
# Define function create
|
||||
def create(
|
||||
cls,
|
||||
*,
|
||||
# Entry: mitre_id
|
||||
mitre_id: str,
|
||||
# Entry: name
|
||||
name: str,
|
||||
# Entry: tactic
|
||||
tactic: str | None = None,
|
||||
# Entry: description
|
||||
description: str | None = None,
|
||||
# Entry: platforms
|
||||
platforms: list[str] | None = None,
|
||||
) -> TechniqueEntity:
|
||||
"""Create a new technique, validating the MITRE ID format.
|
||||
|
||||
Args:
|
||||
mitre_id (str): MITRE ATT&CK identifier (e.g. ``"T1059"`` or ``"T1059.001"``).
|
||||
name (str): Human-readable name of the technique.
|
||||
tactic (str | None): MITRE tactic category the technique belongs to.
|
||||
description (str | None): Optional free-text description.
|
||||
platforms (list[str] | None): List of platform strings the technique applies to.
|
||||
|
||||
Returns:
|
||||
TechniqueEntity: A new entity with a freshly generated UUID and
|
||||
``status_global`` set to ``not_evaluated``.
|
||||
"""
|
||||
# Assign validated_id = MitreId(mitre_id)
|
||||
validated_id = MitreId(mitre_id)
|
||||
# Return cls(
|
||||
return cls(
|
||||
# Keyword argument: id
|
||||
id=uuid.uuid4(),
|
||||
# Keyword argument: mitre_id
|
||||
mitre_id=validated_id.value,
|
||||
# Keyword argument: name
|
||||
name=name,
|
||||
# Keyword argument: tactic
|
||||
tactic=tactic,
|
||||
# Keyword argument: description
|
||||
description=description,
|
||||
# Keyword argument: platforms
|
||||
platforms=platforms or [],
|
||||
# Keyword argument: is_subtechnique
|
||||
is_subtechnique=validated_id.is_subtechnique,
|
||||
# Keyword argument: parent_mitre_id
|
||||
parent_mitre_id=validated_id.parent_id,
|
||||
# Keyword argument: status_global
|
||||
status_global=TechniqueStatus.not_evaluated,
|
||||
)
|
||||
|
||||
# Apply the @classmethod decorator
|
||||
@classmethod
|
||||
# Define function from_orm
|
||||
def from_orm(cls, model: TechniqueORM) -> TechniqueEntity:
|
||||
"""Build a TechniqueEntity from a SQLAlchemy Technique model.
|
||||
|
||||
Args:
|
||||
model (TechniqueORM): The ORM model instance to convert.
|
||||
|
||||
Returns:
|
||||
TechniqueEntity: A fully populated domain entity reflecting the ORM state.
|
||||
"""
|
||||
# Assign raw_status = model.status_global
|
||||
raw_status = model.status_global
|
||||
# Check: raw_status is None
|
||||
if raw_status is None:
|
||||
# Assign status = TechniqueStatus.not_evaluated
|
||||
status = TechniqueStatus.not_evaluated
|
||||
# Alternative: isinstance(raw_status, TechniqueStatus)
|
||||
elif isinstance(raw_status, TechniqueStatus):
|
||||
# Assign status = raw_status
|
||||
status = raw_status
|
||||
# Fallback: handle remaining cases
|
||||
else:
|
||||
# Assign status = TechniqueStatus(raw_status)
|
||||
status = TechniqueStatus(raw_status)
|
||||
# Return cls(
|
||||
return cls(
|
||||
# Keyword argument: id
|
||||
id=model.id,
|
||||
# Keyword argument: mitre_id
|
||||
mitre_id=model.mitre_id,
|
||||
# Keyword argument: name
|
||||
name=model.name,
|
||||
# Keyword argument: tactic
|
||||
tactic=model.tactic,
|
||||
# Keyword argument: description
|
||||
description=model.description,
|
||||
# Keyword argument: platforms
|
||||
platforms=model.platforms or [],
|
||||
# Keyword argument: is_subtechnique
|
||||
is_subtechnique=model.is_subtechnique or False,
|
||||
# Keyword argument: parent_mitre_id
|
||||
parent_mitre_id=model.parent_mitre_id,
|
||||
# Keyword argument: status_global
|
||||
status_global=status,
|
||||
# Keyword argument: review_required
|
||||
review_required=model.review_required or False,
|
||||
# Keyword argument: last_review_date
|
||||
last_review_date=model.last_review_date,
|
||||
# Keyword argument: mitre_version
|
||||
mitre_version=getattr(model, "mitre_version", None),
|
||||
# Keyword argument: mitre_last_modified
|
||||
mitre_last_modified=getattr(model, "mitre_last_modified", None),
|
||||
)
|
||||
|
||||
# Define function apply_to
|
||||
def apply_to(self, model: TechniqueORM) -> None:
|
||||
"""Copy mutable fields back onto the ORM model.
|
||||
|
||||
Args:
|
||||
model (TechniqueORM): The ORM model to update in-place.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# Assign model.status_global = self.status_global
|
||||
model.status_global = self.status_global
|
||||
# Assign model.review_required = self.review_required
|
||||
model.review_required = self.review_required
|
||||
# Assign model.last_review_date = self.last_review_date
|
||||
model.last_review_date = self.last_review_date
|
||||
|
||||
# -- Business logic ----------------------------------------------------
|
||||
|
||||
def recalculate_status(
|
||||
self,
|
||||
# Entry: test_snapshots
|
||||
test_snapshots: list[tuple[str, str | None]],
|
||||
) -> TechniqueStatus:
|
||||
"""Recompute ``status_global`` from a list of (state, detection_result) pairs.
|
||||
|
||||
Rules (v3):
|
||||
1. No tests -> not_evaluated
|
||||
2. All tests validated -> inspect detection results:
|
||||
a. All detected AND ≥ 1 validated test -> validated
|
||||
b. Any partially_detected -> partial
|
||||
d. Otherwise (no detected results) -> not_covered
|
||||
3. Some validated, others in intermediate states -> partial
|
||||
4. All tests in intermediate states (draft/executing/evaluating/review/rejected)
|
||||
-> in_progress
|
||||
|
||||
Minimum validated count for "validated": 1 test.
|
||||
|
||||
Args:
|
||||
test_snapshots (list[tuple[str, str | None]]): Each element is a
|
||||
``(state, detection_result)`` pair where *state* is a
|
||||
:class:`TestState` value string and *detection_result* is a
|
||||
:class:`TestResult` value string or ``None``.
|
||||
|
||||
Returns:
|
||||
TechniqueStatus: The newly computed status, which is also stored on
|
||||
the entity's ``status_global`` field.
|
||||
"""
|
||||
min_validated_for_full = 1 # require ≥ N validated tests for "validated"
|
||||
|
||||
tests = [
|
||||
_TestSnapshot(
|
||||
# Keyword argument: state
|
||||
state=s if isinstance(s, TestState) else TestState(s),
|
||||
# Keyword argument: detection_result
|
||||
detection_result=dr,
|
||||
)
|
||||
for s, dr in test_snapshots
|
||||
]
|
||||
|
||||
# Check: not tests
|
||||
if not tests:
|
||||
# Assign self.status_global = TechniqueStatus.not_evaluated
|
||||
self.status_global = TechniqueStatus.not_evaluated
|
||||
# Alternative: all(t.state == TestState.validated for t in tests)
|
||||
elif all(t.state == TestState.validated for t in tests):
|
||||
validated_count = len(tests)
|
||||
results = [t.detection_result for t in tests if t.detection_result]
|
||||
# Check: results and all(r == TestResult.detected or r == "detected" for r i...
|
||||
if results and all(r == TestResult.detected or r == "detected" for r in results):
|
||||
# Need at least min_validated_for_full tests for "validated"
|
||||
if validated_count >= min_validated_for_full:
|
||||
self.status_global = TechniqueStatus.validated
|
||||
else:
|
||||
self.status_global = TechniqueStatus.partial
|
||||
elif any(
|
||||
# Keyword argument: r
|
||||
r == TestResult.partially_detected or r == "partially_detected"
|
||||
for r in results
|
||||
):
|
||||
# Assign self.status_global = TechniqueStatus.partial
|
||||
self.status_global = TechniqueStatus.partial
|
||||
# Fallback: handle remaining cases
|
||||
else:
|
||||
# Assign self.status_global = TechniqueStatus.not_covered
|
||||
self.status_global = TechniqueStatus.not_covered
|
||||
# Alternative: any(t.state == TestState.validated for t in tests)
|
||||
elif any(t.state == TestState.validated for t in tests):
|
||||
# Assign self.status_global = TechniqueStatus.partial
|
||||
self.status_global = TechniqueStatus.partial
|
||||
# Fallback: handle remaining cases
|
||||
else:
|
||||
# Assign self.status_global = TechniqueStatus.in_progress
|
||||
self.status_global = TechniqueStatus.in_progress
|
||||
|
||||
# Return self.status_global
|
||||
return self.status_global
|
||||
|
||||
# Define function mark_reviewed
|
||||
def mark_reviewed(self) -> None:
|
||||
"""Mark the technique as reviewed, clearing the review flag.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# Assign self.review_required = False
|
||||
self.review_required = False
|
||||
# Assign self.last_review_date = datetime.utcnow()
|
||||
self.last_review_date = datetime.utcnow()
|
||||
|
||||
# Define function flag_for_review
|
||||
def flag_for_review(self) -> None:
|
||||
"""Flag the technique as needing review.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# Assign self.review_required = True
|
||||
self.review_required = True
|
||||
@@ -1,206 +0,0 @@
|
||||
"""Threat actor domain entity with coverage analysis logic.
|
||||
|
||||
Pure domain logic — no framework imports.
|
||||
"""
|
||||
|
||||
# Enable future language features for compatibility
|
||||
from __future__ import annotations
|
||||
|
||||
# Import uuid
|
||||
import uuid
|
||||
|
||||
# Import dataclass, field from dataclasses
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
# Import TYPE_CHECKING from typing
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
# Check: TYPE_CHECKING
|
||||
if TYPE_CHECKING:
|
||||
# Import ThreatActor as ThreatActorORM from app.models.threat_actor
|
||||
from app.models.threat_actor import ThreatActor as ThreatActorORM
|
||||
|
||||
|
||||
# Apply the @dataclass decorator
|
||||
@dataclass
|
||||
# Define class ThreatActorTechniqueRef
|
||||
class ThreatActorTechniqueRef:
|
||||
"""Lightweight reference to a technique used by an actor."""
|
||||
|
||||
# technique_id: uuid.UUID
|
||||
technique_id: uuid.UUID
|
||||
# Assign mitre_id = None
|
||||
mitre_id: str | None = None
|
||||
# Assign name = None
|
||||
name: str | None = None
|
||||
# Assign status = None
|
||||
status: str | None = None
|
||||
# Assign usage_description = None
|
||||
usage_description: str | None = None
|
||||
|
||||
|
||||
# Apply the @dataclass decorator
|
||||
@dataclass
|
||||
# Define class ThreatActorEntity
|
||||
class ThreatActorEntity:
|
||||
"""Pure domain representation of a MITRE ATT&CK threat actor (group).
|
||||
|
||||
Aggregates references to the techniques the actor is known to use and
|
||||
provides coverage analysis properties.
|
||||
"""
|
||||
|
||||
# name: str
|
||||
name: str
|
||||
# Assign id = None
|
||||
id: uuid.UUID | None = None
|
||||
# Assign mitre_id = None
|
||||
mitre_id: str | None = None
|
||||
# Assign aliases = field(default_factory=list)
|
||||
aliases: list[str] = field(default_factory=list)
|
||||
# Assign description = None
|
||||
description: str | None = None
|
||||
# Assign country = None
|
||||
country: str | None = None
|
||||
# Assign target_sectors = field(default_factory=list)
|
||||
target_sectors: list[str] = field(default_factory=list)
|
||||
# Assign target_regions = field(default_factory=list)
|
||||
target_regions: list[str] = field(default_factory=list)
|
||||
# Assign motivation = None
|
||||
motivation: str | None = None
|
||||
# Assign sophistication = None
|
||||
sophistication: str | None = None
|
||||
# Assign first_seen = None
|
||||
first_seen: str | None = None
|
||||
# Assign last_seen = None
|
||||
last_seen: str | None = None
|
||||
# Assign is_active = True
|
||||
is_active: bool = True
|
||||
# Assign techniques = field(default_factory=list)
|
||||
techniques: list[ThreatActorTechniqueRef] = field(default_factory=list)
|
||||
|
||||
# Apply the @property decorator
|
||||
@property
|
||||
# Define function technique_count
|
||||
def technique_count(self) -> int:
|
||||
"""Return the total number of techniques associated with this actor.
|
||||
|
||||
Returns:
|
||||
int: Count of technique references.
|
||||
"""
|
||||
# Return len(self.techniques)
|
||||
return len(self.techniques)
|
||||
|
||||
# Apply the @property decorator
|
||||
@property
|
||||
# Define function covered_techniques
|
||||
def covered_techniques(self) -> list[ThreatActorTechniqueRef]:
|
||||
"""Return technique references whose coverage status is ``validated`` or ``partial``.
|
||||
|
||||
Returns:
|
||||
list[ThreatActorTechniqueRef]: Subset of techniques considered covered.
|
||||
"""
|
||||
# Return [
|
||||
return [
|
||||
t for t in self.techniques
|
||||
if t.status in ("validated", "partial")
|
||||
]
|
||||
|
||||
# Apply the @property decorator
|
||||
@property
|
||||
# Define function uncovered_techniques
|
||||
def uncovered_techniques(self) -> list[ThreatActorTechniqueRef]:
|
||||
"""Return technique references whose coverage status is neither ``validated`` nor ``partial``.
|
||||
|
||||
Returns:
|
||||
list[ThreatActorTechniqueRef]: Subset of techniques not yet covered.
|
||||
"""
|
||||
# Return [
|
||||
return [
|
||||
t for t in self.techniques
|
||||
if t.status not in ("validated", "partial")
|
||||
]
|
||||
|
||||
# Apply the @property decorator
|
||||
@property
|
||||
# Define function coverage_pct
|
||||
def coverage_pct(self) -> float:
|
||||
"""Return the percentage of the actor's techniques that are covered.
|
||||
|
||||
Returns:
|
||||
float: A value from 0.0 to 100.0, rounded to one decimal place.
|
||||
Returns 0.0 when the actor has no associated techniques.
|
||||
"""
|
||||
# Check: not self.techniques
|
||||
if not self.techniques:
|
||||
# Return 0.0
|
||||
return 0.0
|
||||
# Return round(len(self.covered_techniques) / len(self.techniques) * 100, 1)
|
||||
return round(len(self.covered_techniques) / len(self.techniques) * 100, 1)
|
||||
|
||||
# Apply the @classmethod decorator
|
||||
@classmethod
|
||||
# Define function from_orm
|
||||
def from_orm(cls, orm: ThreatActorORM) -> ThreatActorEntity:
|
||||
"""Build a ThreatActorEntity from a SQLAlchemy ThreatActor model.
|
||||
|
||||
Args:
|
||||
orm (ThreatActorORM): The ORM model instance to convert.
|
||||
|
||||
Returns:
|
||||
ThreatActorEntity: A fully populated domain entity including
|
||||
technique references resolved from the ORM relationship.
|
||||
"""
|
||||
# Assign techs = []
|
||||
techs: list[ThreatActorTechniqueRef] = []
|
||||
# Iterate over getattr(orm, "techniques", None) or []
|
||||
for tat in getattr(orm, "techniques", None) or []:
|
||||
# Assign technique = getattr(tat, "technique", None)
|
||||
technique = getattr(tat, "technique", None)
|
||||
# Call techs.append()
|
||||
techs.append(ThreatActorTechniqueRef(
|
||||
# Keyword argument: technique_id
|
||||
technique_id=tat.technique_id,
|
||||
# Keyword argument: mitre_id
|
||||
mitre_id=getattr(technique, "mitre_id", None) if technique else None,
|
||||
# Keyword argument: name
|
||||
name=getattr(technique, "name", None) if technique else None,
|
||||
# Keyword argument: status
|
||||
status=(
|
||||
technique.status_global.value
|
||||
if technique and hasattr(technique.status_global, "value")
|
||||
else getattr(technique, "status_global", None) if technique else None
|
||||
),
|
||||
# Keyword argument: usage_description
|
||||
usage_description=tat.usage_description,
|
||||
))
|
||||
# Return cls(
|
||||
return cls(
|
||||
# Keyword argument: id
|
||||
id=orm.id,
|
||||
# Keyword argument: name
|
||||
name=orm.name,
|
||||
# Keyword argument: mitre_id
|
||||
mitre_id=orm.mitre_id,
|
||||
# Keyword argument: aliases
|
||||
aliases=orm.aliases or [],
|
||||
# Keyword argument: description
|
||||
description=orm.description,
|
||||
# Keyword argument: country
|
||||
country=orm.country,
|
||||
# Keyword argument: target_sectors
|
||||
target_sectors=orm.target_sectors or [],
|
||||
# Keyword argument: target_regions
|
||||
target_regions=orm.target_regions or [],
|
||||
# Keyword argument: motivation
|
||||
motivation=orm.motivation,
|
||||
# Keyword argument: sophistication
|
||||
sophistication=orm.sophistication,
|
||||
# Keyword argument: first_seen
|
||||
first_seen=orm.first_seen,
|
||||
# Keyword argument: last_seen
|
||||
last_seen=orm.last_seen,
|
||||
# Keyword argument: is_active
|
||||
is_active=orm.is_active if orm.is_active is not None else True,
|
||||
# Keyword argument: techniques
|
||||
techniques=techs,
|
||||
)
|
||||
@@ -1,82 +0,0 @@
|
||||
"""Canonical domain enums for Aegis.
|
||||
|
||||
These enums represent core domain concepts and are the single source of
|
||||
truth. ``models/enums.py`` re-exports them so that existing ORM code
|
||||
continues to work without changes.
|
||||
"""
|
||||
|
||||
# Import enum
|
||||
import enum
|
||||
|
||||
|
||||
# Define class TechniqueStatus
|
||||
class TechniqueStatus(str, enum.Enum):
|
||||
"""Coverage and evaluation status for a MITRE ATT&CK technique."""
|
||||
|
||||
# Assign not_evaluated = "not_evaluated"
|
||||
not_evaluated = "not_evaluated"
|
||||
# Assign in_progress = "in_progress"
|
||||
in_progress = "in_progress"
|
||||
# Assign validated = "validated"
|
||||
validated = "validated"
|
||||
# Assign partial = "partial"
|
||||
partial = "partial"
|
||||
# Assign not_covered = "not_covered"
|
||||
not_covered = "not_covered"
|
||||
# Assign review_required = "review_required"
|
||||
review_required = "review_required"
|
||||
|
||||
|
||||
# Define class TestState
|
||||
class TestState(str, enum.Enum):
|
||||
"""Lifecycle states in the security test state machine."""
|
||||
|
||||
# Assign draft = "draft"
|
||||
draft = "draft"
|
||||
# Assign red_executing = "red_executing"
|
||||
red_executing = "red_executing"
|
||||
# Assign blue_evaluating = "blue_evaluating"
|
||||
blue_evaluating = "blue_evaluating"
|
||||
# Assign in_review = "in_review"
|
||||
in_review = "in_review"
|
||||
# Assign validated = "validated"
|
||||
validated = "validated"
|
||||
# Assign rejected = "rejected"
|
||||
rejected = "rejected"
|
||||
disputed = "disputed" # one lead approved, the other rejected
|
||||
|
||||
|
||||
# Define class TeamSide
|
||||
class TeamSide(str, enum.Enum):
|
||||
"""Identifies which team (red or blue) an action belongs to."""
|
||||
|
||||
# Assign red = "red"
|
||||
red = "red"
|
||||
# Assign blue = "blue"
|
||||
blue = "blue"
|
||||
|
||||
|
||||
# Define class TestResult
|
||||
class TestResult(str, enum.Enum):
|
||||
"""Outcome of a red-team test from a detection perspective."""
|
||||
|
||||
# Assign detected = "detected"
|
||||
detected = "detected"
|
||||
# Assign not_detected = "not_detected"
|
||||
not_detected = "not_detected"
|
||||
# Assign partially_detected = "partially_detected"
|
||||
partially_detected = "partially_detected"
|
||||
|
||||
|
||||
# Define class DataClassification
|
||||
class DataClassification(str, enum.Enum):
|
||||
"""Data sensitivity classification levels for compliance and retention policies."""
|
||||
|
||||
# Assign public = "public"
|
||||
public = "public"
|
||||
# Assign internal = "internal"
|
||||
internal = "internal"
|
||||
# Assign sensitive = "sensitive"
|
||||
sensitive = "sensitive"
|
||||
# Assign restricted = "restricted"
|
||||
restricted = "restricted"
|
||||
@@ -1,192 +0,0 @@
|
||||
"""Canonical domain error hierarchy for Aegis.
|
||||
|
||||
Every service-layer error should be a subclass of :class:`DomainError`.
|
||||
The global exception handler in ``app.middleware.error_handler`` maps
|
||||
each concrete subclass to an appropriate HTTP status code so that
|
||||
services never depend on FastAPI.
|
||||
|
||||
Existing code that imports from ``app.domain.exceptions`` continues to
|
||||
work — that module re-exports everything defined here.
|
||||
"""
|
||||
|
||||
# Enable future language features for compatibility
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
# Define class DomainError
|
||||
class DomainError(Exception):
|
||||
"""Base for all domain errors."""
|
||||
|
||||
# Define function __init__
|
||||
def __init__(self, message: str, *, code: str = "DOMAIN_ERROR") -> None:
|
||||
"""Initialise the domain error with a human-readable message and error code.
|
||||
|
||||
Args:
|
||||
message (str): Human-readable description of the error.
|
||||
code (str): Machine-readable error code used by the HTTP error handler.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# Assign self.message = message
|
||||
self.message = message
|
||||
# Assign self.code = code
|
||||
self.code = code
|
||||
# Call super()
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
# ── Entity lifecycle ──────────────────────────────────────────────────
|
||||
|
||||
|
||||
class EntityNotFoundError(DomainError):
|
||||
"""A requested entity does not exist."""
|
||||
|
||||
# Define function __init__
|
||||
def __init__(self, entity: str, identifier: str) -> None:
|
||||
"""Initialise an entity-not-found error.
|
||||
|
||||
Args:
|
||||
entity (str): Name of the entity type that was not found (e.g. "Technique").
|
||||
identifier (str): The ID or key used in the failed lookup.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# Call super()
|
||||
super().__init__(f"{entity} not found: {identifier}", code="NOT_FOUND")
|
||||
# Assign self.entity = entity
|
||||
self.entity = entity
|
||||
# Assign self.identifier = identifier
|
||||
self.identifier = identifier
|
||||
|
||||
|
||||
# Define class DuplicateEntityError
|
||||
class DuplicateEntityError(DomainError):
|
||||
"""Creating an entity that already exists."""
|
||||
|
||||
# Define function __init__
|
||||
def __init__(self, entity: str, field: str, value: str) -> None:
|
||||
"""Initialise a duplicate-entity error.
|
||||
|
||||
Args:
|
||||
entity (str): Name of the entity type that already exists (e.g. "Campaign").
|
||||
field (str): Name of the field whose value conflicts (e.g. "name").
|
||||
value (str): The conflicting value that is already in use.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# Call super()
|
||||
super().__init__(
|
||||
f"{entity} with {field}='{value}' already exists",
|
||||
# Keyword argument: code
|
||||
code="DUPLICATE",
|
||||
)
|
||||
|
||||
|
||||
# ── State machine ────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class InvalidStateTransition(DomainError): # noqa: N818 — DDD term, renaming would break 96 call sites
|
||||
"""A state-machine transition is not allowed."""
|
||||
|
||||
# Define function __init__
|
||||
def __init__(
|
||||
self,
|
||||
# Entry: current_state
|
||||
current_state: str,
|
||||
# Entry: target_state
|
||||
target_state: str,
|
||||
# Entry: valid_transitions
|
||||
valid_transitions: list[str] | None = None,
|
||||
) -> None:
|
||||
"""Initialise an invalid state-transition error.
|
||||
|
||||
Args:
|
||||
current_state (str): The entity's present state (e.g. "draft").
|
||||
target_state (str): The state that was illegally requested.
|
||||
valid_transitions (list[str] | None): Allowed target states from the
|
||||
current state; included in the error message when provided.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# Assign msg = f"Cannot transition from '{current_state}' to '{target_state}'"
|
||||
msg = f"Cannot transition from '{current_state}' to '{target_state}'"
|
||||
# Check: valid_transitions
|
||||
if valid_transitions:
|
||||
# Assign msg = f". Valid transitions: {valid_transitions}"
|
||||
msg += f". Valid transitions: {valid_transitions}"
|
||||
# Call super()
|
||||
super().__init__(msg, code="INVALID_TRANSITION")
|
||||
# Assign self.current_state = current_state
|
||||
self.current_state = current_state
|
||||
# Assign self.target_state = target_state
|
||||
self.target_state = target_state
|
||||
# Assign self.valid_transitions = valid_transitions or []
|
||||
self.valid_transitions = valid_transitions or []
|
||||
|
||||
|
||||
# ── Business rules ────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class BusinessRuleViolation(DomainError): # noqa: N818 — DDD term, renaming would break 96 call sites
|
||||
"""An operation violates a business invariant."""
|
||||
|
||||
# Define function __init__
|
||||
def __init__(self, message: str) -> None:
|
||||
"""Initialise a business-rule violation error.
|
||||
|
||||
Args:
|
||||
message (str): Human-readable description of the violated rule.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# Call super()
|
||||
super().__init__(message, code="BUSINESS_RULE_VIOLATION")
|
||||
|
||||
|
||||
# Define class InvalidOperationError
|
||||
class InvalidOperationError(BusinessRuleViolation):
|
||||
"""An operation is invalid in the current context.
|
||||
|
||||
Kept for backward compatibility; new code should prefer
|
||||
:class:`BusinessRuleViolation` directly.
|
||||
"""
|
||||
|
||||
# Define function __init__
|
||||
def __init__(self, message: str) -> None:
|
||||
"""Initialise an invalid-operation error.
|
||||
|
||||
Args:
|
||||
message (str): Human-readable description of why the operation is invalid.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# Call super()
|
||||
super().__init__(message)
|
||||
# Assign self.code = "INVALID_OPERATION"
|
||||
self.code = "INVALID_OPERATION"
|
||||
|
||||
|
||||
# ── Authorization ────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class PermissionViolation(DomainError): # noqa: N818 — DDD term, renaming would break 96 call sites
|
||||
"""The user lacks permissions for an action."""
|
||||
|
||||
# Define function __init__
|
||||
def __init__(self, message: str = "Insufficient permissions") -> None:
|
||||
"""Initialise a permission-violation error.
|
||||
|
||||
Args:
|
||||
message (str): Human-readable description of the access denial.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# Call super()
|
||||
super().__init__(message, code="FORBIDDEN")
|
||||
@@ -1,25 +0,0 @@
|
||||
"""Backward-compatible re-exports from :mod:`app.domain.errors`.
|
||||
|
||||
All domain errors now live in ``errors.py``. This module preserves the
|
||||
old import paths so that existing code keeps working without changes::
|
||||
|
||||
from app.domain.exceptions import InvalidTransitionError # still works
|
||||
"""
|
||||
|
||||
# Import # noqa: F401 from app.domain.errors
|
||||
from app.domain.errors import ( # noqa: F401
|
||||
BusinessRuleViolation,
|
||||
DomainError,
|
||||
DuplicateEntityError,
|
||||
EntityNotFoundError,
|
||||
InvalidOperationError,
|
||||
InvalidStateTransition,
|
||||
PermissionViolation,
|
||||
)
|
||||
|
||||
# Legacy aliases — old name → new name
|
||||
DomainException = DomainError
|
||||
# Assign InvalidTransitionError = InvalidStateTransition
|
||||
InvalidTransitionError = InvalidStateTransition
|
||||
# Assign AuthorizationError = PermissionViolation
|
||||
AuthorizationError = PermissionViolation
|
||||
@@ -1 +0,0 @@
|
||||
"""Abstract port interfaces that infrastructure adapters must implement."""
|
||||
@@ -1,165 +0,0 @@
|
||||
"""Port defining the common interface for data import services.
|
||||
|
||||
All import services (Atomic Red Team, Sigma, CALDERA, etc.) follow the
|
||||
same contract: they receive a database session and return a summary dict
|
||||
with import statistics.
|
||||
|
||||
New import sources can be added by:
|
||||
1. Implementing the ``ImportService`` protocol in a new module
|
||||
2. Registering the handler in the ``IMPORT_REGISTRY``
|
||||
|
||||
This satisfies the Open/Closed Principle — the system is open for new
|
||||
import sources without modifying existing code.
|
||||
"""
|
||||
|
||||
# Enable future language features for compatibility
|
||||
from __future__ import annotations
|
||||
|
||||
# Import Any, Protocol, runtime_checkable from typing
|
||||
from typing import Any, Protocol, runtime_checkable
|
||||
|
||||
# Import Session from sqlalchemy.orm
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
|
||||
# Apply the @runtime_checkable decorator
|
||||
@runtime_checkable
|
||||
# Define class ImportService
|
||||
class ImportService(Protocol):
|
||||
"""Contract for any data-import operation.
|
||||
|
||||
Each implementation is a callable ``(Session) -> dict`` that
|
||||
downloads, parses, and upserts records from an external source.
|
||||
"""
|
||||
|
||||
# Define function __call__
|
||||
def __call__(self, db: Session) -> dict[str, Any]:
|
||||
"""Execute the import operation against the given database session.
|
||||
|
||||
Args:
|
||||
db (Session): Active SQLAlchemy session to use for all DB operations.
|
||||
|
||||
Returns:
|
||||
dict[str, Any]: Summary statistics for the import run (e.g. created,
|
||||
updated, skipped counts).
|
||||
"""
|
||||
# ...
|
||||
...
|
||||
|
||||
|
||||
# Define class ImportServiceEntry
|
||||
class ImportServiceEntry:
|
||||
"""Lazy-loading wrapper that resolves a module-level function on first call."""
|
||||
|
||||
# Assign __slots__ = ("_module_path", "_func_name", "_resolved")
|
||||
__slots__ = ("_module_path", "_func_name", "_resolved")
|
||||
|
||||
# Define function __init__
|
||||
def __init__(self, module_path: str, func_name: str) -> None:
|
||||
"""Initialise the lazy entry with the module path and function name to resolve later.
|
||||
|
||||
Args:
|
||||
module_path (str): Dotted Python module path, e.g.
|
||||
``"app.services.atomic_import_service"``.
|
||||
func_name (str): Name of the callable to import from *module_path*.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# Assign self._module_path = module_path
|
||||
self._module_path = module_path
|
||||
# Assign self._func_name = func_name
|
||||
self._func_name = func_name
|
||||
# Assign self._resolved = None
|
||||
self._resolved: ImportService | None = None
|
||||
|
||||
# Define function __call__
|
||||
def __call__(self, db: Session) -> dict[str, Any]:
|
||||
"""Resolve the import function on first call and invoke it with *db*.
|
||||
|
||||
Args:
|
||||
db (Session): SQLAlchemy session passed through to the underlying
|
||||
import function.
|
||||
|
||||
Returns:
|
||||
dict[str, Any]: Import statistics returned by the underlying function
|
||||
(e.g. counts of created/updated/skipped records).
|
||||
"""
|
||||
# Check: self._resolved is None
|
||||
if self._resolved is None:
|
||||
# Import importlib
|
||||
import importlib
|
||||
# Assign mod = importlib.import_module(self._module_path)
|
||||
mod = importlib.import_module(self._module_path)
|
||||
# Assign self._resolved = getattr(mod, self._func_name)
|
||||
self._resolved = getattr(mod, self._func_name)
|
||||
# Return self._resolved(db)
|
||||
return self._resolved(db)
|
||||
|
||||
# Apply the @property decorator
|
||||
@property
|
||||
# Define function source_info
|
||||
def source_info(self) -> str:
|
||||
"""Return a human-readable identifier for this import entry.
|
||||
|
||||
Returns:
|
||||
str: The fully qualified function reference as
|
||||
``"<module_path>.<func_name>"``.
|
||||
"""
|
||||
# Return f"{self._module_path}.{self._func_name}"
|
||||
return f"{self._module_path}.{self._func_name}"
|
||||
|
||||
|
||||
# Assign IMPORT_REGISTRY = {
|
||||
IMPORT_REGISTRY: dict[str, ImportServiceEntry] = {
|
||||
# Literal argument value
|
||||
"atomic_red_team": ImportServiceEntry(
|
||||
# Literal argument value
|
||||
"app.services.atomic_import_service", "import_atomic_red_team",
|
||||
),
|
||||
# Literal argument value
|
||||
"sigma": ImportServiceEntry(
|
||||
# Literal argument value
|
||||
"app.services.sigma_import_service", "sync",
|
||||
),
|
||||
# Literal argument value
|
||||
"lolbas": ImportServiceEntry(
|
||||
# Literal argument value
|
||||
"app.services.lolbas_import_service", "sync",
|
||||
),
|
||||
# Literal argument value
|
||||
"gtfobins": ImportServiceEntry(
|
||||
# Literal argument value
|
||||
"app.services.lolbas_import_service", "sync_gtfobins",
|
||||
),
|
||||
# Literal argument value
|
||||
"caldera": ImportServiceEntry(
|
||||
# Literal argument value
|
||||
"app.services.caldera_import_service", "sync",
|
||||
),
|
||||
# Literal argument value
|
||||
"elastic_rules": ImportServiceEntry(
|
||||
# Literal argument value
|
||||
"app.services.elastic_import_service", "sync",
|
||||
),
|
||||
# Literal argument value
|
||||
"mitre_cti": ImportServiceEntry(
|
||||
# Literal argument value
|
||||
"app.services.threat_actor_import_service", "sync",
|
||||
),
|
||||
# Literal argument value
|
||||
"d3fend": ImportServiceEntry(
|
||||
# Literal argument value
|
||||
"app.services.d3fend_import_service", "sync",
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
# Define function get_import_handler
|
||||
def get_import_handler(source_name: str) -> ImportServiceEntry | None:
|
||||
"""Look up the import handler for *source_name*.
|
||||
|
||||
Returns ``None`` when no handler is registered.
|
||||
"""
|
||||
# Return IMPORT_REGISTRY.get(source_name)
|
||||
return IMPORT_REGISTRY.get(source_name)
|
||||
@@ -1,9 +0,0 @@
|
||||
"""Abstract repository port interfaces for domain entity persistence."""
|
||||
# Import TechniqueRepository from app.domain.ports.repositories.technique_repository
|
||||
from app.domain.ports.repositories.technique_repository import TechniqueRepository
|
||||
|
||||
# Import TestRepository from app.domain.ports.repositories.test_repository
|
||||
from app.domain.ports.repositories.test_repository import TestRepository
|
||||
|
||||
# Assign __all__ = ["TechniqueRepository", "TestRepository"]
|
||||
__all__ = ["TechniqueRepository", "TestRepository"]
|
||||
@@ -1,160 +0,0 @@
|
||||
"""Port defining how the application accesses technique data.
|
||||
|
||||
This is a domain contract — implementations live in infrastructure/.
|
||||
The domain layer NEVER imports the implementation.
|
||||
"""
|
||||
|
||||
# Enable future language features for compatibility
|
||||
from __future__ import annotations
|
||||
|
||||
# Import uuid
|
||||
import uuid
|
||||
|
||||
# Import NamedTuple, Protocol, runtime_checkable from typing
|
||||
from typing import NamedTuple, Protocol, runtime_checkable
|
||||
|
||||
# Import TechniqueEntity from app.domain.entities.technique
|
||||
from app.domain.entities.technique import TechniqueEntity
|
||||
|
||||
# Import TechniqueStatus from app.domain.enums
|
||||
from app.domain.enums import TechniqueStatus
|
||||
|
||||
|
||||
# Define class TechniqueWithCounts
|
||||
class TechniqueWithCounts(NamedTuple):
|
||||
"""Pre-aggregated technique data for heatmap/scoring."""
|
||||
|
||||
# entity: TechniqueEntity
|
||||
entity: TechniqueEntity
|
||||
# test_count: int
|
||||
test_count: int
|
||||
# validated_test_count: int
|
||||
validated_test_count: int
|
||||
# detection_rule_count: int
|
||||
detection_rule_count: int
|
||||
|
||||
|
||||
# Apply the @runtime_checkable decorator
|
||||
@runtime_checkable
|
||||
# Define class TechniqueRepository
|
||||
class TechniqueRepository(Protocol):
|
||||
"""Data access contract for techniques (one per aggregate root)."""
|
||||
|
||||
# -- Single-entity access ----------------------------------------------
|
||||
|
||||
def find_by_id(self, technique_id: uuid.UUID) -> TechniqueEntity | None:
|
||||
"""Return the technique with the given primary key, or None if absent.
|
||||
|
||||
Args:
|
||||
technique_id (uuid.UUID): Primary key of the technique to look up.
|
||||
|
||||
Returns:
|
||||
TechniqueEntity | None: The matching entity, or None if not found.
|
||||
"""
|
||||
# ...
|
||||
...
|
||||
|
||||
# Define function find_by_mitre_id
|
||||
def find_by_mitre_id(self, mitre_id: str) -> TechniqueEntity | None:
|
||||
"""Return the technique matching the given MITRE ATT&CK identifier, or None.
|
||||
|
||||
Args:
|
||||
mitre_id (str): MITRE ATT&CK ID (e.g. ``"T1059"`` or ``"T1059.001"``).
|
||||
|
||||
Returns:
|
||||
TechniqueEntity | None: The matching entity, or None if not found.
|
||||
"""
|
||||
# ...
|
||||
...
|
||||
|
||||
# -- List access -------------------------------------------------------
|
||||
|
||||
def list_all(
|
||||
self,
|
||||
*,
|
||||
# Entry: tactic
|
||||
tactic: str | None = None,
|
||||
# Entry: status
|
||||
status: TechniqueStatus | None = None,
|
||||
# Entry: review_required
|
||||
review_required: bool | None = None,
|
||||
) -> list[TechniqueEntity]:
|
||||
"""Return all techniques, optionally filtered by tactic, status, or review flag.
|
||||
|
||||
Args:
|
||||
tactic (str | None): When provided, restrict results to this tactic category.
|
||||
status (TechniqueStatus | None): When provided, restrict results to this status.
|
||||
review_required (bool | None): When provided, restrict results to techniques
|
||||
whose ``review_required`` flag matches this value.
|
||||
|
||||
Returns:
|
||||
list[TechniqueEntity]: Matching technique entities; may be empty.
|
||||
"""
|
||||
# ...
|
||||
...
|
||||
|
||||
# Define function list_by_ids
|
||||
def list_by_ids(self, ids: list[uuid.UUID]) -> list[TechniqueEntity]:
|
||||
"""Return all techniques whose primary keys are in *ids*.
|
||||
|
||||
Args:
|
||||
ids (list[uuid.UUID]): List of technique UUIDs to retrieve.
|
||||
|
||||
Returns:
|
||||
list[TechniqueEntity]: Entities found for the supplied IDs; order
|
||||
is not guaranteed and missing IDs are silently omitted.
|
||||
"""
|
||||
# ...
|
||||
...
|
||||
|
||||
# -- Batch queries (scoring/heatmap performance) -----------------------
|
||||
|
||||
def count_by_status(self) -> dict[TechniqueStatus, int]:
|
||||
"""Return a count of techniques grouped by their global status.
|
||||
|
||||
Returns:
|
||||
dict[TechniqueStatus, int]: Mapping from each status value to the
|
||||
number of techniques in that state.
|
||||
"""
|
||||
# ...
|
||||
...
|
||||
|
||||
# Define function find_all_with_test_counts
|
||||
def find_all_with_test_counts(self) -> list[TechniqueWithCounts]:
|
||||
"""Return all techniques together with pre-aggregated test and rule counts.
|
||||
|
||||
Returns:
|
||||
list[TechniqueWithCounts]: Each element bundles a TechniqueEntity
|
||||
with its total, validated, and detection-rule counts for use
|
||||
in heatmap and scoring calculations.
|
||||
"""
|
||||
# ...
|
||||
...
|
||||
|
||||
# -- Mutations ---------------------------------------------------------
|
||||
|
||||
def save(self, technique: TechniqueEntity) -> TechniqueEntity:
|
||||
"""Persist a technique entity and return the saved state.
|
||||
|
||||
Args:
|
||||
technique (TechniqueEntity): The entity to create or update.
|
||||
|
||||
Returns:
|
||||
TechniqueEntity: The persisted entity, potentially with updated
|
||||
fields (e.g. server-side timestamps).
|
||||
"""
|
||||
# ...
|
||||
...
|
||||
|
||||
# Define function exists_by_mitre_id
|
||||
def exists_by_mitre_id(self, mitre_id: str) -> bool:
|
||||
"""Return True if a technique with the given MITRE ID exists in the repository.
|
||||
|
||||
Args:
|
||||
mitre_id (str): MITRE ATT&CK ID to check (e.g. ``"T1059"``).
|
||||
|
||||
Returns:
|
||||
bool: True if a matching technique is found, False otherwise.
|
||||
"""
|
||||
# ...
|
||||
...
|
||||
@@ -1,108 +0,0 @@
|
||||
"""Port defining how the application accesses test data.
|
||||
|
||||
This is a domain contract — implementations live in infrastructure/.
|
||||
"""
|
||||
|
||||
# Enable future language features for compatibility
|
||||
from __future__ import annotations
|
||||
|
||||
# Import uuid
|
||||
import uuid
|
||||
|
||||
# Import Protocol from typing
|
||||
from typing import Protocol
|
||||
|
||||
# Import TestState from app.domain.enums
|
||||
from app.domain.enums import TestState
|
||||
|
||||
|
||||
# Define class TestRepository
|
||||
class TestRepository(Protocol):
|
||||
"""Data access contract for tests."""
|
||||
|
||||
# -- Single-entity access ----------------------------------------------
|
||||
|
||||
def find_by_id(self, test_id: uuid.UUID) -> object | None:
|
||||
"""Return a Test ORM model by primary key, or None.
|
||||
|
||||
Returns the ORM model directly (not a domain entity) because
|
||||
the TestEntity is constructed at the service layer via
|
||||
``TestEntity.from_orm()``.
|
||||
|
||||
Args:
|
||||
test_id (uuid.UUID): Primary key of the test to look up.
|
||||
|
||||
Returns:
|
||||
object | None: The ORM model instance, or None if not found.
|
||||
"""
|
||||
# ...
|
||||
...
|
||||
|
||||
# -- List access -------------------------------------------------------
|
||||
|
||||
def list_by_technique(self, technique_id: uuid.UUID) -> list[object]:
|
||||
"""Return all test ORM models associated with the given technique.
|
||||
|
||||
Args:
|
||||
technique_id (uuid.UUID): Primary key of the technique whose tests to retrieve.
|
||||
|
||||
Returns:
|
||||
list[object]: ORM model instances for all tests linked to this technique.
|
||||
"""
|
||||
# ...
|
||||
...
|
||||
|
||||
# Define function list_by_state
|
||||
def list_by_state(self, state: TestState) -> list[object]:
|
||||
"""Return all test ORM models in the given state.
|
||||
|
||||
Args:
|
||||
state (TestState): The state to filter tests by.
|
||||
|
||||
Returns:
|
||||
list[object]: ORM model instances for all tests currently in *state*.
|
||||
"""
|
||||
# ...
|
||||
...
|
||||
|
||||
# Define function count_by_technique_and_state
|
||||
def count_by_technique_and_state(
|
||||
self,
|
||||
# Entry: technique_id
|
||||
technique_id: uuid.UUID,
|
||||
) -> dict[TestState, int]:
|
||||
"""Return test counts grouped by state for a single technique.
|
||||
|
||||
Args:
|
||||
technique_id (uuid.UUID): Primary key of the technique whose test
|
||||
counts to aggregate.
|
||||
|
||||
Returns:
|
||||
dict[TestState, int]: Mapping from each test state to the number of
|
||||
tests in that state for the given technique.
|
||||
"""
|
||||
# ...
|
||||
...
|
||||
|
||||
# -- Batch queries -----------------------------------------------------
|
||||
|
||||
def get_states_and_results_for_technique(
|
||||
self,
|
||||
# Entry: technique_id
|
||||
technique_id: uuid.UUID,
|
||||
) -> list[tuple[str, str | None]]:
|
||||
"""Return (state, detection_result) pairs for all tests of a technique.
|
||||
|
||||
Used by TechniqueEntity.recalculate_status() without loading full
|
||||
test models.
|
||||
|
||||
Args:
|
||||
technique_id (uuid.UUID): Primary key of the technique whose test
|
||||
data to retrieve.
|
||||
|
||||
Returns:
|
||||
list[tuple[str, str | None]]: Each tuple contains the test state
|
||||
string and the detection result string (or None if not yet set).
|
||||
"""
|
||||
# ...
|
||||
...
|
||||
@@ -1,648 +0,0 @@
|
||||
"""TestEntity — pure domain object for the test lifecycle state machine.
|
||||
|
||||
This entity owns ALL state-transition logic and business rules for a
|
||||
security test. It has **no** dependency on FastAPI, SQLAlchemy, or any
|
||||
infrastructure concern.
|
||||
|
||||
Usage::
|
||||
|
||||
entity = TestEntity.from_orm(test_orm_model)
|
||||
entity.start_execution() # draft → red_executing
|
||||
entity.submit_red_evidence() # red_executing → blue_evaluating
|
||||
entity.pause_timer()
|
||||
entity.resume_timer()
|
||||
entity.submit_blue_evidence() # blue_evaluating → in_review
|
||||
entity.validate_red("approved")
|
||||
entity.validate_blue("approved") # triggers dual-validation → validated
|
||||
entity.reopen() # rejected → draft
|
||||
|
||||
After mutations, the service layer copies ``entity.changes`` back onto
|
||||
the ORM model and persists via Unit of Work.
|
||||
"""
|
||||
|
||||
# Enable future language features for compatibility
|
||||
from __future__ import annotations
|
||||
|
||||
# Import enum
|
||||
import enum
|
||||
|
||||
# Import uuid
|
||||
import uuid
|
||||
|
||||
# Import dataclass, field from dataclasses
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
# Import datetime from datetime
|
||||
from datetime import datetime
|
||||
|
||||
# Import TYPE_CHECKING, Any from typing
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
# Import from app.domain.errors
|
||||
from app.domain.errors import (
|
||||
BusinessRuleViolation,
|
||||
InvalidOperationError,
|
||||
InvalidStateTransition,
|
||||
)
|
||||
|
||||
# Check: TYPE_CHECKING
|
||||
if TYPE_CHECKING:
|
||||
# Import Test as TestORM from app.models.test
|
||||
from app.models.test import Test as TestORM
|
||||
|
||||
# ── Value objects ────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestState(str, enum.Enum):
|
||||
"""Ordered lifecycle states for a security test."""
|
||||
|
||||
# Assign draft = "draft"
|
||||
draft = "draft"
|
||||
# Assign red_executing = "red_executing"
|
||||
red_executing = "red_executing"
|
||||
# Assign blue_evaluating = "blue_evaluating"
|
||||
blue_evaluating = "blue_evaluating"
|
||||
# Assign in_review = "in_review"
|
||||
in_review = "in_review"
|
||||
# Assign validated = "validated"
|
||||
validated = "validated"
|
||||
# Assign rejected = "rejected"
|
||||
rejected = "rejected"
|
||||
disputed = "disputed" # one lead approved, the other rejected
|
||||
|
||||
|
||||
# Assign VALID_TRANSITIONS = {
|
||||
VALID_TRANSITIONS: dict[TestState, list[TestState]] = {
|
||||
TestState.draft: [TestState.red_executing],
|
||||
TestState.red_executing: [TestState.blue_evaluating],
|
||||
TestState.blue_evaluating: [TestState.in_review],
|
||||
TestState.in_review: [TestState.validated, TestState.rejected, TestState.disputed],
|
||||
TestState.disputed: [TestState.validated, TestState.rejected],
|
||||
TestState.rejected: [TestState.draft],
|
||||
TestState.validated: [],
|
||||
}
|
||||
|
||||
# Assign _PAUSABLE_STATES = frozenset({TestState.red_executing, TestState.blue_evaluating})
|
||||
_PAUSABLE_STATES = frozenset({TestState.red_executing, TestState.blue_evaluating})
|
||||
|
||||
|
||||
# ── Domain events (lightweight records of what happened) ─────────────
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
# Define class DomainEvent
|
||||
class DomainEvent:
|
||||
"""Immutable record of a domain-level event emitted by the test entity."""
|
||||
|
||||
# name: str
|
||||
name: str
|
||||
# Assign payload = field(default_factory=dict)
|
||||
payload: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
# ── Entity ───────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@dataclass
|
||||
# Define class TestEntity
|
||||
class TestEntity:
|
||||
"""Pure domain representation of a security test."""
|
||||
|
||||
# id: uuid.UUID
|
||||
id: uuid.UUID
|
||||
# state: TestState
|
||||
state: TestState
|
||||
|
||||
# Red validation
|
||||
red_validation_status: str | None = None
|
||||
# Assign red_validated_by = None
|
||||
red_validated_by: uuid.UUID | None = None
|
||||
# Assign red_validated_at = None
|
||||
red_validated_at: datetime | None = None
|
||||
# Assign red_validation_notes = None
|
||||
red_validation_notes: str | None = None
|
||||
|
||||
# Blue validation
|
||||
blue_validation_status: str | None = None
|
||||
# Assign blue_validated_by = None
|
||||
blue_validated_by: uuid.UUID | None = None
|
||||
# Assign blue_validated_at = None
|
||||
blue_validated_at: datetime | None = None
|
||||
# Assign blue_validation_notes = None
|
||||
blue_validation_notes: str | None = None
|
||||
|
||||
# Phase timing
|
||||
execution_date: datetime | None = None
|
||||
# Assign red_started_at = None
|
||||
red_started_at: datetime | None = None
|
||||
# Assign blue_started_at = None
|
||||
blue_started_at: datetime | None = None
|
||||
# Assign paused_at = None
|
||||
paused_at: datetime | None = None
|
||||
# Assign red_paused_seconds = 0
|
||||
red_paused_seconds: int = 0
|
||||
# Assign blue_paused_seconds = 0
|
||||
blue_paused_seconds: int = 0
|
||||
|
||||
# Internal bookkeeping (not persisted as-is)
|
||||
_events: list[DomainEvent] = field(default_factory=list, repr=False)
|
||||
|
||||
# -- Factory --------------------------------------------------------
|
||||
|
||||
@classmethod
|
||||
# Define function from_orm
|
||||
def from_orm(cls, model: TestORM) -> TestEntity:
|
||||
"""Build a TestEntity from a SQLAlchemy ``Test`` model instance.
|
||||
|
||||
Args:
|
||||
model (TestORM): The ORM model whose fields will be copied into the entity.
|
||||
|
||||
Returns:
|
||||
TestEntity: A fully populated domain entity reflecting the ORM state.
|
||||
"""
|
||||
# Assign raw_state = model.state
|
||||
raw_state = model.state
|
||||
# Assign state = raw_state if isinstance(raw_state, TestState) else TestState(raw_st...
|
||||
state = raw_state if isinstance(raw_state, TestState) else TestState(raw_state)
|
||||
# Return cls(
|
||||
return cls(
|
||||
# Keyword argument: id
|
||||
id=model.id,
|
||||
# Keyword argument: state
|
||||
state=state,
|
||||
# Keyword argument: red_validation_status
|
||||
red_validation_status=model.red_validation_status,
|
||||
# Keyword argument: red_validated_by
|
||||
red_validated_by=model.red_validated_by,
|
||||
# Keyword argument: red_validated_at
|
||||
red_validated_at=model.red_validated_at,
|
||||
# Keyword argument: red_validation_notes
|
||||
red_validation_notes=model.red_validation_notes,
|
||||
# Keyword argument: blue_validation_status
|
||||
blue_validation_status=model.blue_validation_status,
|
||||
# Keyword argument: blue_validated_by
|
||||
blue_validated_by=model.blue_validated_by,
|
||||
# Keyword argument: blue_validated_at
|
||||
blue_validated_at=model.blue_validated_at,
|
||||
# Keyword argument: blue_validation_notes
|
||||
blue_validation_notes=model.blue_validation_notes,
|
||||
# Keyword argument: execution_date
|
||||
execution_date=model.execution_date,
|
||||
# Keyword argument: red_started_at
|
||||
red_started_at=model.red_started_at,
|
||||
# Keyword argument: blue_started_at
|
||||
blue_started_at=model.blue_started_at,
|
||||
# Keyword argument: paused_at
|
||||
paused_at=model.paused_at,
|
||||
# Keyword argument: red_paused_seconds
|
||||
red_paused_seconds=model.red_paused_seconds or 0,
|
||||
# Keyword argument: blue_paused_seconds
|
||||
blue_paused_seconds=model.blue_paused_seconds or 0,
|
||||
)
|
||||
|
||||
# Define function apply_to
|
||||
def apply_to(self, model: TestORM) -> None:
|
||||
"""Copy the entity's mutable fields back onto the ORM model.
|
||||
|
||||
Args:
|
||||
model (TestORM): The ORM model to update in-place.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# Assign model.state = self.state
|
||||
model.state = self.state
|
||||
# Assign model.red_validation_status = self.red_validation_status
|
||||
model.red_validation_status = self.red_validation_status
|
||||
# Assign model.red_validated_by = self.red_validated_by
|
||||
model.red_validated_by = self.red_validated_by
|
||||
# Assign model.red_validated_at = self.red_validated_at
|
||||
model.red_validated_at = self.red_validated_at
|
||||
# Assign model.red_validation_notes = self.red_validation_notes
|
||||
model.red_validation_notes = self.red_validation_notes
|
||||
# Assign model.blue_validation_status = self.blue_validation_status
|
||||
model.blue_validation_status = self.blue_validation_status
|
||||
# Assign model.blue_validated_by = self.blue_validated_by
|
||||
model.blue_validated_by = self.blue_validated_by
|
||||
# Assign model.blue_validated_at = self.blue_validated_at
|
||||
model.blue_validated_at = self.blue_validated_at
|
||||
# Assign model.blue_validation_notes = self.blue_validation_notes
|
||||
model.blue_validation_notes = self.blue_validation_notes
|
||||
# Assign model.execution_date = self.execution_date
|
||||
model.execution_date = self.execution_date
|
||||
# Assign model.red_started_at = self.red_started_at
|
||||
model.red_started_at = self.red_started_at
|
||||
# Assign model.blue_started_at = self.blue_started_at
|
||||
model.blue_started_at = self.blue_started_at
|
||||
# Assign model.paused_at = self.paused_at
|
||||
model.paused_at = self.paused_at
|
||||
# Assign model.red_paused_seconds = self.red_paused_seconds
|
||||
model.red_paused_seconds = self.red_paused_seconds
|
||||
# Assign model.blue_paused_seconds = self.blue_paused_seconds
|
||||
model.blue_paused_seconds = self.blue_paused_seconds
|
||||
|
||||
# -- Query helpers --------------------------------------------------
|
||||
|
||||
@property
|
||||
# Define function events
|
||||
def events(self) -> list[DomainEvent]:
|
||||
"""Return a snapshot of all domain events raised on this entity.
|
||||
|
||||
Returns:
|
||||
list[DomainEvent]: Ordered list of events emitted since the entity
|
||||
was constructed or last cleared.
|
||||
"""
|
||||
# Return list(self._events)
|
||||
return list(self._events)
|
||||
|
||||
# Define function can_transition
|
||||
def can_transition(self, target: TestState) -> bool:
|
||||
"""Check whether a transition from the current state to *target* is valid.
|
||||
|
||||
Args:
|
||||
target (TestState): The desired next state.
|
||||
|
||||
Returns:
|
||||
bool: True if the transition is allowed, False otherwise.
|
||||
"""
|
||||
# Return target in VALID_TRANSITIONS.get(self.state, [])
|
||||
return target in VALID_TRANSITIONS.get(self.state, [])
|
||||
|
||||
# Apply the @property decorator
|
||||
@property
|
||||
# Define function is_terminal
|
||||
def is_terminal(self) -> bool:
|
||||
"""Return True if the test has reached its final (validated) state.
|
||||
|
||||
Returns:
|
||||
bool: True when state is ``validated``, False for all other states.
|
||||
"""
|
||||
# Return self.state == TestState.validated
|
||||
return self.state == TestState.validated
|
||||
|
||||
# -- Core transition ------------------------------------------------
|
||||
|
||||
def transition_to(self, target: TestState | str) -> str:
|
||||
"""Validate and apply a state transition.
|
||||
|
||||
Accepts either a :class:`TestState` member or its string value
|
||||
(so callers using ``models.enums.TestState`` work transparently).
|
||||
|
||||
Returns the *previous* state value as a plain string.
|
||||
|
||||
Raises :class:`InvalidStateTransition` when the move is illegal.
|
||||
|
||||
Args:
|
||||
target (TestState | str): The desired next state, as an enum member
|
||||
or its string equivalent.
|
||||
|
||||
Returns:
|
||||
str: The previous state value before the transition.
|
||||
"""
|
||||
# Assign value = target.value if hasattr(target, "value") else str(target)
|
||||
value = target.value if hasattr(target, "value") else str(target)
|
||||
# Assign resolved = target if isinstance(target, TestState) else TestState(value)
|
||||
resolved = target if isinstance(target, TestState) else TestState(value)
|
||||
# Return self._transition(resolved)
|
||||
return self._transition(resolved)
|
||||
|
||||
# Define function _transition
|
||||
def _transition(self, target: TestState) -> str:
|
||||
"""Validate and apply a state transition, returning the previous state value.
|
||||
|
||||
Args:
|
||||
target (TestState): The desired next state enum member.
|
||||
|
||||
Returns:
|
||||
str: The previous state value before the transition was applied.
|
||||
"""
|
||||
# Check: not self.can_transition(target)
|
||||
if not self.can_transition(target):
|
||||
# Assign valid = [s.value for s in VALID_TRANSITIONS.get(self.state, [])]
|
||||
valid = [s.value for s in VALID_TRANSITIONS.get(self.state, [])]
|
||||
# Raise InvalidStateTransition
|
||||
raise InvalidStateTransition(
|
||||
# Keyword argument: current_state
|
||||
current_state=self.state.value,
|
||||
# Keyword argument: target_state
|
||||
target_state=target.value,
|
||||
# Keyword argument: valid_transitions
|
||||
valid_transitions=valid,
|
||||
)
|
||||
# Assign previous = self.state.value
|
||||
previous = self.state.value
|
||||
# Assign self.state = target
|
||||
self.state = target
|
||||
# Call self._events.append()
|
||||
self._events.append(DomainEvent(
|
||||
# Literal argument value
|
||||
"state_changed",
|
||||
{"previous": previous, "new": target.value},
|
||||
))
|
||||
# Return previous
|
||||
return previous
|
||||
|
||||
# -- Lifecycle commands --------------------------------------------
|
||||
|
||||
def start_execution(self) -> None:
|
||||
"""Transition the test from ``draft`` to ``red_executing``.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# Call self._transition()
|
||||
self._transition(TestState.red_executing)
|
||||
# Assign now = datetime.utcnow()
|
||||
now = datetime.utcnow()
|
||||
# Assign self.execution_date = now
|
||||
self.execution_date = now
|
||||
# Assign self.red_started_at = now
|
||||
self.red_started_at = now
|
||||
# Call self._events.append()
|
||||
self._events.append(DomainEvent("execution_started"))
|
||||
|
||||
# Define function submit_red_evidence
|
||||
def submit_red_evidence(self) -> int:
|
||||
"""Transition the test from ``red_executing`` to ``blue_evaluating``.
|
||||
|
||||
Auto-resumes if paused. Returns paused seconds accumulated
|
||||
during this phase (for worklog calculation).
|
||||
|
||||
Returns:
|
||||
int: Total seconds the red phase was paused.
|
||||
"""
|
||||
# Assign paused_extra = self._auto_resume()
|
||||
paused_extra = self._auto_resume()
|
||||
# Call self._transition()
|
||||
self._transition(TestState.blue_evaluating)
|
||||
# Assign total_paused = self.red_paused_seconds + paused_extra
|
||||
total_paused = self.red_paused_seconds + paused_extra
|
||||
# Assign self.blue_started_at = datetime.utcnow()
|
||||
self.blue_started_at = datetime.utcnow()
|
||||
# Assign self.blue_paused_seconds = 0
|
||||
self.blue_paused_seconds = 0
|
||||
# Call self._events.append()
|
||||
self._events.append(DomainEvent(
|
||||
# Literal argument value
|
||||
"red_evidence_submitted",
|
||||
{"red_paused_seconds": total_paused},
|
||||
))
|
||||
# Return total_paused
|
||||
return total_paused
|
||||
|
||||
# Define function submit_blue_evidence
|
||||
def submit_blue_evidence(self) -> int:
|
||||
"""Transition the test from ``blue_evaluating`` to ``in_review``.
|
||||
|
||||
Auto-resumes if paused. Returns paused seconds accumulated
|
||||
during this phase (for worklog calculation).
|
||||
|
||||
Returns:
|
||||
int: Total seconds the blue phase was paused.
|
||||
"""
|
||||
# Assign paused_extra = self._auto_resume()
|
||||
paused_extra = self._auto_resume()
|
||||
# Call self._transition()
|
||||
self._transition(TestState.in_review)
|
||||
# Assign total_paused = self.blue_paused_seconds + paused_extra
|
||||
total_paused = self.blue_paused_seconds + paused_extra
|
||||
# Call self._events.append()
|
||||
self._events.append(DomainEvent(
|
||||
# Literal argument value
|
||||
"blue_evidence_submitted",
|
||||
{"blue_paused_seconds": total_paused},
|
||||
))
|
||||
# Return total_paused
|
||||
return total_paused
|
||||
|
||||
# Define function pause_timer
|
||||
def pause_timer(self) -> None:
|
||||
"""Pause the active phase timer.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# Check: self.state not in _PAUSABLE_STATES
|
||||
if self.state not in _PAUSABLE_STATES:
|
||||
# Raise BusinessRuleViolation
|
||||
raise BusinessRuleViolation(
|
||||
f"Cannot pause timer in '{self.state.value}' state"
|
||||
)
|
||||
# Check: self.paused_at is not None
|
||||
if self.paused_at is not None:
|
||||
# Raise BusinessRuleViolation
|
||||
raise BusinessRuleViolation("Timer is already paused")
|
||||
# Assign self.paused_at = datetime.utcnow()
|
||||
self.paused_at = datetime.utcnow()
|
||||
# Call self._events.append()
|
||||
self._events.append(DomainEvent("timer_paused"))
|
||||
|
||||
# Define function resume_timer
|
||||
def resume_timer(self) -> int:
|
||||
"""Resume a paused timer.
|
||||
|
||||
Returns:
|
||||
int: Number of seconds the timer was paused for.
|
||||
"""
|
||||
# Check: self.paused_at is None
|
||||
if self.paused_at is None:
|
||||
# Raise BusinessRuleViolation
|
||||
raise BusinessRuleViolation("Timer is not paused")
|
||||
# Assign now = datetime.utcnow()
|
||||
now = datetime.utcnow()
|
||||
# Assign paused_seconds = max(int((now - self.paused_at).total_seconds()), 0)
|
||||
paused_seconds = max(int((now - self.paused_at).total_seconds()), 0)
|
||||
# Check: self.state == TestState.red_executing
|
||||
if self.state == TestState.red_executing:
|
||||
# Assign self.red_paused_seconds = paused_seconds
|
||||
self.red_paused_seconds += paused_seconds
|
||||
# Alternative: self.state == TestState.blue_evaluating
|
||||
elif self.state == TestState.blue_evaluating:
|
||||
# Assign self.blue_paused_seconds = paused_seconds
|
||||
self.blue_paused_seconds += paused_seconds
|
||||
# Assign self.paused_at = None
|
||||
self.paused_at = None
|
||||
# Call self._events.append()
|
||||
self._events.append(DomainEvent("timer_resumed", {"paused_seconds": paused_seconds}))
|
||||
# Return paused_seconds
|
||||
return paused_seconds
|
||||
|
||||
# Define function validate_red
|
||||
def validate_red(self, status: str, *, by: uuid.UUID, notes: str | None = None) -> None:
|
||||
"""Record Red Lead's validation decision.
|
||||
|
||||
Args:
|
||||
status (str): Validation outcome; must be ``"approved"`` or ``"rejected"``.
|
||||
by (uuid.UUID): UUID of the Red Lead recording the decision.
|
||||
notes (str | None): Optional free-text notes about the decision.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# Call self._assert_in_review()
|
||||
self._assert_in_review("red")
|
||||
# Call self._assert_valid_vote()
|
||||
self._assert_valid_vote(status)
|
||||
# Assign now = datetime.utcnow()
|
||||
now = datetime.utcnow()
|
||||
# Assign self.red_validation_status = status
|
||||
self.red_validation_status = status
|
||||
# Assign self.red_validated_by = by
|
||||
self.red_validated_by = by
|
||||
# Assign self.red_validated_at = now
|
||||
self.red_validated_at = now
|
||||
# Assign self.red_validation_notes = notes
|
||||
self.red_validation_notes = notes
|
||||
# Call self._events.append()
|
||||
self._events.append(DomainEvent("red_validated", {"status": status}))
|
||||
# Call self._check_dual_validation()
|
||||
self._check_dual_validation()
|
||||
|
||||
# Define function validate_blue
|
||||
def validate_blue(self, status: str, *, by: uuid.UUID, notes: str | None = None) -> None:
|
||||
"""Record Blue Lead's validation decision.
|
||||
|
||||
Args:
|
||||
status (str): Validation outcome; must be ``"approved"`` or ``"rejected"``.
|
||||
by (uuid.UUID): UUID of the Blue Lead recording the decision.
|
||||
notes (str | None): Optional free-text notes about the decision.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# Call self._assert_in_review()
|
||||
self._assert_in_review("blue")
|
||||
# Call self._assert_valid_vote()
|
||||
self._assert_valid_vote(status)
|
||||
# Assign now = datetime.utcnow()
|
||||
now = datetime.utcnow()
|
||||
# Assign self.blue_validation_status = status
|
||||
self.blue_validation_status = status
|
||||
# Assign self.blue_validated_by = by
|
||||
self.blue_validated_by = by
|
||||
# Assign self.blue_validated_at = now
|
||||
self.blue_validated_at = now
|
||||
# Assign self.blue_validation_notes = notes
|
||||
self.blue_validation_notes = notes
|
||||
# Call self._events.append()
|
||||
self._events.append(DomainEvent("blue_validated", {"status": status}))
|
||||
# Call self._check_dual_validation()
|
||||
self._check_dual_validation()
|
||||
|
||||
# Define function reopen
|
||||
def reopen(self) -> None:
|
||||
"""Transition the test from ``rejected`` back to ``draft``, clearing all validation and timing fields.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# Call self._transition()
|
||||
self._transition(TestState.draft)
|
||||
# Assign self.red_validation_status = None
|
||||
self.red_validation_status = None
|
||||
# Assign self.red_validated_by = None
|
||||
self.red_validated_by = None
|
||||
# Assign self.red_validated_at = None
|
||||
self.red_validated_at = None
|
||||
# Assign self.red_validation_notes = None
|
||||
self.red_validation_notes = None
|
||||
# Assign self.blue_validation_status = None
|
||||
self.blue_validation_status = None
|
||||
# Assign self.blue_validated_by = None
|
||||
self.blue_validated_by = None
|
||||
# Assign self.blue_validated_at = None
|
||||
self.blue_validated_at = None
|
||||
# Assign self.blue_validation_notes = None
|
||||
self.blue_validation_notes = None
|
||||
# Assign self.red_started_at = None
|
||||
self.red_started_at = None
|
||||
# Assign self.blue_started_at = None
|
||||
self.blue_started_at = None
|
||||
# Assign self.paused_at = None
|
||||
self.paused_at = None
|
||||
# Assign self.red_paused_seconds = 0
|
||||
self.red_paused_seconds = 0
|
||||
# Assign self.blue_paused_seconds = 0
|
||||
self.blue_paused_seconds = 0
|
||||
# Call self._events.append()
|
||||
self._events.append(DomainEvent("test_reopened"))
|
||||
|
||||
# -- Private -------------------------------------------------------
|
||||
|
||||
def _auto_resume(self) -> int:
|
||||
"""Accumulate pause time and clear the paused timestamp if currently paused.
|
||||
|
||||
Returns:
|
||||
int: Extra seconds that were accumulated from the current pause, or 0
|
||||
if the timer was not paused.
|
||||
"""
|
||||
# Check: self.paused_at is None
|
||||
if self.paused_at is None:
|
||||
# Return 0
|
||||
return 0
|
||||
# Assign now = datetime.utcnow()
|
||||
now = datetime.utcnow()
|
||||
# Assign extra = max(int((now - self.paused_at).total_seconds()), 0)
|
||||
extra = max(int((now - self.paused_at).total_seconds()), 0)
|
||||
# Assign self.paused_at = None
|
||||
self.paused_at = None
|
||||
# Return extra
|
||||
return extra
|
||||
|
||||
# Define function check_dual_validation
|
||||
def check_dual_validation(self) -> None:
|
||||
"""Evaluate both leads' votes and advance state if appropriate.
|
||||
|
||||
Rules (v2 — consensus required):
|
||||
- Both **approved** -> ``validated``
|
||||
- Both **rejected** -> ``rejected``
|
||||
- One approved + one rejected -> ``disputed`` (conflict, needs discussion)
|
||||
- Otherwise (one or both still pending) -> no change
|
||||
|
||||
Called automatically by :meth:`validate_red` and :meth:`validate_blue`.
|
||||
"""
|
||||
# Call self._check_dual_validation()
|
||||
self._check_dual_validation()
|
||||
|
||||
# Define function _assert_in_review
|
||||
def _assert_in_review(self, side: str) -> None:
|
||||
if self.state not in (TestState.in_review, TestState.disputed):
|
||||
raise InvalidOperationError(
|
||||
f"Cannot validate {side} side while test is in "
|
||||
f"'{self.state.value}' state (must be in_review or disputed)"
|
||||
)
|
||||
|
||||
# Apply the @staticmethod decorator
|
||||
@staticmethod
|
||||
# Define function _assert_valid_vote
|
||||
def _assert_valid_vote(status: str) -> None:
|
||||
"""Raise InvalidOperationError if *status* is not a valid vote value.
|
||||
|
||||
Args:
|
||||
status (str): The vote value to validate; must be ``"approved"`` or ``"rejected"``.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# Check: status not in ("approved", "rejected")
|
||||
if status not in ("approved", "rejected"):
|
||||
# Raise InvalidOperationError
|
||||
raise InvalidOperationError(
|
||||
# Literal argument value
|
||||
"validation_status must be 'approved' or 'rejected'"
|
||||
)
|
||||
|
||||
# Define function _check_dual_validation
|
||||
def _check_dual_validation(self) -> None:
|
||||
"""Advance the test state once both leads have voted."""
|
||||
r, b = self.red_validation_status, self.blue_validation_status
|
||||
|
||||
if r == "approved" and b == "approved":
|
||||
self.state = TestState.validated
|
||||
# Call self._events.append()
|
||||
self._events.append(DomainEvent("dual_validation_approved"))
|
||||
|
||||
elif r == "rejected" or b == "rejected":
|
||||
# Any rejection is a veto — one lead can reject without waiting for the other
|
||||
self.state = TestState.rejected
|
||||
self._events.append(DomainEvent("dual_validation_rejected"))
|
||||
@@ -1,103 +0,0 @@
|
||||
"""Unit of Work — wraps a SQLAlchemy session for explicit transaction control.
|
||||
|
||||
Usage in routers::
|
||||
|
||||
with UnitOfWork(db) as uow:
|
||||
service_a(db, ...)
|
||||
service_b(db, ...)
|
||||
uow.commit() # single commit for the entire operation
|
||||
|
||||
If an exception propagates, ``__exit__`` issues a rollback automatically.
|
||||
Services should **never** call ``db.commit()``; they use ``db.add()`` /
|
||||
``db.flush()`` to stage work and let the caller decide when to commit.
|
||||
|
||||
**Documented exceptions** (services that may commit internally):
|
||||
- Import services (atomic_import, sigma_import, etc.) — self-contained sync ops.
|
||||
- Background jobs (campaign_scheduler, intel_service, stale_detection,
|
||||
mitre_sync) — self-contained operations.
|
||||
- Self-contained batch ops (e.g. detection_rule_service.auto_associate_rules,
|
||||
snapshot_service.create_snapshot, campaign_service.generate_campaign_from_*,
|
||||
osint_enrichment_service.enrich_technique_with_cves).
|
||||
"""
|
||||
|
||||
# Enable future language features for compatibility
|
||||
from __future__ import annotations
|
||||
|
||||
# Import TracebackType from types
|
||||
from types import TracebackType
|
||||
|
||||
# Import Session from sqlalchemy.orm
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
|
||||
# Define class UnitOfWork
|
||||
class UnitOfWork:
|
||||
"""Lightweight transaction wrapper around an existing SQLAlchemy session."""
|
||||
|
||||
# Define function __init__
|
||||
def __init__(self, session: Session) -> None:
|
||||
"""Wrap an existing SQLAlchemy session in a Unit of Work.
|
||||
|
||||
Args:
|
||||
session (Session): The active SQLAlchemy session to manage.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# Assign self._session = session
|
||||
self._session = session
|
||||
|
||||
# -- context manager -----------------------------------------------------
|
||||
|
||||
def __enter__(self) -> "UnitOfWork":
|
||||
"""Enter the runtime context, returning this UnitOfWork instance.
|
||||
|
||||
Returns:
|
||||
UnitOfWork: The UnitOfWork itself, for use in ``with`` statements.
|
||||
"""
|
||||
# Return self
|
||||
return self
|
||||
|
||||
# Define function __exit__
|
||||
def __exit__(
|
||||
self,
|
||||
# Entry: exc_type
|
||||
exc_type: type[BaseException] | None,
|
||||
# Entry: exc_val
|
||||
exc_val: BaseException | None,
|
||||
# Entry: exc_tb
|
||||
exc_tb: TracebackType | None,
|
||||
) -> None:
|
||||
"""Exit the runtime context, rolling back if an exception propagated.
|
||||
|
||||
Args:
|
||||
exc_type (type[BaseException] | None): Exception class, if raised.
|
||||
exc_val (BaseException | None): Exception instance, if raised.
|
||||
exc_tb (TracebackType | None): Traceback object, if an exception was raised.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# Check: exc_type is not None
|
||||
if exc_type is not None:
|
||||
# Call self.rollback()
|
||||
self.rollback()
|
||||
|
||||
# -- public API ----------------------------------------------------------
|
||||
|
||||
def commit(self) -> None:
|
||||
"""Flush pending changes and commit the transaction."""
|
||||
# Call self._session.commit()
|
||||
self._session.commit()
|
||||
|
||||
# Define function rollback
|
||||
def rollback(self) -> None:
|
||||
"""Roll back the current transaction."""
|
||||
# Call self._session.rollback()
|
||||
self._session.rollback()
|
||||
|
||||
# Define function flush
|
||||
def flush(self) -> None:
|
||||
"""Flush pending changes without committing (useful for getting IDs)."""
|
||||
# Call self._session.flush()
|
||||
self._session.flush()
|
||||
@@ -1,9 +0,0 @@
|
||||
"""Immutable domain value objects."""
|
||||
# Import MitreId from app.domain.value_objects.mitre_id
|
||||
from app.domain.value_objects.mitre_id import MitreId
|
||||
|
||||
# Import ScoringWeights from app.domain.value_objects.scoring_weights
|
||||
from app.domain.value_objects.scoring_weights import ScoringWeights
|
||||
|
||||
# Assign __all__ = ["MitreId", "ScoringWeights"]
|
||||
__all__ = ["MitreId", "ScoringWeights"]
|
||||
@@ -1,115 +0,0 @@
|
||||
"""MitreId — validated MITRE ATT&CK technique identifier.
|
||||
|
||||
Immutable value object that ensures the identifier follows the ATT&CK
|
||||
format: ``T`` followed by 4 digits, optionally a dot and 3 more digits
|
||||
for sub-techniques (e.g. ``T1059``, ``T1059.001``).
|
||||
"""
|
||||
|
||||
# Enable future language features for compatibility
|
||||
from __future__ import annotations
|
||||
|
||||
# Import re
|
||||
import re
|
||||
|
||||
# Import dataclass from dataclasses
|
||||
from dataclasses import dataclass
|
||||
|
||||
# Assign _MITRE_ID_RE = re.compile(r"^T\d{4}(\.\d{3})?$")
|
||||
_MITRE_ID_RE = re.compile(r"^T\d{4}(\.\d{3})?$")
|
||||
|
||||
|
||||
# Apply the @dataclass decorator
|
||||
@dataclass(frozen=True, slots=True)
|
||||
# Define class MitreId
|
||||
class MitreId:
|
||||
"""Validated MITRE ATT&CK technique identifier."""
|
||||
|
||||
# value: str
|
||||
value: str
|
||||
|
||||
# Define function __post_init__
|
||||
def __post_init__(self) -> None:
|
||||
"""Validate that *value* matches the expected MITRE ATT&CK ID format.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# Check: not _MITRE_ID_RE.match(self.value)
|
||||
if not _MITRE_ID_RE.match(self.value):
|
||||
# Raise ValueError
|
||||
raise ValueError(
|
||||
f"Invalid MITRE ATT&CK ID '{self.value}'. "
|
||||
# Literal argument value
|
||||
"Expected format: T1234 or T1234.001"
|
||||
)
|
||||
|
||||
# Apply the @property decorator
|
||||
@property
|
||||
# Define function is_subtechnique
|
||||
def is_subtechnique(self) -> bool:
|
||||
"""Return True if this identifier represents a sub-technique.
|
||||
|
||||
Returns:
|
||||
bool: True when the ID contains a dot (e.g. ``T1059.001``).
|
||||
"""
|
||||
# Return "." in self.value
|
||||
return "." in self.value
|
||||
|
||||
# Apply the @property decorator
|
||||
@property
|
||||
# Define function parent_id
|
||||
def parent_id(self) -> str | None:
|
||||
"""Return the parent technique ID (e.g. ``T1059`` for ``T1059.001``).
|
||||
|
||||
Returns:
|
||||
str | None: The parent ID string, or None if this is not a sub-technique.
|
||||
"""
|
||||
# Check: not self.is_subtechnique
|
||||
if not self.is_subtechnique:
|
||||
# Return None
|
||||
return None
|
||||
# Return self.value.split(".")[0]
|
||||
return self.value.split(".")[0]
|
||||
|
||||
# Define function __str__
|
||||
def __str__(self) -> str:
|
||||
"""Return the string representation of the MITRE ID.
|
||||
|
||||
Returns:
|
||||
str: The raw identifier string (e.g. ``"T1059.001"``).
|
||||
"""
|
||||
# Return self.value
|
||||
return self.value
|
||||
|
||||
# Define function __eq__
|
||||
def __eq__(self, other: object) -> bool:
|
||||
"""Compare this MitreId to another MitreId or a plain string.
|
||||
|
||||
Args:
|
||||
other (object): The value to compare against; may be a
|
||||
:class:`MitreId` instance or a plain ``str``.
|
||||
|
||||
Returns:
|
||||
bool: True if the identifiers are equal, NotImplemented for
|
||||
unsupported types.
|
||||
"""
|
||||
# Check: isinstance(other, MitreId)
|
||||
if isinstance(other, MitreId):
|
||||
# Return self.value == other.value
|
||||
return self.value == other.value
|
||||
# Check: isinstance(other, str)
|
||||
if isinstance(other, str):
|
||||
# Return self.value == other
|
||||
return self.value == other
|
||||
# Return NotImplemented
|
||||
return NotImplemented
|
||||
|
||||
# Define function __hash__
|
||||
def __hash__(self) -> int:
|
||||
"""Return the hash of the identifier string.
|
||||
|
||||
Returns:
|
||||
int: Hash value derived from the raw identifier string.
|
||||
"""
|
||||
# Return hash(self.value)
|
||||
return hash(self.value)
|
||||
@@ -1,107 +0,0 @@
|
||||
"""ScoringWeights — validated immutable weight set for the scoring engine.
|
||||
|
||||
Enforces that all five weights are non-negative and sum to exactly 100.
|
||||
"""
|
||||
|
||||
# Enable future language features for compatibility
|
||||
from __future__ import annotations
|
||||
|
||||
# Import dataclass from dataclasses
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
# Apply the @dataclass decorator
|
||||
@dataclass(frozen=True, slots=True)
|
||||
# Define class ScoringWeights
|
||||
class ScoringWeights:
|
||||
"""Five scoring dimension weights that must sum to 100."""
|
||||
|
||||
# tests: float
|
||||
tests: float
|
||||
# detection_rules: float
|
||||
detection_rules: float
|
||||
# d3fend: float
|
||||
d3fend: float
|
||||
# recency: float
|
||||
recency: float
|
||||
# severity: float
|
||||
severity: float
|
||||
|
||||
# Define function __post_init__
|
||||
def __post_init__(self) -> None:
|
||||
"""Validate that all weights are non-negative and sum to exactly 100.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# Assign fields = [
|
||||
fields = [
|
||||
self.tests,
|
||||
self.detection_rules,
|
||||
self.d3fend,
|
||||
self.recency,
|
||||
self.severity,
|
||||
]
|
||||
# Iterate over fields
|
||||
for f in fields:
|
||||
# Check: f < 0
|
||||
if f < 0:
|
||||
# Raise ValueError
|
||||
raise ValueError("Scoring weights must be non-negative")
|
||||
|
||||
# Assign total = sum(fields)
|
||||
total = sum(fields)
|
||||
# Check: abs(total - 100) > 0.01
|
||||
if abs(total - 100) > 0.01:
|
||||
# Raise ValueError
|
||||
raise ValueError(
|
||||
f"Scoring weights must sum to 100, got {total}"
|
||||
)
|
||||
|
||||
# Apply the @classmethod decorator
|
||||
@classmethod
|
||||
# Define function default
|
||||
def default(cls) -> ScoringWeights:
|
||||
"""Return the default weight distribution.
|
||||
|
||||
Returns:
|
||||
ScoringWeights: A weight set with tests=40, detection_rules=25,
|
||||
d3fend=15, recency=10, severity=10.
|
||||
"""
|
||||
# Return cls(
|
||||
return cls(
|
||||
# Keyword argument: tests
|
||||
tests=40.0,
|
||||
# Keyword argument: detection_rules
|
||||
detection_rules=25.0,
|
||||
# Keyword argument: d3fend
|
||||
d3fend=15.0,
|
||||
# Keyword argument: recency
|
||||
recency=10.0,
|
||||
# Keyword argument: severity
|
||||
severity=10.0,
|
||||
)
|
||||
|
||||
# Backward-compatible aliases for older API payloads
|
||||
@property
|
||||
# Define function freshness
|
||||
def freshness(self) -> float:
|
||||
"""Return the recency weight (backward-compatible alias).
|
||||
|
||||
Returns:
|
||||
float: The value of the ``recency`` weight.
|
||||
"""
|
||||
# Return self.recency
|
||||
return self.recency
|
||||
|
||||
# Apply the @property decorator
|
||||
@property
|
||||
# Define function platform_diversity
|
||||
def platform_diversity(self) -> float:
|
||||
"""Return the severity weight (backward-compatible alias).
|
||||
|
||||
Returns:
|
||||
float: The value of the ``severity`` weight.
|
||||
"""
|
||||
# Return self.severity
|
||||
return self.severity
|
||||
@@ -1 +0,0 @@
|
||||
"""Infrastructure adapters — persistence, caching, and external services."""
|
||||
@@ -1 +0,0 @@
|
||||
"""SQLAlchemy-based persistence adapters for the domain repository ports."""
|
||||
@@ -1 +0,0 @@
|
||||
"""ORM-to-domain entity mapper functions."""
|
||||
@@ -1,28 +0,0 @@
|
||||
"""Technique ORM model <-> domain entity mapper."""
|
||||
|
||||
# Enable future language features for compatibility
|
||||
from __future__ import annotations
|
||||
|
||||
# Import TechniqueEntity from app.domain.entities.technique
|
||||
from app.domain.entities.technique import TechniqueEntity
|
||||
|
||||
|
||||
# Define class TechniqueMapper
|
||||
class TechniqueMapper:
|
||||
"""Converts between SQLAlchemy Technique model and TechniqueEntity."""
|
||||
|
||||
# Apply the @staticmethod decorator
|
||||
@staticmethod
|
||||
# Define function to_entity
|
||||
def to_entity(model: object) -> TechniqueEntity:
|
||||
"""Convert an ORM Technique model to a domain TechniqueEntity."""
|
||||
# Return TechniqueEntity.from_orm(model)
|
||||
return TechniqueEntity.from_orm(model)
|
||||
|
||||
# Apply the @staticmethod decorator
|
||||
@staticmethod
|
||||
# Define function to_model_updates
|
||||
def to_model_updates(entity: TechniqueEntity, model: object) -> None:
|
||||
"""Apply entity changes back onto an existing ORM model."""
|
||||
# Call entity.apply_to()
|
||||
entity.apply_to(model)
|
||||
@@ -1,13 +0,0 @@
|
||||
"""Concrete SQLAlchemy repository implementations."""
|
||||
# Import from app.infrastructure.persistence.repositories.sa_technique_repository
|
||||
from app.infrastructure.persistence.repositories.sa_technique_repository import (
|
||||
SATechniqueRepository,
|
||||
)
|
||||
|
||||
# Import from app.infrastructure.persistence.repositories.sa_test_repository
|
||||
from app.infrastructure.persistence.repositories.sa_test_repository import (
|
||||
SATestRepository,
|
||||
)
|
||||
|
||||
# Assign __all__ = ["SATechniqueRepository", "SATestRepository"]
|
||||
__all__ = ["SATechniqueRepository", "SATestRepository"]
|
||||
@@ -1,380 +0,0 @@
|
||||
"""SQLAlchemy implementation of TechniqueRepository.
|
||||
|
||||
Receives a Session from the caller — does NOT create its own.
|
||||
Does NOT call commit() — the Unit of Work owns that.
|
||||
"""
|
||||
|
||||
# Enable future language features for compatibility
|
||||
from __future__ import annotations
|
||||
|
||||
# Import uuid
|
||||
import uuid
|
||||
|
||||
# Import func from sqlalchemy
|
||||
from sqlalchemy import func
|
||||
|
||||
# Import Session from sqlalchemy.orm
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
# Import TechniqueEntity from app.domain.entities.technique
|
||||
from app.domain.entities.technique import TechniqueEntity
|
||||
|
||||
# Import TechniqueStatus, TestState from app.domain.enums
|
||||
from app.domain.enums import TechniqueStatus, TestState
|
||||
|
||||
# Import TechniqueWithCounts from app.domain.ports.repositories.technique_repository
|
||||
from app.domain.ports.repositories.technique_repository import TechniqueWithCounts
|
||||
|
||||
# Import TechniqueMapper from app.infrastructure.persistence.mappers.technique_mapper
|
||||
from app.infrastructure.persistence.mappers.technique_mapper import TechniqueMapper
|
||||
|
||||
# Import DetectionRule from app.models.detection_rule
|
||||
from app.models.detection_rule import DetectionRule
|
||||
|
||||
# Import Technique from app.models.technique
|
||||
from app.models.technique import Technique
|
||||
|
||||
# Import Test from app.models.test
|
||||
from app.models.test import Test
|
||||
|
||||
|
||||
# Define class SATechniqueRepository
|
||||
class SATechniqueRepository:
|
||||
"""Concrete repository backed by SQLAlchemy."""
|
||||
|
||||
# Define function __init__
|
||||
def __init__(self, session: Session) -> None:
|
||||
"""Initialise the repository with a caller-provided session.
|
||||
|
||||
Args:
|
||||
session (Session): The SQLAlchemy session to use for all queries.
|
||||
"""
|
||||
# Assign self._session = session
|
||||
self._session = session
|
||||
|
||||
# -- Single-entity access ----------------------------------------------
|
||||
|
||||
def find_by_id(self, technique_id: uuid.UUID) -> TechniqueEntity | None:
|
||||
"""Return a single technique by its primary key.
|
||||
|
||||
Args:
|
||||
technique_id (uuid.UUID): The UUID primary key of the technique.
|
||||
|
||||
Returns:
|
||||
TechniqueEntity | None: The matching entity, or ``None`` if not found.
|
||||
"""
|
||||
# Assign model = (
|
||||
model = (
|
||||
self._session.query(Technique)
|
||||
# Chain .filter() call
|
||||
.filter(Technique.id == technique_id)
|
||||
# Chain .first() call
|
||||
.first()
|
||||
)
|
||||
# Return TechniqueMapper.to_entity(model) if model else None
|
||||
return TechniqueMapper.to_entity(model) if model else None
|
||||
|
||||
# Define function find_by_mitre_id
|
||||
def find_by_mitre_id(self, mitre_id: str) -> TechniqueEntity | None:
|
||||
"""Return a single technique by its MITRE ATT&CK ID (e.g. ``T1059.001``).
|
||||
|
||||
Args:
|
||||
mitre_id (str): The MITRE ATT&CK identifier string.
|
||||
|
||||
Returns:
|
||||
TechniqueEntity | None: The matching entity, or ``None`` if not found.
|
||||
"""
|
||||
# Assign model = (
|
||||
model = (
|
||||
self._session.query(Technique)
|
||||
# Chain .filter() call
|
||||
.filter(Technique.mitre_id == mitre_id)
|
||||
# Chain .first() call
|
||||
.first()
|
||||
)
|
||||
# Return TechniqueMapper.to_entity(model) if model else None
|
||||
return TechniqueMapper.to_entity(model) if model else None
|
||||
|
||||
# -- List access -------------------------------------------------------
|
||||
|
||||
def list_all(
|
||||
self,
|
||||
*,
|
||||
# Entry: tactic
|
||||
tactic: str | None = None,
|
||||
# Entry: status
|
||||
status: TechniqueStatus | None = None,
|
||||
# Entry: review_required
|
||||
review_required: bool | None = None,
|
||||
) -> list[TechniqueEntity]:
|
||||
"""Return all techniques, optionally filtered by tactic, status, or review flag.
|
||||
|
||||
Args:
|
||||
tactic (str | None): Filter to techniques belonging to this tactic name.
|
||||
status (TechniqueStatus | None): Filter to techniques with this coverage status.
|
||||
review_required (bool | None): Filter to techniques where ``review_required`` matches.
|
||||
|
||||
Returns:
|
||||
list[TechniqueEntity]: Ordered list of matching technique entities.
|
||||
"""
|
||||
# Assign query = self._session.query(Technique)
|
||||
query = self._session.query(Technique)
|
||||
# Check: tactic is not None
|
||||
if tactic is not None:
|
||||
# Assign query = query.filter(Technique.tactic == tactic)
|
||||
query = query.filter(Technique.tactic == tactic)
|
||||
# Check: status is not None
|
||||
if status is not None:
|
||||
# Assign query = query.filter(Technique.status_global == status)
|
||||
query = query.filter(Technique.status_global == status)
|
||||
# Check: review_required is not None
|
||||
if review_required is not None:
|
||||
# Assign query = query.filter(Technique.review_required == review_required)
|
||||
query = query.filter(Technique.review_required == review_required)
|
||||
# Assign models = query.order_by(Technique.mitre_id).all()
|
||||
models = query.order_by(Technique.mitre_id).all()
|
||||
# Return [TechniqueMapper.to_entity(m) for m in models]
|
||||
return [TechniqueMapper.to_entity(m) for m in models]
|
||||
|
||||
# Define function list_by_ids
|
||||
def list_by_ids(self, ids: list[uuid.UUID]) -> list[TechniqueEntity]:
|
||||
"""Return techniques matching the provided list of UUIDs.
|
||||
|
||||
Args:
|
||||
ids (list[uuid.UUID]): UUIDs of the techniques to retrieve.
|
||||
|
||||
Returns:
|
||||
list[TechniqueEntity]: Technique entities corresponding to the given IDs.
|
||||
"""
|
||||
# Check: not ids
|
||||
if not ids:
|
||||
# Return []
|
||||
return []
|
||||
# Assign models = (
|
||||
models = (
|
||||
self._session.query(Technique)
|
||||
# Chain .filter() call
|
||||
.filter(Technique.id.in_(ids))
|
||||
# Chain .all() call
|
||||
.all()
|
||||
)
|
||||
# Return [TechniqueMapper.to_entity(m) for m in models]
|
||||
return [TechniqueMapper.to_entity(m) for m in models]
|
||||
|
||||
# -- Batch queries (for scoring/heatmap) -------------------------------
|
||||
|
||||
def count_by_status(self) -> dict[TechniqueStatus, int]:
|
||||
"""Return a count of techniques grouped by their coverage status.
|
||||
|
||||
Returns:
|
||||
dict[TechniqueStatus, int]: Mapping of each status value to its technique count.
|
||||
"""
|
||||
# Assign rows = (
|
||||
rows = (
|
||||
self._session.query(
|
||||
Technique.status_global,
|
||||
func.count(Technique.id),
|
||||
)
|
||||
# Chain .group_by() call
|
||||
.group_by(Technique.status_global)
|
||||
# Chain .all() call
|
||||
.all()
|
||||
)
|
||||
# Assign result = {s: 0 for s in TechniqueStatus}
|
||||
result = {s: 0 for s in TechniqueStatus}
|
||||
# Iterate over rows
|
||||
for status_val, count in rows:
|
||||
# Assign key = (
|
||||
key = (
|
||||
status_val
|
||||
if isinstance(status_val, TechniqueStatus)
|
||||
else TechniqueStatus(status_val)
|
||||
)
|
||||
# Assign result[key] = count
|
||||
result[key] = count
|
||||
# Return result
|
||||
return result
|
||||
|
||||
# Define function find_all_with_test_counts
|
||||
def find_all_with_test_counts(self) -> list[TechniqueWithCounts]:
|
||||
"""Return all techniques with pre-aggregated test and detection rule counts.
|
||||
|
||||
Uses a single query with subqueries to avoid the N+1 pattern.
|
||||
|
||||
Returns:
|
||||
list[TechniqueWithCounts]: All techniques with their associated counts.
|
||||
"""
|
||||
# Assign test_count_sq = (
|
||||
test_count_sq = (
|
||||
self._session.query(
|
||||
Test.technique_id,
|
||||
func.count(Test.id).label("test_count"),
|
||||
func.sum(
|
||||
func.cast(Test.state == TestState.validated, self._int_type())
|
||||
).label("validated_count"),
|
||||
)
|
||||
# Chain .group_by() call
|
||||
.group_by(Test.technique_id)
|
||||
# Chain .subquery() call
|
||||
.subquery()
|
||||
)
|
||||
# Assign rule_count_sq = (
|
||||
rule_count_sq = (
|
||||
self._session.query(
|
||||
DetectionRule.mitre_technique_id,
|
||||
func.count(DetectionRule.id).label("rule_count"),
|
||||
)
|
||||
# Chain .group_by() call
|
||||
.group_by(DetectionRule.mitre_technique_id)
|
||||
# Chain .subquery() call
|
||||
.subquery()
|
||||
)
|
||||
|
||||
# Assign rows = (
|
||||
rows = (
|
||||
self._session.query(
|
||||
Technique,
|
||||
func.coalesce(test_count_sq.c.test_count, 0),
|
||||
func.coalesce(test_count_sq.c.validated_count, 0),
|
||||
func.coalesce(rule_count_sq.c.rule_count, 0),
|
||||
)
|
||||
# Chain .outerjoin() call
|
||||
.outerjoin(test_count_sq, Technique.id == test_count_sq.c.technique_id)
|
||||
# Chain .outerjoin() call
|
||||
.outerjoin(
|
||||
rule_count_sq,
|
||||
Technique.mitre_id == rule_count_sq.c.mitre_technique_id,
|
||||
)
|
||||
# Chain .order_by() call
|
||||
.order_by(Technique.mitre_id)
|
||||
# Chain .all() call
|
||||
.all()
|
||||
)
|
||||
|
||||
# Return [
|
||||
return [
|
||||
TechniqueWithCounts(
|
||||
# Keyword argument: entity
|
||||
entity=TechniqueMapper.to_entity(tech),
|
||||
# Keyword argument: test_count
|
||||
test_count=int(tc),
|
||||
# Keyword argument: validated_test_count
|
||||
validated_test_count=int(vtc),
|
||||
# Keyword argument: detection_rule_count
|
||||
detection_rule_count=int(rc),
|
||||
)
|
||||
for tech, tc, vtc, rc in rows
|
||||
]
|
||||
|
||||
# -- Mutations ---------------------------------------------------------
|
||||
|
||||
def save(self, technique: TechniqueEntity) -> TechniqueEntity:
|
||||
"""Persist a technique entity, inserting or updating as needed.
|
||||
|
||||
Args:
|
||||
technique (TechniqueEntity): The domain entity to persist.
|
||||
|
||||
Returns:
|
||||
TechniqueEntity: The persisted entity reflecting the current DB state.
|
||||
"""
|
||||
# Assign existing = (
|
||||
existing = (
|
||||
self._session.query(Technique)
|
||||
# Chain .filter() call
|
||||
.filter(Technique.id == technique.id)
|
||||
# Chain .first() call
|
||||
.first()
|
||||
)
|
||||
# Check: existing
|
||||
if existing:
|
||||
# Call technique.apply_to()
|
||||
technique.apply_to(existing)
|
||||
# Assign existing.mitre_id = technique.mitre_id
|
||||
existing.mitre_id = technique.mitre_id
|
||||
# Assign existing.name = technique.name
|
||||
existing.name = technique.name
|
||||
# Assign existing.tactic = technique.tactic
|
||||
existing.tactic = technique.tactic
|
||||
# Assign existing.description = technique.description
|
||||
existing.description = technique.description
|
||||
# Assign existing.platforms = technique.platforms
|
||||
existing.platforms = technique.platforms
|
||||
# Assign existing.is_subtechnique = technique.is_subtechnique
|
||||
existing.is_subtechnique = technique.is_subtechnique
|
||||
# Assign existing.parent_mitre_id = technique.parent_mitre_id
|
||||
existing.parent_mitre_id = technique.parent_mitre_id
|
||||
# Assign existing.mitre_version = technique.mitre_version
|
||||
existing.mitre_version = technique.mitre_version
|
||||
# Assign existing.mitre_last_modified = technique.mitre_last_modified
|
||||
existing.mitre_last_modified = technique.mitre_last_modified
|
||||
# Call self._session.flush()
|
||||
self._session.flush()
|
||||
# Return TechniqueMapper.to_entity(existing)
|
||||
return TechniqueMapper.to_entity(existing)
|
||||
# Fallback: handle remaining cases
|
||||
else:
|
||||
# Assign model = Technique(
|
||||
model = Technique(
|
||||
# Keyword argument: id
|
||||
id=technique.id,
|
||||
# Keyword argument: mitre_id
|
||||
mitre_id=technique.mitre_id,
|
||||
# Keyword argument: name
|
||||
name=technique.name,
|
||||
# Keyword argument: tactic
|
||||
tactic=technique.tactic,
|
||||
# Keyword argument: description
|
||||
description=technique.description,
|
||||
# Keyword argument: platforms
|
||||
platforms=technique.platforms,
|
||||
# Keyword argument: is_subtechnique
|
||||
is_subtechnique=technique.is_subtechnique,
|
||||
# Keyword argument: parent_mitre_id
|
||||
parent_mitre_id=technique.parent_mitre_id,
|
||||
# Keyword argument: status_global
|
||||
status_global=technique.status_global,
|
||||
# Keyword argument: review_required
|
||||
review_required=technique.review_required,
|
||||
# Keyword argument: last_review_date
|
||||
last_review_date=technique.last_review_date,
|
||||
# Keyword argument: mitre_version
|
||||
mitre_version=technique.mitre_version,
|
||||
# Keyword argument: mitre_last_modified
|
||||
mitre_last_modified=technique.mitre_last_modified,
|
||||
)
|
||||
# Call self._session.add()
|
||||
self._session.add(model)
|
||||
# Call self._session.flush()
|
||||
self._session.flush()
|
||||
# Return TechniqueMapper.to_entity(model)
|
||||
return TechniqueMapper.to_entity(model)
|
||||
|
||||
# Define function exists_by_mitre_id
|
||||
def exists_by_mitre_id(self, mitre_id: str) -> bool:
|
||||
"""Check whether a technique with the given MITRE ID already exists.
|
||||
|
||||
Args:
|
||||
mitre_id (str): The MITRE ATT&CK identifier to look up.
|
||||
|
||||
Returns:
|
||||
bool: ``True`` if the technique exists, ``False`` otherwise.
|
||||
"""
|
||||
# Return (
|
||||
return (
|
||||
self._session.query(Technique.id)
|
||||
# Chain .filter() call
|
||||
.filter(Technique.mitre_id == mitre_id)
|
||||
# Chain .first() call
|
||||
.first()
|
||||
) is not None
|
||||
|
||||
# -- Internal ----------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
# Define function _int_type
|
||||
def _int_type() -> type:
|
||||
"""Return an Integer type for CAST expressions (SQLite-compatible)."""
|
||||
# Import Integer from sqlalchemy
|
||||
from sqlalchemy import Integer
|
||||
# Return Integer
|
||||
return Integer
|
||||
@@ -1,171 +0,0 @@
|
||||
"""SQLAlchemy implementation of TestRepository."""
|
||||
|
||||
# Enable future language features for compatibility
|
||||
from __future__ import annotations
|
||||
|
||||
# Import uuid
|
||||
import uuid
|
||||
|
||||
# Import func from sqlalchemy
|
||||
from sqlalchemy import func
|
||||
|
||||
# Import Session from sqlalchemy.orm
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
# Import TestState from app.domain.enums
|
||||
from app.domain.enums import TestState
|
||||
|
||||
# Import Test from app.models.test
|
||||
from app.models.test import Test
|
||||
|
||||
|
||||
# Define class SATestRepository
|
||||
class SATestRepository:
|
||||
"""Concrete test repository backed by SQLAlchemy."""
|
||||
|
||||
# Define function __init__
|
||||
def __init__(self, session: Session) -> None:
|
||||
"""Initialise the repository with a caller-provided session.
|
||||
|
||||
Args:
|
||||
session (Session): The SQLAlchemy session to use for all queries.
|
||||
"""
|
||||
# Assign self._session = session
|
||||
self._session = session
|
||||
|
||||
# Define function find_by_id
|
||||
def find_by_id(self, test_id: uuid.UUID) -> Test | None:
|
||||
"""Return a single test by its primary key.
|
||||
|
||||
Args:
|
||||
test_id (uuid.UUID): The UUID primary key of the test.
|
||||
|
||||
Returns:
|
||||
Test | None: The ORM model instance, or ``None`` if not found.
|
||||
"""
|
||||
# Return (
|
||||
return (
|
||||
self._session.query(Test)
|
||||
# Chain .filter() call
|
||||
.filter(Test.id == test_id)
|
||||
# Chain .first() call
|
||||
.first()
|
||||
)
|
||||
|
||||
# Define function list_by_technique
|
||||
def list_by_technique(self, technique_id: uuid.UUID) -> list[Test]:
|
||||
"""Return all tests for a given technique, ordered by creation date.
|
||||
|
||||
Args:
|
||||
technique_id (uuid.UUID): The UUID of the parent technique.
|
||||
|
||||
Returns:
|
||||
list[Test]: ORM model instances ordered by ``created_at`` ascending.
|
||||
"""
|
||||
# Return (
|
||||
return (
|
||||
self._session.query(Test)
|
||||
# Chain .filter() call
|
||||
.filter(Test.technique_id == technique_id)
|
||||
# Chain .order_by() call
|
||||
.order_by(Test.created_at)
|
||||
# Chain .all() call
|
||||
.all()
|
||||
)
|
||||
|
||||
# Define function list_by_state
|
||||
def list_by_state(self, state: TestState) -> list[Test]:
|
||||
"""Return all tests that are currently in the given workflow state.
|
||||
|
||||
Args:
|
||||
state (TestState): The workflow state to filter on.
|
||||
|
||||
Returns:
|
||||
list[Test]: All ORM model instances with the specified state.
|
||||
"""
|
||||
# Return (
|
||||
return (
|
||||
self._session.query(Test)
|
||||
# Chain .filter() call
|
||||
.filter(Test.state == state)
|
||||
# Chain .all() call
|
||||
.all()
|
||||
)
|
||||
|
||||
# Define function count_by_technique_and_state
|
||||
def count_by_technique_and_state(
|
||||
self,
|
||||
# Entry: technique_id
|
||||
technique_id: uuid.UUID,
|
||||
) -> dict[TestState, int]:
|
||||
"""Return per-state test counts for a specific technique.
|
||||
|
||||
Args:
|
||||
technique_id (uuid.UUID): The UUID of the technique to aggregate for.
|
||||
|
||||
Returns:
|
||||
dict[TestState, int]: Mapping of each state to the number of tests in that state.
|
||||
"""
|
||||
# Assign rows = (
|
||||
rows = (
|
||||
self._session.query(Test.state, func.count(Test.id))
|
||||
# Chain .filter() call
|
||||
.filter(Test.technique_id == technique_id)
|
||||
# Chain .group_by() call
|
||||
.group_by(Test.state)
|
||||
# Chain .all() call
|
||||
.all()
|
||||
)
|
||||
# Assign result = {}
|
||||
result: dict[TestState, int] = {}
|
||||
# Iterate over rows
|
||||
for state_val, count in rows:
|
||||
# Assign key = (
|
||||
key = (
|
||||
state_val
|
||||
if isinstance(state_val, TestState)
|
||||
else TestState(state_val)
|
||||
)
|
||||
# Assign result[key] = count
|
||||
result[key] = count
|
||||
# Return result
|
||||
return result
|
||||
|
||||
# Define function get_states_and_results_for_technique
|
||||
def get_states_and_results_for_technique(
|
||||
self,
|
||||
# Entry: technique_id
|
||||
technique_id: uuid.UUID,
|
||||
) -> list[tuple[str, str | None]]:
|
||||
"""Return lightweight ``(state, detection_result)`` pairs for a technique.
|
||||
|
||||
Used by ``TechniqueEntity.recalculate_status()`` to avoid loading full
|
||||
``Test`` models.
|
||||
|
||||
Args:
|
||||
technique_id (uuid.UUID): The UUID of the technique to query.
|
||||
|
||||
Returns:
|
||||
list[tuple[str, str | None]]: Each tuple contains the state string
|
||||
and the detection result string (or ``None``).
|
||||
"""
|
||||
# Assign rows = (
|
||||
rows = (
|
||||
self._session.query(Test.state, Test.detection_result)
|
||||
# Chain .filter() call
|
||||
.filter(Test.technique_id == technique_id)
|
||||
# Chain .all() call
|
||||
.all()
|
||||
)
|
||||
# Return [
|
||||
return [
|
||||
(
|
||||
r.state.value if hasattr(r.state, "value") else str(r.state),
|
||||
(
|
||||
r.detection_result.value
|
||||
if hasattr(r.detection_result, "value")
|
||||
else r.detection_result
|
||||
),
|
||||
)
|
||||
for r in rows
|
||||
]
|
||||
@@ -1,91 +0,0 @@
|
||||
"""Redis client factories.
|
||||
|
||||
``settings.REDIS_URL`` selects the default logical database (usually ``0``).
|
||||
Token blacklist and application cache use separate logical DBs on the same
|
||||
Redis instance (``REDIS_TOKEN_BLACKLIST_DB``, ``REDIS_CACHE_DB``) so keys never
|
||||
collide and TTL policies can differ per workload.
|
||||
|
||||
Usage::
|
||||
|
||||
from app.infrastructure.redis_client import get_redis, get_redis_blacklist
|
||||
|
||||
get_redis().set("key", "value", ex=300)
|
||||
get_redis_blacklist().setex("blacklist:…", ttl, "1")
|
||||
"""
|
||||
|
||||
# Enable future language features for compatibility
|
||||
from __future__ import annotations
|
||||
|
||||
# Import logging
|
||||
import logging
|
||||
|
||||
# Import urlparse, urlunparse from urllib.parse
|
||||
from urllib.parse import urlparse, urlunparse
|
||||
|
||||
# Import redis
|
||||
import redis
|
||||
|
||||
# Import settings from app.config
|
||||
from app.config import settings
|
||||
|
||||
# Assign logger = logging.getLogger(__name__)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Assign _clients = {}
|
||||
_clients: dict[str, redis.Redis] = {}
|
||||
|
||||
|
||||
# Define function _redis_url_with_db
|
||||
def _redis_url_with_db(base_url: str, db_index: int) -> str:
|
||||
"""Return *base_url* with its path replaced by ``/{db_index}``."""
|
||||
# Assign parsed = urlparse(base_url)
|
||||
parsed = urlparse(base_url)
|
||||
# Assign path = f"/{db_index}"
|
||||
path = f"/{db_index}"
|
||||
# Return urlunparse(
|
||||
return urlunparse(
|
||||
(parsed.scheme, parsed.netloc, path, "", "", ""),
|
||||
)
|
||||
|
||||
|
||||
# Define function _get_client
|
||||
def _get_client(url: str) -> redis.Redis:
|
||||
# Check: url not in _clients
|
||||
if url not in _clients:
|
||||
# Assign _clients[url] = redis.from_url(url, decode_responses=True)
|
||||
_clients[url] = redis.from_url(url, decode_responses=True)
|
||||
# Log info: "Redis client connected to %s", url
|
||||
logger.info("Redis client connected to %s", url)
|
||||
# Return _clients[url]
|
||||
return _clients[url]
|
||||
|
||||
|
||||
# Define function get_redis
|
||||
def get_redis() -> redis.Redis:
|
||||
"""Default Redis connection (URL from ``settings.REDIS_URL``)."""
|
||||
# Return _get_client(settings.REDIS_URL)
|
||||
return _get_client(settings.REDIS_URL)
|
||||
|
||||
|
||||
# Define function get_redis_blacklist
|
||||
def get_redis_blacklist() -> redis.Redis:
|
||||
"""Redis DB used for JWT revocation (``jti`` keys with TTL)."""
|
||||
# Assign url = _redis_url_with_db(
|
||||
url = _redis_url_with_db(
|
||||
settings.REDIS_URL,
|
||||
settings.REDIS_TOKEN_BLACKLIST_DB,
|
||||
)
|
||||
# Return _get_client(url)
|
||||
return _get_client(url)
|
||||
|
||||
|
||||
# Define function get_redis_cache
|
||||
def get_redis_cache() -> redis.Redis:
|
||||
"""Redis DB reserved for shared cache (scores, queues, etc.)."""
|
||||
# Assign url = _redis_url_with_db(
|
||||
url = _redis_url_with_db(
|
||||
settings.REDIS_URL,
|
||||
settings.REDIS_CACHE_DB,
|
||||
)
|
||||
# Return _get_client(url)
|
||||
return _get_client(url)
|
||||
@@ -1 +0,0 @@
|
||||
"""Background scheduler jobs (MITRE sync, Jira sync, data retention)."""
|
||||
|
||||
@@ -1,65 +0,0 @@
|
||||
"""Scheduled job — syncs all Jira links hourly."""
|
||||
|
||||
# Import logging
|
||||
import logging
|
||||
|
||||
# Import settings from app.config
|
||||
from app.config import settings
|
||||
|
||||
# Import SessionLocal from app.database
|
||||
from app.database import SessionLocal
|
||||
|
||||
# Import JiraLink from app.models.jira_link
|
||||
from app.models.jira_link import JiraLink
|
||||
|
||||
# Import jira_service from app.services
|
||||
from app.services import jira_service
|
||||
|
||||
# Assign logger = logging.getLogger(__name__)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Define function sync_all_jira_links
|
||||
def sync_all_jira_links() -> None:
|
||||
"""Pull latest status from Jira for every stored link.
|
||||
|
||||
Silently skips if ``JIRA_ENABLED`` is ``False``. Individual link
|
||||
failures are logged but do not abort the rest of the batch.
|
||||
"""
|
||||
# Check: not settings.JIRA_ENABLED
|
||||
if not settings.JIRA_ENABLED:
|
||||
# Return control to caller
|
||||
return
|
||||
|
||||
# Assign db = SessionLocal()
|
||||
db = SessionLocal()
|
||||
# Attempt the following; catch errors below
|
||||
try:
|
||||
# Assign links = db.query(JiraLink).all()
|
||||
links = db.query(JiraLink).all()
|
||||
# Assign synced = 0
|
||||
synced = 0
|
||||
# Iterate over links
|
||||
for link in links:
|
||||
# Attempt the following; catch errors below
|
||||
try:
|
||||
# Call jira_service.sync_jira_to_aegis()
|
||||
jira_service.sync_jira_to_aegis(db, link)
|
||||
# Assign synced = 1
|
||||
synced += 1
|
||||
# Handle Exception
|
||||
except Exception as e:
|
||||
# Log warning: "Jira sync failed for link %s: %s", link.id, e
|
||||
logger.warning("Jira sync failed for link %s: %s", link.id, e)
|
||||
# Commit all pending changes to the database
|
||||
db.commit()
|
||||
# Log info: "Jira sync completed: %d/%d links updated", synced
|
||||
logger.info("Jira sync completed: %d/%d links updated", synced, len(links))
|
||||
# Handle Exception
|
||||
except Exception:
|
||||
# Log exception: "Jira sync batch job failed"
|
||||
logger.exception("Jira sync batch job failed")
|
||||
# Always execute this cleanup block
|
||||
finally:
|
||||
# Close the database session
|
||||
db.close()
|
||||
@@ -10,44 +10,14 @@ Each job manages its own database session (created on entry, closed in
|
||||
sessions.
|
||||
"""
|
||||
|
||||
# Import logging
|
||||
import logging
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
# Import BackgroundScheduler from apscheduler.schedulers.background
|
||||
from apscheduler.schedulers.background import BackgroundScheduler
|
||||
|
||||
# Import SessionLocal from app.database
|
||||
from app.database import SessionLocal
|
||||
|
||||
# Import sync_all_jira_links from app.jobs.jira_sync_job
|
||||
from app.jobs.jira_sync_job import sync_all_jira_links
|
||||
|
||||
# Import run_retention_job from app.jobs.retention_job
|
||||
from app.jobs.retention_job import run_retention_job
|
||||
|
||||
# Import check_and_run_recurring_campaigns from app.services.campaign_scheduler_service
|
||||
from app.services.campaign_scheduler_service import check_and_run_recurring_campaigns
|
||||
|
||||
# Import scan_intel from app.services.intel_service
|
||||
from app.services.mitre_sync_service import sync_mitre
|
||||
from app.services.intel_service import scan_intel
|
||||
|
||||
# Import sync_mitre from app.services.mitre_sync_service
|
||||
from app.services.mitre_sync_service import sync_mitre
|
||||
|
||||
# Import cleanup_old_notifications from app.services.notification_service
|
||||
from app.services.notification_service import cleanup_old_notifications
|
||||
|
||||
# Import enrich_all_techniques from app.services.osint_enrichment_service
|
||||
from app.services.osint_enrichment_service import enrich_all_techniques
|
||||
|
||||
# Import cleanup_old_snapshots, create_snapshot from app.services.snapshot_service
|
||||
from app.services.snapshot_service import cleanup_old_snapshots, create_snapshot
|
||||
|
||||
# Import detect_stale_coverage from app.services.stale_detection_service
|
||||
from app.services.stale_detection_service import detect_stale_coverage
|
||||
|
||||
# Assign logger = logging.getLogger(__name__)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -64,437 +34,26 @@ scheduler = BackgroundScheduler()
|
||||
|
||||
def _run_mitre_sync() -> None:
|
||||
"""Execute a MITRE sync inside its own DB session."""
|
||||
from app.services.webhook_service import dispatch_webhook
|
||||
logger.info("Scheduled MITRE sync job starting...")
|
||||
# Assign db = SessionLocal()
|
||||
db = SessionLocal()
|
||||
# Attempt the following; catch errors below
|
||||
try:
|
||||
# Assign summary = sync_mitre(db)
|
||||
summary = sync_mitre(db)
|
||||
# Log info: "Scheduled MITRE sync job finished — %s", summary
|
||||
logger.info("Scheduled MITRE sync job finished — %s", summary)
|
||||
dispatch_webhook("mitre.synced", {"created": summary.get("created", 0), "updated": summary.get("updated", 0)})
|
||||
except Exception:
|
||||
# Log exception: "Scheduled MITRE sync job failed"
|
||||
logger.exception("Scheduled MITRE sync job failed")
|
||||
# Always execute this cleanup block
|
||||
finally:
|
||||
# Close the database session
|
||||
db.close()
|
||||
|
||||
|
||||
# Define function _run_notification_cleanup
|
||||
def _run_notification_cleanup() -> None:
|
||||
"""Clean up old read notifications."""
|
||||
# Log info: "Scheduled notification cleanup job starting..."
|
||||
logger.info("Scheduled notification cleanup job starting...")
|
||||
# Assign db = SessionLocal()
|
||||
db = SessionLocal()
|
||||
# Attempt the following; catch errors below
|
||||
try:
|
||||
# Assign deleted = cleanup_old_notifications(db, days=90)
|
||||
deleted = cleanup_old_notifications(db, days=90)
|
||||
# Log info: "Notification cleanup finished — deleted %d old no
|
||||
logger.info("Notification cleanup finished — deleted %d old notifications", deleted)
|
||||
# Handle Exception
|
||||
except Exception:
|
||||
# Log exception: "Notification cleanup job failed"
|
||||
logger.exception("Notification cleanup job failed")
|
||||
# Always execute this cleanup block
|
||||
finally:
|
||||
# Close the database session
|
||||
db.close()
|
||||
|
||||
|
||||
# Define function _run_weekly_snapshot
|
||||
def _run_weekly_snapshot() -> None:
|
||||
"""Create a weekly coverage snapshot and clean up old ones."""
|
||||
# Log info: "Scheduled weekly snapshot job starting..."
|
||||
logger.info("Scheduled weekly snapshot job starting...")
|
||||
# Assign db = SessionLocal()
|
||||
db = SessionLocal()
|
||||
# Attempt the following; catch errors below
|
||||
try:
|
||||
# Assign snapshot = create_snapshot(db, name="Auto-weekly")
|
||||
snapshot = create_snapshot(db, name="Auto-weekly")
|
||||
# Log info:
|
||||
logger.info(
|
||||
# Literal argument value
|
||||
"Weekly snapshot created — score %.1f, %d techniques",
|
||||
snapshot.organization_score,
|
||||
snapshot.total_techniques,
|
||||
)
|
||||
# Assign deleted = cleanup_old_snapshots(db, keep_last=52)
|
||||
deleted = cleanup_old_snapshots(db, keep_last=52)
|
||||
# Check: deleted
|
||||
if deleted:
|
||||
# Log info: "Cleaned up %d old snapshots", deleted
|
||||
logger.info("Cleaned up %d old snapshots", deleted)
|
||||
# Handle Exception
|
||||
except Exception:
|
||||
# Log exception: "Weekly snapshot job failed"
|
||||
logger.exception("Weekly snapshot job failed")
|
||||
# Always execute this cleanup block
|
||||
finally:
|
||||
# Close the database session
|
||||
db.close()
|
||||
|
||||
|
||||
# Define function _run_recurring_campaigns
|
||||
def _run_recurring_campaigns() -> None:
|
||||
"""Check and run any due recurring campaigns."""
|
||||
# Log info: "Scheduled recurring campaigns check starting..."
|
||||
logger.info("Scheduled recurring campaigns check starting...")
|
||||
# Assign db = SessionLocal()
|
||||
db = SessionLocal()
|
||||
# Attempt the following; catch errors below
|
||||
try:
|
||||
# Assign spawned = check_and_run_recurring_campaigns(db)
|
||||
spawned = check_and_run_recurring_campaigns(db)
|
||||
# Log info: "Recurring campaigns check finished — spawned %d c
|
||||
logger.info("Recurring campaigns check finished — spawned %d campaigns", spawned)
|
||||
# Handle Exception
|
||||
except Exception:
|
||||
# Log exception: "Recurring campaigns check failed"
|
||||
logger.exception("Recurring campaigns check failed")
|
||||
# Always execute this cleanup block
|
||||
finally:
|
||||
# Close the database session
|
||||
db.close()
|
||||
|
||||
|
||||
def _run_scheduled_campaign_activation() -> None:
|
||||
"""Auto-activate campaigns whose start_date has arrived.
|
||||
|
||||
Finds all campaigns in 'draft' state with a start_date <= now,
|
||||
activates them, creates Jira tickets, and notifies the red_tech team.
|
||||
Runs every hour so campaigns activate within ~1 hour of their scheduled time.
|
||||
"""
|
||||
logger.info("Scheduled campaign auto-activation check starting...")
|
||||
db = SessionLocal()
|
||||
try:
|
||||
from datetime import datetime as _dt
|
||||
from app.models.campaign import Campaign
|
||||
from app.models.user import User
|
||||
from app.services.campaign_crud_service import activate_campaign as _activate
|
||||
from app.services.notification_service import notify_role
|
||||
from app.services.audit_service import log_action
|
||||
|
||||
now = _dt.utcnow()
|
||||
due_campaigns = (
|
||||
db.query(Campaign)
|
||||
.filter(
|
||||
Campaign.status == "draft",
|
||||
Campaign.start_date != None, # noqa: E711
|
||||
Campaign.start_date <= now,
|
||||
)
|
||||
.all()
|
||||
)
|
||||
|
||||
activated = 0
|
||||
for campaign in due_campaigns:
|
||||
try:
|
||||
_activate(db, str(campaign.id))
|
||||
notify_role(
|
||||
db,
|
||||
role="red_tech",
|
||||
type="campaign_activated",
|
||||
title="Campaign auto-activated",
|
||||
message=f'Campaign "{campaign.name}" has been automatically activated on its scheduled start date.',
|
||||
entity_type="campaign",
|
||||
entity_id=campaign.id,
|
||||
)
|
||||
log_action(
|
||||
db,
|
||||
user_id=None,
|
||||
action="auto_activate_campaign",
|
||||
entity_type="campaign",
|
||||
entity_id=campaign.id,
|
||||
details={"name": campaign.name, "start_date": str(campaign.start_date)},
|
||||
)
|
||||
|
||||
# Create Jira tickets non-fatally
|
||||
try:
|
||||
from app.services.jira_service import (
|
||||
auto_create_campaign_issue,
|
||||
auto_create_test_issue,
|
||||
get_campaign_jira_key,
|
||||
get_test_jira_key,
|
||||
)
|
||||
# Use first admin user as actor for Jira auth
|
||||
admin_user = db.query(User).filter(User.role == "admin").first()
|
||||
if admin_user:
|
||||
db.refresh(campaign)
|
||||
campaign_jira_key = get_campaign_jira_key(db, str(campaign.id))
|
||||
if not campaign_jira_key:
|
||||
campaign_jira_key = auto_create_campaign_issue(db, campaign, admin_user)
|
||||
if campaign_jira_key:
|
||||
for ct in campaign.campaign_tests:
|
||||
if ct.test and not get_test_jira_key(db, ct.test.id):
|
||||
auto_create_test_issue(
|
||||
db, ct.test, admin_user,
|
||||
parent_ticket_override=campaign_jira_key,
|
||||
campaign_start_date=campaign.start_date,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Jira auto-create failed for auto-activated campaign %s", campaign.id)
|
||||
|
||||
db.commit()
|
||||
activated += 1
|
||||
logger.info("Auto-activated campaign %s (%s)", campaign.id, campaign.name)
|
||||
except Exception:
|
||||
logger.exception("Failed to auto-activate campaign %s", campaign.id)
|
||||
db.rollback()
|
||||
|
||||
logger.info("Campaign auto-activation check finished — activated %d campaigns", activated)
|
||||
except Exception:
|
||||
logger.exception("Campaign auto-activation job failed")
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
|
||||
def _run_intel_scan() -> None:
|
||||
"""Execute an intel scan inside its own DB session."""
|
||||
# Log info: "Scheduled intel scan job starting..."
|
||||
logger.info("Scheduled intel scan job starting...")
|
||||
# Assign db = SessionLocal()
|
||||
db = SessionLocal()
|
||||
# Attempt the following; catch errors below
|
||||
try:
|
||||
# Assign summary = scan_intel(db)
|
||||
summary = scan_intel(db)
|
||||
# Log info: "Scheduled intel scan job finished — %s", summary
|
||||
logger.info("Scheduled intel scan job finished — %s", summary)
|
||||
# Handle Exception
|
||||
except Exception:
|
||||
# Log exception: "Scheduled intel scan job failed"
|
||||
logger.exception("Scheduled intel scan job failed")
|
||||
# Always execute this cleanup block
|
||||
finally:
|
||||
# Close the database session
|
||||
db.close()
|
||||
|
||||
|
||||
def _run_evaluation_round_check() -> None:
|
||||
"""Weekly job: check if a new ATT&CK Evaluation round is available.
|
||||
|
||||
If a new round is found it is imported automatically and an admin
|
||||
notification is created so the team knows new baseline data is available.
|
||||
"""
|
||||
logger.info("ATT&CK Evaluations new-round check starting...")
|
||||
db = SessionLocal()
|
||||
try:
|
||||
from app.services.attck_evaluations_service import check_for_new_round, import_evaluation_round
|
||||
from app.models.user import User as UserModel
|
||||
|
||||
result = check_for_new_round(db)
|
||||
if result.get("error"):
|
||||
logger.warning("ATT&CK Evaluations check failed: %s", result["error"])
|
||||
return
|
||||
|
||||
if not result.get("new_round_available"):
|
||||
logger.info(
|
||||
"ATT&CK Evaluations check — latest round '%s' already imported",
|
||||
result.get("latest_round", {}).get("display_name", "?"),
|
||||
)
|
||||
return
|
||||
|
||||
latest = result["latest_round"]
|
||||
logger.info(
|
||||
"New ATT&CK Evaluation round detected: %s (round %d) — starting auto-import",
|
||||
latest["display_name"], latest["eval_round"],
|
||||
)
|
||||
|
||||
# Use the first admin user as the importer (system action)
|
||||
admin = db.query(UserModel).filter(UserModel.role == "admin").first()
|
||||
if not admin:
|
||||
logger.warning("ATT&CK Evaluations auto-import: no admin user found — skipping")
|
||||
return
|
||||
|
||||
summary = import_evaluation_round(
|
||||
db,
|
||||
latest["name"],
|
||||
latest["display_name"],
|
||||
latest["eval_round"],
|
||||
admin,
|
||||
)
|
||||
logger.info(
|
||||
"ATT&CK Evaluations auto-import complete — round %d (%s): %d tests created",
|
||||
latest["eval_round"], latest["display_name"], summary["created"],
|
||||
)
|
||||
|
||||
# Notify all admins
|
||||
try:
|
||||
from app.services.notification_service import create_notification
|
||||
admins = db.query(UserModel).filter(UserModel.role == "admin").all()
|
||||
for adm in admins:
|
||||
create_notification(
|
||||
db,
|
||||
user_id=adm.id,
|
||||
title="New ATT&CK Evaluation round imported",
|
||||
message=(
|
||||
f"Round {latest['eval_round']} — {latest['display_name']} — "
|
||||
f"has been automatically imported. "
|
||||
f"{summary['created']} tests created in In Review state. "
|
||||
f"Blue Leads must validate each result before it counts as coverage."
|
||||
),
|
||||
notification_type="eval_import",
|
||||
entity_type="evaluation",
|
||||
entity_id=None,
|
||||
)
|
||||
db.commit()
|
||||
except Exception:
|
||||
logger.warning("Failed to send eval import notifications", exc_info=True)
|
||||
|
||||
except Exception:
|
||||
logger.exception("ATT&CK Evaluations round check job failed")
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
|
||||
def _run_osint_enrichment() -> None:
|
||||
"""Execute weekly OSINT enrichment inside its own DB session."""
|
||||
# Log info: "Scheduled OSINT enrichment job starting..."
|
||||
logger.info("Scheduled OSINT enrichment job starting...")
|
||||
# Assign db = SessionLocal()
|
||||
db = SessionLocal()
|
||||
# Attempt the following; catch errors below
|
||||
try:
|
||||
# Assign total = enrich_all_techniques(db)
|
||||
total = enrich_all_techniques(db)
|
||||
# Log info: "OSINT enrichment finished — %d new items", total
|
||||
logger.info("OSINT enrichment finished — %d new items", total)
|
||||
# Handle Exception
|
||||
except Exception:
|
||||
# Log exception: "OSINT enrichment job failed"
|
||||
logger.exception("OSINT enrichment job failed")
|
||||
# Always execute this cleanup block
|
||||
finally:
|
||||
# Close the database session
|
||||
db.close()
|
||||
|
||||
|
||||
_FREQUENCY_INTERVALS: dict[str, timedelta] = {
|
||||
"daily": timedelta(days=1),
|
||||
"weekly": timedelta(weeks=1),
|
||||
"monthly": timedelta(days=30),
|
||||
}
|
||||
|
||||
|
||||
def _run_data_sources_sync() -> None:
|
||||
"""Check all enabled data sources and sync those that are overdue."""
|
||||
logger.info("Scheduled data sources sync check starting...")
|
||||
db = SessionLocal()
|
||||
try:
|
||||
from app.models.data_source import DataSource
|
||||
from app.services.data_source_service import sync_source
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
sources = (
|
||||
db.query(DataSource)
|
||||
.filter(DataSource.is_enabled == True) # noqa: E712
|
||||
.all()
|
||||
)
|
||||
synced = 0
|
||||
for ds in sources:
|
||||
freq = ds.sync_frequency
|
||||
if not freq or freq == "manual":
|
||||
continue
|
||||
interval = _FREQUENCY_INTERVALS.get(freq)
|
||||
if interval is None:
|
||||
continue
|
||||
last = ds.last_sync_at
|
||||
if last is None:
|
||||
# Never synced — run it now
|
||||
overdue = True
|
||||
else:
|
||||
# Make last timezone-aware if needed
|
||||
if last.tzinfo is None:
|
||||
last = last.replace(tzinfo=timezone.utc)
|
||||
overdue = now - last >= interval
|
||||
if overdue:
|
||||
logger.info(
|
||||
"Data source '%s' is overdue (freq=%s, last=%s) — syncing",
|
||||
ds.name, freq, last,
|
||||
)
|
||||
try:
|
||||
sync_source(db, str(ds.id))
|
||||
synced += 1
|
||||
except Exception:
|
||||
logger.exception("Failed to sync data source '%s'", ds.name)
|
||||
logger.info("Data sources sync check finished — %d source(s) synced", synced)
|
||||
except Exception:
|
||||
logger.exception("Data sources sync check failed")
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
|
||||
def _run_stale_detection() -> None:
|
||||
"""Execute daily stale coverage detection inside its own DB session."""
|
||||
# Log info: "Scheduled stale coverage detection starting..."
|
||||
logger.info("Scheduled stale coverage detection starting...")
|
||||
# Assign db = SessionLocal()
|
||||
db = SessionLocal()
|
||||
# Attempt the following; catch errors below
|
||||
try:
|
||||
# Assign count = detect_stale_coverage(db)
|
||||
count = detect_stale_coverage(db)
|
||||
# Log info: "Stale detection finished — %d techniques flagged"
|
||||
logger.info("Stale detection finished — %d techniques flagged", count)
|
||||
# Handle Exception
|
||||
except Exception:
|
||||
# Log exception: "Stale coverage detection job failed"
|
||||
logger.exception("Stale coverage detection job failed")
|
||||
# Always execute this cleanup block
|
||||
finally:
|
||||
# Close the database session
|
||||
db.close()
|
||||
|
||||
|
||||
def _run_decay_engine() -> None:
|
||||
"""Execute the decay engine inside its own DB session."""
|
||||
logger.info("Scheduled decay engine job starting...")
|
||||
db = SessionLocal()
|
||||
try:
|
||||
from app.services.decay_engine_service import run_decay_engine
|
||||
results = run_decay_engine(db)
|
||||
logger.info("Decay engine job finished — %s", results)
|
||||
except Exception:
|
||||
logger.exception("Decay engine job failed")
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
|
||||
def _run_queue_generation() -> None:
|
||||
"""Generate revalidation queue items for analysts — runs after decay engine."""
|
||||
logger.info("Scheduled revalidation queue generation starting...")
|
||||
db = SessionLocal()
|
||||
try:
|
||||
from app.services.revalidation_queue_service import generate_queue_items
|
||||
results = generate_queue_items(db)
|
||||
logger.info("Queue generation finished — %s", results)
|
||||
except Exception:
|
||||
logger.exception("Queue generation job failed")
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
|
||||
def _run_alert_evaluation() -> None:
|
||||
"""Evaluate all enabled operational alert rules (hourly)."""
|
||||
logger.info("Scheduled alert evaluation job starting...")
|
||||
db = SessionLocal()
|
||||
try:
|
||||
from app.services.operational_alert_service import evaluate_all_rules
|
||||
result = evaluate_all_rules(db)
|
||||
logger.info(
|
||||
"Alert evaluation finished — %d rules, %d alerts fired in %.3fs",
|
||||
result["rules_evaluated"],
|
||||
result["alerts_fired"],
|
||||
result["duration_seconds"],
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Alert evaluation job failed")
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
@@ -514,199 +73,21 @@ def start_scheduler() -> None:
|
||||
|
||||
Neither job fires immediately on startup.
|
||||
"""
|
||||
# Call scheduler.add_job()
|
||||
scheduler.add_job(
|
||||
_run_mitre_sync,
|
||||
# Keyword argument: trigger
|
||||
trigger="interval",
|
||||
# Keyword argument: hours
|
||||
hours=24,
|
||||
# Keyword argument: id
|
||||
id="mitre_sync",
|
||||
# Keyword argument: name
|
||||
name="MITRE ATT&CK sync (every 24h)",
|
||||
# Keyword argument: replace_existing
|
||||
replace_existing=True,
|
||||
)
|
||||
# Call scheduler.add_job()
|
||||
scheduler.add_job(
|
||||
_run_intel_scan,
|
||||
# Keyword argument: trigger
|
||||
trigger="interval",
|
||||
# Keyword argument: weeks
|
||||
weeks=1,
|
||||
# Keyword argument: id
|
||||
id="intel_scan",
|
||||
# Keyword argument: name
|
||||
name="Intel scan (every 7d)",
|
||||
# Keyword argument: replace_existing
|
||||
replace_existing=True,
|
||||
)
|
||||
# Call scheduler.add_job()
|
||||
scheduler.add_job(
|
||||
_run_notification_cleanup,
|
||||
# Keyword argument: trigger
|
||||
trigger="interval",
|
||||
# Keyword argument: hours
|
||||
hours=24,
|
||||
# Keyword argument: id
|
||||
id="notification_cleanup",
|
||||
# Keyword argument: name
|
||||
name="Notification cleanup (daily)",
|
||||
# Keyword argument: replace_existing
|
||||
replace_existing=True,
|
||||
)
|
||||
# Call scheduler.add_job()
|
||||
scheduler.add_job(
|
||||
_run_weekly_snapshot,
|
||||
# Keyword argument: trigger
|
||||
trigger="cron",
|
||||
# Keyword argument: day_of_week
|
||||
day_of_week="sun",
|
||||
# Keyword argument: hour
|
||||
hour=0,
|
||||
# Keyword argument: minute
|
||||
minute=0,
|
||||
# Keyword argument: id
|
||||
id="weekly_snapshot",
|
||||
# Keyword argument: name
|
||||
name="Weekly coverage snapshot (Sundays 00:00)",
|
||||
# Keyword argument: replace_existing
|
||||
replace_existing=True,
|
||||
)
|
||||
# Call scheduler.add_job()
|
||||
scheduler.add_job(
|
||||
_run_scheduled_campaign_activation,
|
||||
trigger="interval",
|
||||
hours=1,
|
||||
id="scheduled_campaign_activation",
|
||||
name="Auto-activate campaigns on start_date (hourly)",
|
||||
replace_existing=True,
|
||||
)
|
||||
scheduler.add_job(
|
||||
_run_recurring_campaigns,
|
||||
# Keyword argument: trigger
|
||||
trigger="interval",
|
||||
# Keyword argument: hours
|
||||
hours=24,
|
||||
# Keyword argument: id
|
||||
id="recurring_campaigns",
|
||||
# Keyword argument: name
|
||||
name="Recurring campaigns check (daily)",
|
||||
# Keyword argument: replace_existing
|
||||
replace_existing=True,
|
||||
)
|
||||
# Call scheduler.add_job()
|
||||
scheduler.add_job(
|
||||
sync_all_jira_links,
|
||||
# Keyword argument: trigger
|
||||
trigger="interval",
|
||||
# Keyword argument: hours
|
||||
hours=1,
|
||||
# Keyword argument: id
|
||||
id="jira_sync",
|
||||
# Keyword argument: name
|
||||
name="Jira link sync (hourly)",
|
||||
# Keyword argument: replace_existing
|
||||
replace_existing=True,
|
||||
)
|
||||
# Call scheduler.add_job()
|
||||
scheduler.add_job(
|
||||
_run_osint_enrichment,
|
||||
# Keyword argument: trigger
|
||||
trigger="interval",
|
||||
# Keyword argument: weeks
|
||||
weeks=1,
|
||||
# Keyword argument: id
|
||||
id="osint_enrichment",
|
||||
# Keyword argument: name
|
||||
name="OSINT enrichment (weekly)",
|
||||
# Keyword argument: replace_existing
|
||||
replace_existing=True,
|
||||
)
|
||||
# Call scheduler.add_job()
|
||||
scheduler.add_job(
|
||||
_run_stale_detection,
|
||||
# Keyword argument: trigger
|
||||
trigger="interval",
|
||||
# Keyword argument: hours
|
||||
hours=24,
|
||||
# Keyword argument: id
|
||||
id="stale_detection",
|
||||
# Keyword argument: name
|
||||
name="Stale coverage detection (daily)",
|
||||
# Keyword argument: replace_existing
|
||||
replace_existing=True,
|
||||
)
|
||||
# Call scheduler.add_job()
|
||||
scheduler.add_job(
|
||||
run_retention_job,
|
||||
# Keyword argument: trigger
|
||||
trigger="interval",
|
||||
# Keyword argument: hours
|
||||
hours=24,
|
||||
# Keyword argument: id
|
||||
id="retention_policies",
|
||||
# Keyword argument: name
|
||||
name="Data retention policies (daily)",
|
||||
# Keyword argument: replace_existing
|
||||
replace_existing=True,
|
||||
)
|
||||
scheduler.add_job(
|
||||
_run_data_sources_sync,
|
||||
trigger="interval",
|
||||
hours=6,
|
||||
id="data_sources_sync",
|
||||
name="Data sources auto-sync (every 6h)",
|
||||
replace_existing=True,
|
||||
)
|
||||
scheduler.add_job(
|
||||
_run_decay_engine,
|
||||
trigger="cron",
|
||||
hour=2,
|
||||
minute=0,
|
||||
id="decay_engine",
|
||||
name="Detection decay engine (daily 02:00)",
|
||||
replace_existing=True,
|
||||
)
|
||||
scheduler.add_job(
|
||||
_run_queue_generation,
|
||||
trigger="cron",
|
||||
hour=2,
|
||||
minute=30,
|
||||
id="queue_generation",
|
||||
name="Revalidation queue generation (daily 02:30)",
|
||||
replace_existing=True,
|
||||
)
|
||||
scheduler.add_job(
|
||||
_run_alert_evaluation,
|
||||
trigger="interval",
|
||||
hours=1,
|
||||
id="alert_evaluation",
|
||||
name="Operational alert evaluation (hourly)",
|
||||
replace_existing=True,
|
||||
)
|
||||
scheduler.add_job(
|
||||
_run_evaluation_round_check,
|
||||
trigger="cron",
|
||||
day_of_week="mon",
|
||||
hour=6,
|
||||
minute=0,
|
||||
id="attck_evaluation_check",
|
||||
name="ATT&CK Evaluations new-round check (Mondays 06:00)",
|
||||
replace_existing=True,
|
||||
)
|
||||
scheduler.start()
|
||||
# Log info:
|
||||
logger.info(
|
||||
# Literal argument value
|
||||
"Background scheduler started — mitre_sync (24h), intel_scan (7d), "
|
||||
# Literal argument value
|
||||
"notification_cleanup (24h), weekly_snapshot (Sundays 00:00), "
|
||||
# Literal argument value
|
||||
"recurring_campaigns (daily), jira_sync (1h), "
|
||||
# Literal argument value
|
||||
"osint_enrichment (weekly), stale_detection (daily), "
|
||||
"retention_policies (daily), data_sources_sync (6h), "
|
||||
"alert_evaluation (1h), attck_evaluation_check (Mondays 06:00)"
|
||||
)
|
||||
logger.info("Background scheduler started — mitre_sync (24h), intel_scan (7d)")
|
||||
|
||||
@@ -1,89 +0,0 @@
|
||||
"""Data retention policies — scheduled cleanup of aged records."""
|
||||
|
||||
# Enable future language features for compatibility
|
||||
from __future__ import annotations
|
||||
|
||||
# Import logging
|
||||
import logging
|
||||
|
||||
# Import datetime, timedelta, timezone from datetime
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
# Import Session from sqlalchemy.orm
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
# Import SessionLocal from app.database
|
||||
from app.database import SessionLocal
|
||||
|
||||
# Import AuditLog from app.models.audit
|
||||
from app.models.audit import AuditLog
|
||||
|
||||
# Import cleanup_old_notifications from app.services.notification_service
|
||||
from app.services.notification_service import cleanup_old_notifications
|
||||
|
||||
# Assign logger = logging.getLogger(__name__)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Assign AUDIT_LOG_RETENTION_DAYS = 730
|
||||
AUDIT_LOG_RETENTION_DAYS = 730
|
||||
|
||||
|
||||
# Define function apply_retention_policies
|
||||
def apply_retention_policies(db: Session) -> dict[str, int]:
|
||||
"""Apply retention rules. Commits the session before returning."""
|
||||
# Assign cutoff = datetime.now(timezone.utc) - timedelta(days=AUDIT_LOG_RETENTION_DAYS)
|
||||
cutoff = datetime.now(timezone.utc) - timedelta(days=AUDIT_LOG_RETENTION_DAYS)
|
||||
# Assign deleted_audit = (
|
||||
deleted_audit = (
|
||||
db.query(AuditLog)
|
||||
# Chain .filter() call
|
||||
.filter(AuditLog.timestamp < cutoff)
|
||||
# Chain .delete() call
|
||||
.delete(synchronize_session=False)
|
||||
)
|
||||
# Check: deleted_audit
|
||||
if deleted_audit:
|
||||
# Log info:
|
||||
logger.info(
|
||||
# Literal argument value
|
||||
"Retention: deleted %d audit logs older than %d days",
|
||||
deleted_audit,
|
||||
AUDIT_LOG_RETENTION_DAYS,
|
||||
)
|
||||
|
||||
# Assign deleted_notifications = cleanup_old_notifications(db, days=90)
|
||||
deleted_notifications = cleanup_old_notifications(db, days=90)
|
||||
# Commit all pending changes to the database
|
||||
db.commit()
|
||||
# Return {
|
||||
return {
|
||||
# Literal argument value
|
||||
"audit_logs_deleted": deleted_audit,
|
||||
# Literal argument value
|
||||
"notifications_deleted": deleted_notifications,
|
||||
}
|
||||
|
||||
|
||||
# Define function run_retention_job
|
||||
def run_retention_job() -> None:
|
||||
"""Entry point for the daily retention scheduler job."""
|
||||
# Log info: "Scheduled retention job starting..."
|
||||
logger.info("Scheduled retention job starting...")
|
||||
# Assign db = SessionLocal()
|
||||
db = SessionLocal()
|
||||
# Attempt the following; catch errors below
|
||||
try:
|
||||
# Assign summary = apply_retention_policies(db)
|
||||
summary = apply_retention_policies(db)
|
||||
# Log info: "Retention job finished — %s", summary
|
||||
logger.info("Retention job finished — %s", summary)
|
||||
# Handle Exception
|
||||
except Exception:
|
||||
# Log exception: "Retention job failed"
|
||||
logger.exception("Retention job failed")
|
||||
# Roll back all uncommitted changes
|
||||
db.rollback()
|
||||
# Always execute this cleanup block
|
||||
finally:
|
||||
# Close the database session
|
||||
db.close()
|
||||
@@ -1,10 +0,0 @@
|
||||
"""Shared SlowAPI rate limiter for all routers."""
|
||||
|
||||
# Import Limiter from slowapi
|
||||
from slowapi import Limiter
|
||||
|
||||
# Import get_remote_address from slowapi.util
|
||||
from slowapi.util import get_remote_address
|
||||
|
||||
# Assign limiter = Limiter(key_func=get_remote_address)
|
||||
limiter = Limiter(key_func=get_remote_address)
|
||||
@@ -1,108 +0,0 @@
|
||||
"""Structured JSON logging configuration.
|
||||
|
||||
In **production** (``AEGIS_ENV=production``), emits one JSON object per
|
||||
line so that log aggregators (ELK, CloudWatch, Datadog) can ingest them
|
||||
without custom parsing.
|
||||
|
||||
In **development** (default), uses a human-readable text format for
|
||||
comfortable local work.
|
||||
"""
|
||||
|
||||
# Enable future language features for compatibility
|
||||
from __future__ import annotations
|
||||
|
||||
# Import json
|
||||
import json
|
||||
|
||||
# Import logging
|
||||
import logging
|
||||
|
||||
# Import os
|
||||
import os
|
||||
|
||||
# Import sys
|
||||
import sys
|
||||
|
||||
# Import datetime, timezone from datetime
|
||||
from datetime import datetime, timezone
|
||||
|
||||
|
||||
# Define class _JSONFormatter
|
||||
class _JSONFormatter(logging.Formatter):
|
||||
"""Emit each log record as a single-line JSON object."""
|
||||
|
||||
# Define function format
|
||||
def format(self, record: logging.LogRecord) -> str:
|
||||
# Assign payload = {
|
||||
payload: dict = {
|
||||
# Literal argument value
|
||||
"timestamp": datetime.fromtimestamp(record.created, tz=timezone.utc).isoformat(),
|
||||
# Literal argument value
|
||||
"level": record.levelname,
|
||||
# Literal argument value
|
||||
"logger": record.name,
|
||||
# Literal argument value
|
||||
"message": record.getMessage(),
|
||||
}
|
||||
|
||||
# Check: record.exc_info and record.exc_info[1] is not None
|
||||
if record.exc_info and record.exc_info[1] is not None:
|
||||
# Assign payload["exception"] = self.formatException(record.exc_info)
|
||||
payload["exception"] = self.formatException(record.exc_info)
|
||||
|
||||
# Assign extra = getattr(record, "_extra", None)
|
||||
extra = getattr(record, "_extra", None)
|
||||
# Check: extra
|
||||
if extra:
|
||||
# Call payload.update()
|
||||
payload.update(extra)
|
||||
|
||||
# Return json.dumps(payload, default=str)
|
||||
return json.dumps(payload, default=str)
|
||||
|
||||
|
||||
# Assign _DEV_FORMAT = "%(asctime)s %(levelname)-8s %(name)s — %(message)s"
|
||||
_DEV_FORMAT = "%(asctime)s %(levelname)-8s %(name)s — %(message)s"
|
||||
|
||||
|
||||
# Define function setup_logging
|
||||
def setup_logging() -> None:
|
||||
"""Configure the root logger based on the environment."""
|
||||
# Assign is_production = os.environ.get("AEGIS_ENV", "").lower() == "production"
|
||||
is_production = os.environ.get("AEGIS_ENV", "").lower() == "production"
|
||||
# Assign level_name = os.environ.get("LOG_LEVEL", "INFO").upper()
|
||||
level_name = os.environ.get("LOG_LEVEL", "INFO").upper()
|
||||
# Assign level = getattr(logging, level_name, logging.INFO)
|
||||
level = getattr(logging, level_name, logging.INFO)
|
||||
|
||||
# Assign root = logging.getLogger()
|
||||
root = logging.getLogger()
|
||||
# Call root.setLevel()
|
||||
root.setLevel(level)
|
||||
|
||||
# Check: root.handlers
|
||||
if root.handlers:
|
||||
# Call root.handlers.clear()
|
||||
root.handlers.clear()
|
||||
|
||||
# Assign handler = logging.StreamHandler(sys.stdout)
|
||||
handler = logging.StreamHandler(sys.stdout)
|
||||
# Call handler.setLevel()
|
||||
handler.setLevel(level)
|
||||
|
||||
# Check: is_production
|
||||
if is_production:
|
||||
# Call handler.setFormatter()
|
||||
handler.setFormatter(_JSONFormatter())
|
||||
# Fallback: handle remaining cases
|
||||
else:
|
||||
# Call handler.setFormatter()
|
||||
handler.setFormatter(logging.Formatter(_DEV_FORMAT))
|
||||
|
||||
# Call root.addHandler()
|
||||
root.addHandler(handler)
|
||||
|
||||
# Call logging.getLogger()
|
||||
logging.getLogger("uvicorn.access").setLevel(logging.WARNING)
|
||||
# Call logging.getLogger()
|
||||
logging.getLogger("sqlalchemy.engine").setLevel(logging.WARNING)
|
||||
+23
-324
@@ -1,41 +1,10 @@
|
||||
"""FastAPI application factory and global middleware/exception configuration.
|
||||
|
||||
Builds the ``app`` instance, wires up CORS, rate limiting, domain-error
|
||||
mapping, all API routers, and async lifespan hooks (MinIO bucket creation,
|
||||
APScheduler startup/shutdown).
|
||||
"""
|
||||
|
||||
# Import logging
|
||||
import logging
|
||||
|
||||
# Import os
|
||||
import os
|
||||
|
||||
# Import AsyncGenerator from collections.abc
|
||||
from collections.abc import AsyncGenerator
|
||||
|
||||
# Import asynccontextmanager from contextlib
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
# Import FastAPI, Request, status from fastapi
|
||||
from fastapi import FastAPI, Request, status
|
||||
|
||||
# Import RequestValidationError from fastapi.exceptions
|
||||
from fastapi.exceptions import RequestValidationError
|
||||
|
||||
# Import CORSMiddleware from fastapi.middleware.cors
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
|
||||
# Import JSONResponse from fastapi.responses
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
# Import _rate_limit_exceeded_handler from slowapi
|
||||
from slowapi import _rate_limit_exceeded_handler
|
||||
|
||||
# Import RateLimitExceeded from slowapi.errors
|
||||
from slowapi.errors import RateLimitExceeded
|
||||
|
||||
# Import SQLAlchemyError from sqlalchemy.exc
|
||||
from fastapi.exceptions import RequestValidationError
|
||||
from sqlalchemy.exc import SQLAlchemyError
|
||||
|
||||
from app.routers import auth as auth_router
|
||||
@@ -47,360 +16,90 @@ from app.routers import system as system_router
|
||||
from app.routers import metrics as metrics_router
|
||||
from app.routers import users as users_router
|
||||
from app.routers import audit as audit_router
|
||||
from app.routers import notifications as notifications_router
|
||||
from app.routers import reports as reports_router
|
||||
from app.routers import data_sources as data_sources_router
|
||||
from app.routers import threat_actors as threat_actors_router
|
||||
from app.routers import d3fend as d3fend_router
|
||||
from app.routers import detection_rules as detection_rules_router
|
||||
from app.routers import campaigns as campaigns_router
|
||||
from app.routers import heatmap as heatmap_router
|
||||
from app.routers import scores as scores_router
|
||||
from app.routers import operational_metrics as operational_metrics_router
|
||||
from app.routers import compliance as compliance_router
|
||||
from app.routers import snapshots as snapshots_router
|
||||
from app.routers import jira as jira_router
|
||||
from app.routers import worklogs as worklogs_router
|
||||
from app.routers import professional_reports as professional_reports_router
|
||||
from app.routers import analytics as analytics_router
|
||||
from app.routers import advanced_metrics as advanced_metrics_router
|
||||
from app.routers import osint as osint_router
|
||||
from app.routers import webhooks as webhooks_router
|
||||
from app.routers import detection_lifecycle as detection_lifecycle_router
|
||||
from app.routers import intel as intel_router
|
||||
from app.routers import admin_config as admin_config_router
|
||||
from app.routers import ownership as ownership_router
|
||||
from app.routers import attack_paths as attack_paths_router
|
||||
from app.routers import knowledge as knowledge_router
|
||||
from app.routers import risk_intelligence as risk_router
|
||||
from app.routers import executive_dashboard as dashboard_router
|
||||
from app.routers import api_keys as api_keys_router
|
||||
from app.routers import sso as sso_router
|
||||
from app.routers import operational_alerts as alerts_router
|
||||
from app.domain.errors import DomainError
|
||||
|
||||
# Import scheduler, start_scheduler from app.jobs.mitre_sync_job
|
||||
from app.jobs.mitre_sync_job import scheduler, start_scheduler
|
||||
|
||||
# Import limiter from app.limiter
|
||||
from app.limiter import limiter
|
||||
|
||||
# Import setup_logging from app.logging_config
|
||||
from app.logging_config import setup_logging
|
||||
|
||||
# Import domain_exception_handler from app.middleware.error_handler
|
||||
from app.middleware.error_handler import domain_exception_handler
|
||||
|
||||
# Import RequestContextMiddleware from app.middleware.request_context
|
||||
from app.middleware.request_context import RequestContextMiddleware
|
||||
from app.storage import ensure_bucket_exists
|
||||
from app.config import settings as _settings
|
||||
from starlette.middleware.base import BaseHTTPMiddleware
|
||||
from app.jobs.mitre_sync_job import start_scheduler, scheduler
|
||||
|
||||
# Configure structured logging before any module initialises its own logger
|
||||
setup_logging()
|
||||
# ── Logging ───────────────────────────────────────────────────────────────
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s %(levelname)-8s %(name)s — %(message)s",
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ── Environment detection ─────────────────────────────────────────────────
|
||||
_IS_PRODUCTION = os.environ.get("AEGIS_ENV", "").lower() == "production"
|
||||
|
||||
# Apply the @asynccontextmanager decorator
|
||||
@asynccontextmanager
|
||||
# Define async function lifespan
|
||||
async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
|
||||
"""Manage application startup and shutdown lifecycle.
|
||||
|
||||
Args:
|
||||
app (FastAPI): The FastAPI application instance.
|
||||
|
||||
Yields:
|
||||
None: Control is yielded to the running application.
|
||||
"""
|
||||
# Call ensure_bucket_exists()
|
||||
async def lifespan(app: FastAPI):
|
||||
"""Startup / shutdown logic."""
|
||||
ensure_bucket_exists()
|
||||
# Call start_scheduler()
|
||||
start_scheduler()
|
||||
# Seed decay policies
|
||||
from app.database import SessionLocal
|
||||
from app.seed_decay_policies import seed_decay_policies
|
||||
db = SessionLocal()
|
||||
try:
|
||||
seed_decay_policies(db)
|
||||
except Exception as e:
|
||||
logger.warning("seed_decay_policies failed at startup: %s", e)
|
||||
finally:
|
||||
db.close()
|
||||
# Seed operational alert system rules
|
||||
db2 = SessionLocal()
|
||||
try:
|
||||
from app.services.operational_alert_service import seed_system_rules
|
||||
seed_system_rules(db2)
|
||||
except Exception as e:
|
||||
logger.warning("seed_system_rules failed at startup: %s", e)
|
||||
finally:
|
||||
db2.close()
|
||||
yield
|
||||
# Graceful shutdown of the background scheduler
|
||||
scheduler.shutdown(wait=False)
|
||||
|
||||
|
||||
# ── In production, disable Swagger UI and ReDoc to hide API surface ──────
|
||||
app = FastAPI(
|
||||
# Keyword argument: title
|
||||
title="Attack Coverage Platform",
|
||||
# Keyword argument: lifespan
|
||||
lifespan=lifespan,
|
||||
# Keyword argument: docs_url
|
||||
docs_url=None if _IS_PRODUCTION else "/docs",
|
||||
# Keyword argument: redoc_url
|
||||
redoc_url=None if _IS_PRODUCTION else "/redoc",
|
||||
# Keyword argument: openapi_url
|
||||
openapi_url=None if _IS_PRODUCTION else "/openapi.json",
|
||||
)
|
||||
|
||||
# ── Rate Limiter ──────────────────────────────────────────────────────────
|
||||
app.state.limiter = limiter
|
||||
# Call app.add_exception_handler()
|
||||
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
|
||||
|
||||
# Call app.add_middleware()
|
||||
app.add_middleware(RequestContextMiddleware)
|
||||
|
||||
|
||||
# ── No-cache middleware for all /api/ responses ───────────────────────────
|
||||
# Prevents Cloudflare and browser caches from storing API responses,
|
||||
# which would cause stale/empty data to be served after backend restarts.
|
||||
class NoCacheAPIMiddleware(BaseHTTPMiddleware):
|
||||
async def dispatch(self, request: Request, call_next):
|
||||
response = await call_next(request)
|
||||
if request.url.path.startswith("/api/"):
|
||||
response.headers["Cache-Control"] = "no-store, no-cache, must-revalidate"
|
||||
response.headers["Pragma"] = "no-cache"
|
||||
return response
|
||||
|
||||
app.add_middleware(NoCacheAPIMiddleware)
|
||||
|
||||
|
||||
# ── Domain exception → HTTP mapping ──────────────────────────────────────
|
||||
app.add_exception_handler(DomainError, domain_exception_handler)
|
||||
app = FastAPI(title="Attack Coverage Platform", lifespan=lifespan)
|
||||
|
||||
# ── CORS ──────────────────────────────────────────────────────────────────
|
||||
_cors_origins: list[str] = [
|
||||
o.strip() for o in _settings.CORS_ORIGINS.split(",") if o.strip()
|
||||
]
|
||||
|
||||
# Call app.add_middleware()
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
# Keyword argument: allow_origins
|
||||
allow_origins=_cors_origins,
|
||||
# Keyword argument: allow_credentials
|
||||
allow_origins=["http://localhost:3000", "http://localhost:5173"],
|
||||
allow_credentials=True,
|
||||
# Keyword argument: allow_methods
|
||||
allow_methods=["GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS"],
|
||||
# Keyword argument: allow_headers
|
||||
allow_headers=["Authorization", "Content-Type"],
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
# ── Routers ──────────────────────────────────────────────────────────────
|
||||
app.include_router(auth_router.router, prefix="/api/v1")
|
||||
# Call app.include_router()
|
||||
app.include_router(techniques_router.router, prefix="/api/v1")
|
||||
# Call app.include_router()
|
||||
app.include_router(tests_router.router, prefix="/api/v1")
|
||||
# Call app.include_router()
|
||||
app.include_router(evidence_router.router, prefix="/api/v1")
|
||||
# Call app.include_router()
|
||||
app.include_router(test_templates_router.router, prefix="/api/v1")
|
||||
# Call app.include_router()
|
||||
app.include_router(system_router.router, prefix="/api/v1")
|
||||
# Call app.include_router()
|
||||
app.include_router(metrics_router.router, prefix="/api/v1")
|
||||
# Call app.include_router()
|
||||
app.include_router(users_router.router, prefix="/api/v1")
|
||||
# Call app.include_router()
|
||||
app.include_router(audit_router.router, prefix="/api/v1")
|
||||
# Call app.include_router()
|
||||
app.include_router(notifications_router.router, prefix="/api/v1")
|
||||
# Call app.include_router()
|
||||
app.include_router(reports_router.router, prefix="/api/v1")
|
||||
# Call app.include_router()
|
||||
app.include_router(data_sources_router.router, prefix="/api/v1")
|
||||
# Call app.include_router()
|
||||
app.include_router(threat_actors_router.router, prefix="/api/v1")
|
||||
# Call app.include_router()
|
||||
app.include_router(d3fend_router.router, prefix="/api/v1")
|
||||
# Call app.include_router()
|
||||
app.include_router(detection_rules_router.router, prefix="/api/v1")
|
||||
# Call app.include_router()
|
||||
app.include_router(campaigns_router.router, prefix="/api/v1")
|
||||
# Call app.include_router()
|
||||
app.include_router(heatmap_router.router, prefix="/api/v1")
|
||||
# Call app.include_router()
|
||||
app.include_router(scores_router.router, prefix="/api/v1")
|
||||
# Call app.include_router()
|
||||
app.include_router(operational_metrics_router.router, prefix="/api/v1")
|
||||
# Call app.include_router()
|
||||
app.include_router(compliance_router.router, prefix="/api/v1")
|
||||
app.include_router(intel_router.router, prefix="/api/v1")
|
||||
app.include_router(admin_config_router.router, prefix="/api/v1")
|
||||
app.include_router(snapshots_router.router, prefix="/api/v1")
|
||||
# Call app.include_router()
|
||||
app.include_router(jira_router.router, prefix="/api/v1")
|
||||
# Call app.include_router()
|
||||
app.include_router(worklogs_router.router, prefix="/api/v1")
|
||||
# Call app.include_router()
|
||||
app.include_router(professional_reports_router.router, prefix="/api/v1")
|
||||
# Call app.include_router()
|
||||
app.include_router(analytics_router.router, prefix="/api/v1")
|
||||
# Call app.include_router()
|
||||
app.include_router(advanced_metrics_router.router, prefix="/api/v1")
|
||||
# Call app.include_router()
|
||||
app.include_router(osint_router.router, prefix="/api/v1")
|
||||
app.include_router(webhooks_router.router, prefix="/api/v1")
|
||||
app.include_router(detection_lifecycle_router.router, prefix="/api/v1")
|
||||
app.include_router(ownership_router.router, prefix="/api/v1")
|
||||
app.include_router(attack_paths_router.router, prefix="/api/v1")
|
||||
app.include_router(knowledge_router.router, prefix="/api/v1")
|
||||
app.include_router(risk_router.router, prefix="/api/v1")
|
||||
app.include_router(dashboard_router.router, prefix="/api/v1")
|
||||
app.include_router(api_keys_router.router, prefix="/api/v1")
|
||||
app.include_router(sso_router.router, prefix="/api/v1")
|
||||
app.include_router(alerts_router.router, prefix="/api/v1")
|
||||
|
||||
|
||||
# Apply the @app.get decorator
|
||||
@app.get("/health", include_in_schema=False)
|
||||
# Define function health
|
||||
def health() -> dict[str, str]:
|
||||
"""Return a minimal liveness probe response.
|
||||
|
||||
Access is restricted to internal networks at the Nginx level
|
||||
(see ``frontend/nginx.conf``).
|
||||
|
||||
Returns:
|
||||
dict[str, str]: A dict with ``{"status": "ok"}``.
|
||||
"""
|
||||
# Return {"status": "ok"}
|
||||
@app.get("/health")
|
||||
def health():
|
||||
return {"status": "ok"}
|
||||
|
||||
|
||||
# ── Exception Handlers ────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _serialize_validation_errors(exc: RequestValidationError) -> list[dict]:
|
||||
"""Return validation errors safe for JSON serialization.
|
||||
|
||||
Converts non-serializable values inside ``ctx`` dictionaries to strings
|
||||
so the response body can be safely encoded.
|
||||
|
||||
Args:
|
||||
exc (RequestValidationError): The Pydantic validation exception.
|
||||
|
||||
Returns:
|
||||
list[dict]: A list of sanitised error detail dictionaries.
|
||||
"""
|
||||
# Assign serialized = []
|
||||
serialized: list[dict] = []
|
||||
# Iterate over exc.errors()
|
||||
for err in exc.errors():
|
||||
# Assign item = dict(err)
|
||||
item = dict(err)
|
||||
# Assign ctx = item.get("ctx")
|
||||
ctx = item.get("ctx")
|
||||
# Check: isinstance(ctx, dict)
|
||||
if isinstance(ctx, dict):
|
||||
# Assign item["ctx"] = {key: str(value) for key, value in ctx.items()}
|
||||
item["ctx"] = {key: str(value) for key, value in ctx.items()}
|
||||
# Call serialized.append()
|
||||
serialized.append(item)
|
||||
# Return serialized
|
||||
return serialized
|
||||
|
||||
|
||||
# Apply the @app.exception_handler decorator
|
||||
@app.exception_handler(RequestValidationError)
|
||||
# Define async function validation_exception_handler
|
||||
async def validation_exception_handler(request: Request, exc: RequestValidationError) -> JSONResponse:
|
||||
"""Handle Pydantic validation errors and return a structured 422 response.
|
||||
|
||||
Args:
|
||||
request (Request): The incoming HTTP request.
|
||||
exc (RequestValidationError): The caught validation exception.
|
||||
|
||||
Returns:
|
||||
JSONResponse: A 422 response with a ``VALIDATION_ERROR`` code and error details.
|
||||
"""
|
||||
# Return JSONResponse(
|
||||
async def validation_exception_handler(request: Request, exc: RequestValidationError):
|
||||
"""Handle validation errors with consistent format."""
|
||||
return JSONResponse(
|
||||
# Keyword argument: status_code
|
||||
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
|
||||
# Keyword argument: content
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
content={
|
||||
# Literal argument value
|
||||
"detail": "Validation error",
|
||||
# Literal argument value
|
||||
"code": "VALIDATION_ERROR",
|
||||
# Literal argument value
|
||||
"errors": _serialize_validation_errors(exc),
|
||||
"errors": exc.errors(),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
# Apply the @app.exception_handler decorator
|
||||
@app.exception_handler(SQLAlchemyError)
|
||||
# Define async function sqlalchemy_exception_handler
|
||||
async def sqlalchemy_exception_handler(request: Request, exc: SQLAlchemyError) -> JSONResponse:
|
||||
"""Handle SQLAlchemy database errors and return a structured 500 response.
|
||||
|
||||
Args:
|
||||
request (Request): The incoming HTTP request.
|
||||
exc (SQLAlchemyError): The caught SQLAlchemy exception.
|
||||
|
||||
Returns:
|
||||
JSONResponse: A 500 response with a ``DATABASE_ERROR`` code.
|
||||
"""
|
||||
# Log error: f"Database error: {exc}"
|
||||
async def sqlalchemy_exception_handler(request: Request, exc: SQLAlchemyError):
|
||||
"""Handle database errors."""
|
||||
logging.error(f"Database error: {exc}")
|
||||
# Return JSONResponse(
|
||||
return JSONResponse(
|
||||
# Keyword argument: status_code
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
# Keyword argument: content
|
||||
content={
|
||||
# Literal argument value
|
||||
"detail": "Database error occurred",
|
||||
# Literal argument value
|
||||
"code": "DATABASE_ERROR",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
# Apply the @app.exception_handler decorator
|
||||
@app.exception_handler(Exception)
|
||||
# Define async function general_exception_handler
|
||||
async def general_exception_handler(request: Request, exc: Exception) -> JSONResponse:
|
||||
"""Handle all otherwise-unhandled exceptions and return a structured 500 response.
|
||||
|
||||
Args:
|
||||
request (Request): The incoming HTTP request.
|
||||
exc (Exception): The unhandled exception.
|
||||
|
||||
Returns:
|
||||
JSONResponse: A 500 response with an ``INTERNAL_ERROR`` code.
|
||||
"""
|
||||
# Log error: f"Unhandled exception: {exc}"
|
||||
async def general_exception_handler(request: Request, exc: Exception):
|
||||
"""Handle all unhandled exceptions."""
|
||||
logging.error(f"Unhandled exception: {exc}")
|
||||
# Return JSONResponse(
|
||||
return JSONResponse(
|
||||
# Keyword argument: status_code
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
# Keyword argument: content
|
||||
content={
|
||||
# Literal argument value
|
||||
"detail": "An internal server error occurred",
|
||||
# Literal argument value
|
||||
"code": "INTERNAL_ERROR",
|
||||
},
|
||||
)
|
||||
|
||||
@@ -1 +0,0 @@
|
||||
"""ASGI middleware components for request context, error handling, and rate limiting."""
|
||||
@@ -1,66 +0,0 @@
|
||||
"""Domain error → HTTP response mapping.
|
||||
|
||||
This module provides a single exception handler that converts
|
||||
domain-layer errors into structured JSON responses, keeping
|
||||
the service layer free from FastAPI's ``HTTPException``.
|
||||
"""
|
||||
|
||||
# Import Request from fastapi
|
||||
from fastapi import Request
|
||||
|
||||
# Import JSONResponse from fastapi.responses
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
# Import from app.domain.errors
|
||||
from app.domain.errors import (
|
||||
BusinessRuleViolation,
|
||||
DomainError,
|
||||
DuplicateEntityError,
|
||||
EntityNotFoundError,
|
||||
InvalidOperationError,
|
||||
InvalidStateTransition,
|
||||
PermissionViolation,
|
||||
)
|
||||
|
||||
# Assign EXCEPTION_STATUS_MAP = {
|
||||
EXCEPTION_STATUS_MAP: dict[type[DomainError], int] = {
|
||||
# Entry: EntityNotFoundError
|
||||
EntityNotFoundError: 404,
|
||||
# Entry: DuplicateEntityError
|
||||
DuplicateEntityError: 409,
|
||||
# Entry: InvalidStateTransition
|
||||
InvalidStateTransition: 400,
|
||||
# Entry: InvalidOperationError
|
||||
InvalidOperationError: 400,
|
||||
# Entry: BusinessRuleViolation
|
||||
BusinessRuleViolation: 400,
|
||||
# Entry: PermissionViolation
|
||||
PermissionViolation: 403,
|
||||
}
|
||||
|
||||
|
||||
# Define async function domain_exception_handler
|
||||
async def domain_exception_handler(
|
||||
# Entry: request
|
||||
request: Request,
|
||||
# Entry: exc
|
||||
exc: DomainError,
|
||||
) -> JSONResponse:
|
||||
"""Convert a :class:`DomainError` into a JSON error response."""
|
||||
# Assign status_code = EXCEPTION_STATUS_MAP.get(type(exc), 400)
|
||||
status_code = EXCEPTION_STATUS_MAP.get(type(exc), 400)
|
||||
|
||||
# Assign content = {"detail": exc.message, "code": exc.code}
|
||||
content: dict = {"detail": exc.message, "code": exc.code}
|
||||
|
||||
# Check: isinstance(exc, InvalidStateTransition)
|
||||
if isinstance(exc, InvalidStateTransition):
|
||||
# Assign content["current_state"] = exc.current_state
|
||||
content["current_state"] = exc.current_state
|
||||
# Assign content["target_state"] = exc.target_state
|
||||
content["target_state"] = exc.target_state
|
||||
# Assign content["valid_transitions"] = exc.valid_transitions
|
||||
content["valid_transitions"] = exc.valid_transitions
|
||||
|
||||
# Return JSONResponse(status_code=status_code, content=content)
|
||||
return JSONResponse(status_code=status_code, content=content)
|
||||
@@ -1,74 +0,0 @@
|
||||
"""Request context middleware — captures client IP and User-Agent per request."""
|
||||
|
||||
# Import Awaitable, Callable from collections.abc
|
||||
from collections.abc import Awaitable, Callable
|
||||
|
||||
# Import ContextVar from contextvars
|
||||
from contextvars import ContextVar
|
||||
|
||||
# Import Request from fastapi
|
||||
from fastapi import Request
|
||||
|
||||
# Import BaseHTTPMiddleware from starlette.middleware.base
|
||||
from starlette.middleware.base import BaseHTTPMiddleware
|
||||
|
||||
# Import Response from starlette.responses
|
||||
from starlette.responses import Response
|
||||
|
||||
# Assign request_ip = ContextVar("request_ip", default="")
|
||||
request_ip: ContextVar[str] = ContextVar("request_ip", default="")
|
||||
# Assign request_user_agent = ContextVar("request_user_agent", default="")
|
||||
request_user_agent: ContextVar[str] = ContextVar("request_user_agent", default="")
|
||||
|
||||
|
||||
# Define function resolve_client_ip
|
||||
def resolve_client_ip(request: Request) -> str:
|
||||
"""Extract the real client IP, honouring ``X-Forwarded-For`` when present.
|
||||
|
||||
Args:
|
||||
request (Request): The incoming Starlette/FastAPI request.
|
||||
|
||||
Returns:
|
||||
str: The resolved client IP address, or ``"unknown"`` when unavailable.
|
||||
"""
|
||||
# Assign forwarded = request.headers.get("X-Forwarded-For")
|
||||
forwarded = request.headers.get("X-Forwarded-For")
|
||||
# Check: forwarded
|
||||
if forwarded:
|
||||
# Return forwarded.split(",")[0].strip()
|
||||
return forwarded.split(",")[0].strip()
|
||||
# Check: request.client
|
||||
if request.client:
|
||||
# Return request.client.host
|
||||
return request.client.host
|
||||
# Return "unknown"
|
||||
return "unknown"
|
||||
|
||||
|
||||
# Define class RequestContextMiddleware
|
||||
class RequestContextMiddleware(BaseHTTPMiddleware):
|
||||
"""Middleware that captures client IP and User-Agent into context variables."""
|
||||
|
||||
# Define async function dispatch
|
||||
async def dispatch(
|
||||
self,
|
||||
# Entry: request
|
||||
request: Request,
|
||||
# Entry: call_next
|
||||
call_next: Callable[[Request], Awaitable[Response]],
|
||||
) -> Response:
|
||||
"""Store client IP and User-Agent in context vars for the current request.
|
||||
|
||||
Args:
|
||||
request (Request): The incoming HTTP request.
|
||||
call_next (Callable[[Request], Awaitable[Response]]): The next middleware or route handler.
|
||||
|
||||
Returns:
|
||||
Response: The HTTP response produced by the downstream handler.
|
||||
"""
|
||||
# Call request_ip.set()
|
||||
request_ip.set(resolve_client_ip(request))
|
||||
# Call request_user_agent.set()
|
||||
request_user_agent.set(request.headers.get("User-Agent", ""))
|
||||
# Return await call_next(request)
|
||||
return await call_next(request)
|
||||
@@ -1,89 +1,15 @@
|
||||
"""SQLAlchemy ORM model definitions for all database tables."""
|
||||
# Import all models here so Alembic can detect them
|
||||
from app.models.audit import AuditLog
|
||||
from app.models.notification import Notification
|
||||
from app.models.data_source import DataSource
|
||||
from app.models.detection_rule import DetectionRule
|
||||
from app.models.threat_actor import ThreatActor, ThreatActorTechnique
|
||||
from app.models.defensive_technique import DefensiveTechnique, DefensiveTechniqueMapping
|
||||
from app.models.test_template_detection_rule import TestTemplateDetectionRule
|
||||
from app.models.test_detection_result import TestDetectionResult
|
||||
from app.models.campaign import Campaign, CampaignTest
|
||||
from app.models.compliance import ComplianceFramework, ComplianceControl, ComplianceControlMapping
|
||||
from app.models.coverage_snapshot import CoverageSnapshot, SnapshotTechniqueState
|
||||
from app.models.jira_link import JiraLink, JiraLinkEntityType, JiraSyncDirection
|
||||
from app.models.worklog import Worklog
|
||||
from app.models.osint_item import OsintItem
|
||||
from app.models.scoring_config import ScoringConfig
|
||||
from app.models.enums import TechniqueStatus, TestState, TestResult, TeamSide
|
||||
from app.models.webhook_config import WebhookConfig
|
||||
from app.models.system_config import SystemConfig
|
||||
from app.models.detection_lifecycle import (
|
||||
DetectionAsset, DetectionTechniqueMapping, DetectionValidation,
|
||||
TechniqueConfidenceScore, InfrastructureChangeLog,
|
||||
DetectionConfidence, DetectionHealthStatus, InvalidationReason,
|
||||
)
|
||||
from app.models.decay_policy import DecayPolicy
|
||||
from app.models.ownership_queue import (
|
||||
TechniqueOwnership, RevalidationQueueItem,
|
||||
QueuePriority, QueueStatus, QueueReason,
|
||||
)
|
||||
from app.models.attack_path import (
|
||||
AttackPath, AttackPathStep, AttackPathExecution,
|
||||
AttackPathStepResult, TimelineEntry,
|
||||
ExecutionStatus, StepResultStatus, TimelineActorSide, TimelineEntryType,
|
||||
)
|
||||
from app.models.knowledge import Playbook, PlaybookVersion, LessonLearned
|
||||
from app.models.risk_intelligence import TechniqueRiskProfile
|
||||
from app.models.executive_dashboard import PostureSnapshot
|
||||
from app.models.api_key import ApiKey
|
||||
from app.models.sso_config import SsoConfig
|
||||
from app.models.operational_alert import AlertRule, AlertInstance
|
||||
from app.models.evidence import Evidence
|
||||
from app.models.intel import IntelItem
|
||||
from app.models.user import User
|
||||
from app.models.technique import Technique
|
||||
from app.models.test import Test
|
||||
from app.models.test_template import TestTemplate
|
||||
from app.models.user import User
|
||||
from app.models.evidence import Evidence
|
||||
from app.models.intel import IntelItem
|
||||
from app.models.audit import AuditLog
|
||||
from app.models.enums import TechniqueStatus, TestState, TestResult, TeamSide
|
||||
|
||||
# Assign __all__ = [
|
||||
__all__ = [
|
||||
# Literal argument value
|
||||
"User", "Technique", "Test", "TestTemplate", "Evidence",
|
||||
# Literal argument value
|
||||
"IntelItem", "AuditLog", "Notification", "DataSource",
|
||||
# Literal argument value
|
||||
"DetectionRule", "ThreatActor", "ThreatActorTechnique",
|
||||
# Literal argument value
|
||||
"DefensiveTechnique", "DefensiveTechniqueMapping",
|
||||
# Literal argument value
|
||||
"TestTemplateDetectionRule", "TestDetectionResult",
|
||||
# Literal argument value
|
||||
"Campaign", "CampaignTest",
|
||||
# Literal argument value
|
||||
"ComplianceFramework", "ComplianceControl", "ComplianceControlMapping",
|
||||
# Literal argument value
|
||||
"CoverageSnapshot", "SnapshotTechniqueState",
|
||||
# Literal argument value
|
||||
"JiraLink", "JiraLinkEntityType", "JiraSyncDirection",
|
||||
# Literal argument value
|
||||
"Worklog", "OsintItem", "ScoringConfig",
|
||||
# Literal argument value
|
||||
"IntelItem", "AuditLog",
|
||||
"TechniqueStatus", "TestState", "TestResult", "TeamSide",
|
||||
"WebhookConfig", "SystemConfig",
|
||||
"DetectionAsset", "DetectionTechniqueMapping", "DetectionValidation",
|
||||
"TechniqueConfidenceScore", "InfrastructureChangeLog",
|
||||
"DetectionConfidence", "DetectionHealthStatus", "InvalidationReason", "DecayPolicy",
|
||||
"TechniqueOwnership", "RevalidationQueueItem",
|
||||
"QueuePriority", "QueueStatus", "QueueReason",
|
||||
"AttackPath", "AttackPathStep", "AttackPathExecution",
|
||||
"AttackPathStepResult", "TimelineEntry",
|
||||
"ExecutionStatus", "StepResultStatus", "TimelineActorSide", "TimelineEntryType",
|
||||
"Playbook", "PlaybookVersion", "LessonLearned",
|
||||
"TechniqueRiskProfile",
|
||||
"PostureSnapshot",
|
||||
"ApiKey",
|
||||
"SsoConfig",
|
||||
"AlertRule",
|
||||
"AlertInstance",
|
||||
]
|
||||
|
||||
@@ -1,81 +0,0 @@
|
||||
"""Phase 14: API Key model for programmatic access."""
|
||||
|
||||
import hashlib
|
||||
import secrets
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
|
||||
from sqlalchemy import Boolean, Column, DateTime, ForeignKey, Index, String, Text
|
||||
from sqlalchemy.dialects.postgresql import JSONB, UUID
|
||||
from sqlalchemy.orm import relationship
|
||||
|
||||
from app.database import Base
|
||||
|
||||
# ── Key generation constants ──────────────────────────────────────────────────
|
||||
KEY_PREFIX = "aegis_"
|
||||
KEY_BYTES = 32 # 32 random bytes → 64 hex chars → 70-char key total
|
||||
DISPLAY_LEN = 12 # chars stored as prefix for UI display
|
||||
|
||||
|
||||
def generate_raw_key() -> str:
|
||||
"""Generate a fresh raw API key (must be shown to user only once)."""
|
||||
return KEY_PREFIX + secrets.token_hex(KEY_BYTES)
|
||||
|
||||
|
||||
def hash_key(raw_key: str) -> str:
|
||||
"""SHA-256 hash of a raw API key for secure storage."""
|
||||
return hashlib.sha256(raw_key.encode()).hexdigest()
|
||||
|
||||
|
||||
def key_prefix_display(raw_key: str) -> str:
|
||||
"""First DISPLAY_LEN characters of the raw key (safe for UI)."""
|
||||
return raw_key[:DISPLAY_LEN]
|
||||
|
||||
|
||||
# ── Valid scopes ──────────────────────────────────────────────────────────────
|
||||
VALID_SCOPES = {"read", "write", "admin"}
|
||||
|
||||
|
||||
class ApiKey(Base):
|
||||
"""
|
||||
Scoped API key for programmatic / BI / SOAR access.
|
||||
|
||||
The full raw key is **never stored** — only a SHA-256 hash.
|
||||
The first 12 characters (``key_prefix``) are retained for display.
|
||||
"""
|
||||
|
||||
__tablename__ = "api_keys"
|
||||
|
||||
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
|
||||
name = Column(String(200), nullable=False)
|
||||
description = Column(Text, nullable=True)
|
||||
|
||||
# Display only — never use for auth
|
||||
key_prefix = Column(String(DISPLAY_LEN + 1), nullable=False)
|
||||
|
||||
# Auth token — SHA-256 of the full raw key
|
||||
key_hash = Column(String(64), nullable=False, unique=True)
|
||||
|
||||
# Owner
|
||||
user_id = Column(
|
||||
UUID(as_uuid=True),
|
||||
ForeignKey("users.id", ondelete="CASCADE"),
|
||||
nullable=False,
|
||||
)
|
||||
|
||||
# Permissions
|
||||
scopes = Column(JSONB, nullable=False, default=["read"]) # ["read","write","admin"]
|
||||
|
||||
# Lifecycle
|
||||
last_used_at = Column(DateTime, nullable=True)
|
||||
expires_at = Column(DateTime, nullable=True)
|
||||
is_active = Column(Boolean, nullable=False, default=True)
|
||||
created_at = Column(DateTime, default=datetime.utcnow)
|
||||
|
||||
user = relationship("User", foreign_keys=[user_id])
|
||||
|
||||
__table_args__ = (
|
||||
Index("ix_api_keys_user_id", "user_id"),
|
||||
Index("ix_api_keys_key_hash", "key_hash"),
|
||||
Index("ix_api_keys_active", "is_active"),
|
||||
)
|
||||
@@ -1,253 +0,0 @@
|
||||
"""Phase 10: Attack Paths & Advanced Purple Team models."""
|
||||
|
||||
import enum
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
|
||||
from sqlalchemy import (
|
||||
Boolean, Column, DateTime, Enum, Float, ForeignKey,
|
||||
Index, Integer, String, Text,
|
||||
)
|
||||
from sqlalchemy.dialects.postgresql import UUID, JSONB
|
||||
from sqlalchemy.orm import relationship
|
||||
|
||||
from app.database import Base
|
||||
|
||||
|
||||
class ExecutionStatus(str, enum.Enum):
|
||||
planned = "planned"
|
||||
in_progress = "in_progress"
|
||||
completed = "completed"
|
||||
aborted = "aborted"
|
||||
|
||||
|
||||
class StepResultStatus(str, enum.Enum):
|
||||
pending = "pending"
|
||||
executing = "executing"
|
||||
detected = "detected"
|
||||
not_detected = "not_detected"
|
||||
skipped = "skipped"
|
||||
|
||||
|
||||
class TimelineActorSide(str, enum.Enum):
|
||||
red = "red"
|
||||
blue = "blue"
|
||||
system = "system"
|
||||
|
||||
|
||||
class TimelineEntryType(str, enum.Enum):
|
||||
action = "action"
|
||||
detection = "detection"
|
||||
note = "note"
|
||||
phase_transition = "phase_transition"
|
||||
flag = "flag"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class AttackPath(Base):
|
||||
"""
|
||||
A reusable attack scenario composed of ordered kill-chain steps.
|
||||
Can be a template (shared) or a one-off scenario.
|
||||
"""
|
||||
__tablename__ = "attack_paths"
|
||||
|
||||
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
|
||||
name = Column(String(300), nullable=False)
|
||||
description = Column(Text, nullable=True)
|
||||
objective = Column(Text, nullable=True) # what the attacker aims to achieve
|
||||
is_template = Column(Boolean, default=False) # reusable template flag
|
||||
threat_actor_id = Column(
|
||||
UUID(as_uuid=True), ForeignKey("threat_actors.id", ondelete="SET NULL"), nullable=True
|
||||
)
|
||||
created_by = Column(
|
||||
UUID(as_uuid=True), ForeignKey("users.id", ondelete="SET NULL"), nullable=True
|
||||
)
|
||||
tags = Column(JSONB, nullable=True, default=list)
|
||||
is_active = Column(Boolean, default=True)
|
||||
created_at = Column(DateTime, default=datetime.utcnow)
|
||||
updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
|
||||
|
||||
steps = relationship(
|
||||
"AttackPathStep", back_populates="attack_path",
|
||||
cascade="all, delete-orphan",
|
||||
order_by="AttackPathStep.order_index",
|
||||
)
|
||||
executions = relationship("AttackPathExecution", back_populates="attack_path")
|
||||
creator = relationship("User", foreign_keys=[created_by])
|
||||
threat_actor = relationship("ThreatActor", foreign_keys=[threat_actor_id])
|
||||
|
||||
__table_args__ = (
|
||||
Index("ix_attack_paths_created_by", "created_by"),
|
||||
Index("ix_attack_paths_is_template", "is_template"),
|
||||
)
|
||||
|
||||
|
||||
class AttackPathStep(Base):
|
||||
"""One step in an attack path — maps to a kill-chain phase + technique."""
|
||||
|
||||
__tablename__ = "attack_path_steps"
|
||||
|
||||
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
|
||||
attack_path_id = Column(
|
||||
UUID(as_uuid=True), ForeignKey("attack_paths.id", ondelete="CASCADE"), nullable=False
|
||||
)
|
||||
order_index = Column(Integer, nullable=False, default=0)
|
||||
kill_chain_phase = Column(String(60), nullable=True) # initial_access, execution, …
|
||||
technique_id = Column(
|
||||
UUID(as_uuid=True), ForeignKey("techniques.id", ondelete="SET NULL"), nullable=True
|
||||
)
|
||||
test_id = Column(
|
||||
UUID(as_uuid=True), ForeignKey("tests.id", ondelete="SET NULL"), nullable=True
|
||||
)
|
||||
name = Column(String(300), nullable=True) # human label for the step
|
||||
description = Column(Text, nullable=True)
|
||||
expected_detection = Column(Boolean, default=True) # do we expect blue to detect this?
|
||||
notes = Column(Text, nullable=True)
|
||||
|
||||
attack_path = relationship("AttackPath", back_populates="steps")
|
||||
technique = relationship("Technique", foreign_keys=[technique_id])
|
||||
test = relationship("Test", foreign_keys=[test_id])
|
||||
|
||||
__table_args__ = (
|
||||
Index("ix_ap_steps_path_id", "attack_path_id"),
|
||||
Index("ix_ap_steps_technique_id", "technique_id"),
|
||||
)
|
||||
|
||||
|
||||
class AttackPathExecution(Base):
|
||||
"""
|
||||
A single run of an attack path.
|
||||
Tracks Red/Blue participants, timing, and aggregated kill-chain metrics.
|
||||
"""
|
||||
__tablename__ = "attack_path_executions"
|
||||
|
||||
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
|
||||
attack_path_id = Column(
|
||||
UUID(as_uuid=True), ForeignKey("attack_paths.id", ondelete="CASCADE"), nullable=False
|
||||
)
|
||||
status = Column(
|
||||
Enum(ExecutionStatus, name="execution_status"), nullable=False,
|
||||
default=ExecutionStatus.planned,
|
||||
)
|
||||
environment = Column(String(100), nullable=True) # prod, staging, lab
|
||||
red_team_lead = Column(
|
||||
UUID(as_uuid=True), ForeignKey("users.id", ondelete="SET NULL"), nullable=True
|
||||
)
|
||||
blue_team_lead = Column(
|
||||
UUID(as_uuid=True), ForeignKey("users.id", ondelete="SET NULL"), nullable=True
|
||||
)
|
||||
started_by = Column(
|
||||
UUID(as_uuid=True), ForeignKey("users.id", ondelete="SET NULL"), nullable=True
|
||||
)
|
||||
started_at = Column(DateTime, nullable=True)
|
||||
completed_at = Column(DateTime, nullable=True)
|
||||
notes = Column(Text, nullable=True)
|
||||
created_at = Column(DateTime, default=datetime.utcnow)
|
||||
|
||||
# ── Computed kill-chain metrics (written on complete) ─────────────────
|
||||
total_steps = Column(Integer, nullable=True)
|
||||
detected_steps = Column(Integer, nullable=True)
|
||||
not_detected_steps = Column(Integer, nullable=True)
|
||||
skipped_steps = Column(Integer, nullable=True)
|
||||
detection_rate = Column(Float, nullable=True) # 0.0–1.0
|
||||
mttd_seconds = Column(Float, nullable=True) # mean time to detect (avg across detected)
|
||||
furthest_undetected_step = Column(Integer, nullable=True) # order_index of deepest undetected step
|
||||
|
||||
attack_path = relationship("AttackPath", back_populates="executions")
|
||||
step_results = relationship(
|
||||
"AttackPathStepResult", back_populates="execution",
|
||||
cascade="all, delete-orphan",
|
||||
order_by="AttackPathStepResult.step_order",
|
||||
)
|
||||
timeline = relationship(
|
||||
"TimelineEntry", back_populates="execution",
|
||||
cascade="all, delete-orphan",
|
||||
order_by="TimelineEntry.timestamp",
|
||||
)
|
||||
red_lead_user = relationship("User", foreign_keys=[red_team_lead])
|
||||
blue_lead_user = relationship("User", foreign_keys=[blue_team_lead])
|
||||
|
||||
__table_args__ = (
|
||||
Index("ix_ap_exec_path_id", "attack_path_id"),
|
||||
Index("ix_ap_exec_status", "status"),
|
||||
)
|
||||
|
||||
|
||||
class AttackPathStepResult(Base):
|
||||
"""Result of executing one step in an attack path execution."""
|
||||
|
||||
__tablename__ = "attack_path_step_results"
|
||||
|
||||
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
|
||||
execution_id = Column(
|
||||
UUID(as_uuid=True), ForeignKey("attack_path_executions.id", ondelete="CASCADE"),
|
||||
nullable=False,
|
||||
)
|
||||
step_id = Column(
|
||||
UUID(as_uuid=True), ForeignKey("attack_path_steps.id", ondelete="CASCADE"),
|
||||
nullable=False,
|
||||
)
|
||||
step_order = Column(Integer, nullable=False, default=0) # denormalized for sorting
|
||||
status = Column(
|
||||
Enum(StepResultStatus, name="step_result_status"), nullable=False,
|
||||
default=StepResultStatus.pending,
|
||||
)
|
||||
executed_by = Column(
|
||||
UUID(as_uuid=True), ForeignKey("users.id", ondelete="SET NULL"), nullable=True
|
||||
)
|
||||
executed_at = Column(DateTime, nullable=True)
|
||||
detected_at = Column(DateTime, nullable=True)
|
||||
time_to_detect_seconds = Column(Float, nullable=True)
|
||||
detection_asset_id = Column(
|
||||
UUID(as_uuid=True),
|
||||
ForeignKey("detection_assets.id", ondelete="SET NULL"), nullable=True
|
||||
)
|
||||
notes = Column(Text, nullable=True)
|
||||
evidence_ids = Column(JSONB, nullable=True, default=list)
|
||||
|
||||
execution = relationship("AttackPathExecution", back_populates="step_results")
|
||||
step = relationship("AttackPathStep")
|
||||
detection_asset = relationship("DetectionAsset", foreign_keys=[detection_asset_id])
|
||||
executor = relationship("User", foreign_keys=[executed_by])
|
||||
|
||||
__table_args__ = (
|
||||
Index("ix_ap_stepres_execution_id", "execution_id"),
|
||||
Index("ix_ap_stepres_step_id", "step_id"),
|
||||
)
|
||||
|
||||
|
||||
class TimelineEntry(Base):
|
||||
"""Timestamped Red/Blue action during an execution — used for MTTD/MTTR."""
|
||||
|
||||
__tablename__ = "attack_path_timeline"
|
||||
|
||||
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
|
||||
execution_id = Column(
|
||||
UUID(as_uuid=True), ForeignKey("attack_path_executions.id", ondelete="CASCADE"),
|
||||
nullable=False,
|
||||
)
|
||||
step_id = Column(
|
||||
UUID(as_uuid=True), ForeignKey("attack_path_steps.id", ondelete="SET NULL"),
|
||||
nullable=True,
|
||||
)
|
||||
timestamp = Column(DateTime, nullable=False, default=datetime.utcnow)
|
||||
actor_side = Column(
|
||||
Enum(TimelineActorSide, name="timeline_actor_side"), nullable=False,
|
||||
)
|
||||
actor_id = Column(
|
||||
UUID(as_uuid=True), ForeignKey("users.id", ondelete="SET NULL"), nullable=True
|
||||
)
|
||||
entry_type = Column(
|
||||
Enum(TimelineEntryType, name="timeline_entry_type"), nullable=False,
|
||||
)
|
||||
content = Column(Text, nullable=False)
|
||||
extra = Column(JSONB, nullable=True)
|
||||
|
||||
execution = relationship("AttackPathExecution", back_populates="timeline")
|
||||
actor = relationship("User", foreign_keys=[actor_id])
|
||||
|
||||
__table_args__ = (
|
||||
Index("ix_timeline_execution_id", "execution_id"),
|
||||
Index("ix_timeline_timestamp", "timestamp"),
|
||||
)
|
||||
@@ -1,60 +1,29 @@
|
||||
"""SQLAlchemy model for the audit log table."""
|
||||
|
||||
# Import uuid
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
|
||||
# Import Column, DateTime, ForeignKey, Index, String, func from sqlalchemy
|
||||
from sqlalchemy import Column, DateTime, ForeignKey, Index, String, func
|
||||
|
||||
# Import JSONB, UUID from sqlalchemy.dialects.postgresql
|
||||
from sqlalchemy.dialects.postgresql import JSONB, UUID
|
||||
|
||||
# Import relationship from sqlalchemy.orm
|
||||
from sqlalchemy import Column, String, DateTime, ForeignKey
|
||||
from sqlalchemy.dialects.postgresql import UUID, JSONB
|
||||
from sqlalchemy.orm import relationship
|
||||
|
||||
# Import Base from app.database
|
||||
from app.database import Base
|
||||
|
||||
|
||||
# Define class AuditLog
|
||||
class AuditLog(Base):
|
||||
"""Audit log model for tracking all system actions.
|
||||
|
||||
"""
|
||||
Audit log model for tracking all system actions.
|
||||
|
||||
Records user actions, entity changes, and system events
|
||||
for security auditing and compliance purposes.
|
||||
"""
|
||||
# Assign __tablename__ = "audit_logs"
|
||||
__tablename__ = "audit_logs"
|
||||
|
||||
# Assign id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
|
||||
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
|
||||
# Assign user_id = Column(UUID(as_uuid=True), ForeignKey("users.id"), nullable=True)
|
||||
user_id = Column(UUID(as_uuid=True), ForeignKey("users.id"), nullable=True)
|
||||
# Assign action = Column(String, nullable=False)
|
||||
action = Column(String, nullable=False)
|
||||
# Assign entity_type = Column(String, nullable=True)
|
||||
entity_type = Column(String, nullable=True)
|
||||
# Assign entity_id = Column(String, nullable=True)
|
||||
entity_id = Column(String, nullable=True)
|
||||
# Assign timestamp = Column(DateTime(timezone=True), server_default=func.now())
|
||||
timestamp = Column(DateTime(timezone=True), server_default=func.now())
|
||||
# Assign details = Column(JSONB, nullable=True)
|
||||
timestamp = Column(DateTime, default=datetime.utcnow)
|
||||
details = Column(JSONB, nullable=True)
|
||||
# Assign ip_address = Column(String(45), nullable=True)
|
||||
ip_address = Column(String(45), nullable=True)
|
||||
# Assign user_agent = Column(String(500), nullable=True)
|
||||
user_agent = Column(String(500), nullable=True)
|
||||
# Assign integrity_hash = Column(String(64), nullable=True)
|
||||
integrity_hash = Column(String(64), nullable=True)
|
||||
# Assign session_id = Column(String(100), nullable=True)
|
||||
session_id = Column(String(100), nullable=True)
|
||||
|
||||
# Relationships
|
||||
user = relationship("User")
|
||||
|
||||
# Assign __table_args__ = (
|
||||
__table_args__ = (
|
||||
Index("ix_audit_logs_entity", "entity_type", "entity_id"),
|
||||
Index("ix_audit_logs_timestamp", "timestamp"),
|
||||
Index("ix_audit_logs_entity_type_entity_id_action", "entity_type", "entity_id", "action"),
|
||||
)
|
||||
|
||||
@@ -1,231 +0,0 @@
|
||||
"""Campaign and CampaignTest models.
|
||||
|
||||
Campaigns group multiple tests into a kill chain sequence,
|
||||
enabling simulation of complete attack chains and APT emulations.
|
||||
"""
|
||||
|
||||
# Import uuid
|
||||
import uuid
|
||||
|
||||
# Import from sqlalchemy
|
||||
from sqlalchemy import (
|
||||
Boolean,
|
||||
Column,
|
||||
DateTime,
|
||||
ForeignKey,
|
||||
Index,
|
||||
Integer,
|
||||
String,
|
||||
Text,
|
||||
func,
|
||||
)
|
||||
|
||||
# Import JSONB, UUID from sqlalchemy.dialects.postgresql
|
||||
from sqlalchemy.dialects.postgresql import JSONB, UUID
|
||||
|
||||
# Import relationship from sqlalchemy.orm
|
||||
from sqlalchemy.orm import relationship
|
||||
|
||||
# Import Base from app.database
|
||||
from app.database import Base
|
||||
|
||||
|
||||
# Define class Campaign
|
||||
class Campaign(Base):
|
||||
"""A campaign groups multiple tests into a sequenced attack chain.
|
||||
|
||||
Types:
|
||||
- custom: manually created campaign
|
||||
- apt_emulation: generated from a threat actor profile
|
||||
- kill_chain: structured around kill chain phases
|
||||
- compliance: targeting specific compliance requirements
|
||||
|
||||
Status:
|
||||
- draft: being configured
|
||||
- active: tests are being executed
|
||||
- completed: all tests done
|
||||
- archived: historical record
|
||||
"""
|
||||
# Assign __tablename__ = "campaigns"
|
||||
__tablename__ = "campaigns"
|
||||
|
||||
# Assign id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
|
||||
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
|
||||
# Assign name = Column(String, nullable=False)
|
||||
name = Column(String, nullable=False)
|
||||
# Assign description = Column(Text, nullable=True)
|
||||
description = Column(Text, nullable=True)
|
||||
# Assign type = Column(String, nullable=False, default="custom") # custom, ap...
|
||||
type = Column(String, nullable=False, default="custom") # custom, apt_emulation, kill_chain, compliance
|
||||
# Assign threat_actor_id = Column(
|
||||
threat_actor_id = Column(
|
||||
UUID(as_uuid=True),
|
||||
ForeignKey("threat_actors.id", ondelete="SET NULL"),
|
||||
# Keyword argument: nullable
|
||||
nullable=True,
|
||||
)
|
||||
# Assign status = Column(String, nullable=False, default="draft") # draft, activ...
|
||||
status = Column(String, nullable=False, default="draft") # draft, active, completed, archived
|
||||
# Assign created_by = Column(
|
||||
created_by = Column(
|
||||
UUID(as_uuid=True),
|
||||
ForeignKey("users.id", ondelete="SET NULL"),
|
||||
# Keyword argument: nullable
|
||||
nullable=True,
|
||||
)
|
||||
start_date = Column(DateTime, nullable=True) # campaign won't activate before this date
|
||||
scheduled_at = Column(DateTime, nullable=True)
|
||||
# Assign completed_at = Column(DateTime, nullable=True)
|
||||
completed_at = Column(DateTime, nullable=True)
|
||||
# Assign target_platform = Column(String, nullable=True)
|
||||
target_platform = Column(String, nullable=True)
|
||||
# Assign tags = Column(JSONB, nullable=True, default=[])
|
||||
tags = Column(JSONB, nullable=True, default=[])
|
||||
# Assign created_at = Column(DateTime(timezone=True), server_default=func.now())
|
||||
created_at = Column(DateTime(timezone=True), server_default=func.now())
|
||||
# Assign data_classification = Column(String(20), nullable=False, server_default="internal")
|
||||
data_classification = Column(String(20), nullable=False, server_default="internal")
|
||||
|
||||
# Recurring scheduling fields
|
||||
is_recurring = Column(Boolean, default=False)
|
||||
# Assign recurrence_pattern = Column(String, nullable=True) # weekly, monthly, quarterly
|
||||
recurrence_pattern = Column(String, nullable=True) # weekly, monthly, quarterly
|
||||
# Assign next_run_at = Column(DateTime, nullable=True)
|
||||
next_run_at = Column(DateTime, nullable=True)
|
||||
# Assign last_run_at = Column(DateTime, nullable=True)
|
||||
last_run_at = Column(DateTime, nullable=True)
|
||||
# Assign parent_campaign_id = Column(
|
||||
parent_campaign_id = Column(
|
||||
UUID(as_uuid=True),
|
||||
ForeignKey("campaigns.id", ondelete="SET NULL"),
|
||||
# Keyword argument: nullable
|
||||
nullable=True,
|
||||
)
|
||||
|
||||
# Relationships
|
||||
threat_actor = relationship("ThreatActor")
|
||||
# Assign creator = relationship("User", foreign_keys=[created_by])
|
||||
creator = relationship("User", foreign_keys=[created_by])
|
||||
# Assign campaign_tests = relationship(
|
||||
campaign_tests = relationship(
|
||||
# Literal argument value
|
||||
"CampaignTest",
|
||||
# Keyword argument: back_populates
|
||||
back_populates="campaign",
|
||||
# Keyword argument: cascade
|
||||
cascade="all, delete-orphan",
|
||||
# Keyword argument: order_by
|
||||
order_by="CampaignTest.order_index",
|
||||
)
|
||||
# Assign parent_campaign = relationship(
|
||||
parent_campaign = relationship(
|
||||
# Literal argument value
|
||||
"Campaign",
|
||||
# Keyword argument: remote_side
|
||||
remote_side="Campaign.id",
|
||||
# Keyword argument: foreign_keys
|
||||
foreign_keys=[parent_campaign_id],
|
||||
)
|
||||
# Assign child_campaigns = relationship(
|
||||
child_campaigns = relationship(
|
||||
# Literal argument value
|
||||
"Campaign",
|
||||
# Keyword argument: foreign_keys
|
||||
foreign_keys=[parent_campaign_id],
|
||||
# Keyword argument: back_populates
|
||||
back_populates="parent_campaign",
|
||||
)
|
||||
|
||||
# Assign __table_args__ = (
|
||||
__table_args__ = (
|
||||
Index('ix_campaigns_status', 'status'),
|
||||
Index('ix_campaigns_type', 'type'),
|
||||
Index('ix_campaigns_threat_actor', 'threat_actor_id'),
|
||||
Index('ix_campaigns_created_by', 'created_by'),
|
||||
Index('ix_campaigns_next_run', 'next_run_at'),
|
||||
)
|
||||
|
||||
|
||||
# Kill chain phases in order (for sorting and validation)
|
||||
KILL_CHAIN_PHASES = [
|
||||
# Literal argument value
|
||||
"reconnaissance",
|
||||
# Literal argument value
|
||||
"resource_development",
|
||||
# Literal argument value
|
||||
"initial_access",
|
||||
# Literal argument value
|
||||
"execution",
|
||||
# Literal argument value
|
||||
"persistence",
|
||||
# Literal argument value
|
||||
"privilege_escalation",
|
||||
# Literal argument value
|
||||
"defense_evasion",
|
||||
# Literal argument value
|
||||
"credential_access",
|
||||
# Literal argument value
|
||||
"discovery",
|
||||
# Literal argument value
|
||||
"lateral_movement",
|
||||
# Literal argument value
|
||||
"collection",
|
||||
# Literal argument value
|
||||
"command_and_control",
|
||||
# Literal argument value
|
||||
"exfiltration",
|
||||
# Literal argument value
|
||||
"impact",
|
||||
]
|
||||
|
||||
|
||||
# Define class CampaignTest
|
||||
class CampaignTest(Base):
|
||||
"""A test within a campaign, with ordering and dependency information.
|
||||
|
||||
``depends_on`` creates a self-referential chain (A -> B -> C).
|
||||
Circular dependencies are validated at the service layer.
|
||||
"""
|
||||
# Assign __tablename__ = "campaign_tests"
|
||||
__tablename__ = "campaign_tests"
|
||||
|
||||
# Assign id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
|
||||
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
|
||||
# Assign campaign_id = Column(
|
||||
campaign_id = Column(
|
||||
UUID(as_uuid=True),
|
||||
ForeignKey("campaigns.id", ondelete="CASCADE"),
|
||||
# Keyword argument: nullable
|
||||
nullable=False,
|
||||
)
|
||||
# Assign test_id = Column(
|
||||
test_id = Column(
|
||||
UUID(as_uuid=True),
|
||||
ForeignKey("tests.id", ondelete="CASCADE"),
|
||||
# Keyword argument: nullable
|
||||
nullable=False,
|
||||
)
|
||||
# Assign order_index = Column(Integer, nullable=False, default=0)
|
||||
order_index = Column(Integer, nullable=False, default=0)
|
||||
# Assign depends_on = Column(
|
||||
depends_on = Column(
|
||||
UUID(as_uuid=True),
|
||||
ForeignKey("campaign_tests.id", ondelete="SET NULL"),
|
||||
# Keyword argument: nullable
|
||||
nullable=True,
|
||||
)
|
||||
# Assign phase = Column(String, nullable=True) # kill chain phase
|
||||
phase = Column(String, nullable=True) # kill chain phase
|
||||
|
||||
# Relationships
|
||||
campaign = relationship("Campaign", back_populates="campaign_tests")
|
||||
# Assign test = relationship("Test")
|
||||
test = relationship("Test")
|
||||
# Assign dependency = relationship("CampaignTest", remote_side="CampaignTest.id")
|
||||
dependency = relationship("CampaignTest", remote_side="CampaignTest.id")
|
||||
|
||||
# Assign __table_args__ = (
|
||||
__table_args__ = (
|
||||
Index('ix_campaign_tests_campaign', 'campaign_id'),
|
||||
Index('ix_campaign_tests_test', 'test_id'),
|
||||
)
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user