Compare commits

3 Commits

Author SHA1 Message Date
kitos 03e06591c6 feat: Phase 2 - Authentication and authorization (T-010 to T-013) 2026-02-06 13:15:25 +01:00
kitos 966aad689c feat: Phase 1 - Data models and migrations (T-004 to T-009)
Implements all database models for the Aegis platform with full
Alembic migration support.

Models created:
- User: Authentication with role-based access control
- Technique: MITRE ATT&CK techniques with coverage status tracking
- Test: Security tests with validation workflow (draft/review/validated)
- Evidence: File metadata for test evidence (stored in MinIO)
- IntelItem: Threat intelligence items linked to techniques
- AuditLog: System-wide audit trail with JSONB details

Enumerations:
- TechniqueStatus: not_evaluated, in_progress, validated, partial, etc.
- TestState: draft, in_review, validated, rejected
- TestResult: detected, not_detected, partially_detected

Services:
- audit_service.py: log_action() helper for audit logging

All models include proper foreign key relationships and PostgreSQL
enum types are managed correctly in migrations (create/drop).
2026-02-06 12:26:26 +01:00
kitos d392d41bf8 feat: Phase 0 - Infrastructure and scaffolding (T-001 to T-003)
This commit establishes the foundational infrastructure for the Aegis
MITRE ATT&CK Coverage Platform.

T-001: Initialize project and Docker Compose
- Set up Docker Compose with PostgreSQL 15, MinIO, and FastAPI backend
- Create basic FastAPI application with health endpoint
- Configure persistent volumes for data storage

T-002: Configuration and database connection
- Add centralized configuration using pydantic-settings
- Implement SQLAlchemy database connection with session management
- Configure MinIO and JWT settings

T-003: Initialize Alembic for migrations
- Set up Alembic with PostgreSQL connection from settings
- Create initial empty migration
- Configure autogenerate support for future models

Also includes:
- Professional README with setup instructions
- Comprehensive .gitignore for Python/Node/Docker
- Project task plan (AegisTestPlan.md)
2026-02-06 11:28:30 +01:00
477 changed files with 1539 additions and 104985 deletions
-44
View File
@@ -1,44 +0,0 @@
# =============================================================================
# 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)
# =============================================================================
# ── Database ─────────────────────────────────────────────────────────────────
DB_USER=postgres
DB_PASSWORD= # REQUIRED — set a strong password
DB_NAME=attackdb
# ── Security ─────────────────────────────────────────────────────────────────
# REQUIRED in production — the app will refuse to start without it.
# Generate with: openssl rand -hex 32
SECRET_KEY=
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_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
-64
View File
@@ -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
-71
View File
@@ -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
-9
View File
@@ -60,12 +60,3 @@ Thumbs.db
# Local development # Local development
*.local *.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
View File
File diff suppressed because it is too large Load Diff
+130 -370
View File
@@ -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.
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.
## Features ## Features
### Core (V1) - **MITRE ATT&CK Integration**: Automatic synchronization with the MITRE ATT&CK framework via TAXII
- **MITRE ATT&CK Integration** — Automatic synchronization via TAXII 2.0 (with GitHub fallback), scheduled every 24h - **Coverage Tracking**: Track validation status for each technique (validated, partial, not covered, in progress)
- **Red/Blue Validation Workflow** — Structured dual-validation lifecycle: draft → red_executing → blue_evaluating → in_review → validated/rejected - **Test Management**: Document and manage security tests with full audit trail
- **Dual Validation** — Independent Red Lead / Blue Lead approval before finalization - **Evidence Storage**: Secure evidence file storage with SHA256 integrity verification
- **Coverage Tracking** — Per-technique status (validated, partial, not covered, in progress) - **Role-Based Access Control**: Granular permissions for red team, blue team, and leadership roles
- **Evidence Storage** — Secure evidence with SHA256 integrity, separated by team (red/blue) - **Intel Monitoring**: Automated scanning for new threat intelligence related to techniques
- **Role-Based Access Control** — Granular permissions for 6 roles (admin, red_tech, blue_tech, red_lead, blue_lead, viewer) - **Metrics Dashboard**: Real-time coverage metrics and reporting by tactic
### 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** — 0100 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 |
## Tech Stack ## Tech Stack
- **Backend**: FastAPI (Python 3.11) — Clean Modular Monolith with domain entities, services, and repository pattern - **Backend**: FastAPI (Python 3.11)
- **Database**: PostgreSQL 16 with UUID primary keys and JSONB columns - **Database**: PostgreSQL 15
- **Object Storage**: MinIO (S3-compatible) - **Object Storage**: MinIO (S3-compatible)
- **ORM**: SQLAlchemy 2.x with Alembic migrations - **ORM**: SQLAlchemy with Alembic migrations
- **Frontend**: React 19 + TypeScript + Vite 7 + Tailwind CSS v4 + TanStack Query + TanStack Virtual - **Frontend**: React + TypeScript + Vite (coming soon)
- **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
## Quick Start ## Quick Start
@@ -99,341 +26,174 @@ Both Red Lead and Blue Lead must independently vote:
- Docker and Docker Compose - Docker and Docker Compose
- Git - Git
- Linux / macOS (or WSL on Windows)
### Production Deployment ### Installation
The recommended way to deploy Aegis in production:
1. Clone the repository:
```bash ```bash
git clone <repository-url> git clone <repository-url>
cd Aegis cd Aegis
chmod +x scripts/install.sh
./scripts/install.sh
``` ```
The interactive install wizard will guide you through: 2. Start all services:
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:
```bash ```bash
git clone <repository-url>
cd Aegis
docker-compose up -d 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. Verify the installation:
```bash
# Check backend health
curl http://localhost:8000/health
# Expected: {"status":"ok"}
```
### Authentication ### Authentication
JWT-based authentication with HttpOnly cookies. Admin credentials are configured during installation: The platform uses JWT-based authentication. After seeding, log in with the default admin credentials:
- 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:
```bash ```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. > **Important:** Change the default `admin123` password and `SECRET_KEY` in production.
### 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)
## Services ## Services
| Service | Port | Description | | Service | Port | Description |
|---------|------|-------------| |----------|------|-------------|
| Frontend | 5173 | React dev server (Vite) | | Backend | 8000 | FastAPI REST API |
| Backend | 8000 | FastAPI REST API | | PostgreSQL | 5433 | Database (mapped to 5433 to avoid conflicts) |
| PostgreSQL | 5433 | Database |
| MinIO API | 9000 | S3-compatible object storage | | MinIO API | 9000 | S3-compatible object storage |
| MinIO Console | 9001 | MinIO web interface | | 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 ## 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 - **Swagger UI**: http://localhost:8000/docs
- **ReDoc**: http://localhost:8000/redoc - **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
| 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) |
See [docs/API.md](docs/API.md) for the full endpoint reference.
## Configuration
All variables are configured automatically by `scripts/install.sh`. For manual setup, copy `.env.example` to `.env` and fill in the values.
### Required (production)
| 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 |
### 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.
## Project Structure ## Project Structure
``` ```
Aegis/ Aegis/
├── docker-compose.yml ├── docker-compose.yml # Docker services configuration
├── 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)
├── backend/ ├── backend/
│ ├── Dockerfile │ ├── Dockerfile # Backend container definition
│ ├── requirements.txt │ ├── requirements.txt # Python dependencies
│ ├── alembic.ini │ ├── alembic.ini # Alembic configuration
│ ├── alembic/versions/ # Database migration files │ ├── alembic/ # Database migrations
│ ├── pytest.ini │ ├── env.py
│ ├── tests/ # 367+ pytest tests (domain, service, API) │ ├── versions/ # Migration files
│ │ └── ...
│ └── app/ │ └── app/
│ ├── main.py # FastAPI app with all routers + lifespan │ ├── __init__.py
│ ├── config.py # Pydantic Settings from environment │ ├── main.py # FastAPI application entry point
│ ├── database.py # SQLAlchemy engine + session (lazy init) │ ├── config.py # Application settings
│ ├── storage.py # MinIO/S3 helpers │ ├── database.py # SQLAlchemy configuration
│ ├── auth.py # Password hashing + JWT tokens │ ├── auth.py # Password hashing & JWT utilities
│ ├── domain/ # Pure business logic (zero framework imports) │ ├── seed.py # Admin seed script (python -m app.seed)
│ ├── entities/ # Rich domain entities (Technique, Campaign, etc.) ├── models/ # SQLAlchemy models
│ │ ├── ports/ # Protocol interfaces (repos, ImportService) │ │ ├── user.py # User authentication model
│ │ ├── value_objects/ # Immutable types (MitreId, ScoringWeights) │ │ ├── technique.py # MITRE ATT&CK techniques
│ │ ├── errors.py # Domain exception hierarchy │ │ ├── test.py # Security tests
│ │ ── unit_of_work.py # Transaction management │ │ ── evidence.py # Test evidence files
│ ├── infrastructure/ # SQLAlchemy repos, Redis, mappers ├── intel.py # Threat intelligence items
├── models/ # SQLAlchemy ORM models │ ├── audit.py # Audit logging
├── schemas/ # Pydantic request/response schemas │ └── enums.py # Shared enumerations
│ ├── routers/ # 27 thin HTTP adapter routers │ ├── schemas/ # Pydantic request/response schemas
├── services/ # 46 framework-agnostic business services │ └── auth.py # LoginRequest, TokenResponse, UserOut
│ ├── middleware/ # Error handler (domain exceptions → HTTP) │ ├── routers/ # API endpoint routers
├── dependencies/ # FastAPI dependency injection (auth, repos) │ └── auth.py # POST /auth/login, GET /auth/me
── jobs/ # APScheduler background jobs ── dependencies/ # FastAPI dependencies (DI)
└── frontend/src/ │ │ └── auth.py # get_current_user, require_role (RBAC)
├── App.tsx # Routes with lazy loading + role protection │ └── services/ # Business logic services
├── api/ # API client modules (Axios + TanStack Query) │ └── audit_service.py
├── components/ # Reusable UI components └── frontend/ # React frontend (coming soon)
├── hooks/ # Custom hooks (useDebounce, etc.)
├── context/ # Auth state management
└── pages/ # Page components
``` ```
## 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 ## Development
### Running Migrations ### Running Migrations
```bash ```bash
docker exec aegis-backend alembic upgrade head # Generate a new migration after model changes
docker exec aegis-backend alembic revision --autogenerate -m "description" docker exec -w /app aegis-backend-1 alembic revision --autogenerate -m "description"
docker exec aegis-backend alembic downgrade -1
# 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
``` ```
### Running Tests ### Accessing Services
```bash - **MinIO Console**: http://localhost:9001 (login: `minioadmin` / `minioadmin`)
# Run all V3 tests inside the container (recommended) - **PostgreSQL**: `psql -h localhost -p 5433 -U postgres -d attackdb`
docker exec aegis-backend python -m pytest tests/ -v --tb=short
# Run specific test suites ## User Roles
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) | Role | Description |
docker exec aegis-backend python -m pytest tests/ -v -m "not integration" |------|-------------|
``` | `admin` | Full system access |
| `red_tech` | Red team technician - can create and edit tests |
### Generating Reports | `blue_tech` | Blue team technician - can create and edit tests |
| `red_lead` | Red team lead - can validate tests |
```bash | `blue_lead` | Blue team lead - can validate tests |
# Coverage summary (JSON) | `viewer` | Read-only access |
GET /api/v1/reports/coverage-summary
# Coverage CSV export
GET /api/v1/reports/coverage-csv
# Compliance gap analysis
GET /api/v1/compliance/{framework_id}/gaps
```
## Further Documentation
- **[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
## License ## License
This project is proprietary software. All rights reserved. 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.
-2
View File
@@ -1,2 +0,0 @@
skips:
- B311
-11
View File
@@ -1,11 +0,0 @@
__pycache__
*.pyc
*.pyo
*.egg-info
.pytest_cache
.mypy_cache
.venv
venv
env
.env
*.log
+3 -17
View File
@@ -3,14 +3,9 @@ FROM python:3.11-slim
WORKDIR /app WORKDIR /app
# Install system dependencies # Install system dependencies
RUN apt-get update && apt-get upgrade -y && apt-get install -y \ RUN apt-get update && apt-get install -y \
gcc \ gcc \
libpq-dev \ libpq-dev \
curl \
pkg-config \
libxml2-dev \
libxmlsec1-dev \
libxmlsec1-openssl \
&& rm -rf /var/lib/apt/lists/* && rm -rf /var/lib/apt/lists/*
# Copy requirements first for better caching # Copy requirements first for better caching
@@ -20,17 +15,8 @@ RUN pip install --no-cache-dir -r requirements.txt
# Copy application code # Copy application code
COPY . . 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
# Expose port # Expose port
EXPOSE 8000 EXPOSE 8000
# Default command (migrations + seed + uvicorn) # Default command
CMD ["sh", "/app/entrypoint.sh"] CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8000", "--reload"]
@@ -1,32 +0,0 @@
"""add_new_test_states
Revision ID: b001add0test
Revises: a1412d1ef337
Create Date: 2026-02-09 10:00:00.000000
"""
from typing import Sequence, Union
from alembic import op
# revision identifiers, used by Alembic.
revision: str = 'b001add0test'
down_revision: Union[str, Sequence[str], None] = 'a1412d1ef337'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Add red_executing and blue_evaluating values to the teststate enum."""
op.execute("ALTER TYPE teststate ADD VALUE IF NOT EXISTS 'red_executing' AFTER 'draft'")
op.execute("ALTER TYPE teststate ADD VALUE IF NOT EXISTS 'blue_evaluating' AFTER 'red_executing'")
def downgrade() -> None:
"""Downgrade: removing enum values in PostgreSQL requires recreating the type.
This is intentionally left as a no-op because dropping enum values is
destructive and rarely needed in practice.
"""
pass
@@ -1,46 +0,0 @@
"""add_evidence_team_and_notes
Revision ID: b002evidteam
Revises: b001add0test
Create Date: 2026-02-09 10:01: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 = 'b002evidteam'
down_revision: Union[str, Sequence[str], None] = 'b001add0test'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Create teamside enum and add team/notes columns to evidences."""
# Create the new enum type
teamside_enum = postgresql.ENUM('red', 'blue', name='teamside', create_type=False)
op.execute("CREATE TYPE teamside AS ENUM ('red', 'blue')")
# Add columns
op.add_column('evidences', sa.Column(
'team',
teamside_enum,
nullable=False,
server_default='red',
))
op.add_column('evidences', sa.Column(
'notes',
sa.Text(),
nullable=True,
))
def downgrade() -> None:
"""Remove team/notes columns and drop teamside enum."""
op.drop_column('evidences', 'notes')
op.drop_column('evidences', 'team')
op.execute("DROP TYPE IF EXISTS teamside")
@@ -1,87 +0,0 @@
"""add_dual_validation_fields_to_tests
Revision ID: b003dualvalid
Revises: b002evidteam
Create Date: 2026-02-09 10:02: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 = 'b003dualvalid'
down_revision: Union[str, Sequence[str], None] = 'b002evidteam'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Drop legacy validated_by/validated_at and add dual validation columns."""
# Drop legacy single-validation columns
op.drop_constraint('tests_validated_by_fkey', 'tests', type_='foreignkey')
op.drop_column('tests', 'validated_by')
op.drop_column('tests', 'validated_at')
# ── Red Team fields ─────────────────────────────────────────
op.add_column('tests', sa.Column('red_summary', sa.Text(), nullable=True))
op.add_column('tests', sa.Column('attack_success', sa.Boolean(), nullable=True))
op.add_column('tests', sa.Column('red_validated_by', sa.UUID(), nullable=True))
op.add_column('tests', sa.Column('red_validated_at', sa.DateTime(), nullable=True))
op.add_column('tests', sa.Column('red_validation_status', sa.String(), nullable=True))
op.add_column('tests', sa.Column('red_validation_notes', sa.Text(), nullable=True))
# ── Blue Team fields ────────────────────────────────────────
op.add_column('tests', sa.Column('blue_summary', sa.Text(), nullable=True))
op.add_column('tests', sa.Column(
'detection_result',
postgresql.ENUM('detected', 'not_detected', 'partially_detected',
name='testresult', create_type=False),
nullable=True,
))
op.add_column('tests', sa.Column('blue_validated_by', sa.UUID(), nullable=True))
op.add_column('tests', sa.Column('blue_validated_at', sa.DateTime(), nullable=True))
op.add_column('tests', sa.Column('blue_validation_status', sa.String(), nullable=True))
op.add_column('tests', sa.Column('blue_validation_notes', sa.Text(), nullable=True))
# ── Foreign keys ────────────────────────────────────────────
op.create_foreign_key(
'fk_tests_red_validated_by', 'tests', 'users',
['red_validated_by'], ['id'],
)
op.create_foreign_key(
'fk_tests_blue_validated_by', 'tests', 'users',
['blue_validated_by'], ['id'],
)
def downgrade() -> None:
"""Reverse: drop dual validation columns and restore legacy columns."""
# Drop FKs
op.drop_constraint('fk_tests_blue_validated_by', 'tests', type_='foreignkey')
op.drop_constraint('fk_tests_red_validated_by', 'tests', type_='foreignkey')
# Drop new columns
op.drop_column('tests', 'blue_validation_notes')
op.drop_column('tests', 'blue_validation_status')
op.drop_column('tests', 'blue_validated_at')
op.drop_column('tests', 'blue_validated_by')
op.drop_column('tests', 'detection_result')
op.drop_column('tests', 'blue_summary')
op.drop_column('tests', 'red_validation_notes')
op.drop_column('tests', 'red_validation_status')
op.drop_column('tests', 'red_validated_at')
op.drop_column('tests', 'red_validated_by')
op.drop_column('tests', 'attack_success')
op.drop_column('tests', 'red_summary')
# Restore legacy columns
op.add_column('tests', sa.Column('validated_by', sa.UUID(), nullable=True))
op.add_column('tests', sa.Column('validated_at', sa.DateTime(), nullable=True))
op.create_foreign_key(
'tests_validated_by_fkey', 'tests', 'users',
['validated_by'], ['id'],
)
@@ -1,54 +0,0 @@
"""add_test_templates_table
Revision ID: b004templates
Revises: b003dualvalid
Create Date: 2026-02-09 10:03:00.000000
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision: str = 'b004templates'
down_revision: Union[str, Sequence[str], None] = 'b003dualvalid'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Create the test_templates table with indexes."""
op.create_table(
'test_templates',
sa.Column('id', sa.UUID(), nullable=False, default=sa.text('gen_random_uuid()')),
sa.Column('mitre_technique_id', sa.String(), nullable=False),
sa.Column('name', sa.String(), nullable=False),
sa.Column('description', sa.Text(), nullable=True),
sa.Column('source', sa.String(), nullable=False),
sa.Column('source_url', sa.String(), nullable=True),
sa.Column('attack_procedure', sa.Text(), nullable=True),
sa.Column('expected_detection', sa.Text(), nullable=True),
sa.Column('platform', sa.String(), nullable=True),
sa.Column('tool_suggested', sa.String(), nullable=True),
sa.Column('severity', sa.String(), nullable=True),
sa.Column('atomic_test_id', sa.String(), nullable=True),
sa.Column('is_active', sa.Boolean(), nullable=True, server_default=sa.text('true')),
sa.Column('created_at', sa.DateTime(), nullable=True),
sa.PrimaryKeyConstraint('id'),
)
op.create_index('ix_test_templates_mitre_technique_id', 'test_templates', ['mitre_technique_id'])
op.create_index('ix_test_templates_source', 'test_templates', ['source'])
op.create_index('ix_test_templates_platform', 'test_templates', ['platform'])
op.create_index('ix_test_templates_severity', 'test_templates', ['severity'])
def downgrade() -> None:
"""Drop the test_templates table and its indexes."""
op.drop_index('ix_test_templates_severity', table_name='test_templates')
op.drop_index('ix_test_templates_platform', table_name='test_templates')
op.drop_index('ix_test_templates_source', table_name='test_templates')
op.drop_index('ix_test_templates_mitre_technique_id', table_name='test_templates')
op.drop_table('test_templates')
@@ -1,55 +0,0 @@
"""add_v2_indexes
Revision ID: b005v2indexes
Revises: b004templates
Create Date: 2026-02-09 10:04:00.000000
"""
from typing import Sequence, Union
from alembic import op
# revision identifiers, used by Alembic.
revision: str = 'b005v2indexes'
down_revision: Union[str, Sequence[str], None] = 'b004templates'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Create performance indexes for V2 queries."""
# ── Tests ───────────────────────────────────────────────────
op.create_index('ix_tests_state', 'tests', ['state'])
op.create_index('ix_tests_technique_id', 'tests', ['technique_id'])
op.create_index('ix_tests_created_by', 'tests', ['created_by'])
op.create_index('ix_tests_red_validation_status', 'tests', ['red_validation_status'])
op.create_index('ix_tests_blue_validation_status', 'tests', ['blue_validation_status'])
# ── Evidences ───────────────────────────────────────────────
op.create_index('ix_evidences_test_id', 'evidences', ['test_id'])
op.create_index('ix_evidences_team', 'evidences', ['team'])
# ── Techniques (if not already present from MVP) ────────────
op.create_index('ix_techniques_tactic', 'techniques', ['tactic'])
op.create_index('ix_techniques_status_global', 'techniques', ['status_global'])
op.create_index('ix_techniques_review_required', 'techniques', ['review_required'])
def downgrade() -> None:
"""Drop all V2 indexes."""
# Techniques
op.drop_index('ix_techniques_review_required', table_name='techniques')
op.drop_index('ix_techniques_status_global', table_name='techniques')
op.drop_index('ix_techniques_tactic', table_name='techniques')
# Evidences
op.drop_index('ix_evidences_team', table_name='evidences')
op.drop_index('ix_evidences_test_id', table_name='evidences')
# Tests
op.drop_index('ix_tests_blue_validation_status', table_name='tests')
op.drop_index('ix_tests_red_validation_status', table_name='tests')
op.drop_index('ix_tests_created_by', table_name='tests')
op.drop_index('ix_tests_technique_id', table_name='tests')
op.drop_index('ix_tests_state', table_name='tests')
@@ -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"))
-106
View File
@@ -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
View File
@@ -1 +0,0 @@
"""Aegis — MITRE ATT&CK Coverage Platform application package."""
+7 -95
View File
@@ -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: This module provides pure functions for:
- Hashing and verifying passwords using bcrypt via passlib. - Hashing and verifying passwords using bcrypt via passlib.
- Creating JWT access tokens using PyJWT. - Creating JWT access tokens using python-jose.
- Managing a Redis-backed token blacklist for revocation.
No endpoints are defined here. 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 from datetime import datetime, timedelta, timezone
# Import jwt (PyJWT) from jose import jwt
import jwt
# Import CryptContext from passlib.context
from passlib.context import CryptContext from passlib.context import CryptContext
# Import settings from app.config
from app.config import settings from app.config import settings
# Assign logger = logging.getLogger(__name__)
logger = logging.getLogger(__name__)
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# Password hashing # Password hashing
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -36,17 +22,13 @@ logger = logging.getLogger(__name__)
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto") pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
# Define function hash_password
def hash_password(password: str) -> str: def hash_password(password: str) -> str:
"""Return a bcrypt hash of *password*.""" """Return a bcrypt hash of *password*."""
# Return pwd_context.hash(password)
return pwd_context.hash(password) return pwd_context.hash(password)
# Define function verify_password
def verify_password(plain: str, hashed: str) -> bool: def verify_password(plain: str, hashed: str) -> bool:
"""Return ``True`` if *plain* matches the bcrypt *hashed* value.""" """Return ``True`` if *plain* matches the bcrypt *hashed* value."""
# Return pwd_context.verify(plain, hashed)
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: 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. The token expires after ``ACCESS_TOKEN_EXPIRE_MINUTES`` (from settings).
- ``exp``: expiration timestamp based on ``ACCESS_TOKEN_EXPIRE_MINUTES``.
""" """
# Assign to_encode = data.copy()
to_encode = data.copy() to_encode = data.copy()
# Assign expire = datetime.now(timezone.utc) + timedelta(
expire = datetime.now(timezone.utc) + timedelta( expire = datetime.now(timezone.utc) + timedelta(
# Keyword argument: minutes
minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES, minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES,
) )
# Call to_encode.update() to_encode.update({"exp": expire})
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...
return jwt.encode(to_encode, settings.SECRET_KEY, algorithm=settings.ALGORITHM) 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
View File
@@ -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 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): 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" DATABASE_URL: str = "postgresql://postgres:postgres@postgres:5432/attackdb"
SECRET_KEY: str = "change-me-in-production"
# ── 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"
ALGORITHM: str = "HS256" ALGORITHM: str = "HS256"
ACCESS_TOKEN_EXPIRE_MINUTES: int = 480 # 8 hours — /auth/refresh extends active sessions ACCESS_TOKEN_EXPIRE_MINUTES: int = 60
# ── 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 ───────────────────────────────────────────────────
MINIO_ENDPOINT: str = "minio:9000" 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" MINIO_ACCESS_KEY: str = "minioadmin"
# Assign MINIO_SECRET_KEY = "minioadmin"
MINIO_SECRET_KEY: str = "minioadmin" MINIO_SECRET_KEY: str = "minioadmin"
# Assign MINIO_BUCKET = "evidence"
MINIO_BUCKET: str = "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: class Config:
"""Pydantic BaseSettings configuration — load from .env file."""
# Assign env_file = ".env"
env_file = ".env" env_file = ".env"
# Assign settings = Settings()
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
View File
@@ -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 import create_engine
from sqlalchemy.orm import sessionmaker, declarative_base
# Import Engine from sqlalchemy.engine from app.config import settings
from sqlalchemy.engine import Engine
# Import Session, declarative_base, sessionmaker from sqlalchemy.orm engine = create_engine(settings.DATABASE_URL)
from sqlalchemy.orm import Session, declarative_base, sessionmaker SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
# Assign Base = declarative_base()
Base = declarative_base() 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
def get_db():
# 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()
db = SessionLocal() db = SessionLocal()
# Attempt the following; catch errors below
try: try:
# Yield db
yield db yield db
# Always execute this cleanup block
finally: finally:
# Close the database session
db.close() db.close()
-1
View File
@@ -1 +0,0 @@
"""FastAPI dependency injection helpers for auth, DB, and shared state."""
+11 -194
View File
@@ -1,52 +1,26 @@
"""Authentication and RBAC dependencies for FastAPI. """
Authentication and RBAC dependencies for FastAPI.
Provides: Provides:
- ``get_current_user``: decodes JWT from HttpOnly cookie (preferred) or - ``get_current_user``: decodes JWT, fetches user from DB, raises 401 on failure.
Authorization header (fallback), fetches user from DB, raises 401 on failure.
Also accepts Aegis API keys (``aegis_…`` prefix) as Bearer tokens.
- ``require_role``: factory that returns a dependency enforcing a specific role - ``require_role``: factory that returns a dependency enforcing a specific role
(admins always pass). (admins always pass).
""" """
# Import Callable from collections.abc from fastapi import Depends, HTTPException, status
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.security import OAuth2PasswordBearer from fastapi.security import OAuth2PasswordBearer
from jose import JWTError, jwt
# Import jwt (PyJWT)
import jwt
# Import Session from sqlalchemy.orm
from sqlalchemy.orm import Session 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 from app.config import settings
# Import get_db from app.database
from app.database import get_db from app.database import get_db
# Import User from app.models.user
from app.models.user import 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) oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/api/v1/auth/login")
# Cookie name — must match the one set in the auth router
_COOKIE_NAME = "aegis_token"
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# Current-user dependency # Current-user dependency
@@ -54,94 +28,38 @@ _COOKIE_NAME = "aegis_token"
async def get_current_user( async def get_current_user(
# Entry: aegis_token token: str = Depends(oauth2_scheme),
aegis_token: Optional[str] = Cookie(None),
# Entry: bearer_token
bearer_token: Optional[str] = Depends(oauth2_scheme),
# Entry: db
db: Session = Depends(get_db), db: Session = Depends(get_db),
) -> User: ) -> User:
"""Decode the JWT, look up the user in *db*, and return it. """Decode the JWT *token*, 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).
Raises :class:`~fastapi.HTTPException` **401** when: Raises :class:`~fastapi.HTTPException` **401** when:
- no token is found in either location,
- the token cannot be decoded, - the token cannot be decoded,
- the ``sub`` claim is missing, or - the ``sub`` claim is missing, or
- no matching active user exists in the database. - no matching active user exists in the database.
""" """
# Assign credentials_exception = HTTPException(
credentials_exception = HTTPException( credentials_exception = HTTPException(
# Keyword argument: status_code
status_code=status.HTTP_401_UNAUTHORIZED, status_code=status.HTTP_401_UNAUTHORIZED,
# Keyword argument: detail
detail="Could not validate credentials", 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"}, 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: try:
# Assign payload = jwt.decode(
payload = jwt.decode( payload = jwt.decode(
token, token,
settings.SECRET_KEY, settings.SECRET_KEY,
# Keyword argument: algorithms
algorithms=[settings.ALGORITHM], algorithms=[settings.ALGORITHM],
) )
# Assign username = payload.get("sub")
username: str | None = payload.get("sub") username: str | None = payload.get("sub")
# Check: username is None
if username is None: if username is None:
# Raise credentials_exception
raise credentials_exception raise credentials_exception
# Check token blacklist (revoked tokens) except JWTError:
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
raise credentials_exception raise credentials_exception
# Assign user = db.query(User).filter(User.username == username).first()
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:
if user is None or not user.is_active:
# Raise credentials_exception
raise credentials_exception raise credentials_exception
# Return user
return user return user
@@ -150,123 +68,22 @@ 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): def require_role(required_role: str):
"""Return a FastAPI dependency that enforces *required_role*. """Return a FastAPI dependency that enforces *required_role*.
The dependency allows the request to proceed when The dependency allows the request to proceed when
``user.role == required_role`` **or** ``user.role == "admin"``. ``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**. Otherwise it raises :class:`~fastapi.HTTPException` **403**.
""" """
# Define async function role_checker
async def role_checker( async def role_checker(
# Entry: current_user
current_user: User = Depends(get_current_user), current_user: User = Depends(get_current_user),
) -> User: ) -> User:
# Check: current_user.role != required_role and current_user.role != "admin"
if 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( raise HTTPException(
# Keyword argument: status_code
status_code=status.HTTP_403_FORBIDDEN, status_code=status.HTTP_403_FORBIDDEN,
# Keyword argument: detail
detail="Not enough permissions", detail="Not enough permissions",
) )
scope = "admin" if required_role == "admin" else "write"
_check_api_key_scope(current_user, scope)
return current_user return current_user
# Return role_checker
return role_checker return role_checker
# Define function require_any_role
def require_any_role(*roles: str) -> Callable[..., object]:
"""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::
@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
-44
View File
@@ -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
View File
@@ -1 +0,0 @@
"""Domain layer — entities, value objects, errors, and repository ports."""
-34
View File
@@ -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",
]
-219
View File
@@ -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,
)
-164
View File
@@ -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
]
-316
View File
@@ -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
-206
View File
@@ -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,
)
-82
View File
@@ -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"
-192
View File
@@ -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")
-25
View File
@@ -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
View File
@@ -1 +0,0 @@
"""Abstract port interfaces that infrastructure adapters must implement."""
-165
View File
@@ -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).
"""
# ...
...
-648
View File
@@ -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"))
-103
View File
@@ -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
View File
@@ -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
View File
@@ -1 +0,0 @@
"""Background scheduler jobs (MITRE sync, Jira sync, data retention)."""
-65
View File
@@ -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()
-712
View File
@@ -1,712 +0,0 @@
"""Scheduled background jobs.
Registers periodic tasks on an APScheduler ``BackgroundScheduler``:
* **MITRE sync** — every 24 hours (see :func:`sync_mitre`)
* **Intel scan** — every 7 days (see :func:`scan_intel`)
Each job manages its own database session (created on entry, closed in
``finally``) so it is fully independent from FastAPI's request-scoped
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.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__)
# ---------------------------------------------------------------------------
# Module-level scheduler instance
# ---------------------------------------------------------------------------
scheduler = BackgroundScheduler()
# ---------------------------------------------------------------------------
# Job functions
# ---------------------------------------------------------------------------
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()
# ---------------------------------------------------------------------------
# Scheduler bootstrap
# ---------------------------------------------------------------------------
def start_scheduler() -> None:
"""Register all periodic jobs and start the background scheduler.
Jobs registered:
* ``mitre_sync`` — every **24 hours**
* ``intel_scan`` — every **7 days**
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)"
)
-89
View File
@@ -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()
-10
View File
@@ -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)
-108
View File
@@ -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)
+7 -390
View File
@@ -1,406 +1,23 @@
"""FastAPI application factory and global middleware/exception configuration. from fastapi import FastAPI
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 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 sqlalchemy.exc import SQLAlchemyError
from app.routers import auth as auth_router from app.routers import auth as auth_router
from app.routers import techniques as techniques_router
from app.routers import tests as tests_router
from app.routers import evidence as evidence_router
from app.routers import test_templates as test_templates_router
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 app = FastAPI(title="Attack Coverage Platform")
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
# Configure structured logging before any module initialises its own logger
setup_logging()
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()
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)
# ── CORS ────────────────────────────────────────────────────────────────── # ── CORS ──────────────────────────────────────────────────────────────────
_cors_origins: list[str] = [
o.strip() for o in _settings.CORS_ORIGINS.split(",") if o.strip()
]
# Call app.add_middleware()
app.add_middleware( app.add_middleware(
CORSMiddleware, CORSMiddleware,
# Keyword argument: allow_origins allow_origins=["http://localhost:3000", "http://localhost:5173"],
allow_origins=_cors_origins,
# Keyword argument: allow_credentials
allow_credentials=True, allow_credentials=True,
# Keyword argument: allow_methods allow_methods=["*"],
allow_methods=["GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS"], allow_headers=["*"],
# Keyword argument: allow_headers
allow_headers=["Authorization", "Content-Type"],
) )
# ── Routers ────────────────────────────────────────────────────────────── # ── Routers ──────────────────────────────────────────────────────────────
app.include_router(auth_router.router, prefix="/api/v1") 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")
@app.get("/health", include_in_schema=False) def health():
# 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"}
return {"status": "ok"} 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(
return JSONResponse(
# Keyword argument: status_code
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
# Keyword argument: content
content={
# Literal argument value
"detail": "Validation error",
# Literal argument value
"code": "VALIDATION_ERROR",
# Literal argument value
"errors": _serialize_validation_errors(exc),
},
)
# 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}"
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}"
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
View File
@@ -1 +0,0 @@
"""ASGI middleware components for request context, error handling, and rate limiting."""
-66
View File
@@ -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)

Some files were not shown because too many files have changed in this diff Show More