diff --git a/web-agent/app/worker.py b/web-agent/app/worker.py index 2847324..e8578ab 100644 --- a/web-agent/app/worker.py +++ b/web-agent/app/worker.py @@ -1,10 +1,14 @@ #!/usr/bin/env python3 import argparse import base64 +import gzip import json +import logging import os import secrets import string +import tempfile +import time import uuid from collections import deque from logging.handlers import TimedRotatingFileHandler @@ -12,11 +16,7 @@ from typing import Optional, Tuple, Any, Dict import requests -import logging -import time -import gzip -from urllib.parse import unquote -import tempfile +from gevent.pool import Pool # Global variables __version__ = "1.1.3" @@ -152,6 +152,7 @@ def main() -> None: def process() -> None: headers: Dict[str, str] = _get_headers() thread_backoff_time: int = min_backoff_time + pool = Pool(5) while True: try: # Get the next task for the agent @@ -176,11 +177,8 @@ def process() -> None: logger.info("Received task: %s", task['taskId']) task["version"] = __version__ # Process the task - - result: Dict[str, Any] = process_task(task) - - # Update the task status - update_task(result) + pool.wait_available() # Wait if the pool is full + pool.spawn(process_task_async, task) # Submit the task when free elif get_task_response.status_code == 204: logger.info("No task available. Waiting...") time.sleep(5) @@ -200,6 +198,20 @@ def process() -> None: time.sleep(5) +def process_task_async(task: Dict[str, Any]) -> None: + url: str = task.get('url') + taskId: str = task.get('taskId') + method: str = task.get('method').upper() + + try: + result: Dict[str, Any] = process_task(task) + # Update the task status + update_task(result) + except Exception as e: + logger.info("Unexpected error while processing task id: %s, method: %s url: %s, error: %s", taskId, method, url, e) + time.sleep(5) + + def update_task(task: Optional[Dict[str, Any]], count: int = 0) -> None: if task is None: return @@ -528,4 +540,4 @@ def _clean_temp_output_files() -> None: _createFolder(armorcode_folder) # create parent level folder for logs anf files _createFolder(log_folder) # create folder to store log files _createFolder(output_file_folder) # create folder to store output files - main() + main() \ No newline at end of file diff --git a/web-agent/requirements.txt b/web-agent/requirements.txt index 2c24336..124b15d 100644 --- a/web-agent/requirements.txt +++ b/web-agent/requirements.txt @@ -1 +1,8 @@ requests==2.31.0 +certifi==2025.1.31 +gevent==24.11.1 +greenlet==3.1.1 +idna==3.10 +urllib3==2.3.0 +zope.event==5.0 +zope.interface==7.2 \ No newline at end of file