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 &

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_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_description.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,17 @@
         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(
         'github/tool/get_latest_releases/',
         GithubToolCheckGetLatestRelease.as_view(),

web/api/views.py

@@ -5,6 +5,7 @@
 import socket
 from ipaddress import IPv4Network
 from collections import defaultdict
+from datetime import datetime
 
 import requests
 import validators
@@ -30,23 +31,27 @@
 	get_interesting_endpoints,
 	get_interesting_subdomains,
 	get_lookup_keywords,
-	safe_int_cast
+	safe_int_cast,
+	get_open_ai_key,
 )
 from reNgine.definitions import (
 	ABORTED_TASK,
-	OLLAMA_INSTANCE,
 	NUCLEI_SEVERITY_MAP,
-	DEFAULT_GPT_MODELS,
 	RUNNING_TASK,
 	SUCCESS_TASK
 )
+from reNgine.llm.config import (
+	OLLAMA_INSTANCE,
+	DEFAULT_GPT_MODELS,
+	MODEL_REQUIREMENTS
+)
 from reNgine.settings import (
 	RENGINE_CURRENT_VERSION,
 	RENGINE_TOOL_GITHUB_PATH
 )
 from reNgine.tasks import (
 	create_scan_activity,
-	gpt_vulnerability_description,
+	llm_vulnerability_description,
 	initiate_subscan,
 	query_ip_history,
 	query_reverse_whois,
@@ -57,7 +62,7 @@
 	run_wafw00f,
 	send_hackerone_report
 )
-from reNgine.gpt import GPTAttackSuggestionGenerator
+from reNgine.llm.llm import LLMAttackSuggestionGenerator
 from reNgine.utilities import is_safe_path, remove_lead_and_trail_slash
 from scanEngine.models import EngineType, InstalledExternalTool
 from startScan.models import (
@@ -169,17 +174,16 @@
                 defaults={
                     'selected_model': model_name,
                     'use_ollama': use_ollama,
-                    'selected': True
                 }
             )
             return Response({'status': True})
         except Exception as e:
             logger.error(f"Error in OllamaManager PUT: {str(e)}")
             return Response({'status': False, 'message': 'An error occurred while updating Ollama settings.'}, status=500)
 
-class GPTAttackSuggestion(APIView):
+class LLMAttackSuggestion(APIView):
 	def get(self, request):
-		req = self.request
+		req = request
 		subdomain_id = safe_int_cast(req.query_params.get('subdomain_id'))
 		if not subdomain_id:
 			return Response({
@@ -195,11 +199,11 @@
 			})
 
 		if subdomain.attack_surface:
 			return Response({
 				'status': True,
 				'subdomain_name': subdomain.name,
 				'description': subdomain.attack_surface
 			})
 
 		ip_addrs = subdomain.ip_addresses.all()
 		open_ports = ', '.join(f'{port.number}/{port.service_name}' for ip in ip_addrs for port in ip.ports.all())
@@ -216,18 +220,18 @@
 			Page Content Length: {subdomain.content_length}
 		'''
 
-		gpt = GPTAttackSuggestionGenerator()
-		response = gpt.get_attack_suggestion(input_data)
+		llm = LLMAttackSuggestionGenerator()
+		response = llm.get_attack_suggestion(input_data)
 		response['subdomain_name'] = subdomain.name
 
 		if response.get('status'):
 			subdomain.attack_surface = response.get('description')
 			subdomain.save()
 
 		return Response(response)
 
 
-class GPTVulnerabilityReportGenerator(APIView):
+class LLMVulnerabilityReportGenerator(APIView):
 	def get(self, request):
 		req = self.request
 		vulnerability_id = safe_int_cast(req.query_params.get('id'))
@@ -236,7 +240,7 @@
 				'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)
 
@@ -2892,3 +2896,55 @@
 					print(e)
 
 		return qs
+
+class LLMModelsManager(APIView):
+    def get(self, request):
+        """Get all available LLM models (GPT and Ollama) and currently selected model"""
+        try:
+            # Get default GPT models
+            all_models = DEFAULT_GPT_MODELS.copy()
+
+            # Get Ollama models
+            try:
+                response = requests.get(f'{OLLAMA_INSTANCE}/api/tags')
+                if response.status_code == 200:
+                    ollama_models = response.json().get('models', [])
+                    date_format = "%Y-%m-%dT%H:%M:%S"
+                    all_models.extend([{
+                        **model,
+                        'modified_at': datetime.strptime(model['modified_at'].split('.')[0], date_format),
+                        'is_local': True,
+                    } for model in ollama_models])
+            except Exception as e:
+                logger.error(f"Error fetching Ollama models: {str(e)}")
+
+            # Get currently selected model
+            selected_model = OllamaSettings.objects.first()
+            selected_model_name = selected_model.selected_model if selected_model else 'gpt-3.5-turbo'
+
+            # Mark selected model
+            for model in all_models:
+                if model['name'] == selected_model_name:
+                    model['selected'] = True
+
+            # Add model capabilities
+            for model in all_models:
+                # Strip tags from model name (e.g., "llama2:latest" -> "llama2")
+                base_model_name = model['name'].split(':')[0]
+                if base_model_name in MODEL_REQUIREMENTS:
+                    model['capabilities'] = MODEL_REQUIREMENTS[base_model_name]
+
+            return Response({
+                'status': True,
+                'models': all_models,
+                'selected_model': selected_model_name,
+                'openai_key_error': not get_open_ai_key() and 'gpt' in selected_model_name
+            })
+
+        except Exception as e:
+            logger.error(f"Error in LLMModelsManager GET: {str(e)}")
+            return Response({
+                'status': False,
+                'error': 'Failed to fetch LLM models',
+                'message': str(e)
+            }, status=500)

web/config/default_yaml_config.yaml

@@ -124,7 +124,7 @@ vulnerability_scan: {
   'rate_limit': 150,
   'retries': 1,
   'timeout': 5,
-  'fetch_gpt_report': true,
+  'fetch_llm_report': true,
   'nuclei': {
     'use_nuclei_config': false,
     'severities': ['unknown', 'info', 'low', 'medium', 'high', 'critical'],

web/dashboard/templates/dashboard/onboarding.html

@@ -58,7 +58,7 @@ <h4 class="mt-3 mt-lg-0">Default API Keys</h4>
                         <p class="text-muted mb-4">If you have API keys for these services, please enter them here.</p>
                         <div class="mb-3">
                           <label for="key_openai" class="form-label">OpenAI <span class="ms-1 badge bg-soft-danger text-danger">🔥 Recommended</span><span class="ms-1 badge bg-soft-primary text-primary">Experimental</span></label>
-                          <p class="text-muted">OpenAI keys will be used to generate vulnerability description, remediation, impact and vulnerability report writing using ChatGPT.</p>
+                          <p class="text-muted">OpenAI keys will be used to generate vulnerability description, remediation, impact and vulnerability report writing using LLM.</p>
                           {% if openai_key %}
                             <input class="form-control" type="text" id="key_openai" name="key_openai" placeholder="Enter OpenAI Key" value="{{openai_key}}">
                           {% else %}

web/fixtures/auth.json

@@ -1875,36 +1875,36 @@
   "model": "auth.permission",
   "pk": 209,
   "fields": {
-    "name": "Can add gpt vulnerability report",
+    "name": "Can add llm vulnerability report",
     "content_type": 53,
-    "codename": "add_gptvulnerabilityreport"
+    "codename": "add_llmvulnerabilityreport"
   }
 },
 {
   "model": "auth.permission",
   "pk": 210,
   "fields": {
-    "name": "Can change gpt vulnerability report",
+    "name": "Can change llm vulnerability report",
     "content_type": 53,
-    "codename": "change_gptvulnerabilityreport"
+    "codename": "change_llmvulnerabilityreport"
   }
 },
 {
   "model": "auth.permission",
   "pk": 211,
   "fields": {
-    "name": "Can delete gpt vulnerability report",
+    "name": "Can delete llm vulnerability report",
     "content_type": 53,
-    "codename": "delete_gptvulnerabilityreport"
+    "codename": "delete_llmvulnerabilityreport"
   }
 },
 {
   "model": "auth.permission",
   "pk": 212,
   "fields": {
-    "name": "Can view gpt vulnerability report",
+    "name": "Can view llm vulnerability report",
     "content_type": 53,
-    "codename": "view_gptvulnerabilityreport"
+    "codename": "view_llmvulnerabilityreport"
   }
 },
 {

web/fixtures/default_scan_engines.yaml

@@ -28,7 +28,7 @@
       'jpg', 'jpeg', 'gif', 'mp4', 'mpeg', 'mp3'],\r\n  'threads': 30\r\n}\r\nvulnerability_scan: {\r\n
       \ 'run_nuclei': true,\r\n  'run_dalfox': true,\r\n  'run_crlfuzz': true,\r\n
       \ 'enable_http_crawl': true,\r\n  'concurrency': 50,\r\n  'intensity': 'normal',\r\n
-      \ 'rate_limit': 150,\r\n  'retries': 1,\r\n  'timeout': 5,\r\n  'fetch_gpt_report':
+      \ 'rate_limit': 150,\r\n  'retries': 1,\r\n  'timeout': 5,\r\n  'fetch_llm_report':
       true,\r\n  'nuclei': {\r\n    'use_nuclei_config': false,\r\n    'severities': ['unknown',
       'info', 'low', 'medium', 'high', 'critical']\r\n  }\r\n}\r\nwaf_detection: {\r\n\r\n}\r\nscreenshot:
       {\r\n  'enable_http_crawl': true,\r\n  'intensity': 'normal',\r\n  'timeout':
@@ -70,7 +70,7 @@
       \ ],\r\n  'intensity': 'normal',\r\n  'documents_limit': 50\r\n}\r\nvulnerability_scan:
       {\r\n  'run_nuclei': true,\r\n  'run_dalfox': true,\r\n  'run_crlfuzz': true,\r\n
       \ 'enable_http_crawl': true,\r\n  'concurrency': 50,\r\n  'intensity': 'normal',\r\n
-      \ 'rate_limit': 150,\r\n  'retries': 1,\r\n  'timeout': 5,\r\n  'fetch_gpt_report':
+      \ 'rate_limit': 150,\r\n  'retries': 1,\r\n  'timeout': 5,\r\n  'fetch_llm_report':
       true,\r\n  'nuclei': {\r\n    'use_nuclei_config': false,\r\n    'severities': ['unknown',
       'info', 'low', 'medium', 'high', 'critical']\r\n  }\r\n}"
     default_engine: true
@@ -97,7 +97,7 @@
       \ ],\r\n  'intensity': 'normal',\r\n  'documents_limit': 50\r\n}\r\nvulnerability_scan:
       {\r\n  'run_nuclei': true,\r\n  'run_dalfox': true,\r\n  'run_crlfuzz': true,\r\n
       \ 'enable_http_crawl': false,\r\n  'concurrency': 50,\r\n  'intensity': 'normal',\r\n
-      \ 'rate_limit': 150,\r\n  'retries': 1,\r\n  'timeout': 5,\r\n  'fetch_gpt_report':
+      \ 'rate_limit': 150,\r\n  'retries': 1,\r\n  'timeout': 5,\r\n  'fetch_llm_report':
       true,\r\n  'nuclei': {\r\n    'use_nuclei_config': false,\r\n    'severities': ['low',
       'medium', 'high', 'critical']\r\n  }\r\n}"
     default_engine: true