master_workflow.py

from datetime import timedelta

from temporalio import workflow
from temporalio.exceptions import ApplicationError

from fast_scan_workflow import FastScanWorkflow
from complete_scan_workflow import CompleteScanWorkflow
from asset_port_scanning_workflow import AssetScanWorkflow
from utility_activities import build_absolute_root_path
from config import ScanConfig



@workflow.defn
class MasterWorkflow:
    @workflow.run
    async def run(self, config: ScanConfig) -> int:

        if config.scan_type in ("fast", "complete") and not config.domains:
            raise ApplicationError(
                "domains must be provided for fast/complete scans",
            type="InvalidScanConfig",
            non_retryable=True,
            )

        if config.scan_type == "asset" and not config.ips:
            raise ApplicationError(
                "ips must be provided for asset scans",
                type="InvalidScanConfig",
                non_retryable=True,
            )


        timestamp = workflow.now().strftime("%Y%m%d_%H%M%S")
        workflow_id = workflow.info().workflow_id
        config.output_dir = await workflow.execute_activity(
            build_absolute_root_path,
            args=[config, timestamp, workflow_id],
            start_to_close_timeout=timedelta(seconds=5)
        )

        match config.scan_type:
            case "fast":
                await workflow.execute_child_workflow(
                    FastScanWorkflow.run,
                    args=[config],
                    id=f"fast-{workflow.info().workflow_id}",
                    task_queue="scan-task-queue"
                )
            case "complete":
                await workflow.execute_child_workflow(
                    CompleteScanWorkflow.run,
                    args=[config],
                    id=f"complete-{workflow.info().workflow_id}",
                    task_queue="scan-task-queue"
                )
            case "asset":
                await workflow.execute_child_workflow(
                    AssetScanWorkflow.run,
                    args=[config],
                    id=f"asset-{workflow.info().workflow_id}",
                    task_queue="scan-task-queue"
                )
            case _:
                raise ApplicationError(
                    f"Unknown scan kind: {config.scan_type!r}",
                    type="InvalidScanKind",
                    non_retryable=True,
                )


        return 0