refactor: improve tasks code structure and efficiency

.gitignore

@@ -43,6 +43,9 @@ MANIFEST
 .vscode/settings.json
 .vscode/tasks-and-contexts.json
 .idea/
+.cursorrules
+.cursorignore
+.cursor/
 
 *.manifest
 *.spec

docker/celery/Dockerfile

@@ -85,10 +85,10 @@ RUN ARCH=$(dpkg --print-architecture) && \
     fi
 
 # Install python 3.10
-RUN cd /root && wget https://www.python.org/ftp/python/3.10.0/Python-3.10.0.tgz && \
-    tar -xvf Python-3.10.0.tgz && \
-    rm Python-3.10.0.tgz && \
-    cd Python-3.10.0 && \
+RUN cd /root && wget https://www.python.org/ftp/python/3.10.16/Python-3.10.16.tgz && \
+    tar -xvf Python-3.10.16.tgz && \
+    rm Python-3.10.16.tgz && \
+    cd Python-3.10.16 && \
     ./configure --enable-optimizations && \
     make -j4 && \
     make altinstall

docker/celery/entrypoint.sh

@@ -8,6 +8,9 @@ print_msg() {
 }
 
 RENGINE_FOLDER="/home/$USERNAME/rengine"
+MAX_CONCURRENCY=${MAX_CONCURRENCY:-20}
+MIN_CONCURRENCY=${MIN_CONCURRENCY:-5}
+CELERY_LOGLEVEL=${CELERY_LOGLEVEL:-info}
 
 print_msg "Generate Django migrations files"
 poetry run -C $RENGINE_FOLDER python3 manage.py makemigrations
@@ -24,31 +27,42 @@ poetry run -C $RENGINE_FOLDER python3 manage.py loaddata fixtures/default_keywor
 print_msg "Load default external tools"
 poetry run -C $RENGINE_FOLDER python3 manage.py loaddata fixtures/external_tools.yaml --app scanEngine.InstalledExternalTool
 
-if [ ! "$CELERY_LOGLEVEL" ]; then
-  export CELERY_LOGLEVEL='info'
-fi
-
-print_msg "Start celery workers"
-watchmedo auto-restart --recursive --pattern="*.py" --directory="$RENGINE_FOLDER" -- poetry run -C $RENGINE_FOLDER celery -A reNgine.tasks worker --loglevel=$CELERY_LOGLEVEL --autoscale=$MAX_CONCURRENCY,$MIN_CONCURRENCY -Q main_scan_queue &
-watchmedo auto-restart --recursive --pattern="*.py" --directory="$RENGINE_FOLDER" -- poetry run -C $RENGINE_FOLDER celery -A reNgine.tasks worker --pool=gevent --concurrency=30 --loglevel=$CELERY_LOGLEVEL -Q initiate_scan_queue -n initiate_scan_worker &
-watchmedo auto-restart --recursive --pattern="*.py" --directory="$RENGINE_FOLDER" -- poetry run -C $RENGINE_FOLDER celery -A reNgine.tasks worker --pool=gevent --concurrency=30 --loglevel=$CELERY_LOGLEVEL -Q subscan_queue -n subscan_worker &
-watchmedo auto-restart --recursive --pattern="*.py" --directory="$RENGINE_FOLDER" -- poetry run -C $RENGINE_FOLDER celery -A reNgine.tasks worker --pool=gevent --concurrency=20 --loglevel=$CELERY_LOGLEVEL -Q report_queue -n report_worker &
-watchmedo auto-restart --recursive --pattern="*.py" --directory="$RENGINE_FOLDER" -- poetry run -C $RENGINE_FOLDER celery -A reNgine.tasks worker --pool=gevent --concurrency=10 --loglevel=$CELERY_LOGLEVEL -Q send_notif_queue -n send_notif_worker &
-watchmedo auto-restart --recursive --pattern="*.py" --directory="$RENGINE_FOLDER" -- poetry run -C $RENGINE_FOLDER celery -A reNgine.tasks worker --pool=gevent --concurrency=10 --loglevel=$CELERY_LOGLEVEL -Q send_scan_notif_queue -n send_scan_notif_worker &
-watchmedo auto-restart --recursive --pattern="*.py" --directory="$RENGINE_FOLDER" -- poetry run -C $RENGINE_FOLDER celery -A reNgine.tasks worker --pool=gevent --concurrency=10 --loglevel=$CELERY_LOGLEVEL -Q send_task_notif_queue -n send_task_notif_worker &
-watchmedo auto-restart --recursive --pattern="*.py" --directory="$RENGINE_FOLDER" -- poetry run -C $RENGINE_FOLDER celery -A reNgine.tasks worker --pool=gevent --concurrency=5 --loglevel=$CELERY_LOGLEVEL -Q send_file_to_discord_queue -n send_file_to_discord_worker &
-watchmedo auto-restart --recursive --pattern="*.py" --directory="$RENGINE_FOLDER" -- poetry run -C $RENGINE_FOLDER celery -A reNgine.tasks worker --pool=gevent --concurrency=5 --loglevel=$CELERY_LOGLEVEL -Q send_hackerone_report_queue -n send_hackerone_report_worker &
-watchmedo auto-restart --recursive --pattern="*.py" --directory="$RENGINE_FOLDER" -- poetry run -C $RENGINE_FOLDER celery -A reNgine.tasks worker --pool=gevent --concurrency=10 --loglevel=$CELERY_LOGLEVEL -Q parse_nmap_results_queue -n parse_nmap_results_worker &
-watchmedo auto-restart --recursive --pattern="*.py" --directory="$RENGINE_FOLDER" -- poetry run -C $RENGINE_FOLDER celery -A reNgine.tasks worker --pool=gevent --concurrency=20 --loglevel=$CELERY_LOGLEVEL -Q geo_localize_queue -n geo_localize_worker &
-watchmedo auto-restart --recursive --pattern="*.py" --directory="$RENGINE_FOLDER" -- poetry run -C $RENGINE_FOLDER celery -A reNgine.tasks worker --pool=gevent --concurrency=10 --loglevel=$CELERY_LOGLEVEL -Q query_whois_queue -n query_whois_worker &
-watchmedo auto-restart --recursive --pattern="*.py" --directory="$RENGINE_FOLDER" -- poetry run -C $RENGINE_FOLDER celery -A reNgine.tasks worker --pool=gevent --concurrency=30 --loglevel=$CELERY_LOGLEVEL -Q remove_duplicate_endpoints_queue -n remove_duplicate_endpoints_worker &
-watchmedo auto-restart --recursive --pattern="*.py" --directory="$RENGINE_FOLDER" -- poetry run -C $RENGINE_FOLDER celery -A reNgine.tasks worker --pool=gevent --concurrency=50 --loglevel=$CELERY_LOGLEVEL -Q run_command_queue -n run_command_worker &
-watchmedo auto-restart --recursive --pattern="*.py" --directory="$RENGINE_FOLDER" -- poetry run -C $RENGINE_FOLDER celery -A reNgine.tasks worker --pool=gevent --concurrency=10 --loglevel=$CELERY_LOGLEVEL -Q query_reverse_whois_queue -n query_reverse_whois_worker &
-watchmedo auto-restart --recursive --pattern="*.py" --directory="$RENGINE_FOLDER" -- poetry run -C $RENGINE_FOLDER celery -A reNgine.tasks worker --pool=gevent --concurrency=10 --loglevel=$CELERY_LOGLEVEL -Q query_ip_history_queue -n query_ip_history_worker &
-watchmedo auto-restart --recursive --pattern="*.py" --directory="$RENGINE_FOLDER" -- poetry run -C $RENGINE_FOLDER celery -A reNgine.tasks worker --pool=gevent --concurrency=30 --loglevel=$CELERY_LOGLEVEL -Q gpt_queue -n gpt_worker &
-watchmedo auto-restart --recursive --pattern="*.py" --directory="$RENGINE_FOLDER" -- poetry run -C $RENGINE_FOLDER celery -A reNgine.tasks worker --pool=gevent --concurrency=10 --loglevel=$CELERY_LOGLEVEL -Q dorking_queue -n dorking_worker &
-watchmedo auto-restart --recursive --pattern="*.py" --directory="$RENGINE_FOLDER" -- poetry run -C $RENGINE_FOLDER celery -A reNgine.tasks worker --pool=gevent --concurrency=10 --loglevel=$CELERY_LOGLEVEL -Q osint_discovery_queue -n osint_discovery_worker &
-watchmedo auto-restart --recursive --pattern="*.py" --directory="$RENGINE_FOLDER" -- poetry run -C $RENGINE_FOLDER celery -A reNgine.tasks worker --pool=gevent --concurrency=10 --loglevel=$CELERY_LOGLEVEL -Q h8mail_queue -n h8mail_worker &
-watchmedo auto-restart --recursive --pattern="*.py" --directory="$RENGINE_FOLDER" -- poetry run -C $RENGINE_FOLDER celery -A reNgine.tasks worker --pool=gevent --concurrency=10 --loglevel=$CELERY_LOGLEVEL -Q theHarvester_queue -n theHarvester_worker
+worker_command() {
+    local queue=$1
+    local worker_name=$2
+
+    if [ "$CELERY_DEBUG" = "1" ]; then
+        echo "watchmedo auto-restart --recursive --pattern=\"*.py\" --directory=\"$RENGINE_FOLDER\" -- \
+            poetry run -C $RENGINE_FOLDER celery -A reNgine worker \
+            --pool=solo \
+            --loglevel=$CELERY_LOGLEVEL \
+            -Q $queue -n $worker_name"
+    else
+        echo "poetry run -C $RENGINE_FOLDER celery -A reNgine worker \
+            --pool=gevent \
+            --loglevel=$CELERY_LOGLEVEL \
+            --autoscale=$MAX_CONCURRENCY,$MIN_CONCURRENCY \
+            -Q $queue -n $worker_name"
+    fi
+}
+
+queues=(
+    "orchestrator_queue:orchestrator_worker"
+    "io_queue:io_worker"
+    "run_command_queue:run_command_worker"
+    "cpu_queue:cpu_worker"
+    "report_queue:report_worker"
+    "send_notif_queue:send_notif_worker"
+)
+
+commands=""
+for queue in "${queues[@]}"; do
+    IFS=':' read -r queue worker_name <<< "$queue"
+    commands+="$(worker_command "$queue" "$worker_name") &"$'\n'
+done
+
+eval "$commands"
+
+wait
 
 exec "$@"

docker/celery/pyproject.toml

@@ -11,6 +11,7 @@ aiodns = "3.0.0"
 argh = "0.26.2"
 beautifulsoup4 = "4.9.3"
 celery = "5.4.0"
+colorama = "0.4.6"
 debugpy = "1.8.5"
 discord-webhook = "1.3.0"
 django = "3.2.25"
@@ -27,6 +28,7 @@ djangorestframework = "3.14.0"
 djangorestframework-datatables = "0.7.2"
 dotted-dict = "1.1.3"
 drf-yasg = "1.21.5"
+flower = "2.0.1"
 gunicorn = "23.0.0"
 gevent = "24.2.1"
 humanize = "4.3.0"

docker/docker-compose.dev.yml

@@ -13,10 +13,14 @@ services:
       - CELERY_DEBUG=1
       - CELERY_REMOTE_DEBUG=0
       - CELERY_REMOTE_DEBUG_PORT=5679
+      - CELERY_TASK_ALWAYS_EAGER=0
+      - CELERY_TASK_EAGER_PROPAGATES=0
+      - COMMAND_EXECUTOR_DRY_RUN=1
     volumes:
       - ./celery/entrypoint-dev.sh:/entrypoint-dev.sh:ro
     ports:
       - "127.0.0.1:5679:5679"
+      - "127.0.0.1:5555:5555"
 
   celery-beat:
     entrypoint: /entrypoint-dev.sh

docker/docker-compose.yml

@@ -16,6 +16,7 @@ services:
       retries: 5
     networks:
       - rengine_network
+    tty: true
 
   redis:
     image: ghcr.io/security-tools-alliance/rengine-ng:rengine-redis-v${RENGINE_VERSION}
@@ -31,6 +32,7 @@ services:
       retries: 5
     networks:
       - rengine_network
+    tty: true
 
   celery:
     image: ghcr.io/security-tools-alliance/rengine-ng:rengine-celery-v${RENGINE_VERSION}
@@ -64,6 +66,7 @@ services:
         condition: service_healthy
     networks:
       - rengine_network
+    tty: true
 
   celery-beat:
     image: ghcr.io/security-tools-alliance/rengine-ng:rengine-celery-v${RENGINE_VERSION}
@@ -90,6 +93,7 @@ services:
       - wordlist:/home/rengine/wordlists
     networks:
       - rengine_network
+    tty: true
 
   web:
     image: ghcr.io/security-tools-alliance/rengine-ng:rengine-web-v${RENGINE_VERSION}
@@ -106,6 +110,7 @@ services:
     volumes:
       - ../web:/home/rengine/rengine:rw,z
       - ./web/entrypoint.sh:/entrypoint.sh:ro
+      - scan_results:/home/rengine/scan_results
       - tool_config:/home/rengine/.config
       - nuclei_templates:/home/rengine/nuclei-templates
       - gf_patterns:/home/rengine/.gf
@@ -126,6 +131,7 @@ services:
       rengine_network:
         aliases:
           - rengine
+    tty: true
 
   proxy:
     image: ghcr.io/security-tools-alliance/rengine-ng:rengine-proxy-v${RENGINE_VERSION}
@@ -156,6 +162,7 @@ services:
     ports:
       - 8082:8082/tcp
       - 443:443/tcp
+    tty: true
 
   ollama:
     image: ghcr.io/security-tools-alliance/rengine-ng:rengine-ollama-v${RENGINE_VERSION}

docker/web/pyproject.toml

@@ -11,6 +11,7 @@ aiodns = "3.0.0"
 argh = "0.26.2"
 beautifulsoup4 = "4.9.3"
 celery = "5.4.0"
+colorama = "0.4.6"
 debugpy = "1.8.5"
 discord-webhook = "1.3.0"
 django = "3.2.25"
@@ -27,6 +28,7 @@ djangorestframework = "3.14.0"
 djangorestframework-datatables = "0.7.2"
 dotted-dict = "1.1.3"
 drf-yasg = "1.21.5"
+flower = "2.0.1"
 gunicorn = "23.0.0"
 gevent = "24.2.1"
 humanize = "4.3.0"

web/api/serializers.py

@@ -3,14 +3,15 @@
 from django.contrib.humanize.templatetags.humanize import (naturalday, naturaltime)
 from django.db.models import F, JSONField, Value
 from recon_note.models import *
-from reNgine.common_func import *
 from rest_framework import serializers
 from scanEngine.models import *
 from startScan.models import *
 from targetApp.models import *
 from dashboard.models import *
 import yaml
 
+from reNgine.definitions import ENGINE_NAMES
+
 
 class SearchHistorySerializer(serializers.ModelSerializer):
 	class Meta:
@@ -359,6 +360,7 @@ def get_description(self, subdomain):
 		return subdomain.name
 
 	def get_title(self, subdomain):
+		from reNgine.utils.db import get_interesting_subdomains
 		if get_interesting_subdomains(subdomain.scan_history.id).filter(name=subdomain.name).exists():
 			return "Interesting"
 
@@ -724,6 +726,7 @@ def get_change(self, Subdomain):
 		return Subdomain.change
 
 	def get_is_interesting(self, Subdomain):
+		from reNgine.utils.db import get_interesting_subdomains
 		return (
 			get_interesting_subdomains(Subdomain.scan_history.id)
 			.filter(name=Subdomain.name)
@@ -905,6 +908,7 @@ class Meta:
 		fields = '__all__'
 
 	def get_is_interesting(self, subdomain):
+		from reNgine.utils.db import get_interesting_subdomains
 		scan_id = subdomain.scan_history.id if subdomain.scan_history else None
 		return (
 			get_interesting_subdomains(scan_id)

web/api/tests/test_vulnerability.py

@@ -88,7 +88,7 @@ def setUp(self):
         self.data_generator.create_endpoint()
         self.data_generator.create_vulnerability()
 
-    @patch("reNgine.tasks.gpt_vulnerability_description.apply_async")
+    @patch("reNgine.tasks.llm_vulnerability_description.apply_async")
     def test_get_vulnerability_report(self, mock_apply_async):
         """Test generating a vulnerability report."""
         mock_task = MagicMock()

web/api/views.py

@@ -1,4 +1,3 @@
-import logging
 import re
 import os.path
 from pathlib import Path
@@ -25,11 +24,11 @@
 
 from recon_note.models import TodoNote
 from reNgine.celery import app
-from reNgine.common_func import (
-	get_data_from_post_request,
-	get_interesting_endpoints,
-	get_interesting_subdomains,
+from reNgine.utils.logger import default_logger as logger
+from reNgine.utils.db import (
 	get_lookup_keywords,
+)
+from reNgine.utils.utils import (
 	safe_int_cast
 )
 from reNgine.definitions import (
@@ -39,27 +38,31 @@
 	DEFAULT_GPT_MODELS,
 	RUNNING_TASK,
 	SUCCESS_TASK,
-	ENGINE_NAMES
 )
 from reNgine.settings import (
 	RENGINE_CURRENT_VERSION,
 	RENGINE_TOOL_GITHUB_PATH
 )
-from reNgine.tasks import (
-	create_scan_activity,
-	gpt_vulnerability_description,
+from reNgine.tasks.scan import (
 	initiate_subscan,
+)
+from reNgine.tasks.command import run_command_line
+from reNgine.tasks.detect import (
+	run_cmseek,
+	run_wafw00f,
+)
+from reNgine.tasks.dns import (
 	query_ip_history,
 	query_reverse_whois,
 	query_whois,
-	run_cmseek,
-	run_command,
+)
+from reNgine.tasks.llm import llm_vulnerability_description
+from reNgine.tasks.notification import send_hackerone_report
+from reNgine.tasks.url import (
 	run_gf_list,
-	run_wafw00f,
-	send_hackerone_report
 )
 from reNgine.gpt import GPTAttackSuggestionGenerator
-from reNgine.utilities import is_safe_path, remove_lead_and_trail_slash
+from reNgine.utils.utils import is_safe_path, remove_lead_and_trail_slash
 from scanEngine.models import EngineType, InstalledExternalTool
 from startScan.models import (
 	Command,
@@ -80,6 +83,13 @@
 	Vulnerability,
 )
 from targetApp.models import Domain, Organization
+from reNgine.utils.command_executor import run_command
+from reNgine.utils.http import get_data_from_post_request
+from reNgine.utils.db import (
+	create_scan_activity,
+	get_interesting_endpoints,
+	get_interesting_subdomains
+)
 
 from .serializers import (
 	CommandSerializer,
@@ -103,7 +113,6 @@
 	OnlySubdomainNameSerializer,
 	OrganizationSerializer,
 	OrganizationTargetsSerializer,
-	PortSerializer,
  	ProjectSerializer,
 	ReconNoteSerializer,
 	ScanHistorySerializer,
@@ -117,8 +126,6 @@
 	VulnerabilitySerializer
 )
 
-logger = logging.getLogger(__name__)
-
 
 class OllamaManager(APIView):
     def get(self, request):
@@ -240,7 +247,7 @@ def get(self, request):
 				'status': False,
 				'error': 'Missing GET param Vulnerability `id`'
 			})
-		task = gpt_vulnerability_description.apply_async(args=(vulnerability_id,))
+		task = llm_vulnerability_description.apply_async(args=(vulnerability_id,))
 		response = task.wait()
 		return Response(response)
 
@@ -839,7 +846,7 @@ def post(self, request):
 					SUCCESS_TASK)
 				response['status'] = True
 			except Exception as e:
-				logging.error(e)
+				logger.error(e)
 				response = {'status': False, 'message': str(e)}
 		elif scan_id:
 			try:
@@ -855,7 +862,7 @@ def post(self, request):
 					SUCCESS_TASK)
 				response['status'] = True
 			except Exception as e:
-				logging.error(e)
+				logger.error(e)
 				response = {'status': False, 'message': str(e)}
 
 		logger.warning(f'Revoking tasks {task_ids}')
@@ -1015,8 +1022,8 @@ def get(self, request):
 		else:
 			return Response({'status': False, 'message': 'Cannot uninstall tool!'})
 
-		run_command(uninstall_command)
-		run_command.apply_async(args=(uninstall_command,))
+		run_command_line(uninstall_command)
+		run_command_line.apply_async(args=(uninstall_command,))
 
 		tool.delete()
 
@@ -1047,7 +1054,6 @@ def get(self, request):
 			update_command = 'cd ' + str(Path(RENGINE_TOOL_GITHUB_PATH) / tool_name) + ' && git pull && cd -'
 
 		run_command(update_command)
-		run_command.apply_async(args=(update_command,))
 		return Response({'status': True, 'message': tool.name + ' updated successfully.'})