Files
Aegis/backend/app/services/campaign_scheduler_service.py
T
kitos 8f98bdd273 refactor(pep8): enforce full PEP8 compliance across backend Python codebase
- ruff.toml: select E/W/F/I/N rules, line-length=120, drop legacy ignores
- Auto-fix: sort 82 import blocks (isort), remove 29 unused imports,
  strip 6 trailing-whitespace blank lines in docstrings
- main.py: move setup_logging and settings imports to top (E402)
- errors.py: noqa N818 on DDD exception names (96 call sites, safe)
- intel_service.py: noqa N817 for universal ET alias
- atomic/elastic/sigma import services: move _MAX_UNCOMPRESSED_SIZE and
  _MAX_ENTRIES to module level (N806)
- compliance_import_service.py: move SAMPLE_CONTROLS / CIS_CONTROLS to
  module level; wrap long description strings (N806 + E501)
- snapshot_service.py: move STATUS_ORDER dict to module level (N806)
- sigma_import_service.py: remove dead dedup_key expression (F841)
- threat_actor_import_service.py: remove dead stix_to_actor expression (F841)
- data_source.py, seed_demo.py, campaign_scheduler_service.py,
  lolbas_import_service.py: wrap lines exceeding 120 chars (E501)
- d3fend_import_service.py: per-file E501 ignore (data file with long strings)

All 439 unit tests pass. ruff check app/ → All checks passed!
2026-06-11 11:06:54 +02:00

197 lines
6.1 KiB
Python

"""Campaign scheduler service — recurring campaign execution.
Handles checking which recurring campaigns are due, cloning them with
fresh tests, and computing the next run date.
"""
import logging
from datetime import datetime, timedelta
from sqlalchemy.orm import Session
from app.models.campaign import Campaign, CampaignTest
from app.models.enums import TestState
from app.models.test import Test
from app.models.user import User
from app.services.audit_service import log_action
from app.services.notification_service import create_notification
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# Next-run calculation
# ---------------------------------------------------------------------------
def calculate_next_run(current_date: datetime, pattern: str) -> datetime:
"""Compute the next run date from *current_date* and a recurrence pattern.
Supported patterns:
- ``weekly`` : +7 days
- ``monthly`` : +30 days
- ``quarterly``: +90 days
"""
offsets = {
"weekly": timedelta(days=7),
"monthly": timedelta(days=30),
"quarterly": timedelta(days=90),
}
return current_date + offsets.get(pattern, timedelta(days=30))
# ---------------------------------------------------------------------------
# Clone a campaign
# ---------------------------------------------------------------------------
def _clone_campaign(db: Session, original: Campaign) -> Campaign:
"""Create a new child campaign from a recurring template.
1. Clone the campaign with a date-stamped name.
2. For each ``CampaignTest`` in the original, create a new ``Test``
with the same base data (in ``draft`` state) and link it.
3. Activate the new campaign.
"""
now = datetime.utcnow()
run_label = now.strftime("%Y-%m-%d")
child = Campaign(
name=f"{original.name} (Run {run_label})",
description=original.description,
type=original.type,
threat_actor_id=original.threat_actor_id,
status="active",
created_by=original.created_by,
target_platform=original.target_platform,
tags=original.tags or [],
parent_campaign_id=original.id,
)
db.add(child)
db.flush() # get child.id
# Clone each campaign_test with a fresh Test
original_cts = (
db.query(CampaignTest)
.filter(CampaignTest.campaign_id == original.id)
.order_by(CampaignTest.order_index)
.all()
)
for ct in original_cts:
src_test = ct.test
if not src_test:
continue
new_test = Test(
technique_id=src_test.technique_id,
name=src_test.name,
description=src_test.description,
platform=src_test.platform,
procedure_text=src_test.procedure_text,
tool_used=src_test.tool_used,
created_by=original.created_by,
state=TestState.draft,
)
db.add(new_test)
db.flush() # get new_test.id
new_ct = CampaignTest(
campaign_id=child.id,
test_id=new_test.id,
order_index=ct.order_index,
phase=ct.phase,
# depends_on is not copied — would need ID remapping
)
db.add(new_ct)
db.flush()
return child
# ---------------------------------------------------------------------------
# Check and run recurring campaigns (daily job)
# ---------------------------------------------------------------------------
def check_and_run_recurring_campaigns(db: Session) -> int:
"""Check all recurring campaigns and clone any that are due.
Returns the number of campaigns spawned.
"""
now = datetime.utcnow()
due_campaigns = (
db.query(Campaign)
.filter(
Campaign.is_recurring == True, # noqa: E712
Campaign.next_run_at <= now,
)
.all()
)
spawned = 0
for campaign in due_campaigns:
try:
child = _clone_campaign(db, campaign)
# Update the original's scheduling fields
campaign.last_run_at = now
campaign.next_run_at = calculate_next_run(now, campaign.recurrence_pattern or "monthly")
db.commit()
db.refresh(child)
# Audit
log_action(
db,
user_id=campaign.created_by,
action="recurring_campaign_run",
entity_type="campaign",
entity_id=child.id,
details={
"parent_campaign_id": str(campaign.id),
"child_campaign_name": child.name,
"pattern": campaign.recurrence_pattern,
},
)
db.commit()
# Notify
if campaign.created_by:
create_notification(
db,
user_id=campaign.created_by,
type="recurring_campaign_run",
title="Recurring campaign executed",
message=(
f'Campaign "{child.name}" was automatically created '
f'from recurring template "{campaign.name}".'
),
entity_type="campaign",
entity_id=child.id,
)
# Notify red_tech users
red_techs = db.query(User).filter(User.role == "red_tech", User.is_active == True).all() # noqa: E712
for user in red_techs:
create_notification(
db,
user_id=user.id,
type="campaign_activated",
title="New recurring campaign active",
message=f'Campaign "{child.name}" is now active and ready for execution.',
entity_type="campaign",
entity_id=child.id,
)
spawned += 1
logger.info("Spawned child campaign '%s' from parent '%s'", child.name, campaign.name)
except Exception:
db.rollback()
logger.exception("Failed to run recurring campaign '%s'", campaign.name)
return spawned