feat: Complete migration from Prefect to Temporal

BREAKING CHANGE: Replaces Prefect workflow orchestration with Temporal

## Major Changes
- Replace Prefect with Temporal for workflow orchestration
- Implement vertical worker architecture (rust, android)
- Replace Docker registry with MinIO for unified storage
- Refactor activities to be co-located with workflows
- Update all API endpoints for Temporal compatibility

## Infrastructure
- New: docker-compose.temporal.yaml (Temporal + MinIO + workers)
- New: workers/ directory with rust and android vertical workers
- New: backend/src/temporal/ (manager, discovery)
- New: backend/src/storage/ (S3-cached storage with MinIO)
- New: backend/toolbox/common/ (shared storage activities)
- Deleted: docker-compose.yaml (old Prefect setup)
- Deleted: backend/src/core/prefect_manager.py
- Deleted: backend/src/services/prefect_stats_monitor.py
- Deleted: Docker registry and insecure-registries requirement

## Workflows
- Migrated: security_assessment workflow to Temporal
- New: rust_test workflow (example/test workflow)
- Deleted: secret_detection_scan (Prefect-based, to be reimplemented)
- Activities now co-located with workflows for independent testing

## API Changes
- Updated: backend/src/api/workflows.py (Temporal submission)
- Updated: backend/src/api/runs.py (Temporal status/results)
- Updated: backend/src/main.py (727 lines, TemporalManager integration)
- Updated: All 16 MCP tools to use TemporalManager

## Testing
-  All services healthy (Temporal, PostgreSQL, MinIO, workers, backend)
-  All API endpoints functional
-  End-to-end workflow test passed (72 findings from vulnerable_app)
-  MinIO storage integration working (target upload/download, results)
-  Worker activity discovery working (6 activities registered)
-  Tarball extraction working
-  SARIF report generation working

## Documentation
- ARCHITECTURE.md: Complete Temporal architecture documentation
- QUICKSTART_TEMPORAL.md: Getting started guide
- MIGRATION_DECISION.md: Why we chose Temporal over Prefect
- IMPLEMENTATION_STATUS.md: Migration progress tracking
- workers/README.md: Worker development guide

## Dependencies
- Added: temporalio>=1.6.0
- Added: boto3>=1.34.0 (MinIO S3 client)
- Removed: prefect>=3.4.18
This commit is contained in:
Tanguy Duhamel
2025-10-01 15:11:24 +02:00
parent 987c49569c
commit 0680f14df6
48 changed files with 7504 additions and 3101 deletions
+5 -9
View File
@@ -17,25 +17,21 @@ RUN apt-get update && apt-get install -y \
# Docker client configuration removed - localhost:5001 doesn't require insecure registry config
# Install uv for faster package management
RUN pip install uv
# Copy project files
COPY pyproject.toml ./
COPY uv.lock ./
# Install dependencies
RUN uv sync --no-dev
# Install dependencies with pip
RUN pip install --no-cache-dir -e .
# Copy source code
COPY . .
# Expose port
EXPOSE 8000
# Expose ports (API on 8000, MCP on 8010)
EXPOSE 8000 8010
# Health check
HEALTHCHECK --interval=30s --timeout=10s --start-period=5s --retries=3 \
CMD curl -f http://localhost:8000/health || exit 1
# Start the application
CMD ["uv", "run", "uvicorn", "src.main:app", "--host", "0.0.0.0", "--port", "8000"]
CMD ["uvicorn", "src.main:app", "--host", "0.0.0.0", "--port", "8000"]
+2 -1
View File
@@ -7,7 +7,8 @@ readme = "README.md"
requires-python = ">=3.11"
dependencies = [
"fastapi>=0.116.1",
"prefect>=3.4.18",
"temporalio>=1.6.0",
"boto3>=1.34.0",
"pydantic>=2.0.0",
"pyyaml>=6.0",
"docker>=7.0.0",
-1
View File
@@ -25,7 +25,6 @@ from src.models.findings import (
FuzzingStats,
CrashReport
)
from src.core.workflow_discovery import WorkflowDiscovery
logger = logging.getLogger(__name__)
+49 -55
View File
@@ -24,22 +24,22 @@ logger = logging.getLogger(__name__)
router = APIRouter(prefix="/runs", tags=["runs"])
def get_prefect_manager():
"""Dependency to get the Prefect manager instance"""
from src.main import prefect_mgr
return prefect_mgr
def get_temporal_manager():
"""Dependency to get the Temporal manager instance"""
from src.main import temporal_mgr
return temporal_mgr
@router.get("/{run_id}/status", response_model=WorkflowStatus)
async def get_run_status(
run_id: str,
prefect_mgr=Depends(get_prefect_manager)
temporal_mgr=Depends(get_temporal_manager)
) -> WorkflowStatus:
"""
Get the current status of a workflow run.
Args:
run_id: The flow run ID
run_id: The workflow run ID
Returns:
Status information including state, timestamps, and completion flags
@@ -48,25 +48,23 @@ async def get_run_status(
HTTPException: 404 if run not found
"""
try:
status = await prefect_mgr.get_flow_run_status(run_id)
status = await temporal_mgr.get_workflow_status(run_id)
# Find workflow name from deployment
workflow_name = "unknown"
workflow_deployment_id = status.get("workflow", "")
for name, deployment_id in prefect_mgr.deployments.items():
if str(deployment_id) == str(workflow_deployment_id):
workflow_name = name
break
# Map Temporal status to response format
workflow_status = status.get("status", "UNKNOWN")
is_completed = workflow_status in ["COMPLETED", "FAILED", "CANCELLED"]
is_failed = workflow_status == "FAILED"
is_running = workflow_status == "RUNNING"
return WorkflowStatus(
run_id=status["run_id"],
workflow=workflow_name,
status=status["status"],
is_completed=status["is_completed"],
is_failed=status["is_failed"],
is_running=status["is_running"],
created_at=status["created_at"],
updated_at=status["updated_at"]
run_id=run_id,
workflow="unknown", # Temporal doesn't track workflow name in status
status=workflow_status,
is_completed=is_completed,
is_failed=is_failed,
is_running=is_running,
created_at=status.get("start_time"),
updated_at=status.get("close_time") or status.get("execution_time")
)
except Exception as e:
@@ -80,13 +78,13 @@ async def get_run_status(
@router.get("/{run_id}/findings", response_model=WorkflowFindings)
async def get_run_findings(
run_id: str,
prefect_mgr=Depends(get_prefect_manager)
temporal_mgr=Depends(get_temporal_manager)
) -> WorkflowFindings:
"""
Get the findings from a completed workflow run.
Args:
run_id: The flow run ID
run_id: The workflow run ID
Returns:
SARIF-formatted findings from the workflow execution
@@ -96,50 +94,46 @@ async def get_run_findings(
"""
try:
# Get run status first
status = await prefect_mgr.get_flow_run_status(run_id)
status = await temporal_mgr.get_workflow_status(run_id)
workflow_status = status.get("status", "UNKNOWN")
if not status["is_completed"]:
if status["is_running"]:
if workflow_status not in ["COMPLETED", "FAILED", "CANCELLED"]:
if workflow_status == "RUNNING":
raise HTTPException(
status_code=400,
detail=f"Run {run_id} is still running. Current status: {status['status']}"
)
elif status["is_failed"]:
raise HTTPException(
status_code=400,
detail=f"Run {run_id} failed. Status: {status['status']}"
detail=f"Run {run_id} is still running. Current status: {workflow_status}"
)
else:
raise HTTPException(
status_code=400,
detail=f"Run {run_id} not completed. Status: {status['status']}"
detail=f"Run {run_id} not completed. Status: {workflow_status}"
)
# Get the findings
findings = await prefect_mgr.get_flow_run_findings(run_id)
if workflow_status == "FAILED":
raise HTTPException(
status_code=400,
detail=f"Run {run_id} failed. Status: {workflow_status}"
)
# Find workflow name
workflow_name = "unknown"
workflow_deployment_id = status.get("workflow", "")
for name, deployment_id in prefect_mgr.deployments.items():
if str(deployment_id) == str(workflow_deployment_id):
workflow_name = name
break
# Get the workflow result
result = await temporal_mgr.get_workflow_result(run_id)
# Get workflow version if available
# Extract SARIF from result
if isinstance(result, dict):
sarif = result.get("sarif", {})
else:
sarif = {}
# Metadata
metadata = {
"completion_time": status["updated_at"],
"completion_time": status.get("close_time"),
"workflow_version": "unknown"
}
if workflow_name in prefect_mgr.workflows:
workflow_info = prefect_mgr.workflows[workflow_name]
metadata["workflow_version"] = workflow_info.metadata.get("version", "unknown")
return WorkflowFindings(
workflow=workflow_name,
workflow="unknown",
run_id=run_id,
sarif=findings,
sarif=sarif,
metadata=metadata
)
@@ -157,7 +151,7 @@ async def get_run_findings(
async def get_workflow_findings(
workflow_name: str,
run_id: str,
prefect_mgr=Depends(get_prefect_manager)
temporal_mgr=Depends(get_temporal_manager)
) -> WorkflowFindings:
"""
Get findings for a specific workflow run.
@@ -166,7 +160,7 @@ async def get_workflow_findings(
Args:
workflow_name: Name of the workflow
run_id: The flow run ID
run_id: The workflow run ID
Returns:
SARIF-formatted findings from the workflow execution
@@ -174,11 +168,11 @@ async def get_workflow_findings(
Raises:
HTTPException: 404 if workflow or run not found, 400 if run not completed
"""
if workflow_name not in prefect_mgr.workflows:
if workflow_name not in temporal_mgr.workflows:
raise HTTPException(
status_code=404,
detail=f"Workflow not found: {workflow_name}"
)
# Delegate to the main findings endpoint
return await get_run_findings(run_id, prefect_mgr)
return await get_run_findings(run_id, temporal_mgr)
+49 -52
View File
@@ -25,7 +25,7 @@ from src.models.findings import (
WorkflowListItem,
RunSubmissionResponse
)
from src.core.workflow_discovery import WorkflowDiscovery
from src.temporal.discovery import WorkflowDiscovery
logger = logging.getLogger(__name__)
@@ -68,15 +68,15 @@ def create_structured_error_response(
return error_response
def get_prefect_manager():
"""Dependency to get the Prefect manager instance"""
from src.main import prefect_mgr
return prefect_mgr
def get_temporal_manager():
"""Dependency to get the Temporal manager instance"""
from src.main import temporal_mgr
return temporal_mgr
@router.get("/", response_model=List[WorkflowListItem])
async def list_workflows(
prefect_mgr=Depends(get_prefect_manager)
temporal_mgr=Depends(get_temporal_manager)
) -> List[WorkflowListItem]:
"""
List all discovered workflows with their metadata.
@@ -85,7 +85,7 @@ async def list_workflows(
author, and tags.
"""
workflows = []
for name, info in prefect_mgr.workflows.items():
for name, info in temporal_mgr.workflows.items():
workflows.append(WorkflowListItem(
name=name,
version=info.metadata.get("version", "0.6.0"),
@@ -111,7 +111,7 @@ async def get_metadata_schema() -> Dict[str, Any]:
@router.get("/{workflow_name}/metadata", response_model=WorkflowMetadata)
async def get_workflow_metadata(
workflow_name: str,
prefect_mgr=Depends(get_prefect_manager)
temporal_mgr=Depends(get_temporal_manager)
) -> WorkflowMetadata:
"""
Get complete metadata for a specific workflow.
@@ -126,8 +126,8 @@ async def get_workflow_metadata(
Raises:
HTTPException: 404 if workflow not found
"""
if workflow_name not in prefect_mgr.workflows:
available_workflows = list(prefect_mgr.workflows.keys())
if workflow_name not in temporal_mgr.workflows:
available_workflows = list(temporal_mgr.workflows.keys())
error_response = create_structured_error_response(
error_type="WorkflowNotFound",
message=f"Workflow '{workflow_name}' not found",
@@ -143,7 +143,7 @@ async def get_workflow_metadata(
detail=error_response
)
info = prefect_mgr.workflows[workflow_name]
info = temporal_mgr.workflows[workflow_name]
metadata = info.metadata
return WorkflowMetadata(
@@ -156,7 +156,7 @@ async def get_workflow_metadata(
default_parameters=metadata.get("default_parameters", {}),
required_modules=metadata.get("required_modules", []),
supported_volume_modes=metadata.get("supported_volume_modes", ["ro", "rw"]),
has_custom_docker=info.has_docker
has_custom_docker=metadata.get("has_docker", False)
)
@@ -164,14 +164,14 @@ async def get_workflow_metadata(
async def submit_workflow(
workflow_name: str,
submission: WorkflowSubmission,
prefect_mgr=Depends(get_prefect_manager)
temporal_mgr=Depends(get_temporal_manager)
) -> RunSubmissionResponse:
"""
Submit a workflow for execution with volume mounting.
Submit a workflow for execution.
Args:
workflow_name: Name of the workflow to execute
submission: Submission parameters including target path and volume mode
submission: Submission parameters including target path and parameters
Returns:
Run submission response with run_id and initial status
@@ -179,8 +179,8 @@ async def submit_workflow(
Raises:
HTTPException: 404 if workflow not found, 400 for invalid parameters
"""
if workflow_name not in prefect_mgr.workflows:
available_workflows = list(prefect_mgr.workflows.keys())
if workflow_name not in temporal_mgr.workflows:
available_workflows = list(temporal_mgr.workflows.keys())
error_response = create_structured_error_response(
error_type="WorkflowNotFound",
message=f"Workflow '{workflow_name}' not found",
@@ -197,31 +197,32 @@ async def submit_workflow(
)
try:
# Convert ResourceLimits to dict if provided
resource_limits_dict = None
if submission.resource_limits:
resource_limits_dict = {
"cpu_limit": submission.resource_limits.cpu_limit,
"memory_limit": submission.resource_limits.memory_limit,
"cpu_request": submission.resource_limits.cpu_request,
"memory_request": submission.resource_limits.memory_request
}
# Upload target file to MinIO and get target_id
target_path = Path(submission.target_path)
if not target_path.exists():
raise ValueError(f"Target path does not exist: {submission.target_path}")
# Submit the workflow with enhanced parameters
flow_run = await prefect_mgr.submit_workflow(
workflow_name=workflow_name,
target_path=submission.target_path,
volume_mode=submission.volume_mode,
parameters=submission.parameters,
resource_limits=resource_limits_dict,
additional_volumes=submission.additional_volumes,
timeout=submission.timeout
# Upload target (using anonymous user for now)
target_id = await temporal_mgr.upload_target(
file_path=target_path,
user_id="api-user",
metadata={"workflow": workflow_name}
)
run_id = str(flow_run.id)
# Prepare workflow parameters
workflow_params = submission.parameters or {}
# Start workflow execution
handle = await temporal_mgr.run_workflow(
workflow_name=workflow_name,
target_id=target_id,
workflow_params=workflow_params
)
run_id = handle.id
# Initialize fuzzing tracking if this looks like a fuzzing workflow
workflow_info = prefect_mgr.workflows.get(workflow_name, {})
workflow_info = temporal_mgr.workflows.get(workflow_name, {})
workflow_tags = workflow_info.metadata.get("tags", []) if hasattr(workflow_info, 'metadata') else []
if "fuzzing" in workflow_tags or "fuzz" in workflow_name.lower():
from src.api.fuzzing import initialize_fuzzing_tracking
@@ -229,7 +230,7 @@ async def submit_workflow(
return RunSubmissionResponse(
run_id=run_id,
status=flow_run.state.name if flow_run.state else "PENDING",
status="RUNNING",
workflow=workflow_name,
message=f"Workflow '{workflow_name}' submitted successfully"
)
@@ -261,17 +262,13 @@ async def submit_workflow(
error_type = "WorkflowSubmissionError"
# Detect specific error patterns
if "deployment" in error_message.lower():
error_type = "DeploymentError"
deployment_info = {
"status": "failed",
"error": error_message
}
if "workflow" in error_message.lower() and "not found" in error_message.lower():
error_type = "WorkflowError"
suggestions.extend([
"Check if Prefect server is running and accessible",
"Verify Docker is running and has sufficient resources",
"Check container image availability",
"Ensure volume paths exist and are accessible"
"Check if Temporal server is running and accessible",
"Verify workflow workers are running",
"Check if workflow is registered with correct vertical",
"Ensure Docker is running and has sufficient resources"
])
elif "volume" in error_message.lower() or "mount" in error_message.lower():
@@ -327,7 +324,7 @@ async def submit_workflow(
@router.get("/{workflow_name}/parameters")
async def get_workflow_parameters(
workflow_name: str,
prefect_mgr=Depends(get_prefect_manager)
temporal_mgr=Depends(get_temporal_manager)
) -> Dict[str, Any]:
"""
Get the parameters schema for a workflow.
@@ -341,8 +338,8 @@ async def get_workflow_parameters(
Raises:
HTTPException: 404 if workflow not found
"""
if workflow_name not in prefect_mgr.workflows:
available_workflows = list(prefect_mgr.workflows.keys())
if workflow_name not in temporal_mgr.workflows:
available_workflows = list(temporal_mgr.workflows.keys())
error_response = create_structured_error_response(
error_type="WorkflowNotFound",
message=f"Workflow '{workflow_name}' not found",
@@ -357,7 +354,7 @@ async def get_workflow_parameters(
detail=error_response
)
info = prefect_mgr.workflows[workflow_name]
info = temporal_mgr.workflows[workflow_name]
metadata = info.metadata
# Return parameters with enhanced schema information
-770
View File
@@ -1,770 +0,0 @@
"""
Prefect Manager - Core orchestration for workflow deployment and execution
"""
# Copyright (c) 2025 FuzzingLabs
#
# Licensed under the Business Source License 1.1 (BSL). See the LICENSE file
# at the root of this repository for details.
#
# After the Change Date (four years from publication), this version of the
# Licensed Work will be made available under the Apache License, Version 2.0.
# See the LICENSE-APACHE file or http://www.apache.org/licenses/LICENSE-2.0
#
# Additional attribution and requirements are provided in the NOTICE file.
import logging
import os
import platform
import re
from pathlib import Path
from typing import Dict, Optional, Any
from prefect import get_client
from prefect.docker import DockerImage
from prefect.client.schemas import FlowRun
from src.core.workflow_discovery import WorkflowDiscovery, WorkflowInfo
logger = logging.getLogger(__name__)
def get_registry_url(context: str = "default") -> str:
"""
Get the container registry URL to use for a given operation context.
Goals:
- Work reliably across Linux and macOS Docker Desktop
- Prefer in-network service discovery when running inside containers
- Allow full override via env vars from docker-compose
Env overrides:
- FUZZFORGE_REGISTRY_PUSH_URL: used for image builds/pushes
- FUZZFORGE_REGISTRY_PULL_URL: used for workers to pull images
"""
# Normalize context
ctx = (context or "default").lower()
# Always honor explicit overrides first
if ctx in ("push", "build"):
push_url = os.getenv("FUZZFORGE_REGISTRY_PUSH_URL")
if push_url:
logger.debug("Using FUZZFORGE_REGISTRY_PUSH_URL: %s", push_url)
return push_url
# Default to host-published registry for Docker daemon operations
return "localhost:5001"
if ctx == "pull":
pull_url = os.getenv("FUZZFORGE_REGISTRY_PULL_URL")
if pull_url:
logger.debug("Using FUZZFORGE_REGISTRY_PULL_URL: %s", pull_url)
return pull_url
# Prefect worker pulls via host Docker daemon as well
return "localhost:5001"
# Default/fallback
return os.getenv("FUZZFORGE_REGISTRY_PULL_URL", os.getenv("FUZZFORGE_REGISTRY_PUSH_URL", "localhost:5001"))
def _compose_project_name(default: str = "fuzzforge") -> str:
"""Return the docker-compose project name used for network/volume naming.
Always returns 'fuzzforge' regardless of environment variables.
"""
return "fuzzforge"
class PrefectManager:
"""
Manages Prefect deployments and flow runs for discovered workflows.
This class handles:
- Workflow discovery and registration
- Docker image building through Prefect
- Deployment creation and management
- Flow run submission with volume mounting
- Findings retrieval from completed runs
"""
def __init__(self, workflows_dir: Path = None):
"""
Initialize the Prefect manager.
Args:
workflows_dir: Path to the workflows directory (default: toolbox/workflows)
"""
if workflows_dir is None:
workflows_dir = Path("toolbox/workflows")
self.discovery = WorkflowDiscovery(workflows_dir)
self.workflows: Dict[str, WorkflowInfo] = {}
self.deployments: Dict[str, str] = {} # workflow_name -> deployment_id
# Security: Define allowed and forbidden paths for host mounting
self.allowed_base_paths = [
"/tmp",
"/home",
"/Users", # macOS users
"/opt",
"/var/tmp",
"/workspace", # Common container workspace
"/app" # Container application directory (for test projects)
]
self.forbidden_paths = [
"/etc",
"/root",
"/var/run",
"/sys",
"/proc",
"/dev",
"/boot",
"/var/lib/docker", # Critical Docker data
"/var/log", # System logs
"/usr/bin", # System binaries
"/usr/sbin",
"/sbin",
"/bin"
]
@staticmethod
def _parse_memory_to_bytes(memory_str: str) -> int:
"""
Parse memory string (like '512Mi', '1Gi') to bytes.
Args:
memory_str: Memory string with unit suffix
Returns:
Memory in bytes
Raises:
ValueError: If format is invalid
"""
if not memory_str:
return 0
match = re.match(r'^(\d+(?:\.\d+)?)\s*([GMK]i?)$', memory_str.strip())
if not match:
raise ValueError(f"Invalid memory format: {memory_str}. Expected format like '512Mi', '1Gi'")
value, unit = match.groups()
value = float(value)
# Convert to bytes based on unit (binary units: Ki, Mi, Gi)
if unit in ['K', 'Ki']:
multiplier = 1024
elif unit in ['M', 'Mi']:
multiplier = 1024 * 1024
elif unit in ['G', 'Gi']:
multiplier = 1024 * 1024 * 1024
else:
raise ValueError(f"Unsupported memory unit: {unit}")
return int(value * multiplier)
@staticmethod
def _parse_cpu_to_millicores(cpu_str: str) -> int:
"""
Parse CPU string (like '500m', '1', '2.5') to millicores.
Args:
cpu_str: CPU string
Returns:
CPU in millicores (1 core = 1000 millicores)
Raises:
ValueError: If format is invalid
"""
if not cpu_str:
return 0
cpu_str = cpu_str.strip()
# Handle millicores format (e.g., '500m')
if cpu_str.endswith('m'):
try:
return int(cpu_str[:-1])
except ValueError:
raise ValueError(f"Invalid CPU format: {cpu_str}")
# Handle core format (e.g., '1', '2.5')
try:
cores = float(cpu_str)
return int(cores * 1000) # Convert to millicores
except ValueError:
raise ValueError(f"Invalid CPU format: {cpu_str}")
def _extract_resource_requirements(self, workflow_info: WorkflowInfo) -> Dict[str, str]:
"""
Extract resource requirements from workflow metadata.
Args:
workflow_info: Workflow information with metadata
Returns:
Dictionary with resource requirements in Docker format
"""
metadata = workflow_info.metadata
requirements = metadata.get("requirements", {})
resources = requirements.get("resources", {})
resource_config = {}
# Extract memory requirement
memory = resources.get("memory")
if memory:
try:
# Validate memory format and store original string for Docker
self._parse_memory_to_bytes(memory)
resource_config["memory"] = memory
except ValueError as e:
logger.warning(f"Invalid memory requirement in {workflow_info.name}: {e}")
# Extract CPU requirement
cpu = resources.get("cpu")
if cpu:
try:
# Validate CPU format and store original string for Docker
self._parse_cpu_to_millicores(cpu)
resource_config["cpus"] = cpu
except ValueError as e:
logger.warning(f"Invalid CPU requirement in {workflow_info.name}: {e}")
# Extract timeout
timeout = resources.get("timeout")
if timeout and isinstance(timeout, int):
resource_config["timeout"] = str(timeout)
return resource_config
async def initialize(self):
"""
Initialize the manager by discovering and deploying all workflows.
This method:
1. Discovers all valid workflows in the workflows directory
2. Validates their metadata
3. Deploys each workflow to Prefect with Docker images
"""
try:
# Discover workflows
self.workflows = await self.discovery.discover_workflows()
if not self.workflows:
logger.warning("No workflows discovered")
return
logger.info(f"Discovered {len(self.workflows)} workflows: {list(self.workflows.keys())}")
# Deploy each workflow
for name, info in self.workflows.items():
try:
await self._deploy_workflow(name, info)
except Exception as e:
logger.error(f"Failed to deploy workflow '{name}': {e}")
except Exception as e:
logger.error(f"Failed to initialize Prefect manager: {e}")
raise
async def _deploy_workflow(self, name: str, info: WorkflowInfo):
"""
Deploy a single workflow to Prefect with Docker image.
Args:
name: Workflow name
info: Workflow information including metadata and paths
"""
logger.info(f"Deploying workflow '{name}'...")
# Get the flow function from registry
flow_func = self.discovery.get_flow_function(name)
if not flow_func:
logger.error(
f"Failed to get flow function for '{name}' from registry. "
f"Ensure the workflow is properly registered in toolbox/workflows/registry.py"
)
return
# Use the mandatory Dockerfile with absolute paths for Docker Compose
# Get absolute paths for build context and dockerfile
toolbox_path = info.path.parent.parent.resolve()
dockerfile_abs_path = info.dockerfile.resolve()
# Calculate relative dockerfile path from toolbox context
try:
dockerfile_rel_path = dockerfile_abs_path.relative_to(toolbox_path)
except ValueError:
# If relative path fails, use the workflow-specific path
dockerfile_rel_path = Path("workflows") / name / "Dockerfile"
# Determine deployment strategy based on Dockerfile presence
base_image = "prefecthq/prefect:3-python3.11"
has_custom_dockerfile = info.has_docker and info.dockerfile.exists()
logger.info(f"=== DEPLOYMENT DEBUG for '{name}' ===")
logger.info(f"info.has_docker: {info.has_docker}")
logger.info(f"info.dockerfile: {info.dockerfile}")
logger.info(f"info.dockerfile.exists(): {info.dockerfile.exists()}")
logger.info(f"has_custom_dockerfile: {has_custom_dockerfile}")
logger.info(f"toolbox_path: {toolbox_path}")
logger.info(f"dockerfile_rel_path: {dockerfile_rel_path}")
if has_custom_dockerfile:
logger.info(f"Workflow '{name}' has custom Dockerfile - building custom image")
# Decide whether to use registry or keep images local to host engine
import os
# Default to using the local registry; set FUZZFORGE_USE_REGISTRY=false to bypass (not recommended)
use_registry = os.getenv("FUZZFORGE_USE_REGISTRY", "true").lower() == "true"
if use_registry:
registry_url = get_registry_url(context="push")
image_spec = DockerImage(
name=f"{registry_url}/fuzzforge/{name}",
tag="latest",
dockerfile=str(dockerfile_rel_path),
context=str(toolbox_path)
)
deploy_image = f"{registry_url}/fuzzforge/{name}:latest"
build_custom = True
push_custom = True
logger.info(f"Using registry: {registry_url} for '{name}'")
else:
# Single-host mode: build into host engine cache; no push required
image_spec = DockerImage(
name=f"fuzzforge/{name}",
tag="latest",
dockerfile=str(dockerfile_rel_path),
context=str(toolbox_path)
)
deploy_image = f"fuzzforge/{name}:latest"
build_custom = True
push_custom = False
logger.info("Using single-host image (no registry push): %s", deploy_image)
else:
logger.info(f"Workflow '{name}' using base image - no custom dependencies needed")
deploy_image = base_image
build_custom = False
push_custom = False
# Pre-validate registry connectivity when pushing
if push_custom:
try:
from .setup import validate_registry_connectivity
await validate_registry_connectivity(registry_url)
logger.info(f"Registry connectivity validated for {registry_url}")
except Exception as e:
logger.error(f"Registry connectivity validation failed for {registry_url}: {e}")
raise RuntimeError(f"Cannot deploy workflow '{name}': Registry {registry_url} is not accessible. {e}")
# Deploy the workflow
try:
# Ensure any previous deployment is removed so job variables are updated
try:
async with get_client() as client:
existing = await client.read_deployment_by_name(
f"{name}/{name}-deployment"
)
if existing:
logger.info(f"Removing existing deployment for '{name}' to refresh settings...")
await client.delete_deployment(existing.id)
except Exception:
# If not found or deletion fails, continue with deployment
pass
# Extract resource requirements from metadata
workflow_resource_requirements = self._extract_resource_requirements(info)
logger.info(f"Workflow '{name}' resource requirements: {workflow_resource_requirements}")
# Build job variables with resource requirements
job_variables = {
"image": deploy_image, # Use the worker-accessible registry name
"volumes": [], # Populated at run submission with toolbox mount
"env": {
"PYTHONPATH": "/opt/prefect/toolbox:/opt/prefect",
"WORKFLOW_NAME": name
}
}
# Add resource requirements to job variables if present
if workflow_resource_requirements:
job_variables["resources"] = workflow_resource_requirements
# Prepare deployment parameters
deploy_params = {
"name": f"{name}-deployment",
"work_pool_name": "docker-pool",
"image": image_spec if has_custom_dockerfile else deploy_image,
"push": push_custom,
"build": build_custom,
"job_variables": job_variables
}
deployment = await flow_func.deploy(**deploy_params)
self.deployments[name] = str(deployment.id) if hasattr(deployment, 'id') else name
logger.info(f"Successfully deployed workflow '{name}'")
except Exception as e:
# Enhanced error reporting with more context
import traceback
logger.error(f"Failed to deploy workflow '{name}': {e}")
logger.error(f"Deployment traceback: {traceback.format_exc()}")
# Try to capture Docker-specific context
error_context = {
"workflow_name": name,
"has_dockerfile": has_custom_dockerfile,
"image_name": deploy_image if 'deploy_image' in locals() else "unknown",
"registry_url": registry_url if 'registry_url' in locals() else "unknown",
"error_type": type(e).__name__,
"error_message": str(e)
}
# Check for specific error patterns with detailed categorization
error_msg_lower = str(e).lower()
if "registry" in error_msg_lower and ("no such host" in error_msg_lower or "connection" in error_msg_lower):
error_context["category"] = "registry_connectivity_error"
error_context["solution"] = f"Cannot reach registry at {error_context['registry_url']}. Check Docker network and registry service."
elif "docker" in error_msg_lower:
error_context["category"] = "docker_error"
if "build" in error_msg_lower:
error_context["subcategory"] = "image_build_failed"
error_context["solution"] = "Check Dockerfile syntax and dependencies."
elif "pull" in error_msg_lower:
error_context["subcategory"] = "image_pull_failed"
error_context["solution"] = "Check if image exists in registry and network connectivity."
elif "push" in error_msg_lower:
error_context["subcategory"] = "image_push_failed"
error_context["solution"] = f"Check registry connectivity and push permissions to {error_context['registry_url']}."
elif "registry" in error_msg_lower:
error_context["category"] = "registry_error"
error_context["solution"] = "Check registry configuration and accessibility."
elif "prefect" in error_msg_lower:
error_context["category"] = "prefect_error"
error_context["solution"] = "Check Prefect server connectivity and deployment configuration."
else:
error_context["category"] = "unknown_deployment_error"
error_context["solution"] = "Check logs for more specific error details."
logger.error(f"Deployment error context: {error_context}")
# Raise enhanced exception with context
enhanced_error = Exception(f"Deployment failed for workflow '{name}': {str(e)} | Context: {error_context}")
enhanced_error.original_error = e
enhanced_error.context = error_context
raise enhanced_error
async def submit_workflow(
self,
workflow_name: str,
target_path: str,
volume_mode: str = "ro",
parameters: Dict[str, Any] = None,
resource_limits: Dict[str, str] = None,
additional_volumes: list = None,
timeout: int = None
) -> FlowRun:
"""
Submit a workflow for execution with volume mounting.
Args:
workflow_name: Name of the workflow to execute
target_path: Host path to mount as volume
volume_mode: Volume mount mode ("ro" for read-only, "rw" for read-write)
parameters: Workflow-specific parameters
resource_limits: CPU/memory limits for container
additional_volumes: List of additional volume mounts
timeout: Timeout in seconds
Returns:
FlowRun object with run information
Raises:
ValueError: If workflow not found or volume mode not supported
"""
if workflow_name not in self.workflows:
raise ValueError(f"Unknown workflow: {workflow_name}")
# Validate volume mode
workflow_info = self.workflows[workflow_name]
supported_modes = workflow_info.metadata.get("supported_volume_modes", ["ro", "rw"])
if volume_mode not in supported_modes:
raise ValueError(
f"Workflow '{workflow_name}' doesn't support volume mode '{volume_mode}'. "
f"Supported modes: {supported_modes}"
)
# Validate target path with security checks
self._validate_target_path(target_path)
# Validate additional volumes if provided
if additional_volumes:
for volume in additional_volumes:
self._validate_target_path(volume.host_path)
async with get_client() as client:
# Get the deployment, auto-redeploy once if missing
try:
deployment = await client.read_deployment_by_name(
f"{workflow_name}/{workflow_name}-deployment"
)
except Exception as e:
import traceback
logger.error(f"Failed to find deployment for workflow '{workflow_name}': {e}")
logger.error(f"Deployment lookup traceback: {traceback.format_exc()}")
# Attempt a one-time auto-deploy to recover from startup races
try:
logger.info(f"Auto-deploying missing workflow '{workflow_name}' and retrying...")
await self._deploy_workflow(workflow_name, workflow_info)
deployment = await client.read_deployment_by_name(
f"{workflow_name}/{workflow_name}-deployment"
)
except Exception as redeploy_exc:
# Enhanced error with context
error_context = {
"workflow_name": workflow_name,
"error_type": type(e).__name__,
"error_message": str(e),
"redeploy_error": str(redeploy_exc),
"available_deployments": list(self.deployments.keys()),
}
enhanced_error = ValueError(
f"Deployment not found and redeploy failed for workflow '{workflow_name}': {e} | Context: {error_context}"
)
enhanced_error.context = error_context
raise enhanced_error
# Determine the Docker Compose network name and volume names
# Hardcoded to 'fuzzforge' to avoid directory name dependencies
import os
compose_project = "fuzzforge"
docker_network = "fuzzforge_default"
# Build volume mounts
# Add toolbox volume mount for workflow code access
backend_toolbox_path = "/app/toolbox" # Path in backend container
# Hardcoded volume names
prefect_storage_volume = "fuzzforge_prefect_storage"
toolbox_code_volume = "fuzzforge_toolbox_code"
volumes = [
f"{target_path}:/workspace:{volume_mode}",
f"{prefect_storage_volume}:/prefect-storage", # Shared storage for results
f"{toolbox_code_volume}:/opt/prefect/toolbox:ro" # Mount workflow code
]
# Add additional volumes if provided
if additional_volumes:
for volume in additional_volumes:
volume_spec = f"{volume.host_path}:{volume.container_path}:{volume.mode}"
volumes.append(volume_spec)
# Build environment variables
env_vars = {
"PREFECT_API_URL": "http://prefect-server:4200/api", # Use internal network hostname
"PREFECT_LOGGING_LEVEL": "INFO",
"PREFECT_LOCAL_STORAGE_PATH": "/prefect-storage", # Use shared storage
"PREFECT_RESULTS_PERSIST_BY_DEFAULT": "true", # Enable result persistence
"PREFECT_DEFAULT_RESULT_STORAGE_BLOCK": "local-file-system/fuzzforge-results", # Use our storage block
"WORKSPACE_PATH": "/workspace",
"VOLUME_MODE": volume_mode,
"WORKFLOW_NAME": workflow_name
}
# Add additional volume paths to environment for easy access
if additional_volumes:
for i, volume in enumerate(additional_volumes):
env_vars[f"ADDITIONAL_VOLUME_{i}_PATH"] = volume.container_path
# Determine which image to use based on workflow configuration
workflow_info = self.workflows[workflow_name]
has_custom_dockerfile = workflow_info.has_docker and workflow_info.dockerfile.exists()
# Use pull context for worker to pull from registry
registry_url = get_registry_url(context="pull")
workflow_image = f"{registry_url}/fuzzforge/{workflow_name}:latest" if has_custom_dockerfile else "prefecthq/prefect:3-python3.11"
logger.debug(f"Worker will pull image: {workflow_image} (Registry: {registry_url})")
# Configure job variables with volume mounting and network access
job_variables = {
# Use custom image if available, otherwise base Prefect image
"image": workflow_image,
"volumes": volumes,
"networks": [docker_network], # Connect to Docker Compose network
"env": {
**env_vars,
"PYTHONPATH": "/opt/prefect/toolbox:/opt/prefect/toolbox/workflows",
"WORKFLOW_NAME": workflow_name
}
}
# Apply resource requirements from workflow metadata and user overrides
workflow_resource_requirements = self._extract_resource_requirements(workflow_info)
final_resource_config = {}
# Start with workflow requirements as base
if workflow_resource_requirements:
final_resource_config.update(workflow_resource_requirements)
# Apply user-provided resource limits (overrides workflow defaults)
if resource_limits:
user_resource_config = {}
if resource_limits.get("cpu_limit"):
user_resource_config["cpus"] = resource_limits["cpu_limit"]
if resource_limits.get("memory_limit"):
user_resource_config["memory"] = resource_limits["memory_limit"]
# Note: cpu_request and memory_request are not directly supported by Docker
# but could be used for Kubernetes in the future
# User overrides take precedence
final_resource_config.update(user_resource_config)
# Apply final resource configuration
if final_resource_config:
job_variables["resources"] = final_resource_config
logger.info(f"Applied resource limits: {final_resource_config}")
# Merge parameters with defaults from metadata
default_params = workflow_info.metadata.get("default_parameters", {})
final_params = {**default_params, **(parameters or {})}
# Set flow parameters that match the flow signature
final_params["target_path"] = "/workspace" # Container path where volume is mounted
final_params["volume_mode"] = volume_mode
# Create and submit the flow run
# Pass job_variables to ensure network, volumes, and environment are configured
logger.info(f"Submitting flow with job_variables: {job_variables}")
logger.info(f"Submitting flow with parameters: {final_params}")
# Prepare flow run creation parameters
flow_run_params = {
"deployment_id": deployment.id,
"parameters": final_params,
"job_variables": job_variables
}
# Note: Timeout is handled through workflow-level configuration
# Additional timeout configuration can be added to deployment metadata if needed
flow_run = await client.create_flow_run_from_deployment(**flow_run_params)
logger.info(
f"Submitted workflow '{workflow_name}' with run_id: {flow_run.id}, "
f"target: {target_path}, mode: {volume_mode}"
)
return flow_run
async def get_flow_run_findings(self, run_id: str) -> Dict[str, Any]:
"""
Retrieve findings from a completed flow run.
Args:
run_id: The flow run ID
Returns:
Dictionary containing SARIF-formatted findings
Raises:
ValueError: If run not completed or not found
"""
async with get_client() as client:
flow_run = await client.read_flow_run(run_id)
if not flow_run.state.is_completed():
raise ValueError(
f"Flow run {run_id} not completed. Current status: {flow_run.state.name}"
)
# Get the findings from the flow run result
try:
findings = await flow_run.state.result()
return findings
except Exception as e:
logger.error(f"Failed to retrieve findings for run {run_id}: {e}")
raise ValueError(f"Failed to retrieve findings: {e}")
async def get_flow_run_status(self, run_id: str) -> Dict[str, Any]:
"""
Get the current status of a flow run.
Args:
run_id: The flow run ID
Returns:
Dictionary with status information
"""
async with get_client() as client:
flow_run = await client.read_flow_run(run_id)
return {
"run_id": str(flow_run.id),
"workflow": flow_run.deployment_id,
"status": flow_run.state.name,
"is_completed": flow_run.state.is_completed(),
"is_failed": flow_run.state.is_failed(),
"is_running": flow_run.state.is_running(),
"created_at": flow_run.created,
"updated_at": flow_run.updated
}
def _validate_target_path(self, target_path: str) -> None:
"""
Validate target path for security before mounting as volume.
Args:
target_path: Host path to validate
Raises:
ValueError: If path is not allowed for security reasons
"""
target = Path(target_path)
# Path must be absolute
if not target.is_absolute():
raise ValueError(f"Target path must be absolute: {target_path}")
# Resolve path to handle symlinks and relative components
try:
resolved_path = target.resolve()
except (OSError, RuntimeError) as e:
raise ValueError(f"Cannot resolve target path: {target_path} - {e}")
resolved_str = str(resolved_path)
# Check against forbidden paths first (more restrictive)
for forbidden in self.forbidden_paths:
if resolved_str.startswith(forbidden):
raise ValueError(
f"Access denied: Path '{target_path}' resolves to forbidden directory '{forbidden}'. "
f"This path contains sensitive system files and cannot be mounted."
)
# Check if path starts with any allowed base path
path_allowed = False
for allowed in self.allowed_base_paths:
if resolved_str.startswith(allowed):
path_allowed = True
break
if not path_allowed:
allowed_list = ", ".join(self.allowed_base_paths)
raise ValueError(
f"Access denied: Path '{target_path}' is not in allowed directories. "
f"Allowed base paths: {allowed_list}"
)
# Additional security checks
if resolved_str == "/":
raise ValueError("Cannot mount root filesystem")
# Warn if path doesn't exist (but don't block - it might be created later)
if not resolved_path.exists():
logger.warning(f"Target path does not exist: {target_path}")
logger.info(f"Path validation passed for: {target_path} -> {resolved_str}")
+10 -282
View File
@@ -1,5 +1,5 @@
"""
Setup utilities for Prefect infrastructure
Setup utilities for FuzzForge infrastructure
"""
# Copyright (c) 2025 FuzzingLabs
@@ -14,234 +14,21 @@ Setup utilities for Prefect infrastructure
# Additional attribution and requirements are provided in the NOTICE file.
import logging
from prefect import get_client
from prefect.client.schemas.actions import WorkPoolCreate
from prefect.client.schemas.objects import WorkPool
from .prefect_manager import get_registry_url
logger = logging.getLogger(__name__)
async def setup_docker_pool():
"""
Create or update the Docker work pool for container execution.
This work pool is configured to:
- Connect to the local Docker daemon
- Support volume mounting at runtime
- Clean up containers after execution
- Use bridge networking by default
"""
import os
async with get_client() as client:
pool_name = "docker-pool"
# Add force recreation flag for debugging fresh install issues
force_recreate = os.getenv('FORCE_RECREATE_WORK_POOL', 'false').lower() == 'true'
debug_setup = os.getenv('DEBUG_WORK_POOL_SETUP', 'false').lower() == 'true'
if force_recreate:
logger.warning(f"FORCE_RECREATE_WORK_POOL=true - Will recreate work pool regardless of existing configuration")
if debug_setup:
logger.warning(f"DEBUG_WORK_POOL_SETUP=true - Enhanced logging enabled")
# Temporarily set logging level to DEBUG for this function
original_level = logger.level
logger.setLevel(logging.DEBUG)
try:
# Check if pool already exists and supports custom images
existing_pools = await client.read_work_pools()
existing_pool = None
for pool in existing_pools:
if pool.name == pool_name:
existing_pool = pool
break
if existing_pool and not force_recreate:
logger.info(f"Found existing work pool '{pool_name}' - validating configuration...")
# Check if the existing pool has the correct configuration
base_template = existing_pool.base_job_template or {}
logger.debug(f"Base template keys: {list(base_template.keys())}")
job_config = base_template.get("job_configuration", {})
logger.debug(f"Job config keys: {list(job_config.keys())}")
image_config = job_config.get("image", "")
has_image_variable = "{{ image }}" in str(image_config)
logger.debug(f"Image config: '{image_config}' -> has_image_variable: {has_image_variable}")
# Check if volume defaults include toolbox mount
variables = base_template.get("variables", {})
properties = variables.get("properties", {})
volume_config = properties.get("volumes", {})
volume_defaults = volume_config.get("default", [])
has_toolbox_volume = any("toolbox_code" in str(vol) for vol in volume_defaults) if volume_defaults else False
logger.debug(f"Volume defaults: {volume_defaults}")
logger.debug(f"Has toolbox volume: {has_toolbox_volume}")
# Check if environment defaults include required settings
env_config = properties.get("env", {})
env_defaults = env_config.get("default", {})
has_api_url = "PREFECT_API_URL" in env_defaults
has_storage_path = "PREFECT_LOCAL_STORAGE_PATH" in env_defaults
has_results_persist = "PREFECT_RESULTS_PERSIST_BY_DEFAULT" in env_defaults
has_required_env = has_api_url and has_storage_path and has_results_persist
logger.debug(f"Environment defaults: {env_defaults}")
logger.debug(f"Has API URL: {has_api_url}, Has storage path: {has_storage_path}, Has results persist: {has_results_persist}")
logger.debug(f"Has required env: {has_required_env}")
# Log the full validation result
logger.info(f"Work pool validation - Image: {has_image_variable}, Toolbox: {has_toolbox_volume}, Environment: {has_required_env}")
if has_image_variable and has_toolbox_volume and has_required_env:
logger.info(f"Docker work pool '{pool_name}' already exists with correct configuration")
return
else:
reasons = []
if not has_image_variable:
reasons.append("missing image template")
if not has_toolbox_volume:
reasons.append("missing toolbox volume mount")
if not has_required_env:
if not has_api_url:
reasons.append("missing PREFECT_API_URL")
if not has_storage_path:
reasons.append("missing PREFECT_LOCAL_STORAGE_PATH")
if not has_results_persist:
reasons.append("missing PREFECT_RESULTS_PERSIST_BY_DEFAULT")
logger.warning(f"Docker work pool '{pool_name}' exists but lacks: {', '.join(reasons)}. Recreating...")
# Delete the old pool and recreate it
try:
await client.delete_work_pool(pool_name)
logger.info(f"Deleted old work pool '{pool_name}'")
except Exception as e:
logger.warning(f"Failed to delete old work pool: {e}")
elif force_recreate and existing_pool:
logger.warning(f"Force recreation enabled - deleting existing work pool '{pool_name}'")
try:
await client.delete_work_pool(pool_name)
logger.info(f"Deleted existing work pool for force recreation")
except Exception as e:
logger.warning(f"Failed to delete work pool for force recreation: {e}")
logger.info(f"Creating Docker work pool '{pool_name}' with custom image support...")
# Create the work pool with proper Docker configuration
work_pool = WorkPoolCreate(
name=pool_name,
type="docker",
description="Docker work pool for FuzzForge workflows with custom image support",
base_job_template={
"job_configuration": {
"image": "{{ image }}", # Template variable for custom images
"volumes": "{{ volumes }}", # List of volume mounts
"env": "{{ env }}", # Environment variables
"networks": "{{ networks }}", # Docker networks
"stream_output": True,
"auto_remove": True,
"privileged": False,
"network_mode": None, # Use networks instead
"labels": {},
"command": None # Let the image's CMD/ENTRYPOINT run
},
"variables": {
"type": "object",
"properties": {
"image": {
"type": "string",
"title": "Docker Image",
"default": "prefecthq/prefect:3-python3.11",
"description": "Docker image for the flow run"
},
"volumes": {
"type": "array",
"title": "Volume Mounts",
"default": [
"fuzzforge_prefect_storage:/prefect-storage",
"fuzzforge_toolbox_code:/opt/prefect/toolbox:ro"
],
"description": "Volume mounts in format 'host:container:mode'",
"items": {
"type": "string"
}
},
"networks": {
"type": "array",
"title": "Docker Networks",
"default": ["fuzzforge_default"],
"description": "Docker networks to connect container to",
"items": {
"type": "string"
}
},
"env": {
"type": "object",
"title": "Environment Variables",
"default": {
"PREFECT_API_URL": "http://prefect-server:4200/api",
"PREFECT_LOCAL_STORAGE_PATH": "/prefect-storage",
"PREFECT_RESULTS_PERSIST_BY_DEFAULT": "true"
},
"description": "Environment variables for the container",
"additionalProperties": {
"type": "string"
}
}
}
}
}
)
await client.create_work_pool(work_pool)
logger.info(f"Created Docker work pool '{pool_name}'")
except Exception as e:
logger.error(f"Failed to setup Docker work pool: {e}")
raise
finally:
# Restore original logging level if debug mode was enabled
if debug_setup and 'original_level' in locals():
logger.setLevel(original_level)
def get_actual_compose_project_name():
"""
Return the hardcoded compose project name for FuzzForge.
Always returns 'fuzzforge' as per system requirements.
"""
logger.info("Using hardcoded compose project name: fuzzforge")
return "fuzzforge"
async def setup_result_storage():
"""
Create or update Prefect result storage block for findings persistence.
Setup result storage (MinIO).
This sets up a LocalFileSystem storage block pointing to the shared
/prefect-storage volume for result persistence.
MinIO is used for both target upload and result storage.
This is a placeholder for any MinIO-specific setup if needed.
"""
from prefect.filesystems import LocalFileSystem
storage_name = "fuzzforge-results"
try:
# Create the storage block, overwrite if it exists
logger.info(f"Setting up storage block '{storage_name}'...")
storage = LocalFileSystem(basepath="/prefect-storage")
block_doc_id = await storage.save(name=storage_name, overwrite=True)
logger.info(f"Storage block '{storage_name}' configured successfully")
return str(block_doc_id)
except Exception as e:
logger.error(f"Failed to setup result storage: {e}")
# Don't raise the exception - continue without storage block
logger.warning("Continuing without result storage block - findings may not persist")
return None
logger.info("Result storage (MinIO) configured")
# MinIO is configured via environment variables in docker-compose
# No additional setup needed here
return True
async def validate_docker_connection():
@@ -274,60 +61,6 @@ async def validate_docker_connection():
)
async def validate_registry_connectivity(registry_url: str = None):
"""
Validate that the Docker registry is accessible.
Args:
registry_url: URL of the Docker registry to validate (auto-detected if None)
Raises:
RuntimeError: If registry is not accessible
"""
# Resolve a reachable test URL from within this process
if registry_url is None:
# If not specified, prefer internal service name in containers, host port on host
import os
if os.path.exists('/.dockerenv'):
registry_url = "registry:5000"
else:
registry_url = "localhost:5001"
# If we're running inside a container and asked to probe localhost:PORT,
# the probe would hit the container, not the host. Use host.docker.internal instead.
import os
try:
host_part, port_part = registry_url.split(":", 1)
except ValueError:
host_part, port_part = registry_url, "80"
if os.path.exists('/.dockerenv') and host_part in ("localhost", "127.0.0.1"):
test_host = "host.docker.internal"
else:
test_host = host_part
test_url = f"http://{test_host}:{port_part}/v2/"
import aiohttp
import asyncio
logger.info(f"Validating registry connectivity to {registry_url}...")
try:
async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=10)) as session:
async with session.get(test_url) as response:
if response.status == 200:
logger.info(f"Registry at {registry_url} is accessible (tested via {test_host})")
return
else:
raise RuntimeError(f"Registry returned status {response.status}")
except asyncio.TimeoutError:
raise RuntimeError(f"Registry at {registry_url} is not responding (timeout)")
except aiohttp.ClientError as e:
raise RuntimeError(f"Registry at {registry_url} is not accessible: {e}")
except Exception as e:
raise RuntimeError(f"Failed to validate registry connectivity: {e}")
async def validate_docker_network(network_name: str):
"""
Validate that the specified Docker network exists.
@@ -385,18 +118,13 @@ async def validate_infrastructure():
# Validate Docker connection
await validate_docker_connection()
# Validate registry connectivity for custom image building
await validate_registry_connectivity()
# Validate network (hardcoded to avoid directory name dependencies)
import os
compose_project = "fuzzforge"
# Validate network (hardcoded to fuzzforge for Temporal deployment)
docker_network = "fuzzforge_default"
try:
await validate_docker_network(docker_network)
except RuntimeError as e:
logger.warning(f"Network validation failed: {e}")
logger.warning("Workflows may not be able to connect to Prefect services")
logger.warning("Workflows may not be able to connect to Temporal services")
logger.info("Infrastructure validation completed")
+169 -306
View File
@@ -12,7 +12,6 @@
import asyncio
import logging
import os
from uuid import UUID
from contextlib import AsyncExitStack, asynccontextmanager, suppress
from typing import Any, Dict, Optional, List
@@ -23,31 +22,20 @@ from starlette.routing import Mount
from fastmcp.server.http import create_sse_app
from src.core.prefect_manager import PrefectManager
from src.core.setup import setup_docker_pool, setup_result_storage, validate_infrastructure
from src.core.workflow_discovery import WorkflowDiscovery
from src.temporal.manager import TemporalManager
from src.core.setup import setup_result_storage, validate_infrastructure
from src.api import workflows, runs, fuzzing
from src.services.prefect_stats_monitor import prefect_stats_monitor
from fastmcp import FastMCP
from prefect.client.orchestration import get_client
from prefect.client.schemas.filters import (
FlowRunFilter,
FlowRunFilterDeploymentId,
FlowRunFilterState,
FlowRunFilterStateType,
)
from prefect.client.schemas.sorting import FlowRunSort
from prefect.states import StateType
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
prefect_mgr = PrefectManager()
temporal_mgr = TemporalManager()
class PrefectBootstrapState:
"""Tracks Prefect initialization progress for API and MCP consumers."""
class TemporalBootstrapState:
"""Tracks Temporal initialization progress for API and MCP consumers."""
def __init__(self) -> None:
self.ready: bool = False
@@ -64,19 +52,19 @@ class PrefectBootstrapState:
}
prefect_bootstrap_state = PrefectBootstrapState()
temporal_bootstrap_state = TemporalBootstrapState()
# Configure retry strategy for bootstrapping Prefect + infrastructure
# Configure retry strategy for bootstrapping Temporal + infrastructure
STARTUP_RETRY_SECONDS = max(1, int(os.getenv("FUZZFORGE_STARTUP_RETRY_SECONDS", "5")))
STARTUP_RETRY_MAX_SECONDS = max(
STARTUP_RETRY_SECONDS,
int(os.getenv("FUZZFORGE_STARTUP_RETRY_MAX_SECONDS", "60")),
)
prefect_bootstrap_task: Optional[asyncio.Task] = None
temporal_bootstrap_task: Optional[asyncio.Task] = None
# ---------------------------------------------------------------------------
# FastAPI application (REST API remains unchanged)
# FastAPI application (REST API)
# ---------------------------------------------------------------------------
app = FastAPI(
@@ -90,20 +78,19 @@ app.include_router(runs.router)
app.include_router(fuzzing.router)
def get_prefect_status() -> Dict[str, Any]:
"""Return a snapshot of Prefect bootstrap state for diagnostics."""
status = prefect_bootstrap_state.as_dict()
status["workflows_loaded"] = len(prefect_mgr.workflows)
status["deployments_tracked"] = len(prefect_mgr.deployments)
def get_temporal_status() -> Dict[str, Any]:
"""Return a snapshot of Temporal bootstrap state for diagnostics."""
status = temporal_bootstrap_state.as_dict()
status["workflows_loaded"] = len(temporal_mgr.workflows)
status["bootstrap_task_running"] = (
prefect_bootstrap_task is not None and not prefect_bootstrap_task.done()
temporal_bootstrap_task is not None and not temporal_bootstrap_task.done()
)
return status
def _prefect_not_ready_status() -> Optional[Dict[str, Any]]:
"""Return status details if Prefect is not ready yet."""
status = get_prefect_status()
def _temporal_not_ready_status() -> Optional[Dict[str, Any]]:
"""Return status details if Temporal is not ready yet."""
status = get_temporal_status()
if status.get("ready"):
return None
return status
@@ -111,19 +98,19 @@ def _prefect_not_ready_status() -> Optional[Dict[str, Any]]:
@app.get("/")
async def root() -> Dict[str, Any]:
status = get_prefect_status()
status = get_temporal_status()
return {
"name": "FuzzForge API",
"version": "0.6.0",
"status": "ready" if status.get("ready") else "initializing",
"workflows_loaded": status.get("workflows_loaded", 0),
"prefect": status,
"temporal": status,
}
@app.get("/health")
async def health() -> Dict[str, str]:
status = get_prefect_status()
status = get_temporal_status()
health_status = "healthy" if status.get("ready") else "initializing"
return {"status": health_status}
@@ -165,65 +152,61 @@ _fastapi_mcp_imported = False
mcp = FastMCP(name="FuzzForge MCP")
async def _bootstrap_prefect_with_retries() -> None:
"""Initialize Prefect infrastructure with exponential backoff retries."""
async def _bootstrap_temporal_with_retries() -> None:
"""Initialize Temporal infrastructure with exponential backoff retries."""
attempt = 0
while True:
attempt += 1
prefect_bootstrap_state.task_running = True
prefect_bootstrap_state.status = "starting"
prefect_bootstrap_state.ready = False
prefect_bootstrap_state.last_error = None
temporal_bootstrap_state.task_running = True
temporal_bootstrap_state.status = "starting"
temporal_bootstrap_state.ready = False
temporal_bootstrap_state.last_error = None
try:
logger.info("Bootstrapping Prefect infrastructure...")
logger.info("Bootstrapping Temporal infrastructure...")
await validate_infrastructure()
await setup_docker_pool()
await setup_result_storage()
await prefect_mgr.initialize()
await prefect_stats_monitor.start_monitoring()
await temporal_mgr.initialize()
prefect_bootstrap_state.ready = True
prefect_bootstrap_state.status = "ready"
prefect_bootstrap_state.task_running = False
logger.info("Prefect infrastructure ready")
temporal_bootstrap_state.ready = True
temporal_bootstrap_state.status = "ready"
temporal_bootstrap_state.task_running = False
logger.info("Temporal infrastructure ready")
return
except asyncio.CancelledError:
prefect_bootstrap_state.status = "cancelled"
prefect_bootstrap_state.task_running = False
logger.info("Prefect bootstrap task cancelled")
temporal_bootstrap_state.status = "cancelled"
temporal_bootstrap_state.task_running = False
logger.info("Temporal bootstrap task cancelled")
raise
except Exception as exc: # pragma: no cover - defensive logging on infra startup
logger.exception("Prefect bootstrap failed")
prefect_bootstrap_state.ready = False
prefect_bootstrap_state.status = "error"
prefect_bootstrap_state.last_error = str(exc)
logger.exception("Temporal bootstrap failed")
temporal_bootstrap_state.ready = False
temporal_bootstrap_state.status = "error"
temporal_bootstrap_state.last_error = str(exc)
# Ensure partial initialization does not leave stale state behind
prefect_mgr.workflows.clear()
prefect_mgr.deployments.clear()
await prefect_stats_monitor.stop_monitoring()
temporal_mgr.workflows.clear()
wait_time = min(
STARTUP_RETRY_SECONDS * (2 ** (attempt - 1)),
STARTUP_RETRY_MAX_SECONDS,
)
logger.info("Retrying Prefect bootstrap in %s second(s)", wait_time)
logger.info("Retrying Temporal bootstrap in %s second(s)", wait_time)
try:
await asyncio.sleep(wait_time)
except asyncio.CancelledError:
prefect_bootstrap_state.status = "cancelled"
prefect_bootstrap_state.task_running = False
temporal_bootstrap_state.status = "cancelled"
temporal_bootstrap_state.task_running = False
raise
def _lookup_workflow(workflow_name: str):
info = prefect_mgr.workflows.get(workflow_name)
info = temporal_mgr.workflows.get(workflow_name)
if not info:
return None
metadata = info.metadata
@@ -256,16 +239,16 @@ def _lookup_workflow(workflow_name: str):
@mcp.tool
async def list_workflows_mcp() -> Dict[str, Any]:
"""List all discovered workflows and their metadata summary."""
not_ready = _prefect_not_ready_status()
not_ready = _temporal_not_ready_status()
if not_ready:
return {
"workflows": [],
"prefect": not_ready,
"message": "Prefect infrastructure is still initializing",
"temporal": not_ready,
"message": "Temporal infrastructure is still initializing",
}
workflows_summary = []
for name, info in prefect_mgr.workflows.items():
for name, info in temporal_mgr.workflows.items():
metadata = info.metadata
defaults = metadata.get("default_parameters", {})
workflows_summary.append({
@@ -282,17 +265,17 @@ async def list_workflows_mcp() -> Dict[str, Any]:
or defaults.get("target_path"),
"has_custom_docker": bool(info.has_docker),
})
return {"workflows": workflows_summary, "prefect": get_prefect_status()}
return {"workflows": workflows_summary, "temporal": get_temporal_status()}
@mcp.tool
async def get_workflow_metadata_mcp(workflow_name: str) -> Dict[str, Any]:
"""Fetch detailed metadata for a workflow."""
not_ready = _prefect_not_ready_status()
not_ready = _temporal_not_ready_status()
if not_ready:
return {
"error": "Prefect infrastructure not ready",
"prefect": not_ready,
"error": "Temporal infrastructure not ready",
"temporal": not_ready,
}
data = _lookup_workflow(workflow_name)
@@ -304,11 +287,11 @@ async def get_workflow_metadata_mcp(workflow_name: str) -> Dict[str, Any]:
@mcp.tool
async def get_workflow_parameters_mcp(workflow_name: str) -> Dict[str, Any]:
"""Return the parameter schema and defaults for a workflow."""
not_ready = _prefect_not_ready_status()
not_ready = _temporal_not_ready_status()
if not_ready:
return {
"error": "Prefect infrastructure not ready",
"prefect": not_ready,
"error": "Temporal infrastructure not ready",
"temporal": not_ready,
}
data = _lookup_workflow(workflow_name)
@@ -323,72 +306,41 @@ async def get_workflow_parameters_mcp(workflow_name: str) -> Dict[str, Any]:
@mcp.tool
async def get_workflow_metadata_schema_mcp() -> Dict[str, Any]:
"""Return the JSON schema describing workflow metadata files."""
from src.temporal.discovery import WorkflowDiscovery
return WorkflowDiscovery.get_metadata_schema()
@mcp.tool
async def submit_security_scan_mcp(
workflow_name: str,
target_path: str | None = None,
volume_mode: str | None = None,
target_id: str,
parameters: Dict[str, Any] | None = None,
) -> Dict[str, Any] | Dict[str, str]:
"""Submit a Prefect workflow via MCP."""
"""Submit a Temporal workflow via MCP."""
try:
not_ready = _prefect_not_ready_status()
not_ready = _temporal_not_ready_status()
if not_ready:
return {
"error": "Prefect infrastructure not ready",
"prefect": not_ready,
"error": "Temporal infrastructure not ready",
"temporal": not_ready,
}
workflow_info = prefect_mgr.workflows.get(workflow_name)
workflow_info = temporal_mgr.workflows.get(workflow_name)
if not workflow_info:
return {"error": f"Workflow '{workflow_name}' not found"}
metadata = workflow_info.metadata or {}
defaults = metadata.get("default_parameters", {})
resolved_target_path = target_path or metadata.get("default_target_path") or defaults.get("target_path")
if not resolved_target_path:
return {
"error": (
"target_path is required and no default_target_path is defined in metadata"
),
"metadata": {
"workflow": workflow_name,
"default_target_path": metadata.get("default_target_path"),
},
}
requested_volume_mode = volume_mode or metadata.get("default_volume_mode") or defaults.get("volume_mode")
if not requested_volume_mode:
requested_volume_mode = "ro"
normalised_volume_mode = (
str(requested_volume_mode).strip().lower().replace("-", "_")
)
if normalised_volume_mode in {"read_only", "readonly", "ro"}:
normalised_volume_mode = "ro"
elif normalised_volume_mode in {"read_write", "readwrite", "rw"}:
normalised_volume_mode = "rw"
else:
supported_modes = metadata.get("supported_volume_modes", ["ro", "rw"])
if isinstance(supported_modes, list) and normalised_volume_mode in supported_modes:
pass
else:
normalised_volume_mode = "ro"
parameters = parameters or {}
cleaned_parameters: Dict[str, Any] = {**defaults, **parameters}
# Ensure *_config structures default to dicts so Prefect validation passes.
# Ensure *_config structures default to dicts
for key, value in list(cleaned_parameters.items()):
if isinstance(key, str) and key.endswith("_config") and value is None:
cleaned_parameters[key] = {}
# Some workflows expect configuration dictionaries even when omitted.
# Some workflows expect configuration dictionaries even when omitted
parameter_definitions = (
metadata.get("parameters", {}).get("properties", {})
if isinstance(metadata.get("parameters"), dict)
@@ -403,20 +355,19 @@ async def submit_security_scan_mcp(
elif cleaned_parameters[key] is None:
cleaned_parameters[key] = {}
flow_run = await prefect_mgr.submit_workflow(
# Start workflow
handle = await temporal_mgr.run_workflow(
workflow_name=workflow_name,
target_path=resolved_target_path,
volume_mode=normalised_volume_mode,
parameters=cleaned_parameters,
target_id=target_id,
workflow_params=cleaned_parameters,
)
return {
"run_id": str(flow_run.id),
"status": flow_run.state.name if flow_run.state else "PENDING",
"run_id": handle.id,
"status": "RUNNING",
"workflow": workflow_name,
"message": f"Workflow '{workflow_name}' submitted successfully",
"target_path": resolved_target_path,
"volume_mode": normalised_volume_mode,
"target_id": target_id,
"parameters": cleaned_parameters,
"mcp_enabled": True,
}
@@ -427,43 +378,38 @@ async def submit_security_scan_mcp(
@mcp.tool
async def get_comprehensive_scan_summary(run_id: str) -> Dict[str, Any] | Dict[str, str]:
"""Return a summary for the given flow run via MCP."""
"""Return a summary for the given workflow run via MCP."""
try:
not_ready = _prefect_not_ready_status()
not_ready = _temporal_not_ready_status()
if not_ready:
return {
"error": "Prefect infrastructure not ready",
"prefect": not_ready,
"error": "Temporal infrastructure not ready",
"temporal": not_ready,
}
status = await prefect_mgr.get_flow_run_status(run_id)
findings = await prefect_mgr.get_flow_run_findings(run_id)
workflow_name = "unknown"
deployment_id = status.get("workflow", "")
for name, deployment in prefect_mgr.deployments.items():
if str(deployment) == str(deployment_id):
workflow_name = name
break
status = await temporal_mgr.get_workflow_status(run_id)
# Try to get result if completed
total_findings = 0
severity_summary = {"critical": 0, "high": 0, "medium": 0, "low": 0, "info": 0}
if findings and "sarif" in findings:
sarif = findings["sarif"]
if isinstance(sarif, dict):
total_findings = sarif.get("total_findings", 0)
if status.get("status") == "COMPLETED":
try:
result = await temporal_mgr.get_workflow_result(run_id)
if isinstance(result, dict):
summary = result.get("summary", {})
total_findings = summary.get("total_findings", 0)
except Exception as e:
logger.debug(f"Could not retrieve result for {run_id}: {e}")
return {
"run_id": run_id,
"workflow": workflow_name,
"workflow": "unknown", # Temporal doesn't track workflow name in status
"status": status.get("status", "unknown"),
"is_completed": status.get("is_completed", False),
"is_completed": status.get("status") == "COMPLETED",
"total_findings": total_findings,
"severity_summary": severity_summary,
"scan_duration": status.get("updated_at", "")
if status.get("is_completed")
else "In progress",
"scan_duration": status.get("close_time", "In progress"),
"recommendations": (
[
"Review high and critical severity findings first",
@@ -482,32 +428,26 @@ async def get_comprehensive_scan_summary(run_id: str) -> Dict[str, Any] | Dict[s
@mcp.tool
async def get_run_status_mcp(run_id: str) -> Dict[str, Any]:
"""Return current status information for a Prefect run."""
"""Return current status information for a Temporal run."""
try:
not_ready = _prefect_not_ready_status()
not_ready = _temporal_not_ready_status()
if not_ready:
return {
"error": "Prefect infrastructure not ready",
"prefect": not_ready,
"error": "Temporal infrastructure not ready",
"temporal": not_ready,
}
status = await prefect_mgr.get_flow_run_status(run_id)
workflow_name = "unknown"
deployment_id = status.get("workflow", "")
for name, deployment in prefect_mgr.deployments.items():
if str(deployment) == str(deployment_id):
workflow_name = name
break
status = await temporal_mgr.get_workflow_status(run_id)
return {
"run_id": status["run_id"],
"workflow": workflow_name,
"run_id": run_id,
"workflow": "unknown",
"status": status["status"],
"is_completed": status["is_completed"],
"is_failed": status["is_failed"],
"is_running": status["is_running"],
"created_at": status["created_at"],
"updated_at": status["updated_at"],
"is_completed": status["status"] in ["COMPLETED", "FAILED", "CANCELLED"],
"is_failed": status["status"] == "FAILED",
"is_running": status["status"] == "RUNNING",
"created_at": status.get("start_time"),
"updated_at": status.get("close_time") or status.get("execution_time"),
}
except Exception as exc:
logger.exception("MCP run status failed")
@@ -518,38 +458,30 @@ async def get_run_status_mcp(run_id: str) -> Dict[str, Any]:
async def get_run_findings_mcp(run_id: str) -> Dict[str, Any]:
"""Return SARIF findings for a completed run."""
try:
not_ready = _prefect_not_ready_status()
not_ready = _temporal_not_ready_status()
if not_ready:
return {
"error": "Prefect infrastructure not ready",
"prefect": not_ready,
"error": "Temporal infrastructure not ready",
"temporal": not_ready,
}
status = await prefect_mgr.get_flow_run_status(run_id)
if not status.get("is_completed"):
status = await temporal_mgr.get_workflow_status(run_id)
if status.get("status") != "COMPLETED":
return {"error": f"Run {run_id} not completed. Status: {status.get('status')}"}
findings = await prefect_mgr.get_flow_run_findings(run_id)
workflow_name = "unknown"
deployment_id = status.get("workflow", "")
for name, deployment in prefect_mgr.deployments.items():
if str(deployment) == str(deployment_id):
workflow_name = name
break
result = await temporal_mgr.get_workflow_result(run_id)
metadata = {
"completion_time": status.get("updated_at"),
"completion_time": status.get("close_time"),
"workflow_version": "unknown",
}
info = prefect_mgr.workflows.get(workflow_name)
if info:
metadata["workflow_version"] = info.metadata.get("version", "unknown")
sarif = result.get("sarif", {}) if isinstance(result, dict) else {}
return {
"workflow": workflow_name,
"workflow": "unknown",
"run_id": run_id,
"sarif": findings,
"sarif": sarif,
"metadata": metadata,
}
except Exception as exc:
@@ -561,16 +493,15 @@ async def get_run_findings_mcp(run_id: str) -> Dict[str, Any]:
async def list_recent_runs_mcp(
limit: int = 10,
workflow_name: str | None = None,
states: List[str] | None = None,
) -> Dict[str, Any]:
"""List recent Prefect runs with optional workflow/state filters."""
"""List recent Temporal runs with optional workflow filter."""
not_ready = _prefect_not_ready_status()
not_ready = _temporal_not_ready_status()
if not_ready:
return {
"runs": [],
"prefect": not_ready,
"message": "Prefect infrastructure is still initializing",
"temporal": not_ready,
"message": "Temporal infrastructure is still initializing",
}
try:
@@ -579,116 +510,49 @@ async def list_recent_runs_mcp(
limit_value = 10
limit_value = max(1, min(limit_value, 100))
deployment_map = {
str(deployment_id): workflow
for workflow, deployment_id in prefect_mgr.deployments.items()
}
try:
# Build filter query
filter_query = None
if workflow_name:
workflow_info = temporal_mgr.workflows.get(workflow_name)
if workflow_info:
filter_query = f'WorkflowType="{workflow_info.workflow_type}"'
deployment_filter_value = None
if workflow_name:
deployment_id = prefect_mgr.deployments.get(workflow_name)
if not deployment_id:
return {
"runs": [],
"prefect": get_prefect_status(),
"error": f"Workflow '{workflow_name}' has no registered deployment",
}
try:
deployment_filter_value = UUID(str(deployment_id))
except ValueError:
return {
"runs": [],
"prefect": get_prefect_status(),
"error": (
f"Deployment id '{deployment_id}' for workflow '{workflow_name}' is invalid"
),
}
workflows = await temporal_mgr.list_workflows(filter_query, limit_value)
desired_state_types: List[StateType] = []
if states:
for raw_state in states:
if not raw_state:
continue
normalised = raw_state.strip().upper()
if normalised == "ALL":
desired_state_types = []
break
try:
desired_state_types.append(StateType[normalised])
except KeyError:
continue
if not desired_state_types:
desired_state_types = [
StateType.RUNNING,
StateType.COMPLETED,
StateType.FAILED,
StateType.CANCELLED,
]
results: List[Dict[str, Any]] = []
for wf in workflows:
results.append({
"run_id": wf["workflow_id"],
"workflow": workflow_name or "unknown",
"state": wf["status"],
"state_type": wf["status"],
"is_completed": wf["status"] in ["COMPLETED", "FAILED", "CANCELLED"],
"is_running": wf["status"] == "RUNNING",
"is_failed": wf["status"] == "FAILED",
"created_at": wf.get("start_time"),
"updated_at": wf.get("close_time"),
})
flow_filter = FlowRunFilter()
if desired_state_types:
flow_filter.state = FlowRunFilterState(
type=FlowRunFilterStateType(any_=desired_state_types)
)
if deployment_filter_value:
flow_filter.deployment_id = FlowRunFilterDeploymentId(
any_=[deployment_filter_value]
)
return {"runs": results, "temporal": get_temporal_status()}
async with get_client() as client:
flow_runs = await client.read_flow_runs(
limit=limit_value,
flow_run_filter=flow_filter,
sort=FlowRunSort.START_TIME_DESC,
)
results: List[Dict[str, Any]] = []
for flow_run in flow_runs:
deployment_id = getattr(flow_run, "deployment_id", None)
workflow = deployment_map.get(str(deployment_id), "unknown")
state = getattr(flow_run, "state", None)
state_name = getattr(state, "name", None) if state else None
state_type = getattr(state, "type", None) if state else None
results.append(
{
"run_id": str(flow_run.id),
"workflow": workflow,
"deployment_id": str(deployment_id) if deployment_id else None,
"state": state_name or (state_type.name if state_type else None),
"state_type": state_type.name if state_type else None,
"is_completed": bool(getattr(state, "is_completed", lambda: False)()),
"is_running": bool(getattr(state, "is_running", lambda: False)()),
"is_failed": bool(getattr(state, "is_failed", lambda: False)()),
"created_at": getattr(flow_run, "created", None),
"updated_at": getattr(flow_run, "updated", None),
"expected_start_time": getattr(flow_run, "expected_start_time", None),
"start_time": getattr(flow_run, "start_time", None),
}
)
# Normalise datetimes to ISO 8601 strings for serialization
for entry in results:
for key in ("created_at", "updated_at", "expected_start_time", "start_time"):
value = entry.get(key)
if value is None:
continue
try:
entry[key] = value.isoformat()
except AttributeError:
entry[key] = str(value)
return {"runs": results, "prefect": get_prefect_status()}
except Exception as exc:
logger.exception("Failed to list runs")
return {
"runs": [],
"temporal": get_temporal_status(),
"error": str(exc)
}
@mcp.tool
async def get_fuzzing_stats_mcp(run_id: str) -> Dict[str, Any]:
"""Return fuzzing statistics for a run if available."""
not_ready = _prefect_not_ready_status()
not_ready = _temporal_not_ready_status()
if not_ready:
return {
"error": "Prefect infrastructure not ready",
"prefect": not_ready,
"error": "Temporal infrastructure not ready",
"temporal": not_ready,
}
stats = fuzzing.fuzzing_stats.get(run_id)
@@ -708,11 +572,11 @@ async def get_fuzzing_stats_mcp(run_id: str) -> Dict[str, Any]:
@mcp.tool
async def get_fuzzing_crash_reports_mcp(run_id: str) -> Dict[str, Any]:
"""Return crash reports collected for a fuzzing run."""
not_ready = _prefect_not_ready_status()
not_ready = _temporal_not_ready_status()
if not_ready:
return {
"error": "Prefect infrastructure not ready",
"prefect": not_ready,
"error": "Temporal infrastructure not ready",
"temporal": not_ready,
}
reports = fuzzing.crash_reports.get(run_id)
@@ -725,11 +589,11 @@ async def get_fuzzing_crash_reports_mcp(run_id: str) -> Dict[str, Any]:
async def get_backend_status_mcp() -> Dict[str, Any]:
"""Expose backend readiness, workflows, and registered MCP tools."""
status = get_prefect_status()
response: Dict[str, Any] = {"prefect": status}
status = get_temporal_status()
response: Dict[str, Any] = {"temporal": status}
if status.get("ready"):
response["workflows"] = list(prefect_mgr.workflows.keys())
response["workflows"] = list(temporal_mgr.workflows.keys())
try:
tools = await mcp._tool_manager.list_tools()
@@ -775,12 +639,12 @@ def create_mcp_transport_app() -> Starlette:
# ---------------------------------------------------------------------------
# Combined lifespan: Prefect init + dedicated MCP transports
# Combined lifespan: Temporal init + dedicated MCP transports
# ---------------------------------------------------------------------------
@asynccontextmanager
async def combined_lifespan(app: FastAPI):
global prefect_bootstrap_task, _fastapi_mcp_imported
global temporal_bootstrap_task, _fastapi_mcp_imported
logger.info("Starting FuzzForge backend...")
@@ -793,12 +657,12 @@ async def combined_lifespan(app: FastAPI):
except Exception as exc:
logger.exception("Failed to import FastAPI endpoints into MCP", exc_info=exc)
# Kick off Prefect bootstrap in the background if needed
if prefect_bootstrap_task is None or prefect_bootstrap_task.done():
prefect_bootstrap_task = asyncio.create_task(_bootstrap_prefect_with_retries())
logger.info("Prefect bootstrap task started")
# Kick off Temporal bootstrap in the background if needed
if temporal_bootstrap_task is None or temporal_bootstrap_task.done():
temporal_bootstrap_task = asyncio.create_task(_bootstrap_temporal_with_retries())
logger.info("Temporal bootstrap task started")
else:
logger.info("Prefect bootstrap task already running")
logger.info("Temporal bootstrap task already running")
# Start MCP transports on shared port (HTTP + SSE)
mcp_app = create_mcp_transport_app()
@@ -846,18 +710,17 @@ async def combined_lifespan(app: FastAPI):
mcp_server.force_exit = True
await asyncio.gather(mcp_task, return_exceptions=True)
if prefect_bootstrap_task and not prefect_bootstrap_task.done():
prefect_bootstrap_task.cancel()
if temporal_bootstrap_task and not temporal_bootstrap_task.done():
temporal_bootstrap_task.cancel()
with suppress(asyncio.CancelledError):
await prefect_bootstrap_task
prefect_bootstrap_state.task_running = False
if not prefect_bootstrap_state.ready:
prefect_bootstrap_state.status = "stopped"
prefect_bootstrap_state.next_retry_seconds = None
prefect_bootstrap_task = None
await temporal_bootstrap_task
temporal_bootstrap_state.task_running = False
if not temporal_bootstrap_state.ready:
temporal_bootstrap_state.status = "stopped"
temporal_bootstrap_task = None
logger.info("Shutting down Prefect statistics monitor...")
await prefect_stats_monitor.stop_monitoring()
# Close Temporal client
await temporal_mgr.close()
logger.info("Shutting down FuzzForge backend...")
@@ -1,394 +0,0 @@
"""
Generic Prefect Statistics Monitor Service
This service monitors ALL workflows for structured live data logging and
updates the appropriate statistics APIs. Works with any workflow that follows
the standard LIVE_STATS logging pattern.
"""
# Copyright (c) 2025 FuzzingLabs
#
# Licensed under the Business Source License 1.1 (BSL). See the LICENSE file
# at the root of this repository for details.
#
# After the Change Date (four years from publication), this version of the
# Licensed Work will be made available under the Apache License, Version 2.0.
# See the LICENSE-APACHE file or http://www.apache.org/licenses/LICENSE-2.0
#
# Additional attribution and requirements are provided in the NOTICE file.
import asyncio
import json
import logging
from datetime import datetime, timedelta, timezone
from typing import Dict, Any, Optional
from prefect.client.orchestration import get_client
from prefect.client.schemas.objects import FlowRun, TaskRun
from src.models.findings import FuzzingStats
from src.api.fuzzing import fuzzing_stats, initialize_fuzzing_tracking, active_connections
logger = logging.getLogger(__name__)
class PrefectStatsMonitor:
"""Monitors Prefect flows and tasks for live statistics from any workflow"""
def __init__(self):
self.monitoring = False
self.monitor_task = None
self.monitored_runs = set()
self.last_log_ts: Dict[str, datetime] = {}
self._client = None
self._client_refresh_time = None
self._client_refresh_interval = 300 # Refresh connection every 5 minutes
async def start_monitoring(self):
"""Start the Prefect statistics monitoring service"""
if self.monitoring:
logger.warning("Prefect stats monitor already running")
return
self.monitoring = True
self.monitor_task = asyncio.create_task(self._monitor_flows())
logger.info("Started Prefect statistics monitor")
async def stop_monitoring(self):
"""Stop the monitoring service"""
self.monitoring = False
if self.monitor_task:
self.monitor_task.cancel()
try:
await self.monitor_task
except asyncio.CancelledError:
pass
logger.info("Stopped Prefect statistics monitor")
async def _get_or_refresh_client(self):
"""Get or refresh Prefect client with connection pooling."""
now = datetime.now(timezone.utc)
if (self._client is None or
self._client_refresh_time is None or
(now - self._client_refresh_time).total_seconds() > self._client_refresh_interval):
if self._client:
try:
await self._client.aclose()
except Exception:
pass
self._client = get_client()
self._client_refresh_time = now
await self._client.__aenter__()
return self._client
async def _monitor_flows(self):
"""Main monitoring loop that watches Prefect flows"""
try:
while self.monitoring:
try:
# Use connection pooling for better performance
client = await self._get_or_refresh_client()
# Get recent flow runs (limit to reduce load)
flow_runs = await client.read_flow_runs(
limit=50,
sort="START_TIME_DESC",
)
# Only consider runs from the last 15 minutes
recent_cutoff = datetime.now(timezone.utc) - timedelta(minutes=15)
for flow_run in flow_runs:
created = getattr(flow_run, "created", None)
if created is None:
continue
try:
# Ensure timezone-aware comparison
if created.tzinfo is None:
created = created.replace(tzinfo=timezone.utc)
if created >= recent_cutoff:
await self._monitor_flow_run(client, flow_run)
except Exception:
# If comparison fails, attempt monitoring anyway
await self._monitor_flow_run(client, flow_run)
await asyncio.sleep(5) # Check every 5 seconds
except Exception as e:
logger.error(f"Error in Prefect monitoring: {e}")
await asyncio.sleep(10)
except asyncio.CancelledError:
logger.info("Prefect monitoring cancelled")
except Exception as e:
logger.error(f"Fatal error in Prefect monitoring: {e}")
finally:
# Clean up client on exit
if self._client:
try:
await self._client.__aexit__(None, None, None)
except Exception:
pass
self._client = None
async def _monitor_flow_run(self, client, flow_run: FlowRun):
"""Monitor a specific flow run for statistics"""
run_id = str(flow_run.id)
workflow_name = flow_run.name or "unknown"
try:
# Initialize tracking if not exists - only for workflows that might have live stats
if run_id not in fuzzing_stats:
initialize_fuzzing_tracking(run_id, workflow_name)
self.monitored_runs.add(run_id)
# Skip corrupted entries (should not happen after startup cleanup, but defensive)
elif not isinstance(fuzzing_stats[run_id], FuzzingStats):
logger.warning(f"Skipping corrupted stats entry for {run_id}, reinitializing")
initialize_fuzzing_tracking(run_id, workflow_name)
self.monitored_runs.add(run_id)
# Get task runs for this flow
task_runs = await client.read_task_runs(
flow_run_filter={"id": {"any_": [flow_run.id]}},
limit=25,
)
# Check all tasks for live statistics logging
for task_run in task_runs:
await self._extract_stats_from_task(client, run_id, task_run, workflow_name)
# Also scan flow-level logs as a fallback
await self._extract_stats_from_flow_logs(client, run_id, flow_run, workflow_name)
except Exception as e:
logger.warning(f"Error monitoring flow run {run_id}: {e}")
async def _extract_stats_from_task(self, client, run_id: str, task_run: TaskRun, workflow_name: str):
"""Extract statistics from any task that logs live stats"""
try:
# Get task run logs
logs = await client.read_logs(
log_filter={
"task_run_id": {"any_": [task_run.id]}
},
limit=100,
sort="TIMESTAMP_ASC"
)
# Parse logs for LIVE_STATS entries (generic pattern for any workflow)
latest_stats = None
for log in logs:
# Prefer structured extra field if present
extra_data = getattr(log, "extra", None) or getattr(log, "extra_fields", None) or None
if isinstance(extra_data, dict):
stat_type = extra_data.get("stats_type")
if stat_type in ["fuzzing_live_update", "scan_progress", "analysis_update", "live_stats"]:
latest_stats = extra_data
continue
# Fallback to parsing from message text
if ("FUZZ_STATS" in log.message or "LIVE_STATS" in log.message):
stats = self._parse_stats_from_log(log.message)
if stats:
latest_stats = stats
# Update statistics if we found any
if latest_stats:
# Calculate elapsed time from task start
elapsed_time = 0
if task_run.start_time:
# Ensure timezone-aware arithmetic
now = datetime.now(timezone.utc)
try:
elapsed_time = int((now - task_run.start_time).total_seconds())
except Exception:
# Fallback to naive UTC if types mismatch
elapsed_time = int((datetime.utcnow() - task_run.start_time.replace(tzinfo=None)).total_seconds())
updated_stats = FuzzingStats(
run_id=run_id,
workflow=workflow_name,
executions=latest_stats.get("executions", 0),
executions_per_sec=latest_stats.get("executions_per_sec", 0.0),
crashes=latest_stats.get("crashes", 0),
unique_crashes=latest_stats.get("unique_crashes", 0),
corpus_size=latest_stats.get("corpus_size", 0),
elapsed_time=elapsed_time
)
# Update the global stats
previous = fuzzing_stats.get(run_id)
fuzzing_stats[run_id] = updated_stats
# Broadcast to any active WebSocket clients for this run
if active_connections.get(run_id):
# Handle both Pydantic objects and plain dicts
if isinstance(updated_stats, dict):
stats_data = updated_stats
elif hasattr(updated_stats, 'model_dump'):
stats_data = updated_stats.model_dump()
elif hasattr(updated_stats, 'dict'):
stats_data = updated_stats.dict()
else:
stats_data = updated_stats.__dict__
message = {
"type": "stats_update",
"data": stats_data,
}
disconnected = []
for ws in active_connections[run_id]:
try:
await ws.send_text(json.dumps(message))
except Exception:
disconnected.append(ws)
# Clean up disconnected sockets
for ws in disconnected:
try:
active_connections[run_id].remove(ws)
except ValueError:
pass
logger.debug(f"Updated Prefect stats for {run_id}: {updated_stats.executions} execs")
except Exception as e:
logger.warning(f"Error extracting stats from task {task_run.id}: {e}")
async def _extract_stats_from_flow_logs(self, client, run_id: str, flow_run: FlowRun, workflow_name: str):
"""Extract statistics by scanning flow-level logs for LIVE/FUZZ stats"""
try:
logs = await client.read_logs(
log_filter={
"flow_run_id": {"any_": [flow_run.id]}
},
limit=200,
sort="TIMESTAMP_ASC"
)
latest_stats = None
last_seen = self.last_log_ts.get(run_id)
max_ts = last_seen
for log in logs:
# Skip logs we've already processed
ts = getattr(log, "timestamp", None)
if last_seen and ts and ts <= last_seen:
continue
if ts and (max_ts is None or ts > max_ts):
max_ts = ts
# Prefer structured extra field if available
extra_data = getattr(log, "extra", None) or getattr(log, "extra_fields", None) or None
if isinstance(extra_data, dict):
stat_type = extra_data.get("stats_type")
if stat_type in ["fuzzing_live_update", "scan_progress", "analysis_update", "live_stats"]:
latest_stats = extra_data
continue
# Fallback to message parse
if ("FUZZ_STATS" in log.message or "LIVE_STATS" in log.message):
stats = self._parse_stats_from_log(log.message)
if stats:
latest_stats = stats
if max_ts:
self.last_log_ts[run_id] = max_ts
if latest_stats:
# Use flow_run timestamps for elapsed time if available
elapsed_time = 0
start_time = getattr(flow_run, "start_time", None) or getattr(flow_run, "start_time", None)
if start_time:
now = datetime.now(timezone.utc)
try:
if start_time.tzinfo is None:
start_time = start_time.replace(tzinfo=timezone.utc)
elapsed_time = int((now - start_time).total_seconds())
except Exception:
elapsed_time = int((datetime.utcnow() - start_time.replace(tzinfo=None)).total_seconds())
updated_stats = FuzzingStats(
run_id=run_id,
workflow=workflow_name,
executions=latest_stats.get("executions", 0),
executions_per_sec=latest_stats.get("executions_per_sec", 0.0),
crashes=latest_stats.get("crashes", 0),
unique_crashes=latest_stats.get("unique_crashes", 0),
corpus_size=latest_stats.get("corpus_size", 0),
elapsed_time=elapsed_time
)
fuzzing_stats[run_id] = updated_stats
# Broadcast if listeners exist
if active_connections.get(run_id):
# Handle both Pydantic objects and plain dicts
if isinstance(updated_stats, dict):
stats_data = updated_stats
elif hasattr(updated_stats, 'model_dump'):
stats_data = updated_stats.model_dump()
elif hasattr(updated_stats, 'dict'):
stats_data = updated_stats.dict()
else:
stats_data = updated_stats.__dict__
message = {
"type": "stats_update",
"data": stats_data,
}
disconnected = []
for ws in active_connections[run_id]:
try:
await ws.send_text(json.dumps(message))
except Exception:
disconnected.append(ws)
for ws in disconnected:
try:
active_connections[run_id].remove(ws)
except ValueError:
pass
except Exception as e:
logger.warning(f"Error extracting stats from flow logs {run_id}: {e}")
def _parse_stats_from_log(self, log_message: str) -> Optional[Dict[str, Any]]:
"""Parse statistics from a log message"""
try:
import re
# Prefer explicit JSON after marker tokens
m = re.search(r'(?:FUZZ_STATS|LIVE_STATS)\s+(\{.*\})', log_message)
if m:
try:
return json.loads(m.group(1))
except Exception:
pass
# Fallback: Extract the extra= dict and coerce to JSON
stats_match = re.search(r'extra=({.*?})', log_message)
if not stats_match:
return None
extra_str = stats_match.group(1)
extra_str = extra_str.replace("'", '"')
extra_str = extra_str.replace('None', 'null')
extra_str = extra_str.replace('True', 'true')
extra_str = extra_str.replace('False', 'false')
stats_data = json.loads(extra_str)
# Support multiple stat types for different workflows
stat_type = stats_data.get("stats_type")
if stat_type in ["fuzzing_live_update", "scan_progress", "analysis_update", "live_stats"]:
return stats_data
except Exception as e:
logger.debug(f"Error parsing log stats: {e}")
return None
# Global instance
prefect_stats_monitor = PrefectStatsMonitor()
+10
View File
@@ -0,0 +1,10 @@
"""
Storage abstraction layer for FuzzForge.
Provides unified interface for storing and retrieving targets and results.
"""
from .base import StorageBackend
from .s3_cached import S3CachedStorage
__all__ = ["StorageBackend", "S3CachedStorage"]
+153
View File
@@ -0,0 +1,153 @@
"""
Base storage backend interface.
All storage implementations must implement this interface.
"""
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Optional, Dict, Any
class StorageBackend(ABC):
"""
Abstract base class for storage backends.
Implementations handle storage and retrieval of:
- Uploaded targets (code, binaries, etc.)
- Workflow results
- Temporary files
"""
@abstractmethod
async def upload_target(
self,
file_path: Path,
user_id: str,
metadata: Optional[Dict[str, Any]] = None
) -> str:
"""
Upload a target file to storage.
Args:
file_path: Local path to file to upload
user_id: ID of user uploading the file
metadata: Optional metadata to store with file
Returns:
Target ID (unique identifier for retrieval)
Raises:
FileNotFoundError: If file_path doesn't exist
StorageError: If upload fails
"""
pass
@abstractmethod
async def get_target(self, target_id: str) -> Path:
"""
Get target file from storage.
Args:
target_id: Unique identifier from upload_target()
Returns:
Local path to cached file
Raises:
FileNotFoundError: If target doesn't exist
StorageError: If download fails
"""
pass
@abstractmethod
async def delete_target(self, target_id: str) -> None:
"""
Delete target from storage.
Args:
target_id: Unique identifier to delete
Raises:
StorageError: If deletion fails (doesn't raise if not found)
"""
pass
@abstractmethod
async def upload_results(
self,
workflow_id: str,
results: Dict[str, Any],
results_format: str = "json"
) -> str:
"""
Upload workflow results to storage.
Args:
workflow_id: Workflow execution ID
results: Results dictionary
results_format: Format (json, sarif, etc.)
Returns:
URL to uploaded results
Raises:
StorageError: If upload fails
"""
pass
@abstractmethod
async def get_results(self, workflow_id: str) -> Dict[str, Any]:
"""
Get workflow results from storage.
Args:
workflow_id: Workflow execution ID
Returns:
Results dictionary
Raises:
FileNotFoundError: If results don't exist
StorageError: If download fails
"""
pass
@abstractmethod
async def list_targets(
self,
user_id: Optional[str] = None,
limit: int = 100
) -> list[Dict[str, Any]]:
"""
List uploaded targets.
Args:
user_id: Filter by user ID (None = all users)
limit: Maximum number of results
Returns:
List of target metadata dictionaries
Raises:
StorageError: If listing fails
"""
pass
@abstractmethod
async def cleanup_cache(self) -> int:
"""
Clean up local cache (LRU eviction).
Returns:
Number of files removed
Raises:
StorageError: If cleanup fails
"""
pass
class StorageError(Exception):
"""Base exception for storage operations."""
pass
+423
View File
@@ -0,0 +1,423 @@
"""
S3-compatible storage backend with local caching.
Works with MinIO (dev/prod) or AWS S3 (cloud).
"""
import json
import logging
import os
import shutil
from datetime import datetime
from pathlib import Path
from typing import Optional, Dict, Any
from uuid import uuid4
import boto3
from botocore.exceptions import ClientError
from .base import StorageBackend, StorageError
logger = logging.getLogger(__name__)
class S3CachedStorage(StorageBackend):
"""
S3-compatible storage with local caching.
Features:
- Upload targets to S3/MinIO
- Download with local caching (LRU eviction)
- Lifecycle management (auto-cleanup old files)
- Metadata tracking
"""
def __init__(
self,
endpoint_url: Optional[str] = None,
access_key: Optional[str] = None,
secret_key: Optional[str] = None,
bucket: str = "targets",
region: str = "us-east-1",
use_ssl: bool = False,
cache_dir: Optional[Path] = None,
cache_max_size_gb: int = 10
):
"""
Initialize S3 storage backend.
Args:
endpoint_url: S3 endpoint (None = AWS S3, or MinIO URL)
access_key: S3 access key (None = from env)
secret_key: S3 secret key (None = from env)
bucket: S3 bucket name
region: AWS region
use_ssl: Use HTTPS
cache_dir: Local cache directory
cache_max_size_gb: Maximum cache size in GB
"""
# Use environment variables as defaults
self.endpoint_url = endpoint_url or os.getenv('S3_ENDPOINT', 'http://minio:9000')
self.access_key = access_key or os.getenv('S3_ACCESS_KEY', 'fuzzforge')
self.secret_key = secret_key or os.getenv('S3_SECRET_KEY', 'fuzzforge123')
self.bucket = bucket or os.getenv('S3_BUCKET', 'targets')
self.region = region or os.getenv('S3_REGION', 'us-east-1')
self.use_ssl = use_ssl or os.getenv('S3_USE_SSL', 'false').lower() == 'true'
# Cache configuration
self.cache_dir = cache_dir or Path(os.getenv('CACHE_DIR', '/tmp/fuzzforge-cache'))
self.cache_max_size = cache_max_size_gb * (1024 ** 3) # Convert to bytes
# Ensure cache directory exists
self.cache_dir.mkdir(parents=True, exist_ok=True)
# Initialize S3 client
try:
self.s3_client = boto3.client(
's3',
endpoint_url=self.endpoint_url,
aws_access_key_id=self.access_key,
aws_secret_access_key=self.secret_key,
region_name=self.region,
use_ssl=self.use_ssl
)
logger.info(f"Initialized S3 storage: {self.endpoint_url}/{self.bucket}")
except Exception as e:
logger.error(f"Failed to initialize S3 client: {e}")
raise StorageError(f"S3 initialization failed: {e}")
async def upload_target(
self,
file_path: Path,
user_id: str,
metadata: Optional[Dict[str, Any]] = None
) -> str:
"""Upload target file to S3/MinIO."""
if not file_path.exists():
raise FileNotFoundError(f"File not found: {file_path}")
# Generate unique target ID
target_id = str(uuid4())
# Prepare metadata
upload_metadata = {
'user_id': user_id,
'uploaded_at': datetime.now().isoformat(),
'filename': file_path.name,
'size': str(file_path.stat().st_size)
}
if metadata:
upload_metadata.update(metadata)
# Upload to S3
s3_key = f'{target_id}/target'
try:
logger.info(f"Uploading target to s3://{self.bucket}/{s3_key}")
self.s3_client.upload_file(
str(file_path),
self.bucket,
s3_key,
ExtraArgs={
'Metadata': upload_metadata
}
)
file_size_mb = file_path.stat().st_size / (1024 * 1024)
logger.info(
f"✓ Uploaded target {target_id} "
f"({file_path.name}, {file_size_mb:.2f} MB)"
)
return target_id
except ClientError as e:
logger.error(f"S3 upload failed: {e}", exc_info=True)
raise StorageError(f"Failed to upload target: {e}")
except Exception as e:
logger.error(f"Upload failed: {e}", exc_info=True)
raise StorageError(f"Upload error: {e}")
async def get_target(self, target_id: str) -> Path:
"""Get target from cache or download from S3/MinIO."""
# Check cache first
cache_path = self.cache_dir / target_id
cached_file = cache_path / "target"
if cached_file.exists():
# Update access time for LRU
cached_file.touch()
logger.info(f"Cache HIT: {target_id}")
return cached_file
# Cache miss - download from S3
logger.info(f"Cache MISS: {target_id}, downloading from S3...")
try:
# Create cache directory
cache_path.mkdir(parents=True, exist_ok=True)
# Download from S3
s3_key = f'{target_id}/target'
logger.info(f"Downloading s3://{self.bucket}/{s3_key}")
self.s3_client.download_file(
self.bucket,
s3_key,
str(cached_file)
)
# Verify download
if not cached_file.exists():
raise StorageError(f"Downloaded file not found: {cached_file}")
file_size_mb = cached_file.stat().st_size / (1024 * 1024)
logger.info(f"✓ Downloaded target {target_id} ({file_size_mb:.2f} MB)")
return cached_file
except ClientError as e:
error_code = e.response.get('Error', {}).get('Code')
if error_code in ['404', 'NoSuchKey']:
logger.error(f"Target not found: {target_id}")
raise FileNotFoundError(f"Target {target_id} not found in storage")
else:
logger.error(f"S3 download failed: {e}", exc_info=True)
raise StorageError(f"Download failed: {e}")
except Exception as e:
logger.error(f"Download error: {e}", exc_info=True)
# Cleanup partial download
if cache_path.exists():
shutil.rmtree(cache_path, ignore_errors=True)
raise StorageError(f"Download error: {e}")
async def delete_target(self, target_id: str) -> None:
"""Delete target from S3/MinIO."""
try:
s3_key = f'{target_id}/target'
logger.info(f"Deleting s3://{self.bucket}/{s3_key}")
self.s3_client.delete_object(
Bucket=self.bucket,
Key=s3_key
)
# Also delete from cache if present
cache_path = self.cache_dir / target_id
if cache_path.exists():
shutil.rmtree(cache_path, ignore_errors=True)
logger.info(f"✓ Deleted target {target_id} from S3 and cache")
else:
logger.info(f"✓ Deleted target {target_id} from S3")
except ClientError as e:
logger.error(f"S3 delete failed: {e}", exc_info=True)
# Don't raise error if object doesn't exist
if e.response.get('Error', {}).get('Code') not in ['404', 'NoSuchKey']:
raise StorageError(f"Delete failed: {e}")
except Exception as e:
logger.error(f"Delete error: {e}", exc_info=True)
raise StorageError(f"Delete error: {e}")
async def upload_results(
self,
workflow_id: str,
results: Dict[str, Any],
results_format: str = "json"
) -> str:
"""Upload workflow results to S3/MinIO."""
try:
# Prepare results content
if results_format == "json":
content = json.dumps(results, indent=2).encode('utf-8')
content_type = 'application/json'
file_ext = 'json'
elif results_format == "sarif":
content = json.dumps(results, indent=2).encode('utf-8')
content_type = 'application/sarif+json'
file_ext = 'sarif'
else:
content = json.dumps(results, indent=2).encode('utf-8')
content_type = 'application/json'
file_ext = 'json'
# Upload to results bucket
results_bucket = 'results'
s3_key = f'{workflow_id}/results.{file_ext}'
logger.info(f"Uploading results to s3://{results_bucket}/{s3_key}")
self.s3_client.put_object(
Bucket=results_bucket,
Key=s3_key,
Body=content,
ContentType=content_type,
Metadata={
'workflow_id': workflow_id,
'format': results_format,
'uploaded_at': datetime.now().isoformat()
}
)
# Construct URL
results_url = f"{self.endpoint_url}/{results_bucket}/{s3_key}"
logger.info(f"✓ Uploaded results: {results_url}")
return results_url
except Exception as e:
logger.error(f"Results upload failed: {e}", exc_info=True)
raise StorageError(f"Results upload failed: {e}")
async def get_results(self, workflow_id: str) -> Dict[str, Any]:
"""Get workflow results from S3/MinIO."""
try:
results_bucket = 'results'
s3_key = f'{workflow_id}/results.json'
logger.info(f"Downloading results from s3://{results_bucket}/{s3_key}")
response = self.s3_client.get_object(
Bucket=results_bucket,
Key=s3_key
)
content = response['Body'].read().decode('utf-8')
results = json.loads(content)
logger.info(f"✓ Downloaded results for workflow {workflow_id}")
return results
except ClientError as e:
error_code = e.response.get('Error', {}).get('Code')
if error_code in ['404', 'NoSuchKey']:
logger.error(f"Results not found: {workflow_id}")
raise FileNotFoundError(f"Results for workflow {workflow_id} not found")
else:
logger.error(f"Results download failed: {e}", exc_info=True)
raise StorageError(f"Results download failed: {e}")
except Exception as e:
logger.error(f"Results download error: {e}", exc_info=True)
raise StorageError(f"Results download error: {e}")
async def list_targets(
self,
user_id: Optional[str] = None,
limit: int = 100
) -> list[Dict[str, Any]]:
"""List uploaded targets."""
try:
targets = []
paginator = self.s3_client.get_paginator('list_objects_v2')
for page in paginator.paginate(Bucket=self.bucket, PaginationConfig={'MaxItems': limit}):
for obj in page.get('Contents', []):
# Get object metadata
try:
metadata_response = self.s3_client.head_object(
Bucket=self.bucket,
Key=obj['Key']
)
metadata = metadata_response.get('Metadata', {})
# Filter by user_id if specified
if user_id and metadata.get('user_id') != user_id:
continue
targets.append({
'target_id': obj['Key'].split('/')[0],
'key': obj['Key'],
'size': obj['Size'],
'last_modified': obj['LastModified'].isoformat(),
'metadata': metadata
})
except Exception as e:
logger.warning(f"Failed to get metadata for {obj['Key']}: {e}")
continue
logger.info(f"Listed {len(targets)} targets (user_id={user_id})")
return targets
except Exception as e:
logger.error(f"List targets failed: {e}", exc_info=True)
raise StorageError(f"List targets failed: {e}")
async def cleanup_cache(self) -> int:
"""Clean up local cache using LRU eviction."""
try:
cache_files = []
total_size = 0
# Gather all cached files with metadata
for cache_file in self.cache_dir.rglob('*'):
if cache_file.is_file():
try:
stat = cache_file.stat()
cache_files.append({
'path': cache_file,
'size': stat.st_size,
'atime': stat.st_atime # Last access time
})
total_size += stat.st_size
except Exception as e:
logger.warning(f"Failed to stat {cache_file}: {e}")
continue
# Check if cleanup is needed
if total_size <= self.cache_max_size:
logger.info(
f"Cache size OK: {total_size / (1024**3):.2f} GB / "
f"{self.cache_max_size / (1024**3):.2f} GB"
)
return 0
# Sort by access time (oldest first)
cache_files.sort(key=lambda x: x['atime'])
# Remove files until under limit
removed_count = 0
for file_info in cache_files:
if total_size <= self.cache_max_size:
break
try:
file_info['path'].unlink()
total_size -= file_info['size']
removed_count += 1
logger.debug(f"Evicted from cache: {file_info['path']}")
except Exception as e:
logger.warning(f"Failed to delete {file_info['path']}: {e}")
continue
logger.info(
f"✓ Cache cleanup: removed {removed_count} files, "
f"new size: {total_size / (1024**3):.2f} GB"
)
return removed_count
except Exception as e:
logger.error(f"Cache cleanup failed: {e}", exc_info=True)
raise StorageError(f"Cache cleanup failed: {e}")
def get_cache_stats(self) -> Dict[str, Any]:
"""Get cache statistics."""
try:
total_size = 0
file_count = 0
for cache_file in self.cache_dir.rglob('*'):
if cache_file.is_file():
total_size += cache_file.stat().st_size
file_count += 1
return {
'total_size_bytes': total_size,
'total_size_gb': total_size / (1024 ** 3),
'file_count': file_count,
'max_size_gb': self.cache_max_size / (1024 ** 3),
'usage_percent': (total_size / self.cache_max_size) * 100
}
except Exception as e:
logger.error(f"Failed to get cache stats: {e}")
return {'error': str(e)}
+10
View File
@@ -0,0 +1,10 @@
"""
Temporal integration for FuzzForge.
Handles workflow execution, monitoring, and management.
"""
from .manager import TemporalManager
from .discovery import WorkflowDiscovery
__all__ = ["TemporalManager", "WorkflowDiscovery"]
+268
View File
@@ -0,0 +1,268 @@
"""
Workflow Discovery for Temporal
Discovers workflows from the toolbox/workflows directory
and provides metadata about available workflows.
"""
import logging
import yaml
from pathlib import Path
from typing import Dict, Any
from pydantic import BaseModel, Field, ConfigDict
logger = logging.getLogger(__name__)
class WorkflowInfo(BaseModel):
"""Information about a discovered workflow"""
name: str = Field(..., description="Workflow name")
path: Path = Field(..., description="Path to workflow directory")
workflow_file: Path = Field(..., description="Path to workflow.py file")
metadata: Dict[str, Any] = Field(..., description="Workflow metadata from YAML")
workflow_type: str = Field(..., description="Workflow class name")
vertical: str = Field(..., description="Vertical (worker type) for this workflow")
model_config = ConfigDict(arbitrary_types_allowed=True)
class WorkflowDiscovery:
"""
Discovers workflows from the filesystem.
Scans toolbox/workflows/ for directories containing:
- metadata.yaml (required)
- workflow.py (required)
Each workflow declares its vertical (rust, android, web, etc.)
which determines which worker pool will execute it.
"""
def __init__(self, workflows_dir: Path):
"""
Initialize workflow discovery.
Args:
workflows_dir: Path to the workflows directory
"""
self.workflows_dir = workflows_dir
if not self.workflows_dir.exists():
self.workflows_dir.mkdir(parents=True, exist_ok=True)
logger.info(f"Created workflows directory: {self.workflows_dir}")
async def discover_workflows(self) -> Dict[str, WorkflowInfo]:
"""
Discover workflows by scanning the workflows directory.
Returns:
Dictionary mapping workflow names to their information
"""
workflows = {}
logger.info(f"Scanning for workflows in: {self.workflows_dir}")
for workflow_dir in self.workflows_dir.iterdir():
if not workflow_dir.is_dir():
continue
# Skip special directories
if workflow_dir.name.startswith('.') or workflow_dir.name == '__pycache__':
continue
metadata_file = workflow_dir / "metadata.yaml"
if not metadata_file.exists():
logger.debug(f"No metadata.yaml in {workflow_dir.name}, skipping")
continue
workflow_file = workflow_dir / "workflow.py"
if not workflow_file.exists():
logger.warning(
f"Workflow {workflow_dir.name} has metadata but no workflow.py, skipping"
)
continue
try:
# Parse metadata
with open(metadata_file) as f:
metadata = yaml.safe_load(f)
# Validate required fields
if 'name' not in metadata:
logger.warning(f"Workflow {workflow_dir.name} metadata missing 'name' field")
metadata['name'] = workflow_dir.name
if 'vertical' not in metadata:
logger.warning(
f"Workflow {workflow_dir.name} metadata missing 'vertical' field"
)
continue
# Infer workflow class name from metadata or use convention
workflow_type = metadata.get('workflow_class')
if not workflow_type:
# Convention: convert snake_case to PascalCase + Workflow
# e.g., rust_test -> RustTestWorkflow
parts = workflow_dir.name.split('_')
workflow_type = ''.join(part.capitalize() for part in parts) + 'Workflow'
# Create workflow info
info = WorkflowInfo(
name=metadata['name'],
path=workflow_dir,
workflow_file=workflow_file,
metadata=metadata,
workflow_type=workflow_type,
vertical=metadata['vertical']
)
workflows[info.name] = info
logger.info(
f"✓ Discovered workflow: {info.name} "
f"(vertical: {info.vertical}, class: {info.workflow_type})"
)
except Exception as e:
logger.error(
f"Error discovering workflow {workflow_dir.name}: {e}",
exc_info=True
)
continue
logger.info(f"Discovered {len(workflows)} workflows")
return workflows
def get_workflows_by_vertical(
self,
workflows: Dict[str, WorkflowInfo],
vertical: str
) -> Dict[str, WorkflowInfo]:
"""
Filter workflows by vertical.
Args:
workflows: All discovered workflows
vertical: Vertical name to filter by
Returns:
Filtered workflows dictionary
"""
return {
name: info
for name, info in workflows.items()
if info.vertical == vertical
}
def get_available_verticals(self, workflows: Dict[str, WorkflowInfo]) -> list[str]:
"""
Get list of all verticals from discovered workflows.
Args:
workflows: All discovered workflows
Returns:
List of unique vertical names
"""
return list(set(info.vertical for info in workflows.values()))
@staticmethod
def get_metadata_schema() -> Dict[str, Any]:
"""
Get the JSON schema for workflow metadata.
Returns:
JSON schema dictionary
"""
return {
"type": "object",
"required": ["name", "version", "description", "author", "category", "vertical", "parameters", "requirements"],
"properties": {
"name": {
"type": "string",
"description": "Workflow name"
},
"version": {
"type": "string",
"pattern": "^\\d+\\.\\d+\\.\\d+$",
"description": "Semantic version (x.y.z)"
},
"vertical": {
"type": "string",
"description": "Vertical worker type (rust, android, web, etc.)"
},
"description": {
"type": "string",
"description": "Workflow description"
},
"author": {
"type": "string",
"description": "Workflow author"
},
"category": {
"type": "string",
"enum": ["comprehensive", "specialized", "fuzzing", "focused"],
"description": "Workflow category"
},
"tags": {
"type": "array",
"items": {"type": "string"},
"description": "Workflow tags for categorization"
},
"requirements": {
"type": "object",
"required": ["tools", "resources"],
"properties": {
"tools": {
"type": "array",
"items": {"type": "string"},
"description": "Required security tools"
},
"resources": {
"type": "object",
"required": ["memory", "cpu", "timeout"],
"properties": {
"memory": {
"type": "string",
"pattern": "^\\d+[GMK]i$",
"description": "Memory limit (e.g., 1Gi, 512Mi)"
},
"cpu": {
"type": "string",
"pattern": "^\\d+m?$",
"description": "CPU limit (e.g., 1000m, 2)"
},
"timeout": {
"type": "integer",
"minimum": 60,
"maximum": 7200,
"description": "Workflow timeout in seconds"
}
}
}
}
},
"parameters": {
"type": "object",
"description": "Workflow parameters schema"
},
"default_parameters": {
"type": "object",
"description": "Default parameter values"
},
"required_modules": {
"type": "array",
"items": {"type": "string"},
"description": "Required module names"
},
"supported_volume_modes": {
"type": "array",
"items": {"enum": ["ro", "rw"]},
"default": ["ro", "rw"],
"description": "Supported volume mount modes"
},
"has_docker": {
"type": "boolean",
"default": False,
"description": "Whether workflow has custom Docker build"
}
}
}
+369
View File
@@ -0,0 +1,369 @@
"""
Temporal Manager - Workflow execution and management
Handles:
- Workflow discovery from toolbox
- Workflow execution (submit to Temporal)
- Status monitoring
- Results retrieval
"""
import logging
import os
from pathlib import Path
from typing import Dict, Optional, Any
from uuid import uuid4
from temporalio.client import Client, WorkflowHandle
from temporalio.common import RetryPolicy
from datetime import timedelta
from .discovery import WorkflowDiscovery, WorkflowInfo
from src.storage import S3CachedStorage
logger = logging.getLogger(__name__)
class TemporalManager:
"""
Manages Temporal workflow execution for FuzzForge.
This class:
- Discovers available workflows from toolbox
- Submits workflow executions to Temporal
- Monitors workflow status
- Retrieves workflow results
"""
def __init__(
self,
workflows_dir: Optional[Path] = None,
temporal_address: Optional[str] = None,
temporal_namespace: str = "default",
storage: Optional[S3CachedStorage] = None
):
"""
Initialize Temporal manager.
Args:
workflows_dir: Path to workflows directory (default: toolbox/workflows)
temporal_address: Temporal server address (default: from env or localhost:7233)
temporal_namespace: Temporal namespace
storage: Storage backend for file uploads (default: S3CachedStorage)
"""
if workflows_dir is None:
workflows_dir = Path("toolbox/workflows")
self.temporal_address = temporal_address or os.getenv(
'TEMPORAL_ADDRESS',
'localhost:7233'
)
self.temporal_namespace = temporal_namespace
self.discovery = WorkflowDiscovery(workflows_dir)
self.workflows: Dict[str, WorkflowInfo] = {}
self.client: Optional[Client] = None
# Initialize storage backend
self.storage = storage or S3CachedStorage()
logger.info(
f"TemporalManager initialized: {self.temporal_address} "
f"(namespace: {self.temporal_namespace})"
)
async def initialize(self):
"""Initialize the manager by discovering workflows and connecting to Temporal."""
try:
# Discover workflows
self.workflows = await self.discovery.discover_workflows()
if not self.workflows:
logger.warning("No workflows discovered")
else:
logger.info(
f"Discovered {len(self.workflows)} workflows: "
f"{list(self.workflows.keys())}"
)
# Connect to Temporal
self.client = await Client.connect(
self.temporal_address,
namespace=self.temporal_namespace
)
logger.info(f"✓ Connected to Temporal: {self.temporal_address}")
except Exception as e:
logger.error(f"Failed to initialize Temporal manager: {e}", exc_info=True)
raise
async def close(self):
"""Close Temporal client connection."""
if self.client:
# Temporal client doesn't need explicit close in Python SDK
# but we keep this for symmetry with PrefectManager
pass
async def get_workflows(self) -> Dict[str, WorkflowInfo]:
"""
Get all discovered workflows.
Returns:
Dictionary mapping workflow names to their info
"""
return self.workflows
async def get_workflow(self, name: str) -> Optional[WorkflowInfo]:
"""
Get workflow info by name.
Args:
name: Workflow name
Returns:
WorkflowInfo or None if not found
"""
return self.workflows.get(name)
async def upload_target(
self,
file_path: Path,
user_id: str,
metadata: Optional[Dict[str, Any]] = None
) -> str:
"""
Upload target file to storage.
Args:
file_path: Local path to file
user_id: User ID
metadata: Optional metadata
Returns:
Target ID for use in workflow execution
"""
target_id = await self.storage.upload_target(file_path, user_id, metadata)
logger.info(f"Uploaded target: {target_id}")
return target_id
async def run_workflow(
self,
workflow_name: str,
target_id: str,
workflow_params: Optional[Dict[str, Any]] = None,
workflow_id: Optional[str] = None
) -> WorkflowHandle:
"""
Execute a workflow.
Args:
workflow_name: Name of workflow to execute
target_id: Target ID (from upload_target)
workflow_params: Additional workflow parameters
workflow_id: Optional workflow ID (generated if not provided)
Returns:
WorkflowHandle for monitoring/results
Raises:
ValueError: If workflow not found or client not initialized
"""
if not self.client:
raise ValueError("Temporal client not initialized. Call initialize() first.")
# Get workflow info
workflow_info = self.workflows.get(workflow_name)
if not workflow_info:
raise ValueError(f"Workflow not found: {workflow_name}")
# Generate workflow ID if not provided
if not workflow_id:
workflow_id = f"{workflow_name}-{str(uuid4())[:8]}"
# Prepare workflow input arguments in order
# For security_assessment: (target_id, scanner_config, analyzer_config, reporter_config)
workflow_params = workflow_params or {}
workflow_args = [
target_id,
workflow_params.get("scanner_config"),
workflow_params.get("analyzer_config"),
workflow_params.get("reporter_config")
]
# Determine task queue from workflow vertical
vertical = workflow_info.metadata.get("vertical", "default")
task_queue = f"{vertical}-queue"
logger.info(
f"Starting workflow: {workflow_name} "
f"(id={workflow_id}, queue={task_queue}, target={target_id})"
)
try:
# Start workflow execution with positional arguments
handle = await self.client.start_workflow(
workflow=workflow_info.workflow_type, # Workflow class name
args=workflow_args, # Positional arguments
id=workflow_id,
task_queue=task_queue,
retry_policy=RetryPolicy(
initial_interval=timedelta(seconds=1),
maximum_interval=timedelta(minutes=1),
maximum_attempts=3
)
)
logger.info(f"✓ Workflow started: {workflow_id}")
return handle
except Exception as e:
logger.error(f"Failed to start workflow {workflow_name}: {e}", exc_info=True)
raise
async def get_workflow_status(self, workflow_id: str) -> Dict[str, Any]:
"""
Get workflow execution status.
Args:
workflow_id: Workflow execution ID
Returns:
Status dictionary with workflow state
Raises:
ValueError: If client not initialized or workflow not found
"""
if not self.client:
raise ValueError("Temporal client not initialized")
try:
# Get workflow handle
handle = self.client.get_workflow_handle(workflow_id)
# Try to get result (non-blocking describe)
description = await handle.describe()
status = {
"workflow_id": workflow_id,
"status": description.status.name,
"start_time": description.start_time.isoformat() if description.start_time else None,
"execution_time": description.execution_time.isoformat() if description.execution_time else None,
"close_time": description.close_time.isoformat() if description.close_time else None,
"task_queue": description.task_queue,
}
logger.info(f"Workflow {workflow_id} status: {status['status']}")
return status
except Exception as e:
logger.error(f"Failed to get workflow status: {e}", exc_info=True)
raise
async def get_workflow_result(
self,
workflow_id: str,
timeout: Optional[timedelta] = None
) -> Any:
"""
Get workflow execution result (blocking).
Args:
workflow_id: Workflow execution ID
timeout: Maximum time to wait for result
Returns:
Workflow result
Raises:
ValueError: If client not initialized
TimeoutError: If timeout exceeded
"""
if not self.client:
raise ValueError("Temporal client not initialized")
try:
handle = self.client.get_workflow_handle(workflow_id)
logger.info(f"Waiting for workflow result: {workflow_id}")
# Wait for workflow to complete and get result
if timeout:
# Use asyncio timeout if provided
import asyncio
result = await asyncio.wait_for(handle.result(), timeout=timeout.total_seconds())
else:
result = await handle.result()
logger.info(f"✓ Workflow {workflow_id} completed")
return result
except Exception as e:
logger.error(f"Failed to get workflow result: {e}", exc_info=True)
raise
async def cancel_workflow(self, workflow_id: str) -> None:
"""
Cancel a running workflow.
Args:
workflow_id: Workflow execution ID
Raises:
ValueError: If client not initialized
"""
if not self.client:
raise ValueError("Temporal client not initialized")
try:
handle = self.client.get_workflow_handle(workflow_id)
await handle.cancel()
logger.info(f"✓ Workflow cancelled: {workflow_id}")
except Exception as e:
logger.error(f"Failed to cancel workflow: {e}", exc_info=True)
raise
async def list_workflows(
self,
filter_query: Optional[str] = None,
limit: int = 100
) -> list[Dict[str, Any]]:
"""
List workflow executions.
Args:
filter_query: Optional Temporal list filter query
limit: Maximum number of results
Returns:
List of workflow execution info
Raises:
ValueError: If client not initialized
"""
if not self.client:
raise ValueError("Temporal client not initialized")
try:
workflows = []
# Use Temporal's list API
async for workflow in self.client.list_workflows(filter_query):
workflows.append({
"workflow_id": workflow.id,
"workflow_type": workflow.workflow_type,
"status": workflow.status.name,
"start_time": workflow.start_time.isoformat() if workflow.start_time else None,
"close_time": workflow.close_time.isoformat() if workflow.close_time else None,
"task_queue": workflow.task_queue,
})
if len(workflows) >= limit:
break
logger.info(f"Listed {len(workflows)} workflows")
return workflows
except Exception as e:
logger.error(f"Failed to list workflows: {e}", exc_info=True)
raise
@@ -1,82 +0,0 @@
# Copyright (c) 2025 FuzzingLabs
#
# Licensed under the Business Source License 1.1 (BSL). See the LICENSE file
# at the root of this repository for details.
#
# After the Change Date (four years from publication), this version of the
# Licensed Work will be made available under the Apache License, Version 2.0.
# See the LICENSE-APACHE file or http://www.apache.org/licenses/LICENSE-2.0
#
# Additional attribution and requirements are provided in the NOTICE file.
import asyncio
from datetime import datetime, timezone, timedelta
from src.services.prefect_stats_monitor import PrefectStatsMonitor
from src.api import fuzzing
class FakeLog:
def __init__(self, message: str):
self.message = message
class FakeClient:
def __init__(self, logs):
self._logs = logs
async def read_logs(self, log_filter=None, limit=100, sort="TIMESTAMP_ASC"):
return self._logs
class FakeTaskRun:
def __init__(self):
self.id = "task-1"
self.start_time = datetime.now(timezone.utc) - timedelta(seconds=5)
def test_parse_stats_from_log_fuzzing():
mon = PrefectStatsMonitor()
msg = (
"INFO LIVE_STATS extra={'stats_type': 'fuzzing_live_update', "
"'executions': 42, 'executions_per_sec': 3.14, 'crashes': 1, 'unique_crashes': 1, 'corpus_size': 9}"
)
stats = mon._parse_stats_from_log(msg)
assert stats is not None
assert stats["stats_type"] == "fuzzing_live_update"
assert stats["executions"] == 42
def test_extract_stats_updates_and_broadcasts():
mon = PrefectStatsMonitor()
run_id = "run-123"
workflow = "wf"
fuzzing.initialize_fuzzing_tracking(run_id, workflow)
# Prepare a fake websocket to capture messages
sent = []
class FakeWS:
async def send_text(self, text: str):
sent.append(text)
fuzzing.active_connections[run_id] = [FakeWS()]
# Craft a log line the parser understands
msg = (
"INFO LIVE_STATS extra={'stats_type': 'fuzzing_live_update', "
"'executions': 10, 'executions_per_sec': 1.5, 'crashes': 0, 'unique_crashes': 0, 'corpus_size': 2}"
)
fake_client = FakeClient([FakeLog(msg)])
task_run = FakeTaskRun()
asyncio.run(mon._extract_stats_from_task(fake_client, run_id, task_run, workflow))
# Verify stats updated
stats = fuzzing.fuzzing_stats[run_id]
assert stats.executions == 10
assert stats.executions_per_sec == 1.5
# Verify a message was sent to WebSocket
assert sent, "Expected a stats_update message to be sent"
@@ -0,0 +1,238 @@
"""
FuzzForge Common Storage Activities
Activities for interacting with MinIO storage:
- get_target_activity: Download target from MinIO to local cache
- cleanup_cache_activity: Remove target from local cache
- upload_results_activity: Upload workflow results to MinIO
"""
import logging
import os
import shutil
from pathlib import Path
import boto3
from botocore.exceptions import ClientError
from temporalio import activity
# Configure logging
logger = logging.getLogger(__name__)
# Initialize S3 client (MinIO)
s3_client = boto3.client(
's3',
endpoint_url=os.getenv('S3_ENDPOINT', 'http://minio:9000'),
aws_access_key_id=os.getenv('S3_ACCESS_KEY', 'fuzzforge'),
aws_secret_access_key=os.getenv('S3_SECRET_KEY', 'fuzzforge123'),
region_name=os.getenv('S3_REGION', 'us-east-1'),
use_ssl=os.getenv('S3_USE_SSL', 'false').lower() == 'true'
)
# Configuration
S3_BUCKET = os.getenv('S3_BUCKET', 'targets')
CACHE_DIR = Path(os.getenv('CACHE_DIR', '/cache'))
CACHE_MAX_SIZE_GB = int(os.getenv('CACHE_MAX_SIZE', '10').rstrip('GB'))
@activity.defn(name="get_target")
async def get_target_activity(target_id: str) -> str:
"""
Download target from MinIO to local cache.
Args:
target_id: UUID of the uploaded target
Returns:
Local path to the cached target file
Raises:
FileNotFoundError: If target doesn't exist in MinIO
Exception: For other download errors
"""
logger.info(f"Activity: get_target (target_id={target_id})")
# Define cache paths
cache_path = CACHE_DIR / target_id
cached_file = cache_path / "target"
# Check if target is already cached
if cached_file.exists():
# Update access time for LRU
cached_file.touch()
logger.info(f"Cache HIT: {target_id}")
return str(cached_file)
# Cache miss - download from MinIO
logger.info(f"Cache MISS: {target_id}, downloading from MinIO...")
try:
# Create cache directory
cache_path.mkdir(parents=True, exist_ok=True)
# Download from S3/MinIO
s3_key = f'{target_id}/target'
logger.info(f"Downloading s3://{S3_BUCKET}/{s3_key} -> {cached_file}")
s3_client.download_file(
Bucket=S3_BUCKET,
Key=s3_key,
Filename=str(cached_file)
)
# Verify file was downloaded
if not cached_file.exists():
raise FileNotFoundError(f"Downloaded file not found: {cached_file}")
file_size = cached_file.stat().st_size
logger.info(
f"✓ Downloaded target {target_id} "
f"({file_size / 1024 / 1024:.2f} MB)"
)
# Extract tarball if it's an archive
import tarfile
workspace_dir = cache_path / "workspace"
if tarfile.is_tarfile(str(cached_file)):
logger.info(f"Extracting tarball to {workspace_dir}...")
workspace_dir.mkdir(parents=True, exist_ok=True)
with tarfile.open(str(cached_file), 'r:*') as tar:
tar.extractall(path=workspace_dir)
logger.info(f"✓ Extracted tarball to {workspace_dir}")
return str(workspace_dir)
else:
# Not a tarball, return file path
return str(cached_file)
except ClientError as e:
error_code = e.response['Error']['Code']
if error_code == '404' or error_code == 'NoSuchKey':
logger.error(f"Target not found in MinIO: {target_id}")
raise FileNotFoundError(f"Target {target_id} not found in storage")
else:
logger.error(f"S3/MinIO error downloading target: {e}", exc_info=True)
raise
except Exception as e:
logger.error(f"Failed to download target {target_id}: {e}", exc_info=True)
# Cleanup partial download
if cache_path.exists():
shutil.rmtree(cache_path, ignore_errors=True)
raise
@activity.defn(name="cleanup_cache")
async def cleanup_cache_activity(target_path: str) -> None:
"""
Remove target from local cache after workflow completes.
Args:
target_path: Path to the cached target file (from get_target_activity)
"""
logger.info(f"Activity: cleanup_cache (path={target_path})")
try:
cache_file = Path(target_path)
cache_dir = cache_file.parent
if cache_dir.exists() and cache_dir.is_relative_to(CACHE_DIR):
shutil.rmtree(cache_dir)
logger.info(f"✓ Cleaned up cache: {cache_dir}")
else:
logger.warning(f"Cache path not in CACHE_DIR or doesn't exist: {cache_dir}")
except Exception as e:
# Don't fail workflow if cleanup fails
logger.error(f"Failed to cleanup cache {target_path}: {e}", exc_info=True)
@activity.defn(name="upload_results")
async def upload_results_activity(
workflow_id: str,
results: dict,
results_format: str = "json"
) -> str:
"""
Upload workflow results to MinIO.
Args:
workflow_id: Workflow execution ID
results: Results dictionary to upload
results_format: Format for results (json, sarif, etc.)
Returns:
S3 URL to the uploaded results
"""
logger.info(
f"Activity: upload_results "
f"(workflow_id={workflow_id}, format={results_format})"
)
try:
import json
# Prepare results content
if results_format == "json":
content = json.dumps(results, indent=2).encode('utf-8')
content_type = 'application/json'
file_ext = 'json'
elif results_format == "sarif":
content = json.dumps(results, indent=2).encode('utf-8')
content_type = 'application/sarif+json'
file_ext = 'sarif'
else:
# Default to JSON
content = json.dumps(results, indent=2).encode('utf-8')
content_type = 'application/json'
file_ext = 'json'
# Upload to MinIO
s3_key = f'{workflow_id}/results.{file_ext}'
logger.info(f"Uploading results to s3://results/{s3_key}")
s3_client.put_object(
Bucket='results',
Key=s3_key,
Body=content,
ContentType=content_type,
Metadata={
'workflow_id': workflow_id,
'format': results_format
}
)
# Construct S3 URL
s3_endpoint = os.getenv('S3_ENDPOINT', 'http://minio:9000')
s3_url = f"{s3_endpoint}/results/{s3_key}"
logger.info(f"✓ Uploaded results: {s3_url}")
return s3_url
except Exception as e:
logger.error(
f"Failed to upload results for workflow {workflow_id}: {e}",
exc_info=True
)
raise
def _check_cache_size():
"""Check total cache size and log warning if exceeding limit"""
try:
total_size = 0
for item in CACHE_DIR.rglob('*'):
if item.is_file():
total_size += item.stat().st_size
total_size_gb = total_size / (1024 ** 3)
if total_size_gb > CACHE_MAX_SIZE_GB:
logger.warning(
f"Cache size ({total_size_gb:.2f} GB) exceeds "
f"limit ({CACHE_MAX_SIZE_GB} GB). Consider cleanup."
)
except Exception as e:
logger.error(f"Failed to check cache size: {e}")
@@ -1,12 +0,0 @@
# Copyright (c) 2025 FuzzingLabs
#
# Licensed under the Business Source License 1.1 (BSL). See the LICENSE file
# at the root of this repository for details.
#
# After the Change Date (four years from publication), this version of the
# Licensed Work will be made available under the Apache License, Version 2.0.
# See the LICENSE-APACHE file or http://www.apache.org/licenses/LICENSE-2.0
#
# Additional attribution and requirements are provided in the NOTICE file.
@@ -1,47 +0,0 @@
# Secret Detection Workflow Dockerfile
FROM prefecthq/prefect:3-python3.11
# Install system dependencies
RUN apt-get update && apt-get install -y \
curl \
wget \
git \
ca-certificates \
gnupg \
&& rm -rf /var/lib/apt/lists/*
# Install TruffleHog (use direct binary download to avoid install script issues)
RUN curl -sSfL "https://github.com/trufflesecurity/trufflehog/releases/download/v3.63.2/trufflehog_3.63.2_linux_amd64.tar.gz" -o trufflehog.tar.gz \
&& tar -xzf trufflehog.tar.gz \
&& mv trufflehog /usr/local/bin/ \
&& rm trufflehog.tar.gz
# Install Gitleaks (use specific version to avoid API rate limiting)
RUN wget https://github.com/gitleaks/gitleaks/releases/download/v8.18.2/gitleaks_8.18.2_linux_x64.tar.gz \
&& tar -xzf gitleaks_8.18.2_linux_x64.tar.gz \
&& mv gitleaks /usr/local/bin/ \
&& rm gitleaks_8.18.2_linux_x64.tar.gz
# Verify installations
RUN trufflehog --version && gitleaks version
# Set working directory
WORKDIR /opt/prefect
# Create toolbox directory structure
RUN mkdir -p /opt/prefect/toolbox
# Set environment variables
ENV PYTHONPATH=/opt/prefect/toolbox:/opt/prefect/toolbox/workflows
ENV WORKFLOW_NAME=secret_detection_scan
# The toolbox code will be mounted at runtime from the backend container
# This includes:
# - /opt/prefect/toolbox/modules/base.py
# - /opt/prefect/toolbox/modules/secret_detection/ (TruffleHog, Gitleaks modules)
# - /opt/prefect/toolbox/modules/reporter/ (SARIF reporter)
# - /opt/prefect/toolbox/workflows/comprehensive/secret_detection_scan/
VOLUME /opt/prefect/toolbox
# Set working directory for execution
WORKDIR /opt/prefect
@@ -1,58 +0,0 @@
# Secret Detection Workflow Dockerfile - Self-Contained Version
# This version copies all required modules into the image for complete isolation
FROM prefecthq/prefect:3-python3.11
# Install system dependencies
RUN apt-get update && apt-get install -y \
curl \
wget \
git \
ca-certificates \
gnupg \
&& rm -rf /var/lib/apt/lists/*
# Install TruffleHog
RUN curl -sSfL https://raw.githubusercontent.com/trufflesecurity/trufflehog/main/scripts/install.sh | sh -s -- -b /usr/local/bin
# Install Gitleaks
RUN wget https://github.com/gitleaks/gitleaks/releases/latest/download/gitleaks_linux_x64.tar.gz \
&& tar -xzf gitleaks_linux_x64.tar.gz \
&& mv gitleaks /usr/local/bin/ \
&& rm gitleaks_linux_x64.tar.gz
# Verify installations
RUN trufflehog --version && gitleaks version
# Set working directory
WORKDIR /opt/prefect
# Create directory structure
RUN mkdir -p /opt/prefect/toolbox/modules/secret_detection \
/opt/prefect/toolbox/modules/reporter \
/opt/prefect/toolbox/workflows/comprehensive/secret_detection_scan
# Copy the base module and required modules
COPY toolbox/modules/base.py /opt/prefect/toolbox/modules/base.py
COPY toolbox/modules/__init__.py /opt/prefect/toolbox/modules/__init__.py
COPY toolbox/modules/secret_detection/ /opt/prefect/toolbox/modules/secret_detection/
COPY toolbox/modules/reporter/ /opt/prefect/toolbox/modules/reporter/
# Copy the workflow code
COPY toolbox/workflows/comprehensive/secret_detection_scan/ /opt/prefect/toolbox/workflows/comprehensive/secret_detection_scan/
# Copy toolbox init files
COPY toolbox/__init__.py /opt/prefect/toolbox/__init__.py
COPY toolbox/workflows/__init__.py /opt/prefect/toolbox/workflows/__init__.py
COPY toolbox/workflows/comprehensive/__init__.py /opt/prefect/toolbox/workflows/comprehensive/__init__.py
# Install Python dependencies for the modules
RUN pip install --no-cache-dir \
pydantic \
asyncio
# Set environment variables
ENV PYTHONPATH=/opt/prefect/toolbox:/opt/prefect/toolbox/workflows
ENV WORKFLOW_NAME=secret_detection_scan
# Set default command (can be overridden)
CMD ["python", "-m", "toolbox.workflows.comprehensive.secret_detection_scan.workflow"]
@@ -1,130 +0,0 @@
# Secret Detection Scan Workflow
This workflow performs comprehensive secret detection using multiple industry-standard tools:
- **TruffleHog**: Comprehensive secret detection with verification capabilities
- **Gitleaks**: Git-specific secret scanning and leak detection
## Features
- **Parallel Execution**: Runs TruffleHog and Gitleaks concurrently for faster results
- **Deduplication**: Automatically removes duplicate findings across tools
- **SARIF Output**: Generates standardized SARIF reports for integration with security tools
- **Configurable**: Supports extensive configuration for both tools
## Dependencies
### Required Modules
- `toolbox.modules.secret_detection.trufflehog`
- `toolbox.modules.secret_detection.gitleaks`
- `toolbox.modules.reporter` (SARIF reporter)
- `toolbox.modules.base` (Base module interface)
### External Tools
- TruffleHog v3.63.2+
- Gitleaks v8.18.0+
## Docker Deployment
This workflow provides two Docker deployment approaches:
### 1. Volume-Based Approach (Default: `Dockerfile`)
**Advantages:**
- Live code updates without rebuilding images
- Smaller image sizes
- Consistent module versions across workflows
- Faster development iteration
**How it works:**
- Docker image contains only external tools (TruffleHog, Gitleaks)
- Python modules are mounted at runtime from the backend container
- Backend manages code synchronization via shared volumes
### 2. Self-Contained Approach (`Dockerfile.self-contained`)
**Advantages:**
- Complete isolation and reproducibility
- No runtime dependencies on backend code
- Can run independently of FuzzForge platform
- Better for CI/CD integration
**How it works:**
- All required Python modules are copied into the Docker image
- Image is completely self-contained
- Larger image size but fully portable
## Configuration
### TruffleHog Configuration
```json
{
"trufflehog_config": {
"verify": true, // Verify discovered secrets
"concurrency": 10, // Number of concurrent workers
"max_depth": 10, // Maximum directory depth
"include_detectors": [], // Specific detectors to include
"exclude_detectors": [] // Specific detectors to exclude
}
}
```
### Gitleaks Configuration
```json
{
"gitleaks_config": {
"scan_mode": "detect", // "detect" or "protect"
"redact": true, // Redact secrets in output
"max_target_megabytes": 100, // Maximum file size (MB)
"no_git": false, // Scan without Git context
"config_file": "", // Custom Gitleaks config
"baseline_file": "" // Baseline file for known findings
}
}
```
## Usage Example
```bash
curl -X POST "http://localhost:8000/workflows/secret_detection_scan/submit" \
-H "Content-Type: application/json" \
-d '{
"target_path": "/path/to/scan",
"volume_mode": "ro",
"parameters": {
"trufflehog_config": {
"verify": true,
"concurrency": 15
},
"gitleaks_config": {
"scan_mode": "detect",
"max_target_megabytes": 200
}
}
}'
```
## Output Format
The workflow generates a SARIF report containing:
- All unique findings from both tools
- Severity levels mapped to standard scale
- File locations and line numbers
- Detailed descriptions and recommendations
- Tool-specific metadata
## Performance Considerations
- **TruffleHog**: CPU-intensive with verification enabled
- **Gitleaks**: Memory-intensive for large repositories
- **Recommended Resources**: 512Mi memory, 500m CPU
- **Typical Runtime**: 1-5 minutes for small repos, 10-30 minutes for large ones
## Security Notes
- Secrets are redacted in output by default
- Verified secrets are marked with higher severity
- Both tools support custom rules and exclusions
- Consider using baseline files for known false positives
@@ -1,17 +0,0 @@
"""
Secret Detection Scan Workflow
This package contains the comprehensive secret detection workflow that combines
multiple secret detection tools for thorough analysis.
"""
# Copyright (c) 2025 FuzzingLabs
#
# Licensed under the Business Source License 1.1 (BSL). See the LICENSE file
# at the root of this repository for details.
#
# After the Change Date (four years from publication), this version of the
# Licensed Work will be made available under the Apache License, Version 2.0.
# See the LICENSE-APACHE file or http://www.apache.org/licenses/LICENSE-2.0
#
# Additional attribution and requirements are provided in the NOTICE file.
@@ -1,113 +0,0 @@
name: secret_detection_scan
version: "2.0.0"
description: "Comprehensive secret detection using TruffleHog and Gitleaks"
author: "FuzzForge Team"
category: "comprehensive"
tags:
- "secrets"
- "credentials"
- "detection"
- "trufflehog"
- "gitleaks"
- "comprehensive"
supported_volume_modes:
- "ro"
- "rw"
default_volume_mode: "ro"
default_target_path: "/workspace"
requirements:
tools:
- "trufflehog"
- "gitleaks"
resources:
memory: "512Mi"
cpu: "500m"
timeout: 1800
has_docker: true
default_parameters:
target_path: "/workspace"
volume_mode: "ro"
trufflehog_config: {}
gitleaks_config: {}
reporter_config: {}
parameters:
type: object
properties:
target_path:
type: string
default: "/workspace"
description: "Path to analyze"
volume_mode:
type: string
enum: ["ro", "rw"]
default: "ro"
description: "Volume mount mode"
trufflehog_config:
type: object
description: "TruffleHog configuration"
properties:
verify:
type: boolean
description: "Verify discovered secrets"
concurrency:
type: integer
description: "Number of concurrent workers"
max_depth:
type: integer
description: "Maximum directory depth to scan"
include_detectors:
type: array
items:
type: string
description: "Specific detectors to include"
exclude_detectors:
type: array
items:
type: string
description: "Specific detectors to exclude"
gitleaks_config:
type: object
description: "Gitleaks configuration"
properties:
scan_mode:
type: string
enum: ["detect", "protect"]
description: "Scan mode"
redact:
type: boolean
description: "Redact secrets in output"
max_target_megabytes:
type: integer
description: "Maximum file size to scan (MB)"
no_git:
type: boolean
description: "Scan files without Git context"
config_file:
type: string
description: "Path to custom configuration file"
baseline_file:
type: string
description: "Path to baseline file"
reporter_config:
type: object
description: "SARIF reporter configuration"
properties:
output_file:
type: string
description: "Output SARIF file name"
include_code_flows:
type: boolean
description: "Include code flow information"
output_schema:
type: object
properties:
sarif:
type: object
description: "SARIF-formatted security findings"
@@ -1,290 +0,0 @@
"""
Secret Detection Scan Workflow
This workflow performs comprehensive secret detection using multiple tools:
- TruffleHog: Comprehensive secret detection with verification
- Gitleaks: Git-specific secret scanning
"""
# Copyright (c) 2025 FuzzingLabs
#
# Licensed under the Business Source License 1.1 (BSL). See the LICENSE file
# at the root of this repository for details.
#
# After the Change Date (four years from publication), this version of the
# Licensed Work will be made available under the Apache License, Version 2.0.
# See the LICENSE-APACHE file or http://www.apache.org/licenses/LICENSE-2.0
#
# Additional attribution and requirements are provided in the NOTICE file.
import sys
import logging
from pathlib import Path
from typing import Dict, Any, List, Optional
from prefect import flow, task
from prefect.artifacts import create_markdown_artifact, create_table_artifact
import asyncio
import json
# Add modules to path
sys.path.insert(0, '/app')
# Import modules
from toolbox.modules.secret_detection.trufflehog import TruffleHogModule
from toolbox.modules.secret_detection.gitleaks import GitleaksModule
from toolbox.modules.reporter import SARIFReporter
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
@task(name="trufflehog_scan")
async def run_trufflehog_task(workspace: Path, config: Dict[str, Any]) -> Dict[str, Any]:
"""
Task to run TruffleHog secret detection.
Args:
workspace: Path to the workspace
config: TruffleHog configuration
Returns:
TruffleHog results
"""
logger.info("Running TruffleHog secret detection")
module = TruffleHogModule()
result = await module.execute(config, workspace)
logger.info(f"TruffleHog completed: {result.summary.get('total_secrets', 0)} secrets found")
return result.dict()
@task(name="gitleaks_scan")
async def run_gitleaks_task(workspace: Path, config: Dict[str, Any]) -> Dict[str, Any]:
"""
Task to run Gitleaks secret detection.
Args:
workspace: Path to the workspace
config: Gitleaks configuration
Returns:
Gitleaks results
"""
logger.info("Running Gitleaks secret detection")
module = GitleaksModule()
result = await module.execute(config, workspace)
logger.info(f"Gitleaks completed: {result.summary.get('total_leaks', 0)} leaks found")
return result.dict()
@task(name="aggregate_findings")
async def aggregate_findings_task(
trufflehog_results: Dict[str, Any],
gitleaks_results: Dict[str, Any],
config: Dict[str, Any],
workspace: Path
) -> Dict[str, Any]:
"""
Task to aggregate findings from all secret detection tools.
Args:
trufflehog_results: Results from TruffleHog
gitleaks_results: Results from Gitleaks
config: Reporter configuration
workspace: Path to workspace
Returns:
Aggregated SARIF report
"""
logger.info("Aggregating secret detection findings")
# Combine all findings
all_findings = []
# Add TruffleHog findings
trufflehog_findings = trufflehog_results.get("findings", [])
all_findings.extend(trufflehog_findings)
# Add Gitleaks findings
gitleaks_findings = gitleaks_results.get("findings", [])
all_findings.extend(gitleaks_findings)
# Deduplicate findings based on file path and line number
unique_findings = []
seen_signatures = set()
for finding in all_findings:
# Create signature for deduplication
signature = (
finding.get("file_path", ""),
finding.get("line_start", 0),
finding.get("title", "").lower()[:50] # First 50 chars of title
)
if signature not in seen_signatures:
seen_signatures.add(signature)
unique_findings.append(finding)
else:
logger.debug(f"Deduplicated finding: {signature}")
logger.info(f"Aggregated {len(unique_findings)} unique findings from {len(all_findings)} total")
# Generate SARIF report
reporter = SARIFReporter()
reporter_config = {
**config,
"findings": unique_findings,
"tool_name": "FuzzForge Secret Detection",
"tool_version": "1.0.0",
"tool_description": "Comprehensive secret detection using TruffleHog and Gitleaks"
}
result = await reporter.execute(reporter_config, workspace)
return result.dict().get("sarif", {})
@flow(name="secret_detection_scan", log_prints=True)
async def main_flow(
target_path: str = "/workspace",
volume_mode: str = "ro",
trufflehog_config: Optional[Dict[str, Any]] = None,
gitleaks_config: Optional[Dict[str, Any]] = None,
reporter_config: Optional[Dict[str, Any]] = None
) -> Dict[str, Any]:
"""
Main secret detection workflow.
This workflow:
1. Runs TruffleHog for comprehensive secret detection
2. Runs Gitleaks for Git-specific secret detection
3. Aggregates and deduplicates findings
4. Generates a unified SARIF report
Args:
target_path: Path to the mounted workspace (default: /workspace)
volume_mode: Volume mount mode (ro/rw)
trufflehog_config: Configuration for TruffleHog
gitleaks_config: Configuration for Gitleaks
reporter_config: Configuration for SARIF reporter
Returns:
SARIF-formatted findings report
"""
logger.info("Starting comprehensive secret detection workflow")
logger.info(f"Workspace: {target_path}, Mode: {volume_mode}")
# Set workspace path
workspace = Path(target_path)
if not workspace.exists():
logger.error(f"Workspace does not exist: {workspace}")
return {
"error": f"Workspace not found: {workspace}",
"sarif": None
}
# Default configurations - merge with provided configs to ensure defaults are always applied
default_trufflehog_config = {
"verify": False,
"concurrency": 10,
"max_depth": 10,
"no_git": True # Add no_git for filesystem scanning
}
trufflehog_config = {**default_trufflehog_config, **(trufflehog_config or {})}
default_gitleaks_config = {
"scan_mode": "detect",
"redact": True,
"max_target_megabytes": 100,
"no_git": True # Critical for non-git directories
}
gitleaks_config = {**default_gitleaks_config, **(gitleaks_config or {})}
default_reporter_config = {
"include_code_flows": False
}
reporter_config = {**default_reporter_config, **(reporter_config or {})}
try:
# Run secret detection tools in parallel
logger.info("Phase 1: Running secret detection tools")
# Create tasks for parallel execution
trufflehog_task_result = run_trufflehog_task(workspace, trufflehog_config)
gitleaks_task_result = run_gitleaks_task(workspace, gitleaks_config)
# Wait for both to complete
trufflehog_results, gitleaks_results = await asyncio.gather(
trufflehog_task_result,
gitleaks_task_result,
return_exceptions=True
)
# Handle any exceptions
if isinstance(trufflehog_results, Exception):
logger.error(f"TruffleHog failed: {trufflehog_results}")
trufflehog_results = {"findings": [], "status": "failed"}
if isinstance(gitleaks_results, Exception):
logger.error(f"Gitleaks failed: {gitleaks_results}")
gitleaks_results = {"findings": [], "status": "failed"}
# Aggregate findings
logger.info("Phase 2: Aggregating findings")
sarif_report = await aggregate_findings_task(
trufflehog_results,
gitleaks_results,
reporter_config,
workspace
)
# Log summary
if sarif_report and "runs" in sarif_report:
results_count = len(sarif_report["runs"][0].get("results", []))
logger.info(f"Workflow completed successfully with {results_count} unique secret findings")
# Log tool-specific stats
trufflehog_count = len(trufflehog_results.get("findings", []))
gitleaks_count = len(gitleaks_results.get("findings", []))
logger.info(f"Tool results - TruffleHog: {trufflehog_count}, Gitleaks: {gitleaks_count}")
else:
logger.info("Workflow completed successfully with no findings")
return sarif_report
except Exception as e:
logger.error(f"Secret detection workflow failed: {e}")
# Return error in SARIF format
return {
"$schema": "https://raw.githubusercontent.com/oasis-tcs/sarif-spec/master/Schemata/sarif-schema-2.1.0.json",
"version": "2.1.0",
"runs": [
{
"tool": {
"driver": {
"name": "FuzzForge Secret Detection",
"version": "1.0.0"
}
},
"results": [],
"invocations": [
{
"executionSuccessful": False,
"exitCode": 1,
"exitCodeDescription": str(e)
}
]
}
]
}
if __name__ == "__main__":
# For local testing
import asyncio
asyncio.run(main_flow(
target_path="/tmp/test",
trufflehog_config={"verify": True, "max_depth": 5},
gitleaks_config={"scan_mode": "detect"}
))
@@ -0,0 +1,21 @@
name: rust_test
version: 1.0.0
description: "Simple test workflow for Rust vertical worker"
vertical: rust
author: "FuzzForge Team"
tags:
- test
- rust
- example
dependencies:
python: []
parameters:
- name: target_id
type: string
required: true
description: "UUID of the uploaded target in MinIO"
- name: test_message
type: string
required: false
default: "Hello from Rust workflow!"
description: "Test message to include in results"
@@ -0,0 +1,159 @@
"""
Rust Test Workflow
Simple test workflow to verify:
1. Temporal worker discovery works
2. MinIO storage integration works
3. Activities execute correctly
4. Results are properly returned
This workflow:
- Downloads a target from MinIO
- Performs a simple analysis (file inspection)
- Returns results
- Cleans up cache
"""
from datetime import timedelta
from typing import Optional
from temporalio import workflow
from temporalio.common import RetryPolicy
# Import activity interfaces (will be executed by worker)
with workflow.unsafe.imports_passed_through():
import logging
logger = logging.getLogger(__name__)
@workflow.defn
class RustTestWorkflow:
"""
Simple test workflow for Rust vertical.
This demonstrates the basic workflow pattern:
1. Download target from MinIO
2. Execute activities
3. Return results
4. Cleanup
"""
@workflow.run
async def run(
self,
target_id: str,
test_message: Optional[str] = "Hello from Rust workflow!"
) -> dict:
"""
Main workflow execution.
Args:
target_id: UUID of the uploaded target in MinIO
test_message: Optional test message to include in results
Returns:
Dictionary containing workflow results
"""
workflow_id = workflow.info().workflow_id
workflow.logger.info(
f"Starting RustTestWorkflow "
f"(workflow_id={workflow_id}, target_id={target_id})"
)
results = {
"workflow_id": workflow_id,
"target_id": target_id,
"message": test_message,
"steps": []
}
try:
# Step 1: Download target from MinIO
workflow.logger.info("Step 1: Downloading target from MinIO")
target_path = await workflow.execute_activity(
"get_target",
target_id,
start_to_close_timeout=timedelta(minutes=5),
retry_policy=RetryPolicy(
initial_interval=timedelta(seconds=1),
maximum_interval=timedelta(seconds=30),
maximum_attempts=3
)
)
results["steps"].append({
"step": "download_target",
"status": "success",
"target_path": target_path
})
workflow.logger.info(f"✓ Target downloaded to: {target_path}")
# Step 2: Perform simple analysis (inline for testing)
workflow.logger.info("Step 2: Performing simple analysis")
# In a real workflow, this would be an activity that uses
# AFL++, cargo-fuzz, or other Rust tools
analysis_result = {
"file_path": target_path,
"analysis_type": "test",
"findings": [
{
"type": "info",
"message": "Test workflow executed successfully",
"test_message": test_message
}
]
}
results["steps"].append({
"step": "analysis",
"status": "success",
"analysis": analysis_result
})
workflow.logger.info("✓ Analysis completed")
# Step 3: Upload results to MinIO (optional)
workflow.logger.info("Step 3: Uploading results")
try:
results_url = await workflow.execute_activity(
"upload_results",
args=[workflow_id, results, "json"],
start_to_close_timeout=timedelta(minutes=2)
)
results["results_url"] = results_url
workflow.logger.info(f"✓ Results uploaded to: {results_url}")
except Exception as e:
workflow.logger.warning(f"Failed to upload results: {e}")
# Don't fail workflow if upload fails
results["results_url"] = None
# Step 4: Cleanup cache
workflow.logger.info("Step 4: Cleaning up cache")
try:
await workflow.execute_activity(
"cleanup_cache",
target_path,
start_to_close_timeout=timedelta(minutes=1)
)
workflow.logger.info("✓ Cache cleaned up")
except Exception as e:
workflow.logger.warning(f"Cache cleanup failed: {e}")
# Don't fail workflow if cleanup fails
# Mark workflow as successful
results["status"] = "success"
workflow.logger.info(f"✓ Workflow completed successfully: {workflow_id}")
return results
except Exception as e:
workflow.logger.error(f"Workflow failed: {e}")
results["status"] = "error"
results["error"] = str(e)
results["steps"].append({
"step": "error",
"status": "failed",
"error": str(e)
})
raise
@@ -1,30 +0,0 @@
FROM prefecthq/prefect:3-python3.11
WORKDIR /app
# Create toolbox directory structure to match expected import paths
RUN mkdir -p /app/toolbox/workflows /app/toolbox/modules
# Copy base module infrastructure
COPY modules/__init__.py /app/toolbox/modules/
COPY modules/base.py /app/toolbox/modules/
# Copy only required modules (manual selection)
COPY modules/scanner /app/toolbox/modules/scanner
COPY modules/analyzer /app/toolbox/modules/analyzer
COPY modules/reporter /app/toolbox/modules/reporter
# Copy this workflow
COPY workflows/security_assessment /app/toolbox/workflows/security_assessment
# Install workflow-specific requirements if they exist
RUN if [ -f /app/toolbox/workflows/security_assessment/requirements.txt ]; then pip install --no-cache-dir -r /app/toolbox/workflows/security_assessment/requirements.txt; fi
# Install common requirements
RUN pip install --no-cache-dir pyyaml
# Set Python path
ENV PYTHONPATH=/app:$PYTHONPATH
# Create workspace directory
RUN mkdir -p /workspace
@@ -0,0 +1,150 @@
"""
Security Assessment Workflow Activities
Activities specific to the security assessment workflow:
- scan_files_activity: Scan files in the workspace
- analyze_security_activity: Analyze security vulnerabilities
- generate_sarif_report_activity: Generate SARIF report from findings
"""
import logging
import sys
from pathlib import Path
from temporalio import activity
# Configure logging
logger = logging.getLogger(__name__)
# Add toolbox to path for module imports
sys.path.insert(0, '/app/toolbox')
@activity.defn(name="scan_files")
async def scan_files_activity(workspace_path: str, config: dict) -> dict:
"""
Scan files in the workspace.
Args:
workspace_path: Path to the workspace directory
config: Scanner configuration
Returns:
Scanner results dictionary
"""
logger.info(f"Activity: scan_files (workspace={workspace_path})")
try:
from modules.scanner import FileScanner
workspace = Path(workspace_path)
if not workspace.exists():
raise FileNotFoundError(f"Workspace not found: {workspace_path}")
scanner = FileScanner()
result = await scanner.execute(config, workspace)
logger.info(
f"✓ File scanning completed: "
f"{result.summary.get('total_files', 0)} files scanned"
)
return result.dict()
except Exception as e:
logger.error(f"File scanning failed: {e}", exc_info=True)
raise
@activity.defn(name="analyze_security")
async def analyze_security_activity(workspace_path: str, config: dict) -> dict:
"""
Analyze security vulnerabilities in the workspace.
Args:
workspace_path: Path to the workspace directory
config: Analyzer configuration
Returns:
Analysis results dictionary
"""
logger.info(f"Activity: analyze_security (workspace={workspace_path})")
try:
from modules.analyzer import SecurityAnalyzer
workspace = Path(workspace_path)
if not workspace.exists():
raise FileNotFoundError(f"Workspace not found: {workspace_path}")
analyzer = SecurityAnalyzer()
result = await analyzer.execute(config, workspace)
logger.info(
f"✓ Security analysis completed: "
f"{result.summary.get('total_findings', 0)} findings"
)
return result.dict()
except Exception as e:
logger.error(f"Security analysis failed: {e}", exc_info=True)
raise
@activity.defn(name="generate_sarif_report")
async def generate_sarif_report_activity(
scan_results: dict,
analysis_results: dict,
config: dict,
workspace_path: str
) -> dict:
"""
Generate SARIF report from scan and analysis results.
Args:
scan_results: Results from file scanner
analysis_results: Results from security analyzer
config: Reporter configuration
workspace_path: Path to the workspace
Returns:
SARIF report dictionary
"""
logger.info(f"Activity: generate_sarif_report")
try:
from modules.reporter import SARIFReporter
workspace = Path(workspace_path)
# Combine findings from all modules
all_findings = []
# Add scanner findings (only sensitive files, not all files)
scanner_findings = scan_results.get("findings", [])
sensitive_findings = [f for f in scanner_findings if f.get("severity") != "info"]
all_findings.extend(sensitive_findings)
# Add analyzer findings
analyzer_findings = analysis_results.get("findings", [])
all_findings.extend(analyzer_findings)
# Prepare reporter config
reporter_config = {
**config,
"findings": all_findings,
"tool_name": "FuzzForge Security Assessment",
"tool_version": "1.0.0"
}
reporter = SARIFReporter()
result = await reporter.execute(reporter_config, workspace)
# Extract SARIF from result
sarif = result.dict().get("sarif", {})
logger.info(f"✓ SARIF report generated with {len(all_findings)} findings")
return sarif
except Exception as e:
logger.error(f"SARIF report generation failed: {e}", exc_info=True)
raise
@@ -1,5 +1,6 @@
name: security_assessment
version: "2.0.0"
vertical: rust
description: "Comprehensive security assessment workflow that scans files, analyzes code for vulnerabilities, and generates SARIF reports"
author: "FuzzForge Team"
category: "comprehensive"
@@ -1,4 +0,0 @@
# Requirements for security assessment workflow
pydantic>=2.0.0
pyyaml>=6.0
aiofiles>=23.0.0
@@ -1,5 +1,8 @@
"""
Security Assessment Workflow - Comprehensive security analysis using multiple modules
Security Assessment Workflow - Temporal Version
Comprehensive security analysis using multiple modules.
Converted from Prefect to Temporal architecture.
"""
# Copyright (c) 2025 FuzzingLabs
@@ -13,240 +16,217 @@ Security Assessment Workflow - Comprehensive security analysis using multiple mo
#
# Additional attribution and requirements are provided in the NOTICE file.
import sys
import logging
from datetime import timedelta
from pathlib import Path
from typing import Dict, Any, Optional
from prefect import flow, task
import json
# Add modules to path
sys.path.insert(0, '/app')
from temporalio import workflow
from temporalio.common import RetryPolicy
# Import modules
from toolbox.modules.scanner import FileScanner
from toolbox.modules.analyzer import SecurityAnalyzer
from toolbox.modules.reporter import SARIFReporter
# Import activity interfaces (will be executed by worker)
with workflow.unsafe.imports_passed_through():
import logging
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
@task(name="file_scanning")
async def scan_files_task(workspace: Path, config: Dict[str, Any]) -> Dict[str, Any]:
@workflow.defn
class SecurityAssessmentWorkflow:
"""
Task to scan files in the workspace.
Args:
workspace: Path to the workspace
config: Scanner configuration
Returns:
Scanner results
"""
logger.info(f"Starting file scanning in {workspace}")
scanner = FileScanner()
result = await scanner.execute(config, workspace)
logger.info(f"File scanning completed: {result.summary.get('total_files', 0)} files found")
return result.dict()
@task(name="security_analysis")
async def analyze_security_task(workspace: Path, config: Dict[str, Any]) -> Dict[str, Any]:
"""
Task to analyze security vulnerabilities.
Args:
workspace: Path to the workspace
config: Analyzer configuration
Returns:
Analysis results
"""
logger.info("Starting security analysis")
analyzer = SecurityAnalyzer()
result = await analyzer.execute(config, workspace)
logger.info(
f"Security analysis completed: {result.summary.get('total_findings', 0)} findings"
)
return result.dict()
@task(name="report_generation")
async def generate_report_task(
scan_results: Dict[str, Any],
analysis_results: Dict[str, Any],
config: Dict[str, Any],
workspace: Path
) -> Dict[str, Any]:
"""
Task to generate SARIF report from all findings.
Args:
scan_results: Results from scanner
analysis_results: Results from analyzer
config: Reporter configuration
workspace: Path to the workspace
Returns:
SARIF report
"""
logger.info("Generating SARIF report")
reporter = SARIFReporter()
# Combine findings from all modules
all_findings = []
# Add scanner findings (only sensitive files, not all files)
scanner_findings = scan_results.get("findings", [])
sensitive_findings = [f for f in scanner_findings if f.get("severity") != "info"]
all_findings.extend(sensitive_findings)
# Add analyzer findings
analyzer_findings = analysis_results.get("findings", [])
all_findings.extend(analyzer_findings)
# Prepare reporter config
reporter_config = {
**config,
"findings": all_findings,
"tool_name": "FuzzForge Security Assessment",
"tool_version": "1.0.0"
}
result = await reporter.execute(reporter_config, workspace)
# Extract SARIF from result
sarif = result.dict().get("sarif", {})
logger.info(f"Report generated with {len(all_findings)} total findings")
return sarif
@flow(name="security_assessment", log_prints=True)
async def main_flow(
target_path: str = "/workspace",
volume_mode: str = "ro",
scanner_config: Optional[Dict[str, Any]] = None,
analyzer_config: Optional[Dict[str, Any]] = None,
reporter_config: Optional[Dict[str, Any]] = None
) -> Dict[str, Any]:
"""
Main security assessment workflow.
Comprehensive security assessment workflow.
This workflow:
1. Scans files in the workspace
2. Analyzes code for security vulnerabilities
3. Generates a SARIF report with all findings
Args:
target_path: Path to the mounted workspace (default: /workspace)
volume_mode: Volume mount mode (ro/rw)
scanner_config: Configuration for file scanner
analyzer_config: Configuration for security analyzer
reporter_config: Configuration for SARIF reporter
Returns:
SARIF-formatted findings report
1. Downloads target from MinIO
2. Scans files in the workspace
3. Analyzes code for security vulnerabilities
4. Generates a SARIF report with all findings
5. Uploads results to MinIO
6. Cleans up cache
"""
logger.info(f"Starting security assessment workflow")
logger.info(f"Workspace: {target_path}, Mode: {volume_mode}")
# Set workspace path
workspace = Path(target_path)
@workflow.run
async def run(
self,
target_id: str,
scanner_config: Optional[Dict[str, Any]] = None,
analyzer_config: Optional[Dict[str, Any]] = None,
reporter_config: Optional[Dict[str, Any]] = None
) -> Dict[str, Any]:
"""
Main workflow execution.
if not workspace.exists():
logger.error(f"Workspace does not exist: {workspace}")
return {
"error": f"Workspace not found: {workspace}",
"sarif": None
}
Args:
target_id: UUID of the uploaded target in MinIO
scanner_config: Configuration for file scanner
analyzer_config: Configuration for security analyzer
reporter_config: Configuration for SARIF reporter
# Default configurations
if not scanner_config:
scanner_config = {
"patterns": ["*"],
"check_sensitive": True,
"calculate_hashes": False,
"max_file_size": 10485760 # 10MB
}
Returns:
Dictionary containing SARIF report and summary
"""
workflow_id = workflow.info().workflow_id
if not analyzer_config:
analyzer_config = {
"file_extensions": [".py", ".js", ".java", ".php", ".rb", ".go"],
"check_secrets": True,
"check_sql": True,
"check_dangerous_functions": True
}
if not reporter_config:
reporter_config = {
"include_code_flows": False
}
try:
# Execute workflow tasks
logger.info("Phase 1: File scanning")
scan_results = await scan_files_task(workspace, scanner_config)
logger.info("Phase 2: Security analysis")
analysis_results = await analyze_security_task(workspace, analyzer_config)
logger.info("Phase 3: Report generation")
sarif_report = await generate_report_task(
scan_results,
analysis_results,
reporter_config,
workspace
workflow.logger.info(
f"Starting SecurityAssessmentWorkflow "
f"(workflow_id={workflow_id}, target_id={target_id})"
)
# Log summary
if sarif_report and "runs" in sarif_report:
results_count = len(sarif_report["runs"][0].get("results", []))
logger.info(f"Workflow completed successfully with {results_count} findings")
else:
logger.info("Workflow completed successfully")
# Default configurations
if not scanner_config:
scanner_config = {
"patterns": ["*"],
"check_sensitive": True,
"calculate_hashes": False,
"max_file_size": 10485760 # 10MB
}
return sarif_report
if not analyzer_config:
analyzer_config = {
"file_extensions": [".py", ".js", ".java", ".php", ".rb", ".go"],
"check_secrets": True,
"check_sql": True,
"check_dangerous_functions": True
}
except Exception as e:
logger.error(f"Workflow failed: {e}")
# Return error in SARIF format
return {
"$schema": "https://raw.githubusercontent.com/oasis-tcs/sarif-spec/master/Schemata/sarif-schema-2.1.0.json",
"version": "2.1.0",
"runs": [
{
"tool": {
"driver": {
"name": "FuzzForge Security Assessment",
"version": "1.0.0"
}
},
"results": [],
"invocations": [
{
"executionSuccessful": False,
"exitCode": 1,
"exitCodeDescription": str(e)
}
]
}
]
if not reporter_config:
reporter_config = {
"include_code_flows": False
}
results = {
"workflow_id": workflow_id,
"target_id": target_id,
"status": "running",
"steps": []
}
try:
# Step 1: Download target from MinIO
workflow.logger.info("Step 1: Downloading target from MinIO")
target_path = await workflow.execute_activity(
"get_target",
target_id,
start_to_close_timeout=timedelta(minutes=5),
retry_policy=RetryPolicy(
initial_interval=timedelta(seconds=1),
maximum_interval=timedelta(seconds=30),
maximum_attempts=3
)
)
results["steps"].append({
"step": "download_target",
"status": "success",
"target_path": target_path
})
workflow.logger.info(f"✓ Target downloaded to: {target_path}")
if __name__ == "__main__":
# For local testing
import asyncio
# Step 2: File scanning
workflow.logger.info("Step 2: Scanning files")
scan_results = await workflow.execute_activity(
"scan_files",
args=[target_path, scanner_config],
start_to_close_timeout=timedelta(minutes=10),
retry_policy=RetryPolicy(
initial_interval=timedelta(seconds=2),
maximum_interval=timedelta(seconds=60),
maximum_attempts=2
)
)
results["steps"].append({
"step": "file_scanning",
"status": "success",
"files_scanned": scan_results.get("summary", {}).get("total_files", 0)
})
workflow.logger.info(
f"✓ File scanning completed: "
f"{scan_results.get('summary', {}).get('total_files', 0)} files"
)
asyncio.run(main_flow(
target_path="/tmp/test",
scanner_config={"patterns": ["*.py"]},
analyzer_config={"check_secrets": True}
))
# Step 3: Security analysis
workflow.logger.info("Step 3: Analyzing security vulnerabilities")
analysis_results = await workflow.execute_activity(
"analyze_security",
args=[target_path, analyzer_config],
start_to_close_timeout=timedelta(minutes=15),
retry_policy=RetryPolicy(
initial_interval=timedelta(seconds=2),
maximum_interval=timedelta(seconds=60),
maximum_attempts=2
)
)
results["steps"].append({
"step": "security_analysis",
"status": "success",
"findings": analysis_results.get("summary", {}).get("total_findings", 0)
})
workflow.logger.info(
f"✓ Security analysis completed: "
f"{analysis_results.get('summary', {}).get('total_findings', 0)} findings"
)
# Step 4: Generate SARIF report
workflow.logger.info("Step 4: Generating SARIF report")
sarif_report = await workflow.execute_activity(
"generate_sarif_report",
args=[scan_results, analysis_results, reporter_config, target_path],
start_to_close_timeout=timedelta(minutes=5)
)
results["steps"].append({
"step": "report_generation",
"status": "success"
})
# Count total findings in SARIF
total_findings = 0
if sarif_report and "runs" in sarif_report:
total_findings = len(sarif_report["runs"][0].get("results", []))
workflow.logger.info(f"✓ SARIF report generated with {total_findings} findings")
# Step 5: Upload results to MinIO
workflow.logger.info("Step 5: Uploading results")
try:
results_url = await workflow.execute_activity(
"upload_results",
args=[workflow_id, sarif_report, "sarif"],
start_to_close_timeout=timedelta(minutes=2)
)
results["results_url"] = results_url
workflow.logger.info(f"✓ Results uploaded to: {results_url}")
except Exception as e:
workflow.logger.warning(f"Failed to upload results: {e}")
results["results_url"] = None
# Step 6: Cleanup cache
workflow.logger.info("Step 6: Cleaning up cache")
try:
await workflow.execute_activity(
"cleanup_cache",
target_path,
start_to_close_timeout=timedelta(minutes=1)
)
workflow.logger.info("✓ Cache cleaned up")
except Exception as e:
workflow.logger.warning(f"Cache cleanup failed: {e}")
# Mark workflow as successful
results["status"] = "success"
results["sarif"] = sarif_report
results["summary"] = {
"total_findings": total_findings,
"files_scanned": scan_results.get("summary", {}).get("total_files", 0)
}
workflow.logger.info(f"✓ Workflow completed successfully: {workflow_id}")
return results
except Exception as e:
workflow.logger.error(f"Workflow failed: {e}")
results["status"] = "error"
results["error"] = str(e)
results["steps"].append({
"step": "error",
"status": "failed",
"error": str(e)
})
raise