Aegis/backend/app/services/campaign_scheduler_service.py

"""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
import uuid
from datetime import datetime, timedelta

from sqlalchemy.orm import Session

from app.models.campaign import Campaign, CampaignTest
from app.models.test import Test
from app.models.enums import TestState
from app.services.notification_service import create_notification
from app.services.audit_service import log_action
from app.models.user import User

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 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