Skip to content

Commit

Permalink
Removed unneeded attributes of normalizers' tasks
Browse files Browse the repository at this point in the history
  • Loading branch information
Souf149 committed Oct 31, 2024
1 parent aabfeeb commit 5899ed6
Show file tree
Hide file tree
Showing 3 changed files with 13 additions and 8 deletions.
3 changes: 1 addition & 2 deletions boefjes/boefjes/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -272,7 +272,7 @@ def get_runtime_manager(settings: Settings, queue: WorkerManager.Queue, log_leve
plugin_service = PluginService(create_plugin_storage(session), create_config_storage(session), local_repository)

item_handler: Handler
capabilities: list[str]
capabilities: list[str] | None = None
reachable_networks: list[str] | None = None

if queue is WorkerManager.Queue.BOEFJES:
Expand All @@ -283,7 +283,6 @@ def get_runtime_manager(settings: Settings, queue: WorkerManager.Queue, log_leve
item_handler = NormalizerHandler(
LocalNormalizerJobRunner(local_repository), bytes_api_client, settings.scan_profile_whitelist
)
capabilities = [WorkerManager.Queue.NORMALIZERS.value]

return SchedulerWorkerManager(
item_handler,
Expand Down
16 changes: 12 additions & 4 deletions boefjes/boefjes/clients/scheduler_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,9 @@ def push_item(self, p_item: Task) -> None:


class SchedulerAPIClient(SchedulerClientInterface):
def __init__(self, base_url: str, task_capabilities: list[str] = [], reachable_networks: list[str] | None = None):
def __init__(
self, base_url: str, task_capabilities: list[str] | None = None, reachable_networks: list[str] | None = None
):
self._session = Client(base_url=base_url, transport=HTTPTransport(retries=6))
self._task_capabilities = task_capabilities
self._reachable_networks = reachable_networks
Expand All @@ -85,15 +87,21 @@ def get_queues(self) -> list[Queue]:
return TypeAdapter(list[Queue]).validate_json(response.content)

def pop_item(self, queue_id: str) -> Task | None:
filters: list[Filter] = [
Filter(column="data", field="requirements", operator="<@", value=json.dumps(self._task_capabilities))
]
filters: list[Filter] = []

# Client should only pop tasks that lie on a network that the runner is capable of reaching (e.g. the internet)
if self._reachable_networks:
filters.append(
Filter(column="data", field="network", operator="<@", value=json.dumps(self._reachable_networks))
)

# Client should only pop tasks that have requirements that this runner is capable of (e.g. being able
# to handle ipv6 requests)
if self._task_capabilities:
filters.append(
Filter(column="data", field="requirements", operator="<@", value=json.dumps(self._task_capabilities))
)

response = self._session.post(
f"/queues/{queue_id}/pop", data=QueuePopRequest(filters=filters).model_dump_json()
)
Expand Down
2 changes: 0 additions & 2 deletions mula/scheduler/models/task.py
Original file line number Diff line number Diff line change
Expand Up @@ -99,8 +99,6 @@ class NormalizerTask(BaseModel):
normalizer: Normalizer
raw_data: RawData

requirements: list[str] | None = None

@property
def hash(self) -> str:
"""Make NormalizerTask hashable, so that we can de-duplicate it when
Expand Down

0 comments on commit 5899ed6

Please sign in to comment.