Files
Aegis/backend/app/routers/campaigns.py
T
kitos 394d5d9056 refactor(types): add comprehensive type annotations across backend Python codebase
Enable ANN rules in ruff.toml (flake8-annotations) and resolve all 221 violations:

ANN201/ANN202 — return types on 168 public/private functions:
- All 28 FastAPI routers: endpoints annotated with dict/list/specific schema/
  StreamingResponse/FileResponse/JSONResponse as appropriate
- main.py: lifespan→AsyncGenerator[None,None], exception handlers→JSONResponse
- database.py: get_db→Generator[Session,None,None], proxy methods→correct types
- middleware/request_context.py: dispatch→Response with Callable call_next type

ANN001/ANN002/ANN003 — 32 missing argument types:
- seed_demo.py: all db parameters typed as Session
- domain/unit_of_work.py: __aexit__ exc_type/exc_val/exc_tb typed with TracebackType
- services: audit_service user_id→UUID|None, heatmap_service query/model/builder,
  notification_service test→Test, tempo_service test→Test/user→User,
  test_workflow_service test_id→UUID, campaign_crud **fields→object,
  test_crud **fields→object (4 sites)

ANN401 — 16 Any usages resolved:
- Domain entities (campaign/technique/threat_actor/test_entity): replaced Any with
  actual ORM types via TYPE_CHECKING guards to avoid circular imports
- detection_rule_service: test_id/detection_rule_id/evaluator_id→UUID
- score_cache: kept Any with # noqa: ANN401 (genuinely generic cache)
- jira_service/tempo_service: kept Any with # noqa: ANN401 (lazy optional deps)
- d3fend_import_service: _to_str(v: Any) kept with # noqa: ANN401

ANN204/ANN205/ANN206 — special/static/class methods:
- database.py proxy __call__/__getattr__: *args: object/**kwargs: object
- schemas/test.py model_validate: obj→object, **kwargs→object
- sa_technique_repository._int_type→type

All 439 unit tests pass. ruff check app/ → All checks passed!

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-06-09 17:04:51 +02:00

417 lines
14 KiB
Python

"""Campaign endpoints — CRUD, test management, activation, and auto-generation.
Provides comprehensive campaign lifecycle management including
test ordering, progress tracking, and threat actor integration.
"""
import logging
import uuid
from typing import Optional
from fastapi import APIRouter, Depends, Query
from pydantic import BaseModel, Field
from sqlalchemy.orm import Session
from app.database import get_db
from app.dependencies.auth import get_current_user, require_any_role
from app.domain.unit_of_work import UnitOfWork
from app.models.user import User
from app.services.audit_service import log_action
from app.services.campaign_crud_service import (
activate_campaign as crud_activate,
)
from app.services.campaign_crud_service import (
add_test_to_campaign as crud_add_test,
)
from app.services.campaign_crud_service import (
complete_campaign as crud_complete,
)
from app.services.campaign_crud_service import (
create_campaign as crud_create,
)
from app.services.campaign_crud_service import (
get_campaign_detail as crud_get_detail,
)
from app.services.campaign_crud_service import (
get_campaign_history as crud_get_history,
)
from app.services.campaign_crud_service import (
get_campaign_progress_data as crud_get_progress,
)
from app.services.campaign_crud_service import (
list_campaigns as crud_list,
)
from app.services.campaign_crud_service import (
remove_test_from_campaign as crud_remove_test,
)
from app.services.campaign_crud_service import (
schedule_campaign as crud_schedule,
)
from app.services.campaign_crud_service import (
serialize_campaign,
)
from app.services.campaign_crud_service import (
update_campaign as crud_update,
)
from app.services.campaign_service import generate_campaign_from_threat_actor
from app.services.notification_service import notify_role
logger = logging.getLogger(__name__)
router = APIRouter(prefix="/campaigns", tags=["campaigns"])
# ── Pydantic schemas ─────────────────────────────────────────────────
class CampaignCreate(BaseModel):
name: str
description: Optional[str] = None
type: str = "custom"
threat_actor_id: Optional[str] = None
target_platform: Optional[str] = None
tags: Optional[list[str]] = Field(default_factory=list)
scheduled_at: Optional[str] = None
class CampaignUpdate(BaseModel):
name: Optional[str] = None
description: Optional[str] = None
type: Optional[str] = None
target_platform: Optional[str] = None
tags: Optional[list[str]] = None
scheduled_at: Optional[str] = None
class AddTestPayload(BaseModel):
test_id: str
order_index: Optional[int] = None
depends_on: Optional[str] = None
phase: Optional[str] = None
class SchedulePayload(BaseModel):
is_recurring: bool
recurrence_pattern: Optional[str] = None # weekly, monthly, quarterly
next_run_at: Optional[str] = None
# ---------------------------------------------------------------------------
# GET /campaigns — List campaigns with filters
# ---------------------------------------------------------------------------
@router.get("")
def list_campaigns(
type: Optional[str] = Query(None),
status: Optional[str] = Query(None),
threat_actor_id: Optional[str] = Query(None),
search: Optional[str] = Query(None),
offset: int = Query(0, ge=0),
limit: int = Query(50, ge=1, le=200),
db: Session = Depends(get_db),
current_user: User = Depends(get_current_user),
) -> list:
"""List campaigns with optional filters and pagination."""
return crud_list(
db,
type=type,
status=status,
threat_actor_id=threat_actor_id,
search=search,
offset=offset,
limit=limit,
)
# ---------------------------------------------------------------------------
# POST /campaigns — Create campaign
# ---------------------------------------------------------------------------
@router.post("", status_code=201)
def create_campaign(
payload: CampaignCreate,
db: Session = Depends(get_db),
current_user: User = Depends(require_any_role("red_lead", "blue_lead")),
) -> dict:
"""Create a new campaign."""
with UnitOfWork(db) as uow:
result = crud_create(
db,
creator_id=current_user.id,
name=payload.name,
description=payload.description,
type=payload.type,
threat_actor_id=payload.threat_actor_id,
target_platform=payload.target_platform,
tags=payload.tags,
scheduled_at=payload.scheduled_at,
)
log_action(
db,
user_id=current_user.id,
action="create_campaign",
entity_type="campaign",
entity_id=result["id"],
details={"name": payload.name, "type": payload.type},
)
uow.commit()
return result
# ---------------------------------------------------------------------------
# GET /campaigns/{id} — Detail with tests and progress
# ---------------------------------------------------------------------------
@router.get("/{campaign_id}")
def get_campaign(
campaign_id: str,
db: Session = Depends(get_db),
current_user: User = Depends(get_current_user),
) -> dict:
"""Get detailed campaign info including tests and progress."""
return crud_get_detail(db, campaign_id)
# ---------------------------------------------------------------------------
# PATCH /campaigns/{id} — Update campaign
# ---------------------------------------------------------------------------
@router.patch("/{campaign_id}")
def update_campaign(
campaign_id: str,
payload: CampaignUpdate,
db: Session = Depends(get_db),
current_user: User = Depends(require_any_role("red_lead", "blue_lead")),
) -> dict:
"""Update a campaign. Only allowed in draft or active state."""
update_data = payload.model_dump(exclude_unset=True)
with UnitOfWork(db) as uow:
result = crud_update(
db,
campaign_id,
updater_id=current_user.id,
updater_role=current_user.role,
**update_data,
)
log_action(
db,
user_id=current_user.id,
action="update_campaign",
entity_type="campaign",
entity_id=campaign_id,
details={"updated_fields": list(update_data.keys())},
)
uow.commit()
return result
# ---------------------------------------------------------------------------
# POST /campaigns/{id}/tests — Add test to campaign
# ---------------------------------------------------------------------------
@router.post("/{campaign_id}/tests")
def add_test_to_campaign(
campaign_id: str,
payload: AddTestPayload,
db: Session = Depends(get_db),
current_user: User = Depends(require_any_role("red_lead", "blue_lead")),
) -> dict:
"""Add a test to a campaign with optional ordering and dependency."""
with UnitOfWork(db) as uow:
result = crud_add_test(
db,
campaign_id,
test_id=payload.test_id,
order_index=payload.order_index,
depends_on=payload.depends_on,
phase=payload.phase,
)
uow.commit()
return result
# ---------------------------------------------------------------------------
# DELETE /campaigns/{id}/tests/{campaign_test_id} — Remove test from campaign
# ---------------------------------------------------------------------------
@router.delete("/{campaign_id}/tests/{campaign_test_id}")
def remove_test_from_campaign(
campaign_id: str,
campaign_test_id: str,
db: Session = Depends(get_db),
current_user: User = Depends(require_any_role("red_lead", "blue_lead")),
) -> dict:
"""Remove a test from a campaign."""
with UnitOfWork(db) as uow:
crud_remove_test(db, campaign_id, campaign_test_id)
uow.commit()
return {"detail": "Test removed from campaign"}
# ---------------------------------------------------------------------------
# POST /campaigns/{id}/activate — Activate campaign
# ---------------------------------------------------------------------------
@router.post("/{campaign_id}/activate")
def activate_campaign(
campaign_id: str,
db: Session = Depends(get_db),
current_user: User = Depends(require_any_role("red_lead", "blue_lead")),
) -> dict:
"""Activate a campaign, moving it from draft to active."""
with UnitOfWork(db) as uow:
campaign = crud_activate(db, campaign_id)
notify_role(
db,
role="red_tech",
type="campaign_activated",
title="Campaign activated",
message=f'Campaign "{campaign.name}" has been activated.',
entity_type="campaign",
entity_id=campaign.id,
)
log_action(
db,
user_id=current_user.id,
action="activate_campaign",
entity_type="campaign",
entity_id=campaign.id,
details={"name": campaign.name},
)
uow.commit()
db.refresh(campaign)
return serialize_campaign(db, campaign)
# ---------------------------------------------------------------------------
# POST /campaigns/{id}/complete — Mark campaign as completed
# ---------------------------------------------------------------------------
@router.post("/{campaign_id}/complete")
def complete_campaign(
campaign_id: str,
db: Session = Depends(get_db),
current_user: User = Depends(require_any_role("red_lead", "admin")),
) -> dict:
"""Mark a campaign as completed."""
with UnitOfWork(db) as uow:
campaign = crud_complete(db, campaign_id)
log_action(
db,
user_id=current_user.id,
action="complete_campaign",
entity_type="campaign",
entity_id=campaign.id,
details={"name": campaign.name},
)
uow.commit()
db.refresh(campaign)
return serialize_campaign(db, campaign)
# ---------------------------------------------------------------------------
# GET /campaigns/{id}/progress — Campaign progress
# ---------------------------------------------------------------------------
@router.get("/{campaign_id}/progress")
def get_campaign_progress_endpoint(
campaign_id: str,
db: Session = Depends(get_db),
current_user: User = Depends(get_current_user),
) -> dict:
"""Get progress statistics for a campaign."""
return crud_get_progress(db, campaign_id)
# ---------------------------------------------------------------------------
# POST /campaigns/from-threat-actor/{actor_id} — Auto-generate campaign
# ---------------------------------------------------------------------------
@router.post("/from-threat-actor/{actor_id}", status_code=201)
def generate_campaign_from_actor(
actor_id: str,
db: Session = Depends(get_db),
current_user: User = Depends(require_any_role("red_lead", "blue_lead")),
) -> dict:
"""Auto-generate a campaign from a threat actor's uncovered techniques.
Creates tests from the best available templates and orders them
by kill chain phase.
"""
campaign = generate_campaign_from_threat_actor(
db,
uuid.UUID(actor_id),
current_user,
)
with UnitOfWork(db) as uow:
log_action(
db,
user_id=current_user.id,
action="generate_campaign",
entity_type="campaign",
entity_id=campaign.id,
details={"actor_id": actor_id, "campaign_name": campaign.name},
)
uow.commit()
return serialize_campaign(db, campaign)
# ---------------------------------------------------------------------------
# PATCH /campaigns/{id}/schedule — Configure recurrence
# ---------------------------------------------------------------------------
@router.patch("/{campaign_id}/schedule")
def schedule_campaign(
campaign_id: str,
payload: SchedulePayload,
db: Session = Depends(get_db),
current_user: User = Depends(require_any_role("red_lead", "blue_lead")),
) -> dict:
"""Configure or update the recurrence schedule for a campaign.
Only the campaign creator or admin can change scheduling.
"""
with UnitOfWork(db) as uow:
campaign = crud_schedule(
db,
campaign_id,
owner_id=current_user.id,
owner_role=current_user.role,
is_recurring=payload.is_recurring,
recurrence_pattern=payload.recurrence_pattern,
next_run_at=payload.next_run_at,
)
log_action(
db,
user_id=current_user.id,
action="schedule_campaign",
entity_type="campaign",
entity_id=campaign.id,
details={
"is_recurring": campaign.is_recurring,
"recurrence_pattern": campaign.recurrence_pattern,
"next_run_at": campaign.next_run_at.isoformat() if campaign.next_run_at else None,
},
)
uow.commit()
db.refresh(campaign)
return serialize_campaign(db, campaign)
# ---------------------------------------------------------------------------
# GET /campaigns/{id}/history — Execution history (child campaigns)
# ---------------------------------------------------------------------------
@router.get("/{campaign_id}/history")
def get_campaign_history(
campaign_id: str,
db: Session = Depends(get_db),
current_user: User = Depends(get_current_user),
) -> list:
"""List all child campaigns (execution history) of a recurring campaign."""
return crud_get_history(db, campaign_id)