From c30b1707edb5e8d32a80bb3bc2a8a414a46c63ec Mon Sep 17 00:00:00 2001 From: Riccardo Schirone <562321+ret2libc@users.noreply.github.com> Date: Mon, 10 Mar 2025 14:11:29 +0100 Subject: [PATCH] orchestrator/scheduler: make competition api/token configurable (#182) --- env.dev.compose | 26 +++++++++++-------- .../orchestrator/api_client_factory.py | 8 +++--- .../orchestrator/scheduler/__cli__.py | 6 +++++ .../orchestrator/scheduler/config.py | 5 +++- .../orchestrator/scheduler/scheduler.py | 8 ++++-- 5 files changed, 36 insertions(+), 17 deletions(-) diff --git a/env.dev.compose b/env.dev.compose index 7ae5832e..b5c65595 100644 --- a/env.dev.compose +++ b/env.dev.compose @@ -20,6 +20,10 @@ BUTTERCUP_TASK_SERVER_API_TOKEN_HASH='$argon2id$v=19$m=65536,t=3,p=4$Dg1v6NPGTyX # Client Authentication: # CRS_API_KEY_ID=515cc8a0-3019-4c9f-8c1c-72d0b54ae561 # CRS_API_TOKEN=VGuAC8axfOnFXKBB7irpNDOKcDjOlnyB + +# Competition API Authentication +COMPETITION_API_KEY_ID=api_key_id +COMPETITION_API_KEY_TOKEN=api_key_token BUTTERCUP_SCHEDULER_SERVE__REDIS_URL=redis://redis:6379 BUTTERCUP_SCHEDULER_LOG_LEVEL=debug @@ -36,14 +40,14 @@ BUTTERCUP_PATCHER_SCRATCH_DIR=/crs_scratch # Keep these short for local development, there's just one worker for each # service and we want to retry as soon as possible usually. -BUILD_TASK_TIMEOUT_MS = 120000 -BUILD_OUTPUT_TASK_TIMEOUT_MS = 120000 -DOWNLOAD_TASK_TIMEOUT_MS = 120000 -READY_TASK_TIMEOUT_MS = 120000 -DELETE_TASK_TIMEOUT_MS = 120000 -CRASH_TASK_TIMEOUT_MS = 120000 -PATCH_TASK_TIMEOUT_MS = 120000 -CONFIRMED_VULNERABILITIES_TASK_TIMEOUT_MS = 120000 -INDEX_TASK_TIMEOUT_MS = 120000 -INDEX_OUTPUT_TASK_TIMEOUT_MS = 120000 -TRACED_VULNERABILITIES_TASK_TIMEOUT_MS = 120000 +BUILD_TASK_TIMEOUT_MS=120000 +BUILD_OUTPUT_TASK_TIMEOUT_MS=120000 +DOWNLOAD_TASK_TIMEOUT_MS=120000 +READY_TASK_TIMEOUT_MS=120000 +DELETE_TASK_TIMEOUT_MS=120000 +CRASH_TASK_TIMEOUT_MS=120000 +PATCH_TASK_TIMEOUT_MS=120000 +CONFIRMED_VULNERABILITIES_TASK_TIMEOUT_MS=120000 +INDEX_TASK_TIMEOUT_MS=120000 +INDEX_OUTPUT_TASK_TIMEOUT_MS=120000 +TRACED_VULNERABILITIES_TASK_TIMEOUT_MS=120000 diff --git a/orchestrator/src/buttercup/orchestrator/api_client_factory.py b/orchestrator/src/buttercup/orchestrator/api_client_factory.py index dffb5188..19a439c4 100644 --- a/orchestrator/src/buttercup/orchestrator/api_client_factory.py +++ b/orchestrator/src/buttercup/orchestrator/api_client_factory.py @@ -5,7 +5,9 @@ from buttercup.orchestrator.competition_api_client.api_client import ApiClient logger = logging.getLogger(__name__) -def create_api_client(competition_api_url: str) -> ApiClient: +def create_api_client( + competition_api_url: str, competition_api_username: str, competition_api_password: str +) -> ApiClient: """Initialize the competition API client with common configuration. Args: @@ -16,8 +18,8 @@ def create_api_client(competition_api_url: str) -> ApiClient: """ configuration = Configuration( host=competition_api_url, - username="api_key_id", # TODO: Make configurable - password="api_key_token", # TODO: Make configurable + username=competition_api_username, + password=competition_api_password, ) logger.info(f"Initializing API client with URL: {competition_api_url}") return ApiClient(configuration=configuration) diff --git a/orchestrator/src/buttercup/orchestrator/scheduler/__cli__.py b/orchestrator/src/buttercup/orchestrator/scheduler/__cli__.py index f178a65b..768ce613 100644 --- a/orchestrator/src/buttercup/orchestrator/scheduler/__cli__.py +++ b/orchestrator/src/buttercup/orchestrator/scheduler/__cli__.py @@ -9,6 +9,9 @@ from buttercup.common.logger import setup_package_logger from pydantic_settings import get_subcommand from redis import Redis from buttercup.orchestrator.scheduler.scheduler import Task, BuildOutput +import logging + +logger = logging.getLogger(__name__) def _prepare_ready_task(command: ProcessReadyTaskCommand) -> Task: @@ -32,6 +35,7 @@ def _prepare_build_output(command: ProcessBuildOutputCommand) -> BuildOutput: def main(): settings = Settings() setup_package_logger(__name__, settings.log_level) + logger.debug(f"Settings: {settings}") command = get_subcommand(settings) if isinstance(command, ServeCommand): redis = Redis.from_url(command.redis_url, decode_responses=False) @@ -41,6 +45,8 @@ def main(): redis, sleep_time=command.sleep_time, competition_api_url=command.competition_api_url, + competition_api_key_id=command.competition_api_key_id, + competition_api_key_token=command.competition_api_key_token, ) scheduler.serve() elif isinstance(command, ProcessReadyTaskCommand): diff --git a/orchestrator/src/buttercup/orchestrator/scheduler/config.py b/orchestrator/src/buttercup/orchestrator/scheduler/config.py index 99240de8..e765931d 100644 --- a/orchestrator/src/buttercup/orchestrator/scheduler/config.py +++ b/orchestrator/src/buttercup/orchestrator/scheduler/config.py @@ -5,10 +5,12 @@ from pydantic import Field from pathlib import Path -class ServeCommand(BaseModel): +class ServeCommand(BaseSettings): sleep_time: Annotated[float, Field(default=1.0, description="Sleep time between checks in seconds")] redis_url: Annotated[str, Field(default="redis://localhost:6379", description="Redis URL")] competition_api_url: Annotated[str, Field(default="http://competition-api:8080", description="Competition API URL")] + competition_api_key_id: Annotated[str, Field(default="api_key_id", description="Competition API username")] + competition_api_key_token: Annotated[str, Field(default="api_key_token", description="Competition API password")] class Config: nested_model_default_partial_update = True @@ -41,6 +43,7 @@ class Settings(BaseSettings): tasks_storage_dir: Annotated[Path, Field(default="/tmp/task_downloads", description="Directory for Tasks storage")] scratch_dir: Annotated[Path, Field(default="/tmp/crs_scratch", description="Directory for CRS scratch")] log_level: Annotated[str, Field(default="info", description="Log level")] + serve: CliSubCommand[ServeCommand] process_ready_task: CliSubCommand[ProcessReadyTaskCommand] process_build_output: CliSubCommand[ProcessBuildOutputCommand] diff --git a/orchestrator/src/buttercup/orchestrator/scheduler/scheduler.py b/orchestrator/src/buttercup/orchestrator/scheduler/scheduler.py index 49e4f04b..b2275299 100644 --- a/orchestrator/src/buttercup/orchestrator/scheduler/scheduler.py +++ b/orchestrator/src/buttercup/orchestrator/scheduler/scheduler.py @@ -22,7 +22,6 @@ from buttercup.orchestrator.api_client_factory import create_api_client from buttercup.common.utils import serve_loop import random - logger = logging.getLogger(__name__) @@ -33,6 +32,9 @@ class Scheduler: redis: Redis | None = None sleep_time: float = 1.0 competition_api_url: str = "http://competition-api:8080" + competition_api_key_id: str = "api_key_id" + competition_api_key_token: str = "api_key_token" + ready_queue: ReliableQueue | None = field(init=False, default=None) build_requests_queue: ReliableQueue | None = field(init=False, default=None) build_output_queue: ReliableQueue | None = field(init=False, default=None) @@ -47,7 +49,9 @@ class Scheduler: def __post_init__(self): if self.redis is not None: queue_factory = QueueFactory(self.redis) - api_client = create_api_client(self.competition_api_url) + api_client = create_api_client( + self.competition_api_url, self.competition_api_key_id, self.competition_api_key_token + ) # Input queues are non-blocking as we're already sleeping between iterations self.cancellation = Cancellation(redis=self.redis) self.vulnerabilities = Vulnerabilities(redis=self.redis, api_client=api_client)