prowler/api/src/backend/tasks/beat.py

import json
from datetime import datetime, timedelta, timezone

from django_celery_beat.models import IntervalSchedule, PeriodicTask
from tasks.tasks import perform_scheduled_scan_task

from api.db_utils import rls_transaction
from api.exceptions import ConflictException
from api.models import Provider, Scan, StateChoices


def schedule_provider_scan(provider_instance: Provider):
    tenant_id = str(provider_instance.tenant_id)
    provider_id = str(provider_instance.id)

    schedule, _ = IntervalSchedule.objects.get_or_create(
        every=24,
        period=IntervalSchedule.HOURS,
    )

    # Create a unique name for the periodic task
    task_name = f"scan-perform-scheduled-{provider_instance.id}"

    if PeriodicTask.objects.filter(
        interval=schedule, name=task_name, task="scan-perform-scheduled"
    ).exists():
        raise ConflictException(
            detail="There is already a scheduled scan for this provider.",
            pointer="/data/attributes/provider_id",
        )

    with rls_transaction(tenant_id):
        scheduled_scan = Scan.objects.create(
            tenant_id=tenant_id,
            name="Daily scheduled scan",
            provider_id=provider_id,
            trigger=Scan.TriggerChoices.SCHEDULED,
            state=StateChoices.AVAILABLE,
            scheduled_at=datetime.now(timezone.utc),
        )

    # Schedule the task
    periodic_task_instance = PeriodicTask.objects.create(
        interval=schedule,
        name=task_name,
        task="scan-perform-scheduled",
        kwargs=json.dumps(
            {
                "tenant_id": tenant_id,
                "provider_id": provider_id,
            }
        ),
        one_off=False,
        start_time=datetime.now(timezone.utc) + timedelta(hours=24),
    )
    scheduled_scan.scheduler_task_id = periodic_task_instance.id
    scheduled_scan.save()

    return perform_scheduled_scan_task.apply_async(
        kwargs={
            "tenant_id": str(provider_instance.tenant_id),
            "provider_id": provider_id,
        },
    )