feat: refactor LLM model selection and attack surface analysis

docker/celery/entrypoint.sh

@@ -43,7 +43,7 @@ watchmedo auto-restart --recursive --pattern="*.py" --directory="/home/rengine/r
 watchmedo auto-restart --recursive --pattern="*.py" --directory="/home/rengine/rengine/" -- poetry run -C $HOME/ 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="/home/rengine/rengine/" -- poetry run -C $HOME/ 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="/home/rengine/rengine/" -- poetry run -C $HOME/ 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="/home/rengine/rengine/" -- poetry run -C $HOME/ 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="/home/rengine/rengine/" -- poetry run -C $HOME/ celery -A reNgine.tasks worker --pool=gevent --concurrency=30 --loglevel=$CELERY_LOGLEVEL -Q llm_queue -n llm_worker &
 watchmedo auto-restart --recursive --pattern="*.py" --directory="/home/rengine/rengine/" -- poetry run -C $HOME/ 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="/home/rengine/rengine/" -- poetry run -C $HOME/ 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="/home/rengine/rengine/" -- poetry run -C $HOME/ celery -A reNgine.tasks worker --pool=gevent --concurrency=10 --loglevel=$CELERY_LOGLEVEL -Q h8mail_queue -n h8mail_worker &

docker/celery/pyproject.toml

@@ -11,6 +11,8 @@ aiodns = "3.0.0"
 argh = "0.26.2"
 beautifulsoup4 = "4.9.3"
 celery = "5.4.0"
+channels = "3.0.5"
+channels-redis = "3.4.1"
 debugpy = "1.8.5"
 discord-webhook = "1.3.0"
 django = "3.2.25"

docker/proxy/config/rengine.conf

@@ -33,6 +33,29 @@ server {
         alias /home/rengine/rengine/staticfiles/;
     }
 
+
+    # WebSocket support
+    location /ws/ {
+        proxy_pass                                  http://web:8000;
+        proxy_http_version                          1.1;
+        proxy_set_header                            Upgrade $http_upgrade;
+        proxy_set_header                            Connection "upgrade";
+        proxy_set_header                            Host $host;
+        proxy_set_header                            X-Real-IP $remote_addr;
+        proxy_set_header                            X-Forwarded-For $proxy_add_x_forwarded_for;
+        proxy_set_header                            X-Forwarded-Proto $scheme;
+        proxy_read_timeout                          86400;  # 24 hours
+    }
+
+    location /api/tool/ollama/ {
+        proxy_pass                                  http://web:8000;
+        proxy_http_version                          1.1;
+        proxy_set_header                            Connection '';
+        proxy_buffering                             off;
+        proxy_cache                                 off;
+        proxy_read_timeout                          600s;  # 10 minutes timeout
+    }
+
     location /media/ {
         alias /home/rengine/scan_results/;
     }

docker/web/entrypoint.sh

@@ -1,9 +1,29 @@
 #!/bin/bash
 
-# Collect static files
+print_msg() {
+  printf "\r\n"
+  printf "========================================\r\n"
+  printf "$1\r\n"
+  printf "========================================\r\n\r\n"
+}
+
+print_msg "Generate Django migrations files"
+poetry run -C $HOME/ python3 manage.py makemigrations
+
+print_msg "Migrate database"
+poetry run -C $HOME/ python3 manage.py migrate
+
+print_msg "Collect static files"
 poetry run -C $HOME/ python3 manage.py collectstatic --noinput
 
-# Run production server
-poetry run -C $HOME/ gunicorn reNgine.wsgi:application -w 8 --bind 0.0.0.0:8000 --limit-request-line 0
+print_msg "Starting ASGI server with Uvicorn"
+poetry run -C $HOME/ uvicorn reNgine.asgi:application \
+    --host 0.0.0.0 \
+    --port 8000 \
+    --workers 4 \
+    --log-level info \
+    --ws-ping-interval 20 \
+    --ws-ping-timeout 30 \
+    --timeout-keep-alive 65
 
 exec "$@"

docker/web/pyproject.toml

@@ -11,6 +11,8 @@ aiodns = "3.0.0"
 argh = "0.26.2"
 beautifulsoup4 = "4.9.3"
 celery = "5.4.0"
+channels = "3.0.5"
+channels-redis = "3.4.1"
 debugpy = "1.8.5"
 discord-webhook = "1.3.0"
 django = "3.2.25"
@@ -45,6 +47,7 @@ requests = "2.32.2"
 scapy = "2.4.3"
 tldextract = "3.5.0"
 uro = "1.0.0"
+uvicorn = { extras = ["standard"], version = "^0.27.1" }
 validators = "0.18.2"
 watchdog = "4.0.0"
 weasyprint = "53.3"

web/api/consumers.py

@@ -0,0 +1,69 @@
+from channels.generic.websocket import WebsocketConsumer
+from asgiref.sync import async_to_sync
+import json
+import re
+import logging
+
+logger = logging.getLogger(__name__)
+
+class OllamaDownloadConsumer(WebsocketConsumer):
+    def clean_channel_name(self, name):
+        """Clean channel name to only contain valid characters"""
+        return re.sub(r'[^a-zA-Z0-9\-\.]', '-', name)
+
+    def connect(self):
+        try:
+            logger.info(f"WebSocket connection attempt with scope: {self.scope}")
+            self.model_name = self.scope['url_route']['kwargs']['model_name']
+            self.room_group_name = f"ollama-download-{self.clean_channel_name(self.model_name)}"
+
+            logger.info(f"Joining group: {self.room_group_name}")
+
+            # Join room group
+            async_to_sync(self.channel_layer.group_add)(
+                self.room_group_name,
+                self.channel_name
+            )
+
+            logger.info("WebSocket connection accepted")
+            self.accept()
+
+        except Exception as e:
+            logger.error(f"Error in WebSocket connect: {e}")
+            raise
+
+    def disconnect(self, close_code):
+        try:
+            logger.info(f"WebSocket disconnecting with code: {close_code}")
+            # Leave room group
+            async_to_sync(self.channel_layer.group_discard)(
+                self.room_group_name,
+                self.channel_name
+            )
+        except Exception as e:
+            logger.error(f"Error in WebSocket disconnect: {e}")
+
+    def receive(self, text_data):
+        try:
+            logger.info(f"WebSocket received data: {text_data}")
+            text_data_json = json.loads(text_data)
+            message = text_data_json['message']
+
+            # Send message to room group
+            async_to_sync(self.channel_layer.group_send)(
+                self.room_group_name,
+                {
+                    'type': 'download_progress',
+                    'message': message
+                }
+            )
+        except Exception as e:
+            logger.error(f"Error in WebSocket receive: {e}")
+
+    def download_progress(self, event):
+        try:
+            message = event['message']
+            # Send message to WebSocket
+            self.send(text_data=json.dumps(message))
+        except Exception as e:
+            logger.error(f"Error in download_progress: {e}")

web/api/tests/test_project.py

@@ -2,7 +2,6 @@
 This file contains the test cases for the API views.
 """
 
-from unittest.mock import patch
 from django.utils import timezone
 from django.urls import reverse
 from rest_framework import status
@@ -12,7 +11,6 @@
     'TestCreateProjectApi',
     'TestAddReconNote',
     'TestListTodoNotes',
-    'TestGPTAttackSuggestion'
 ]
 
 class TestCreateProjectApi(BaseTestCase):
@@ -108,24 +106,3 @@ def test_list_todo_notes(self):
             self.data_generator.todo_note.scan_history.id,
         )
 
-class TestGPTAttackSuggestion(BaseTestCase):
-    """Tests for the GPT Attack Suggestion API."""
-
-    def setUp(self):
-        super().setUp()
-        self.data_generator.create_project_base()
-
-    @patch("reNgine.gpt.GPTAttackSuggestionGenerator.get_attack_suggestion")
-    def test_get_attack_suggestion(self, mock_get_suggestion):
-        """Test getting an attack suggestion for a subdomain."""
-        mock_get_suggestion.return_value = {
-            "status": True,
-            "description": "Test attack suggestion",
-        }
-        api_url = reverse("api:gpt_get_possible_attacks")
-        response = self.client.get(
-            api_url, {"subdomain_id": self.data_generator.subdomain.id}
-        )
-        self.assertEqual(response.status_code, status.HTTP_200_OK)
-        self.assertTrue(response.data["status"])
-        self.assertEqual(response.data["description"], "Test attack suggestion")

web/api/tests/test_tools.py

@@ -2,11 +2,13 @@
 This file contains the test cases for the API views.
 """
 
-from unittest.mock import patch
+from unittest.mock import patch, AsyncMock, Mock
 from django.urls import reverse
 from rest_framework import status
 from startScan.models import SubScan
 from utils.test_base import BaseTestCase
+from reNgine.llm import config
+from dashboard.models import OllamaSettings
 
 __all__ = [
     'TestOllamaManager',
@@ -25,6 +27,15 @@
 class TestOllamaManager(BaseTestCase):
     """Tests for the OllamaManager API endpoints."""
 
+    def setUp(self):
+        """Set up test environment."""
+        super().setUp()
+        self.ollama_settings = OllamaSettings.objects.create(
+            id=1,
+            selected_model="llama2",
+            use_ollama=True
+        )
+
     @patch("requests.post")
     def test_get_download_model(self, mock_post):
         """Test downloading an Ollama model."""
@@ -34,35 +45,44 @@
         self.assertEqual(response.status_code, status.HTTP_200_OK)
         self.assertTrue(response.data["status"])
 
-    @patch("requests.post")
-    def test_get_download_model_failure(self, mock_post):
-        """Test failed downloading of an Ollama model."""
-        mock_post.return_value.json.return_value = {"error": "pull model manifest: file does not exist"}
-        api_url = reverse("api:ollama_manager")
-        response = self.client.get(api_url, data={"model": "invalid-model"})
-        self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-        self.assertEqual(response.data["message"], "pull model manifest: file does not exist")
-        self.assertFalse(response.data["status"])
-
     @patch("requests.delete")
-    def test_delete_model(self, mock_delete):
+    @patch("requests.get")
+    def test_delete_model(self, mock_get, mock_delete):
         """Test deleting an Ollama model."""
-        mock_delete.return_value.json.return_value = {"status": "success"}
-        api_url = reverse("api:ollama_manager")
-        response = self.client.delete(
-            api_url, data={"model": "gpt-4"}, content_type="application/json"
-        )
+        mock_get.return_value.json.return_value = {
+            "models": [{"name": "llama2"}]
+        }
+        mock_delete.return_value.status_code = 200
+
+        model_name = "llama2"
+        api_url = reverse("api:ollama_detail_manager", kwargs={"model_name": model_name})
+
+        response = self.client.delete(api_url)
+
         self.assertEqual(response.status_code, status.HTTP_200_OK)
         self.assertTrue(response.data["status"])
+        mock_delete.assert_called_once_with(
+            f"{config.OLLAMA_INSTANCE}/api/delete",
+            json={"name": model_name}
+        )
 
-    def test_put_update_model(self):
+    @patch("requests.get")
+    def test_put_update_model(self, mock_get):
         """Test updating the selected Ollama model."""
-        api_url = reverse("api:ollama_manager")
-        response = self.client.put(
-            api_url, data={"model": "gpt-4"}, content_type="application/json"
-        )
+        mock_get.return_value.json.return_value = {
+            "models": [{"name": "gpt-4"}]
+        }
+
+        model_name = "gpt-4"
+        api_url = reverse("api:ollama_detail_manager", kwargs={"model_name": model_name})
+
+        response = self.client.put(api_url)
+
         self.assertEqual(response.status_code, status.HTTP_200_OK)
         self.assertTrue(response.data["status"])
+
+        updated_settings = OllamaSettings.objects.get(id=1)
+        self.assertEqual(updated_settings.selected_model, model_name)
 
 class TestWafDetector(BaseTestCase):
     """Tests for the WAF Detector API."""

web/api/tests/test_vulnerability.py

@@ -10,7 +10,7 @@
 
 __all__ = [
     'TestVulnerabilityViewSet',
-    'TestGPTVulnerabilityReportGenerator',
+    'TestLLMVulnerabilityReportGenerator',
     'TestDeleteVulnerability',
     'TestVulnerabilityReport',
     'TestFetchMostCommonVulnerability',
@@ -79,16 +79,16 @@ def test_list_vulnerabilities_by_severity(self):
             self.data_generator.vulnerabilities[0].name,
         )
 
-class TestGPTVulnerabilityReportGenerator(BaseTestCase):
-    """Tests for the GPT Vulnerability Report Generator API."""
+class TestLLMVulnerabilityReportGenerator(BaseTestCase):
+    """Tests for the LLM Vulnerability Report Generator API."""
 
     def setUp(self):
         super().setUp()
         self.data_generator.create_project_base()
         self.data_generator.create_endpoint()
         self.data_generator.create_vulnerability()
 
-    @patch("reNgine.tasks.gpt_vulnerability_description.apply_async")
+    @patch("reNgine.tasks.llm_vulnerability_report.apply_async")
     def test_get_vulnerability_report(self, mock_apply_async):
         """Test generating a vulnerability report."""
         mock_task = MagicMock()
@@ -97,7 +97,7 @@ def test_get_vulnerability_report(self, mock_apply_async):
             "description": "Test vulnerability report",
         }
         mock_apply_async.return_value = mock_task
-        api_url = reverse("api:gpt_vulnerability_report_generator")
+        api_url = reverse("api:llm_vulnerability_report_generator")
         response = self.client.get(
             api_url, {"id": self.data_generator.vulnerabilities[0].id}
         )

web/api/urls.py

@@ -159,13 +159,21 @@
         GfList.as_view(),
         name='gf_list'),
     path(
-        'tools/gpt_vulnerability_report/',
-        GPTVulnerabilityReportGenerator.as_view(),
-        name='gpt_vulnerability_report_generator'),
+        'tools/llm_vulnerability_report/',
+        LLMVulnerabilityReportGenerator.as_view(),
+        name='llm_vulnerability_report_generator'),
     path(
-        'tools/gpt_get_possible_attacks/',
-        GPTAttackSuggestion.as_view(),
-        name='gpt_get_possible_attacks'),
+        'tools/llm_get_possible_attacks/',
+        LLMAttackSuggestion.as_view(),
+        name='llm_get_possible_attacks'),
+    path(
+        'tools/llm_models/',
+        LLMModelsManager.as_view(),
+        name='llm_models_manager'),
+    path(
+        'tools/available_ollama_models/',
+        AvailableOllamaModels.as_view(),
+        name='available_ollama_models'),
     path(
         'github/tool/get_latest_releases/',
         GithubToolCheckGetLatestRelease.as_view(),
@@ -186,6 +194,10 @@
         'tool/ollama/',
         OllamaManager.as_view(),
         name='ollama_manager'),
+    path(
+        'tool/ollama/<str:model_name>/',
+        OllamaDetailManager.as_view(),
+        name='ollama_detail_manager'),        
     path(
         'rengine/update/',
         RengineUpdateCheck.as_view(),