From 284d83bda76a7952ca4c2a488cb070d82dfdd72c Mon Sep 17 00:00:00 2001 From: psyray Date: Sun, 10 Nov 2024 23:47:15 +0100 Subject: [PATCH 01/21] refactor: replace GPT with LLM for vulnerability reporting The changes involve refactoring the codebase to replace references to "GPT" with "LLM" for generating vulnerability reports and related functionalities. This includes renaming functions, classes, variables, and configuration settings to reflect the new terminology. The update affects various components such as tasks, views, models, templates, and configuration files across the application. --- docker/celery/entrypoint.sh | 2 +- web/api/tests/test_project.py | 10 +- web/api/tests/test_vulnerability.py | 10 +- web/api/urls.py | 12 +- web/api/views.py | 14 +-- web/config/default_yaml_config.yaml | 2 +- .../templates/dashboard/onboarding.html | 2 +- web/fixtures/auth.json | 16 +-- web/fixtures/default_scan_engines.yaml | 6 +- web/reNgine/definitions.py | 6 +- web/reNgine/{gpt.py => llm.py} | 24 ++-- web/reNgine/settings.py | 2 +- web/reNgine/tasks.py | 113 +++++++++--------- web/reNgine/utilities.py | 4 +- web/scanEngine/fixtures/scanEngine.json | 8 +- web/scanEngine/views.py | 2 +- web/startScan/admin.py | 2 +- .../migrations/0058_auto_20241110_2230.py | 22 ++++ web/startScan/models.py | 4 +- .../templates/startScan/detail_scan.html | 4 +- .../templates/startScan/subdomains.html | 2 +- .../templates/startScan/vulnerabilities.html | 2 +- web/static/custom/custom.js | 25 ++-- web/static/custom/vuln_datatables.js | 2 +- web/targetApp/templates/target/summary.html | 4 +- 25 files changed, 160 insertions(+), 140 deletions(-) rename web/reNgine/{gpt.py => llm.py} (85%) create mode 100644 web/startScan/migrations/0058_auto_20241110_2230.py diff --git a/docker/celery/entrypoint.sh b/docker/celery/entrypoint.sh index 2aad31b78..68cc60f71 100755 --- a/docker/celery/entrypoint.sh +++ b/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 & diff --git a/web/api/tests/test_project.py b/web/api/tests/test_project.py index 3522fcadb..d05a96840 100644 --- a/web/api/tests/test_project.py +++ b/web/api/tests/test_project.py @@ -12,7 +12,7 @@ 'TestCreateProjectApi', 'TestAddReconNote', 'TestListTodoNotes', - 'TestGPTAttackSuggestion' + 'TestLLMAttackSuggestion' ] class TestCreateProjectApi(BaseTestCase): @@ -108,21 +108,21 @@ def test_list_todo_notes(self): self.data_generator.todo_note.scan_history.id, ) -class TestGPTAttackSuggestion(BaseTestCase): - """Tests for the GPT Attack Suggestion API.""" +class TestLLMAttackSuggestion(BaseTestCase): + """Tests for the LLM Attack Suggestion API.""" def setUp(self): super().setUp() self.data_generator.create_project_base() - @patch("reNgine.gpt.GPTAttackSuggestionGenerator.get_attack_suggestion") + @patch("reNgine.llm.LLMAttackSuggestionGenerator.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") + api_url = reverse("api:llm_get_possible_attacks") response = self.client.get( api_url, {"subdomain_id": self.data_generator.subdomain.id} ) diff --git a/web/api/tests/test_vulnerability.py b/web/api/tests/test_vulnerability.py index c3cf416a9..a2c78a0b6 100644 --- a/web/api/tests/test_vulnerability.py +++ b/web/api/tests/test_vulnerability.py @@ -10,7 +10,7 @@ __all__ = [ 'TestVulnerabilityViewSet', - 'TestGPTVulnerabilityReportGenerator', + 'TestLLMVulnerabilityReportGenerator', 'TestDeleteVulnerability', 'TestVulnerabilityReport', 'TestFetchMostCommonVulnerability', @@ -79,8 +79,8 @@ 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() @@ -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() @@ -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} ) diff --git a/web/api/urls.py b/web/api/urls.py index ec9e1e108..703120538 100644 --- a/web/api/urls.py +++ b/web/api/urls.py @@ -159,13 +159,13 @@ 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( 'github/tool/get_latest_releases/', GithubToolCheckGetLatestRelease.as_view(), diff --git a/web/api/views.py b/web/api/views.py index 66fe3a8f8..1ee6c3298 100644 --- a/web/api/views.py +++ b/web/api/views.py @@ -46,7 +46,7 @@ ) from reNgine.tasks import ( create_scan_activity, - gpt_vulnerability_description, + llm_vulnerability_description, initiate_subscan, query_ip_history, query_reverse_whois, @@ -57,7 +57,7 @@ run_wafw00f, send_hackerone_report ) -from reNgine.gpt import GPTAttackSuggestionGenerator +from reNgine.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 ( @@ -177,7 +177,7 @@ def put(self, request): 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 subdomain_id = safe_int_cast(req.query_params.get('subdomain_id')) @@ -216,8 +216,8 @@ def get(self, request): 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'): @@ -227,7 +227,7 @@ def get(self, request): 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 +236,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) diff --git a/web/config/default_yaml_config.yaml b/web/config/default_yaml_config.yaml index 4e31bce86..2b70b6415 100644 --- a/web/config/default_yaml_config.yaml +++ b/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'], diff --git a/web/dashboard/templates/dashboard/onboarding.html b/web/dashboard/templates/dashboard/onboarding.html index e34b8dd8c..560340e5a 100644 --- a/web/dashboard/templates/dashboard/onboarding.html +++ b/web/dashboard/templates/dashboard/onboarding.html @@ -58,7 +58,7 @@

Default API Keys

If you have API keys for these services, please enter them here.

-

OpenAI keys will be used to generate vulnerability description, remediation, impact and vulnerability report writing using ChatGPT.

+

OpenAI keys will be used to generate vulnerability description, remediation, impact and vulnerability report writing using LLM.

{% if openai_key %} {% else %} diff --git a/web/fixtures/auth.json b/web/fixtures/auth.json index a22e89ef7..565501b09 100644 --- a/web/fixtures/auth.json +++ b/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" } }, { diff --git a/web/fixtures/default_scan_engines.yaml b/web/fixtures/default_scan_engines.yaml index 744861ba4..f02aa83f5 100644 --- a/web/fixtures/default_scan_engines.yaml +++ b/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 diff --git a/web/reNgine/definitions.py b/web/reNgine/definitions.py index ce1c710c7..ee0fee5a6 100644 --- a/web/reNgine/definitions.py +++ b/web/reNgine/definitions.py @@ -23,7 +23,7 @@ AMASS_WORDLIST = 'amass_wordlist' AUTO_CALIBRATION = 'auto_calibration' CUSTOM_HEADER = 'custom_header' -FETCH_GPT_REPORT = 'fetch_gpt_report' +FETCH_LLM_REPORT = 'fetch_llm_report' RUN_NUCLEI = 'run_nuclei' RUN_CRLFUZZ = 'run_crlfuzz' RUN_DALFOX = 'run_dalfox' @@ -484,7 +484,7 @@ -# GPT Vulnerability Report Generator +# LLM Vulnerability Report Generator VULNERABILITY_DESCRIPTION_SYSTEM_MESSAGE = """ You are a highly skilled penetration tester who has recently completed a penetration testing. You will be given with a @@ -509,7 +509,7 @@ """ -ATTACK_SUGGESTION_GPT_SYSTEM_PROMPT = """ +ATTACK_SUGGESTION_LLM_SYSTEM_PROMPT = """ You are a highly skilled penetration tester who has recently completed a reconnaissance on a target. As a penetration tester, I've conducted a thorough reconnaissance on a specific subdomain. Based on my reconnaissance you will be given with a diff --git a/web/reNgine/gpt.py b/web/reNgine/llm.py similarity index 85% rename from web/reNgine/gpt.py rename to web/reNgine/llm.py index 47af6c43f..d6e40074d 100644 --- a/web/reNgine/gpt.py +++ b/web/reNgine/llm.py @@ -1,7 +1,7 @@ import openai import re from reNgine.common_func import get_open_ai_key, extract_between -from reNgine.definitions import VULNERABILITY_DESCRIPTION_SYSTEM_MESSAGE, ATTACK_SUGGESTION_GPT_SYSTEM_PROMPT, OLLAMA_INSTANCE +from reNgine.definitions import VULNERABILITY_DESCRIPTION_SYSTEM_MESSAGE, ATTACK_SUGGESTION_LLM_SYSTEM_PROMPT, OLLAMA_INSTANCE from langchain_community.llms import Ollama from dashboard.models import OllamaSettings @@ -9,7 +9,7 @@ logger = logging.getLogger(__name__) -class GPTVulnerabilityReportGenerator: +class LLMVulnerabilityReportGenerator: def __init__(self): selected_model = OllamaSettings.objects.first() @@ -19,10 +19,10 @@ def __init__(self): self.ollama = None def get_vulnerability_description(self, description): - """Generate Vulnerability Description using GPT. + """Generate Vulnerability Description using LLM. Args: - description (str): Vulnerability Description message to pass to GPT. + description (str): Vulnerability Description message to pass to LLM. Returns: (dict) of { @@ -49,7 +49,7 @@ def get_vulnerability_description(self, description): } try: openai.api_key = openai_api_key - gpt_response = openai.ChatCompletion.create( + llm_response = openai.ChatCompletion.create( model=self.model_name, messages=[ {'role': 'system', 'content': VULNERABILITY_DESCRIPTION_SYSTEM_MESSAGE}, @@ -57,7 +57,7 @@ def get_vulnerability_description(self, description): ] ) - response_content = gpt_response['choices'][0]['message']['content'] + response_content = llm_response['choices'][0]['message']['content'] except Exception as e: return { 'status': False, @@ -93,7 +93,7 @@ def get_vulnerability_description(self, description): 'references': urls, } -class GPTAttackSuggestionGenerator: +class LLMAttackSuggestionGenerator: def __init__(self): self.api_key = get_open_ai_key() @@ -103,22 +103,22 @@ def __init__(self): def get_attack_suggestion(self, input): ''' - input (str): input for gpt + input (str): input for llm ''' try: if not self.api_key: - prompt = ATTACK_SUGGESTION_GPT_SYSTEM_PROMPT + "\nUser: " + input + prompt = ATTACK_SUGGESTION_LLM_SYSTEM_PROMPT + "\nUser: " + input response_content = self.ollama(prompt) else: openai.api_key = self.api_key - gpt_response = openai.ChatCompletion.create( + llm_response = openai.ChatCompletion.create( model=self.model_name, messages=[ - {'role': 'system', 'content': ATTACK_SUGGESTION_GPT_SYSTEM_PROMPT}, + {'role': 'system', 'content': ATTACK_SUGGESTION_LLM_SYSTEM_PROMPT}, {'role': 'user', 'content': input} ] ) - response_content = gpt_response['choices'][0]['message']['content'] + response_content = llm_response['choices'][0]['message']['content'] return { 'status': True, diff --git a/web/reNgine/settings.py b/web/reNgine/settings.py index 1a63d823e..04e779432 100644 --- a/web/reNgine/settings.py +++ b/web/reNgine/settings.py @@ -50,7 +50,7 @@ DEFAULT_HTTP_TIMEOUT = env.int('DEFAULT_HTTP_TIMEOUT', default=5) # seconds DEFAULT_RETRIES = env.int('DEFAULT_RETRIES', default=1) DEFAULT_THREADS = env.int('DEFAULT_THREADS', default=30) -DEFAULT_GET_GPT_REPORT = env.bool('DEFAULT_GET_GPT_REPORT', default=True) +DEFAULT_GET_LLM_REPORT = env.bool('DEFAULT_GET_LLM_REPORT', default=True) # Globals ALLOWED_HOSTS = ['*'] diff --git a/web/reNgine/tasks.py b/web/reNgine/tasks.py index a3ed329fc..dd1d93dc9 100644 --- a/web/reNgine/tasks.py +++ b/web/reNgine/tasks.py @@ -28,12 +28,12 @@ from metafinder.extractor import extract_metadata_from_google_search from reNgine.celery import app -from reNgine.gpt import GPTVulnerabilityReportGenerator +from reNgine.llm import LLMVulnerabilityReportGenerator from reNgine.celery_custom_task import RengineTask from reNgine.common_func import * from reNgine.definitions import * from reNgine.settings import * -from reNgine.gpt import * +from reNgine.llm import * from reNgine.utilities import * from scanEngine.models import (EngineType, InstalledExternalTool, Notification, Proxy) from startScan.models import * @@ -2108,7 +2108,7 @@ def vulnerability_scan(self, urls=[], ctx={}, description=None): return None @app.task(name='nuclei_individual_severity_module', queue='main_scan_queue', base=RengineTask, bind=True) -def nuclei_individual_severity_module(self, cmd, severity, enable_http_crawl, should_fetch_gpt_report, ctx={}, description=None): +def nuclei_individual_severity_module(self, cmd, severity, enable_http_crawl, should_fetch_llm_report, ctx={}, description=None): ''' This celery task will run vulnerability scan in parallel. All severities supplied should run in parallel as grouped tasks. @@ -2266,11 +2266,10 @@ def nuclei_individual_severity_module(self, cmd, severity, enable_http_crawl, sh } self.notify(fields=fields) - # after vulnerability scan is done, we need to run gpt if - # should_fetch_gpt_report and openapi key exists + # after vulnerability scan is done, we need to run llm if + # should_fetch_llm_report and openapi key exists - if should_fetch_gpt_report and OpenAiAPIKey.objects.all().first(): - logger.info('Getting Vulnerability GPT Report') + if should_fetch_llm_report and OpenAiAPIKey.objects.exists(): vulns = Vulnerability.objects.filter( scan_history__id=self.scan_id ).filter( @@ -2279,7 +2278,7 @@ def nuclei_individual_severity_module(self, cmd, severity, enable_http_crawl, sh severity=0 ) # find all unique vulnerabilities based on path and title - # all unique vulnerability will go thru gpt function and get report + # all unique vulnerability will go thru llm function and get report # once report is got, it will be matched with other vulnerabilities and saved unique_vulns = set() for vuln in vulns: @@ -2288,11 +2287,11 @@ def nuclei_individual_severity_module(self, cmd, severity, enable_http_crawl, sh unique_vulns = list(unique_vulns) with concurrent.futures.ThreadPoolExecutor(max_workers=DEFAULT_THREADS) as executor: - future_to_gpt = {executor.submit(get_vulnerability_gpt_report, vuln): vuln for vuln in unique_vulns} + future_to_llm = {executor.submit(get_vulnerability_llm_report, vuln): vuln for vuln in unique_vulns} # Wait for all tasks to complete - for future in concurrent.futures.as_completed(future_to_gpt): - gpt = future_to_gpt[future] + for future in concurrent.futures.as_completed(future_to_llm): + llm = future_to_llm[future] try: future.result() except Exception as e: @@ -2301,12 +2300,12 @@ def nuclei_individual_severity_module(self, cmd, severity, enable_http_crawl, sh return None -def get_vulnerability_gpt_report(vuln): +def get_vulnerability_llm_report(vuln): title = vuln[0] path = vuln[1] - logger.info(f'Getting GPT Report for {title}, PATH: {path}') + logger.info(f'Getting LLM Report for {title}, PATH: {path}') # check if in db already exists - stored = GPTVulnerabilityReport.objects.filter( + stored = LLMVulnerabilityReport.objects.filter( url_path=path ).filter( title=title @@ -2319,13 +2318,13 @@ def get_vulnerability_gpt_report(vuln): 'references': [url.url for url in stored.references.all()] } else: - report = GPTVulnerabilityReportGenerator() - vulnerability_description = get_gpt_vuln_input_description( + report = LLMVulnerabilityReportGenerator() + vulnerability_description = get_llm_vuln_input_description( title, path ) response = report.get_vulnerability_description(vulnerability_description) - add_gpt_description_db( + add_llm_description_db( title, path, response.get('description'), @@ -2339,7 +2338,7 @@ def get_vulnerability_gpt_report(vuln): vuln.description = response.get('description', vuln.description) vuln.impact = response.get('impact') vuln.remediation = response.get('remediation') - vuln.is_gpt_used = True + vuln.is_llm_used = True vuln.save() for url in response.get('references', []): @@ -2348,19 +2347,19 @@ def get_vulnerability_gpt_report(vuln): vuln.save() -def add_gpt_description_db(title, path, description, impact, remediation, references): - gpt_report = GPTVulnerabilityReport() - gpt_report.url_path = path - gpt_report.title = title - gpt_report.description = description - gpt_report.impact = impact - gpt_report.remediation = remediation - gpt_report.save() +def add_llm_description_db(title, path, description, impact, remediation, references): + llm_report = LLMVulnerabilityReport() + llm_report.url_path = path + llm_report.title = title + llm_report.description = description + llm_report.impact = impact + llm_report.remediation = remediation + llm_report.save() for url in references: ref, created = VulnerabilityReference.objects.get_or_create(url=url) - gpt_report.references.add(ref) - gpt_report.save() + llm_report.references.add(ref) + llm_report.save() @app.task(name='nuclei_scan', queue='main_scan_queue', base=RengineTask, bind=True) def nuclei_scan(self, urls=[], ctx={}, description=None): @@ -2386,7 +2385,7 @@ def nuclei_scan(self, urls=[], ctx={}, description=None): custom_header = config.get(CUSTOM_HEADER) or self.yaml_configuration.get(CUSTOM_HEADER) if custom_header: custom_header = generate_header_param(custom_header, 'common') - should_fetch_gpt_report = config.get(FETCH_GPT_REPORT, DEFAULT_GET_GPT_REPORT) + should_fetch_llm_report = config.get(FETCH_LLM_REPORT, DEFAULT_GET_LLM_REPORT) proxy = get_random_proxy() nuclei_specific_config = config.get('nuclei', {}) use_nuclei_conf = nuclei_specific_config.get(USE_NUCLEI_CONFIG, False) @@ -2474,7 +2473,7 @@ def nuclei_scan(self, urls=[], ctx={}, description=None): cmd, severity, enable_http_crawl, - should_fetch_gpt_report, + should_fetch_llm_report, ctx=custom_ctx, description=f'Nuclei Scan with severity {severity}' ) @@ -2500,7 +2499,7 @@ def dalfox_xss_scan(self, urls=[], ctx={}, description=None): description (str, optional): Task description shown in UI. """ vuln_config = self.yaml_configuration.get(VULNERABILITY_SCAN) or {} - should_fetch_gpt_report = vuln_config.get(FETCH_GPT_REPORT, DEFAULT_GET_GPT_REPORT) + should_fetch_llm_report = vuln_config.get(FETCH_LLM_REPORT, DEFAULT_GET_LLM_REPORT) dalfox_config = vuln_config.get(DALFOX) or {} custom_header = dalfox_config.get(CUSTOM_HEADER) or self.yaml_configuration.get(CUSTOM_HEADER) if custom_header: @@ -2593,11 +2592,11 @@ def dalfox_xss_scan(self, urls=[], ctx={}, description=None): if not vuln: continue - # after vulnerability scan is done, we need to run gpt if - # should_fetch_gpt_report and openapi key exists + # after vulnerability scan is done, we need to run llm if + # should_fetch_llm_report and openapi key exists - if should_fetch_gpt_report and OpenAiAPIKey.objects.all().first(): - logger.info('Getting Dalfox Vulnerability GPT Report') + if should_fetch_llm_report and OpenAiAPIKey.objects.all().first(): + logger.info('Getting Dalfox Vulnerability LLM Report') vulns = Vulnerability.objects.filter( scan_history__id=self.scan_id ).filter( @@ -2611,11 +2610,11 @@ def dalfox_xss_scan(self, urls=[], ctx={}, description=None): _vulns.append((vuln.name, vuln.http_url)) with concurrent.futures.ThreadPoolExecutor(max_workers=DEFAULT_THREADS) as executor: - future_to_gpt = {executor.submit(get_vulnerability_gpt_report, vuln): vuln for vuln in _vulns} + future_to_llm = {executor.submit(get_vulnerability_llm_report, vuln): vuln for vuln in _vulns} # Wait for all tasks to complete - for future in concurrent.futures.as_completed(future_to_gpt): - gpt = future_to_gpt[future] + for future in concurrent.futures.as_completed(future_to_llm): + llm = future_to_llm[future] try: future.result() except Exception as e: @@ -2632,7 +2631,7 @@ def crlfuzz_scan(self, urls=[], ctx={}, description=None): description (str, optional): Task description shown in UI. """ vuln_config = self.yaml_configuration.get(VULNERABILITY_SCAN) or {} - should_fetch_gpt_report = vuln_config.get(FETCH_GPT_REPORT, DEFAULT_GET_GPT_REPORT) + should_fetch_llm_report = vuln_config.get(FETCH_LLM_REPORT, DEFAULT_GET_LLM_REPORT) custom_header = vuln_config.get(CUSTOM_HEADER) or self.yaml_configuration.get(CUSTOM_HEADER) if custom_header: custom_header = generate_header_param(custom_header, 'common') @@ -2719,11 +2718,11 @@ def crlfuzz_scan(self, urls=[], ctx={}, description=None): if not vuln: continue - # after vulnerability scan is done, we need to run gpt if - # should_fetch_gpt_report and openapi key exists + # after vulnerability scan is done, we need to run llm if + # should_fetch_llm_report and openapi key exists - if should_fetch_gpt_report and OpenAiAPIKey.objects.all().first(): - logger.info('Getting CRLFuzz Vulnerability GPT Report') + if should_fetch_llm_report and OpenAiAPIKey.objects.all().first(): + logger.info('Getting CRLFuzz Vulnerability LLM Report') vulns = Vulnerability.objects.filter( scan_history__id=self.scan_id ).filter( @@ -2737,11 +2736,11 @@ def crlfuzz_scan(self, urls=[], ctx={}, description=None): _vulns.append((vuln.name, vuln.http_url)) with concurrent.futures.ThreadPoolExecutor(max_workers=DEFAULT_THREADS) as executor: - future_to_gpt = {executor.submit(get_vulnerability_gpt_report, vuln): vuln for vuln in _vulns} + future_to_llm = {executor.submit(get_vulnerability_llm_report, vuln): vuln for vuln in _vulns} # Wait for all tasks to complete - for future in concurrent.futures.as_completed(future_to_gpt): - gpt = future_to_gpt[future] + for future in concurrent.futures.as_completed(future_to_llm): + llm = future_to_llm[future] try: future.result() except Exception as e: @@ -4836,14 +4835,14 @@ def query_ip_history(domain): return get_domain_historical_ip_address(domain) -@app.task(name='gpt_vulnerability_description', bind=False, queue='gpt_queue') -def gpt_vulnerability_description(vulnerability_id): - """Generate and store Vulnerability Description using GPT. +@app.task(name='llm_vulnerability_description', bind=False, queue='llm_queue') +def llm_vulnerability_description(vulnerability_id): + """Generate and store Vulnerability Description using LLM. Args: vulnerability_id (Vulnerability Model ID): Vulnerability ID to fetch Description. """ - logger.info('Getting GPT Vulnerability Description') + logger.info('Getting LLM Vulnerability Description') try: lookup_vulnerability = Vulnerability.objects.get(id=vulnerability_id) lookup_url = urlparse(lookup_vulnerability.http_url) @@ -4854,8 +4853,8 @@ def gpt_vulnerability_description(vulnerability_id): 'error': str(e) } - # check in db GPTVulnerabilityReport model if vulnerability description and path matches - stored = GPTVulnerabilityReport.objects.filter(url_path=path).filter(title=lookup_vulnerability.name).first() + # check in db LLMVulnerabilityReport model if vulnerability description and path matches + stored = LLMVulnerabilityReport.objects.filter(url_path=path).filter(title=lookup_vulnerability.name).first() if stored: response = { 'status': True, @@ -4865,15 +4864,15 @@ def gpt_vulnerability_description(vulnerability_id): 'references': [url.url for url in stored.references.all()] } else: - vulnerability_description = get_gpt_vuln_input_description( + vulnerability_description = get_llm_vuln_input_description( lookup_vulnerability.name, path ) # one can add more description here later - gpt_generator = GPTVulnerabilityReportGenerator() - response = gpt_generator.get_vulnerability_description(vulnerability_description) - add_gpt_description_db( + llm_generator = LLMVulnerabilityReportGenerator() + response = llm_generator.get_vulnerability_description(vulnerability_description) + add_llm_description_db( lookup_vulnerability.name, path, response.get('description'), @@ -4889,7 +4888,7 @@ def gpt_vulnerability_description(vulnerability_id): vuln.description = response.get('description', vuln.description) vuln.impact = response.get('impact') vuln.remediation = response.get('remediation') - vuln.is_gpt_used = True + vuln.is_llm_used = True vuln.save() for url in response.get('references', []): diff --git a/web/reNgine/utilities.py b/web/reNgine/utilities.py index d985d2e00..d4daafaeb 100644 --- a/web/reNgine/utilities.py +++ b/web/reNgine/utilities.py @@ -68,10 +68,10 @@ def format(self, record): return super().format(record) -def get_gpt_vuln_input_description(title, path): +def get_llm_vuln_input_description(title, path): vulnerability_description = '' vulnerability_description += f'Vulnerability Title: {title}' - # gpt gives concise vulnerability description when a vulnerable URL is provided + # llm gives concise vulnerability description when a vulnerable URL is provided vulnerability_description += f'\nVulnerable URL: {path}' return vulnerability_description diff --git a/web/scanEngine/fixtures/scanEngine.json b/web/scanEngine/fixtures/scanEngine.json index 030b7b148..f5223475e 100644 --- a/web/scanEngine/fixtures/scanEngine.json +++ b/web/scanEngine/fixtures/scanEngine.json @@ -4,7 +4,7 @@ "pk": 1, "fields": { "engine_name": "Full Scan", - "yaml_configuration": "subdomain_discovery: {\r\n 'uses_tools': ['subfinder', 'ctfr', 'sublist3r', 'tlsx', 'oneforall', 'netlas'],\r\n 'enable_http_crawl': true,\r\n 'threads': 30,\r\n 'timeout': 5,\r\n}\r\nhttp_crawl: {}\r\nport_scan: {\r\n 'enable_http_crawl': true,\r\n 'timeout': 5,\r\n # 'exclude_ports': [],\r\n # 'exclude_subdomains': [],\r\n 'ports': ['top-100'],\r\n 'rate_limit': 150,\r\n 'threads': 30,\r\n 'passive': false,\r\n # 'use_naabu_config': false,\r\n # 'enable_nmap': true,\r\n # 'nmap_cmd': '',\r\n # 'nmap_script': '',\r\n # 'nmap_script_args': ''\r\n}\r\nosint: {\r\n 'discover': [\r\n 'emails',\r\n 'metainfo',\r\n 'employees'\r\n ],\r\n 'dorks': [\r\n 'login_pages',\r\n 'admin_panels',\r\n 'dashboard_pages',\r\n 'stackoverflow',\r\n 'social_media',\r\n 'project_management',\r\n 'code_sharing',\r\n 'config_files',\r\n 'jenkins',\r\n 'wordpress_files',\r\n 'php_error',\r\n 'exposed_documents',\r\n 'db_files',\r\n 'git_exposed'\r\n ],\r\n 'intensity': 'normal',\r\n 'documents_limit': 50\r\n}\r\ndir_file_fuzz: {\r\n 'auto_calibration': true,\r\n 'enable_http_crawl': true,\r\n 'rate_limit': 150,\r\n 'extensions': [],\r\n 'follow_redirect': false,\r\n 'max_time': 0,\r\n 'match_http_status': [200, 204],\r\n 'recursive_level': 0,\r\n 'stop_on_error': false,\r\n 'timeout': 5,\r\n 'threads': 30,\r\n 'wordlist_name': 'default', # fuzz-Bo0oM\r\n}\r\nfetch_url: {\r\n 'uses_tools': ['gospider', 'hakrawler', 'waybackurls', 'katana', 'gau'],\r\n 'remove_duplicate_endpoints': true,\r\n 'duplicate_fields': ['content_length', 'page_title'],\r\n 'follow_redirect': false,\r\n 'enable_http_crawl': true,\r\n 'gf_patterns': ['debug_logic', 'idor', 'interestingEXT', 'interestingparams', 'interestingsubs', 'lfi', 'rce', 'redirect', 'sqli', 'ssrf', 'ssti', 'xss'],\r\n 'ignore_file_extensions': ['png', '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': 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': 10,\r\n 'threads': 40\r\n}\r\n\r\n# custom_header: \"Cookie: Test\"", + "yaml_configuration": "subdomain_discovery: {\r\n 'uses_tools': ['subfinder', 'ctfr', 'sublist3r', 'tlsx', 'oneforall', 'netlas'],\r\n 'enable_http_crawl': true,\r\n 'threads': 30,\r\n 'timeout': 5,\r\n}\r\nhttp_crawl: {}\r\nport_scan: {\r\n 'enable_http_crawl': true,\r\n 'timeout': 5,\r\n # 'exclude_ports': [],\r\n # 'exclude_subdomains': [],\r\n 'ports': ['top-100'],\r\n 'rate_limit': 150,\r\n 'threads': 30,\r\n 'passive': false,\r\n # 'use_naabu_config': false,\r\n # 'enable_nmap': true,\r\n # 'nmap_cmd': '',\r\n # 'nmap_script': '',\r\n # 'nmap_script_args': ''\r\n}\r\nosint: {\r\n 'discover': [\r\n 'emails',\r\n 'metainfo',\r\n 'employees'\r\n ],\r\n 'dorks': [\r\n 'login_pages',\r\n 'admin_panels',\r\n 'dashboard_pages',\r\n 'stackoverflow',\r\n 'social_media',\r\n 'project_management',\r\n 'code_sharing',\r\n 'config_files',\r\n 'jenkins',\r\n 'wordpress_files',\r\n 'php_error',\r\n 'exposed_documents',\r\n 'db_files',\r\n 'git_exposed'\r\n ],\r\n 'intensity': 'normal',\r\n 'documents_limit': 50\r\n}\r\ndir_file_fuzz: {\r\n 'auto_calibration': true,\r\n 'enable_http_crawl': true,\r\n 'rate_limit': 150,\r\n 'extensions': [],\r\n 'follow_redirect': false,\r\n 'max_time': 0,\r\n 'match_http_status': [200, 204],\r\n 'recursive_level': 0,\r\n 'stop_on_error': false,\r\n 'timeout': 5,\r\n 'threads': 30,\r\n 'wordlist_name': 'default', # fuzz-Bo0oM\r\n}\r\nfetch_url: {\r\n 'uses_tools': ['gospider', 'hakrawler', 'waybackurls', 'katana', 'gau'],\r\n 'remove_duplicate_endpoints': true,\r\n 'duplicate_fields': ['content_length', 'page_title'],\r\n 'follow_redirect': false,\r\n 'enable_http_crawl': true,\r\n 'gf_patterns': ['debug_logic', 'idor', 'interestingEXT', 'interestingparams', 'interestingsubs', 'lfi', 'rce', 'redirect', 'sqli', 'ssrf', 'ssti', 'xss'],\r\n 'ignore_file_extensions': ['png', '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_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': 10,\r\n 'threads': 40\r\n}\r\n\r\n# custom_header: \"Cookie: Test\"", "default_engine": true } }, @@ -31,7 +31,7 @@ "pk": 4, "fields": { "engine_name": "Vulnerability Scan", - "yaml_configuration": "subdomain_discovery: {\r\n 'uses_tools': ['subfinder', 'ctfr', 'sublist3r', 'tlsx', 'oneforall', 'netlas'],\r\n 'enable_http_crawl': true,\r\n 'threads': 30,\r\n 'timeout': 5,\r\n}\r\nhttp_crawl: {}\r\nosint: {\r\n 'discover': [\r\n 'emails',\r\n 'metainfo',\r\n 'employees'\r\n ],\r\n 'dorks': [\r\n 'login_pages',\r\n 'admin_panels',\r\n 'dashboard_pages',\r\n 'stackoverflow',\r\n 'social_media',\r\n 'project_management',\r\n 'code_sharing',\r\n 'config_files',\r\n 'jenkins',\r\n 'wordpress_files',\r\n 'php_error',\r\n 'exposed_documents',\r\n 'db_files',\r\n 'git_exposed'\r\n ],\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': true,\r\n 'nuclei': {\r\n 'use_nuclei_config': false,\r\n 'severities': ['unknown', 'info', 'low', 'medium', 'high', 'critical']\r\n }\r\n}", + "yaml_configuration": "subdomain_discovery: {\r\n 'uses_tools': ['subfinder', 'ctfr', 'sublist3r', 'tlsx', 'oneforall', 'netlas'],\r\n 'enable_http_crawl': true,\r\n 'threads': 30,\r\n 'timeout': 5,\r\n}\r\nhttp_crawl: {}\r\nosint: {\r\n 'discover': [\r\n 'emails',\r\n 'metainfo',\r\n 'employees'\r\n ],\r\n 'dorks': [\r\n 'login_pages',\r\n 'admin_panels',\r\n 'dashboard_pages',\r\n 'stackoverflow',\r\n 'social_media',\r\n 'project_management',\r\n 'code_sharing',\r\n 'config_files',\r\n 'jenkins',\r\n 'wordpress_files',\r\n 'php_error',\r\n 'exposed_documents',\r\n 'db_files',\r\n 'git_exposed'\r\n ],\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_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 } }, @@ -49,7 +49,7 @@ "pk": 6, "fields": { "engine_name": "reNgine Recommended", - "yaml_configuration": "subdomain_discovery: {\r\n 'uses_tools': ['subfinder', 'ctfr', 'sublist3r', 'tlsx', 'oneforall', 'netlas'],\r\n 'enable_http_crawl': true,\r\n 'threads': 30,\r\n 'timeout': 5,\r\n}\r\nhttp_crawl: {}\r\nosint: {\r\n 'discover': [\r\n 'emails',\r\n 'metainfo'\r\n ],\r\n 'dorks': [\r\n 'login_pages',\r\n 'admin_panels',\r\n 'dashboard_pages',\r\n 'config_files',\r\n 'exposed_documents',\r\n ],\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': true,\r\n 'nuclei': {\r\n 'use_nuclei_config': false,\r\n 'severities': ['low', 'medium', 'high', 'critical']\r\n }\r\n}", + "yaml_configuration": "subdomain_discovery: {\r\n 'uses_tools': ['subfinder', 'ctfr', 'sublist3r', 'tlsx', 'oneforall', 'netlas'],\r\n 'enable_http_crawl': true,\r\n 'threads': 30,\r\n 'timeout': 5,\r\n}\r\nhttp_crawl: {}\r\nosint: {\r\n 'discover': [\r\n 'emails',\r\n 'metainfo'\r\n ],\r\n 'dorks': [\r\n 'login_pages',\r\n 'admin_panels',\r\n 'dashboard_pages',\r\n 'config_files',\r\n 'exposed_documents',\r\n ],\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_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 } }, @@ -58,7 +58,7 @@ "pk": 7, "fields": { "engine_name": "Full (perso)", - "yaml_configuration": "# Global vars for all tools\r\n#\r\n# Custom header - FFUF, Nuclei, Dalfox, CRL Fuzz, HTTPx, Fetch URL (Hakrawler, Katana, Gospider)\r\n# custom_header: {\r\n# 'Cookie':'Test',\r\n# 'User-Agent': 'Mozilla/5.0',\r\n# 'Custom-Header': 'My custom header'\r\n# }\r\n# 'user_agent': '' # Dalfox only\r\n# 'enable_http_crawl': true # All tools\r\n# 'timeout': 10 # Subdomain discovery, Screenshot, Port scan, FFUF, Nuclei \r\n# 'threads': 30 # All tools\r\n# 'rate_limit': 150 # Port scan, FFUF, Nuclei\r\n# 'intensity': 'normal' # Screenshot (grab only the root endpoints of each subdomain), Nuclei (reduce number of endpoints to scan), OSINT (not implemented yet)\r\n# 'retries': 1 # Nuclei\r\n\r\nsubdomain_discovery: {\r\n 'uses_tools': ['subfinder', 'ctfr', 'sublist3r', 'tlsx', 'oneforall', 'netlas'], # amass-passive, amass-active, All\r\n 'enable_http_crawl': true,\r\n 'threads': 30,\r\n 'timeout': 5,\r\n # 'use_subfinder_config': false,\r\n # 'use_amass_config': false,\r\n # 'amass_wordlist': 'deepmagic.com-prefixes-top50000'\r\n}\r\nhttp_crawl: {\r\n # 'custom_header': {\r\n # 'Cookie':'Test',\r\n # 'User-Agent': 'Mozilla/5.0'\r\n # },\r\n # 'threads': 30,\r\n # 'follow_redirect': false\r\n}\r\nport_scan: {\r\n 'enable_http_crawl': true,\r\n 'timeout': 5,\r\n # 'exclude_ports': [],\r\n # 'exclude_subdomains': [],\r\n 'ports': ['top-100'],\r\n 'rate_limit': 150,\r\n 'threads': 30,\r\n 'passive': false,\r\n # 'use_naabu_config': false,\r\n # 'enable_nmap': true,\r\n # 'nmap_cmd': '',\r\n # 'nmap_script': '',\r\n # 'nmap_script_args': ''\r\n}\r\nosint: {\r\n 'discover': [\r\n 'emails',\r\n 'metainfo',\r\n 'employees'\r\n ],\r\n 'dorks': [\r\n 'login_pages',\r\n 'admin_panels',\r\n 'dashboard_pages',\r\n 'stackoverflow',\r\n 'social_media',\r\n 'project_management',\r\n 'code_sharing',\r\n 'config_files',\r\n 'jenkins',\r\n 'wordpress_files',\r\n 'php_error',\r\n 'exposed_documents',\r\n 'db_files',\r\n 'git_exposed'\r\n ],\r\n # 'custom_dorks': [],\r\n 'intensity': 'normal',\r\n 'documents_limit': 50\r\n}\r\ndir_file_fuzz: {\r\n # 'custom_header': {\r\n # 'Cookie':'Test',\r\n # 'User-Agent': 'Mozilla/5.0',\r\n # 'Custom-Header': 'My custom header'\r\n # },\r\n 'auto_calibration': true,\r\n 'enable_http_crawl': true,\r\n 'rate_limit': 150,\r\n 'extensions': [],\r\n 'follow_redirect': false,\r\n 'max_time': 0,\r\n 'match_http_status': [200, 204],\r\n 'recursive_level': 0,\r\n 'stop_on_error': false,\r\n 'timeout': 5,\r\n 'threads': 30,\r\n 'wordlist_name': 'default', # fuzz-Bo0oM,\r\n}\r\nfetch_url: {\r\n # 'custom_header': {\r\n # 'Cookie':'Test',\r\n # 'User-Agent': 'Mozilla/5.0',\r\n # 'Custom-Header': 'My custom header'\r\n # },\r\n 'uses_tools': ['gospider', 'hakrawler', 'waybackurls', 'katana', 'gau'],\r\n 'remove_duplicate_endpoints': true,\r\n 'duplicate_fields': ['content_length', 'page_title'],\r\n 'follow_redirect': false,\r\n 'enable_http_crawl': true,\r\n 'gf_patterns': ['debug_logic', 'idor', 'interestingEXT', 'interestingparams', 'interestingsubs', 'lfi', 'rce', 'redirect', 'sqli', 'ssrf', 'ssti', 'xss'],\r\n 'ignore_file_extensions': ['png', 'jpg', 'jpeg', 'gif', 'mp4', 'mpeg', 'mp3'],\r\n 'threads': 30,\r\n # 'exclude_subdomains': false\r\n}\r\nvulnerability_scan: {\r\n # 'custom_header': {\r\n # 'Cookie':'Test',\r\n # 'User-Agent': 'Mozilla/5.0',\r\n # 'Custom-Header': 'My custom header'\r\n # },\r\n 'run_nuclei': true,\r\n 'run_dalfox': false,\r\n 'run_crlfuzz': false,\r\n 'run_s3scanner': false,\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': true,\r\n 'nuclei': {\r\n 'use_nuclei_config': false,\r\n 'severities': ['unknown', 'info', 'low', 'medium', 'high', 'critical'],\r\n # 'tags': [], # Nuclei tags (https://github.com/projectdiscovery/nuclei-templates)\r\n # 'templates': [], # Nuclei templates (https://github.com/projectdiscovery/nuclei-templates)\r\n # 'custom_templates': [] # Nuclei custom templates uploaded in reNgine\r\n }\r\n}\r\nwaf_detection: {\r\n 'enable_http_crawl': true\r\n}\r\nscreenshot: {\r\n 'enable_http_crawl': true,\r\n 'intensity': 'normal',\r\n 'timeout': 10,\r\n 'threads': 40\r\n}", + "yaml_configuration": "# Global vars for all tools\r\n#\r\n# Custom header - FFUF, Nuclei, Dalfox, CRL Fuzz, HTTPx, Fetch URL (Hakrawler, Katana, Gospider)\r\n# custom_header: {\r\n# 'Cookie':'Test',\r\n# 'User-Agent': 'Mozilla/5.0',\r\n# 'Custom-Header': 'My custom header'\r\n# }\r\n# 'user_agent': '' # Dalfox only\r\n# 'enable_http_crawl': true # All tools\r\n# 'timeout': 10 # Subdomain discovery, Screenshot, Port scan, FFUF, Nuclei \r\n# 'threads': 30 # All tools\r\n# 'rate_limit': 150 # Port scan, FFUF, Nuclei\r\n# 'intensity': 'normal' # Screenshot (grab only the root endpoints of each subdomain), Nuclei (reduce number of endpoints to scan), OSINT (not implemented yet)\r\n# 'retries': 1 # Nuclei\r\n\r\nsubdomain_discovery: {\r\n 'uses_tools': ['subfinder', 'ctfr', 'sublist3r', 'tlsx', 'oneforall', 'netlas'], # amass-passive, amass-active, All\r\n 'enable_http_crawl': true,\r\n 'threads': 30,\r\n 'timeout': 5,\r\n # 'use_subfinder_config': false,\r\n # 'use_amass_config': false,\r\n # 'amass_wordlist': 'deepmagic.com-prefixes-top50000'\r\n}\r\nhttp_crawl: {\r\n # 'custom_header': {\r\n # 'Cookie':'Test',\r\n # 'User-Agent': 'Mozilla/5.0'\r\n # },\r\n # 'threads': 30,\r\n # 'follow_redirect': false\r\n}\r\nport_scan: {\r\n 'enable_http_crawl': true,\r\n 'timeout': 5,\r\n # 'exclude_ports': [],\r\n # 'exclude_subdomains': [],\r\n 'ports': ['top-100'],\r\n 'rate_limit': 150,\r\n 'threads': 30,\r\n 'passive': false,\r\n # 'use_naabu_config': false,\r\n # 'enable_nmap': true,\r\n # 'nmap_cmd': '',\r\n # 'nmap_script': '',\r\n # 'nmap_script_args': ''\r\n}\r\nosint: {\r\n 'discover': [\r\n 'emails',\r\n 'metainfo',\r\n 'employees'\r\n ],\r\n 'dorks': [\r\n 'login_pages',\r\n 'admin_panels',\r\n 'dashboard_pages',\r\n 'stackoverflow',\r\n 'social_media',\r\n 'project_management',\r\n 'code_sharing',\r\n 'config_files',\r\n 'jenkins',\r\n 'wordpress_files',\r\n 'php_error',\r\n 'exposed_documents',\r\n 'db_files',\r\n 'git_exposed'\r\n ],\r\n # 'custom_dorks': [],\r\n 'intensity': 'normal',\r\n 'documents_limit': 50\r\n}\r\ndir_file_fuzz: {\r\n # 'custom_header': {\r\n # 'Cookie':'Test',\r\n # 'User-Agent': 'Mozilla/5.0',\r\n # 'Custom-Header': 'My custom header'\r\n # },\r\n 'auto_calibration': true,\r\n 'enable_http_crawl': true,\r\n 'rate_limit': 150,\r\n 'extensions': [],\r\n 'follow_redirect': false,\r\n 'max_time': 0,\r\n 'match_http_status': [200, 204],\r\n 'recursive_level': 0,\r\n 'stop_on_error': false,\r\n 'timeout': 5,\r\n 'threads': 30,\r\n 'wordlist_name': 'default', # fuzz-Bo0oM,\r\n}\r\nfetch_url: {\r\n # 'custom_header': {\r\n # 'Cookie':'Test',\r\n # 'User-Agent': 'Mozilla/5.0',\r\n # 'Custom-Header': 'My custom header'\r\n # },\r\n 'uses_tools': ['gospider', 'hakrawler', 'waybackurls', 'katana', 'gau'],\r\n 'remove_duplicate_endpoints': true,\r\n 'duplicate_fields': ['content_length', 'page_title'],\r\n 'follow_redirect': false,\r\n 'enable_http_crawl': true,\r\n 'gf_patterns': ['debug_logic', 'idor', 'interestingEXT', 'interestingparams', 'interestingsubs', 'lfi', 'rce', 'redirect', 'sqli', 'ssrf', 'ssti', 'xss'],\r\n 'ignore_file_extensions': ['png', 'jpg', 'jpeg', 'gif', 'mp4', 'mpeg', 'mp3'],\r\n 'threads': 30,\r\n # 'exclude_subdomains': false\r\n}\r\nvulnerability_scan: {\r\n # 'custom_header': {\r\n # 'Cookie':'Test',\r\n # 'User-Agent': 'Mozilla/5.0',\r\n # 'Custom-Header': 'My custom header'\r\n # },\r\n 'run_nuclei': true,\r\n 'run_dalfox': false,\r\n 'run_crlfuzz': false,\r\n 'run_s3scanner': false,\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_llm_report': true,\r\n 'nuclei': {\r\n 'use_nuclei_config': false,\r\n 'severities': ['unknown', 'info', 'low', 'medium', 'high', 'critical'],\r\n # 'tags': [], # Nuclei tags (https://github.com/projectdiscovery/nuclei-templates)\r\n # 'templates': [], # Nuclei templates (https://github.com/projectdiscovery/nuclei-templates)\r\n # 'custom_templates': [] # Nuclei custom templates uploaded in reNgine\r\n }\r\n}\r\nwaf_detection: {\r\n 'enable_http_crawl': true\r\n}\r\nscreenshot: {\r\n 'enable_http_crawl': true,\r\n 'intensity': 'normal',\r\n 'timeout': 10,\r\n 'threads': 40\r\n}", "default_engine": false } }, diff --git a/web/scanEngine/views.py b/web/scanEngine/views.py index f61d47a17..4518db361 100644 --- a/web/scanEngine/views.py +++ b/web/scanEngine/views.py @@ -462,7 +462,7 @@ def api_vault(request): "optional": True, "experimental": True, "name": "OpenAI", - "text": "OpenAI keys will be used to generate vulnerability description, remediation, impact and vulnerability report writing using ChatGPT.", + "text": "OpenAI keys will be used to generate vulnerability description, remediation, impact and vulnerability report writing using LLM.", "hasKey": OpenAiAPIKey.objects.first() is not None }, { diff --git a/web/startScan/admin.py b/web/startScan/admin.py index f65ca8e92..af80ed4ac 100644 --- a/web/startScan/admin.py +++ b/web/startScan/admin.py @@ -23,5 +23,5 @@ admin.site.register(Waf) admin.site.register(CountryISO) admin.site.register(Command) -admin.site.register(GPTVulnerabilityReport) +admin.site.register(LLMVulnerabilityReport) admin.site.register(S3Bucket) diff --git a/web/startScan/migrations/0058_auto_20241110_2230.py b/web/startScan/migrations/0058_auto_20241110_2230.py new file mode 100644 index 000000000..d23087c94 --- /dev/null +++ b/web/startScan/migrations/0058_auto_20241110_2230.py @@ -0,0 +1,22 @@ +# Generated by Django 3.2.25 on 2024-11-10 22:30 + +from django.db import migrations + + +class Migration(migrations.Migration): + + dependencies = [ + ('startScan', '0057_auto_20231201_2354'), + ] + + operations = [ + migrations.RenameModel( + old_name='GPTVulnerabilityReport', + new_name='LLMVulnerabilityReport', + ), + migrations.RenameField( + model_name='vulnerability', + old_name='is_gpt_used', + new_name='is_llm_used', + ), + ] diff --git a/web/startScan/models.py b/web/startScan/models.py index f49c310e0..d1a6dee0c 100644 --- a/web/startScan/models.py +++ b/web/startScan/models.py @@ -411,7 +411,7 @@ def __str__(self): return self.name -class GPTVulnerabilityReport(models.Model): +class LLMVulnerabilityReport(models.Model): url_path = models.CharField(max_length=2000) title = models.CharField(max_length=2500) description = models.TextField(null=True, blank=True) @@ -468,7 +468,7 @@ class Vulnerability(models.Model): hackerone_report_id = models.CharField(max_length=50, null=True, blank=True) request = models.TextField(blank=True, null=True) response = models.TextField(blank=True, null=True) - is_gpt_used = models.BooleanField(null=True, blank=True, default=False) + is_llm_used = models.BooleanField(null=True, blank=True, default=False) # used for subscans vuln_subscan_ids = models.ManyToManyField('SubScan', related_name='vuln_subscan_ids', blank=True) diff --git a/web/startScan/templates/startScan/detail_scan.html b/web/startScan/templates/startScan/detail_scan.html index 67d725c50..b275fa487 100644 --- a/web/startScan/templates/startScan/detail_scan.html +++ b/web/startScan/templates/startScan/detail_scan.html @@ -1730,7 +1730,7 @@

@@ -2373,7 +2373,7 @@

- + diff --git a/web/startScan/templates/startScan/subdomains.html b/web/startScan/templates/startScan/subdomains.html index d764881f3..5494753fa 100644 --- a/web/startScan/templates/startScan/subdomains.html +++ b/web/startScan/templates/startScan/subdomains.html @@ -284,7 +284,7 @@ const cms_detector_http_url = row['http_url'] ? row['http_url'] : 'https://' + row['name']; return `
- + diff --git a/web/startScan/templates/startScan/vulnerabilities.html b/web/startScan/templates/startScan/vulnerabilities.html index 76717a0bd..4b25d8c2b 100644 --- a/web/startScan/templates/startScan/vulnerabilities.html +++ b/web/startScan/templates/startScan/vulnerabilities.html @@ -215,7 +215,7 @@ diff --git a/web/static/custom/custom.js b/web/static/custom/custom.js index 0b36cccf0..8b62bbfb9 100644 --- a/web/static/custom/custom.js +++ b/web/static/custom/custom.js @@ -3118,8 +3118,8 @@ function render_vuln_offcanvas(vuln){
`; } - if (vuln.is_gpt_used) { - body += `(GPT was used to generate vulnerability details.)`; + if (vuln.is_llm_used) { + body += `(LLM was used to generate vulnerability details.)`; } @@ -3143,7 +3143,7 @@ function showSwalLoader(title, text){ }); } -async function send_gpt_api_request(endpoint_url, vuln_id){ +async function send_llm_api_request(endpoint_url, vuln_id){ const api = `${endpoint_url}?format=json&id=${vuln_id}`; try { const response = await fetch(api, { @@ -3156,23 +3156,22 @@ async function send_gpt_api_request(endpoint_url, vuln_id){ if (!response.ok) { throw new Error('Request failed'); } - const data = await response.json(); - return data; + return await response.json(); } catch (error) { throw new Error('Request failed'); } } -async function fetch_gpt_vuln_details(endpoint_url, id, title) { +async function fetch_llm_vuln_details(endpoint_url, id, title) { var loader_title = "Loading..."; - var text = 'Please wait while the GPT is generating vulnerability description.' + var text = 'Please wait while the LLM is generating vulnerability description.' try { showSwalLoader(loader_title, text); - const data = await send_gpt_api_request(endpoint_url, id); + const data = await send_llm_api_request(endpoint_url, id); Swal.close(); if (data.status) { - render_gpt_vuln_modal(data, title); + render_llm_vuln_modal(data, title); } else{ Swal.close(); @@ -3194,7 +3193,7 @@ async function fetch_gpt_vuln_details(endpoint_url, id, title) { } -function render_gpt_vuln_modal(data, title){ +function render_llm_vuln_modal(data, title){ $('#modal_dialog .modal-title').empty(); $('#modal_dialog .modal-text').empty(); $('#modal_dialog .modal-footer').empty(); @@ -3250,7 +3249,7 @@ function endpoint_datatable_col_visibility(endpoint_table){ } -async function send_gpt__attack_surface_api_request(endpoint_url, subdomain_id){ +async function send_llm__attack_surface_api_request(endpoint_url, subdomain_id){ const api = `${endpoint_url}?format=json&subdomain_id=${subdomain_id}`; try { const response = await fetch(api, { @@ -3273,10 +3272,10 @@ async function send_gpt__attack_surface_api_request(endpoint_url, subdomain_id){ async function show_attack_surface_modal(endpoint_url, id){ var loader_title = "Loading..."; - var text = 'Please wait while the GPT is generating attack surface.' + var text = 'Please wait while the LLM is generating attack surface.' try { showSwalLoader(loader_title, text); - const data = await send_gpt__attack_surface_api_request(endpoint_url,id); + const data = await send_llm__attack_surface_api_request(endpoint_url,id); Swal.close(); if (data.status) { $('#modal_dialog .modal-title').html(`Attack Surface Suggestion for ${data.subdomain_name} (BETA)`); diff --git a/web/static/custom/vuln_datatables.js b/web/static/custom/vuln_datatables.js index 10f9ec4c3..de06abd32 100644 --- a/web/static/custom/vuln_datatables.js +++ b/web/static/custom/vuln_datatables.js @@ -33,7 +33,7 @@ const vuln_datatable_columns = [ {'data': 'template_id'}, {'data': 'impact'}, {'data': 'remediation'}, - {'data': 'is_gpt_used'}, + {'data': 'is_llm_used'}, ]; const vuln_datatable_page_length = 50; diff --git a/web/targetApp/templates/target/summary.html b/web/targetApp/templates/target/summary.html index 0eac1dfc3..2807360a9 100644 --- a/web/targetApp/templates/target/summary.html +++ b/web/targetApp/templates/target/summary.html @@ -1150,7 +1150,7 @@

- + @@ -1534,7 +1534,7 @@

From 3ad4ea6cb2bd0a1a1882f5663a859ca1ca089c90 Mon Sep 17 00:00:00 2001 From: psyray Date: Mon, 11 Nov 2024 07:14:42 +0100 Subject: [PATCH 02/21] feat: enhance LLM model selection and attack surface analysis - Introduced a new modal for selecting LLM models for attack surface analysis, allowing users to choose from a list of available models. - Added functionality to update the selected model in the database before proceeding with analysis. - Implemented a new API endpoint to fetch available LLM models and the currently selected model. - Refactored the show_attack_surface_modal function to include model selection and handle errors more robustly. - Removed hardcoded LLM definitions and prompts from definitions.py and moved them to a new configuration file. - Added new classes and methods for managing LLM configurations and generating vulnerability reports and attack suggestions. - Updated tests to cover new functionalities and removed obsolete tests related to LLM attack suggestions. --- web/api/tests/test_project.py | 23 --- web/api/urls.py | 4 + web/api/views.py | 68 +++++++- web/reNgine/definitions.py | 90 ----------- web/reNgine/llm.py | 134 --------------- web/reNgine/llm/config.py | 297 ++++++++++++++++++++++++++++++++++ web/reNgine/llm/llm.py | 213 ++++++++++++++++++++++++ web/reNgine/llm/utils.py | 69 ++++++++ web/reNgine/llm/validators.py | 51 ++++++ web/reNgine/tasks.py | 3 +- web/scanEngine/views.py | 2 +- web/static/custom/custom.js | 289 +++++++++++++++++++++++++++++---- web/tests/test_llm.py | 129 +++++++++++++++ 13 files changed, 1085 insertions(+), 287 deletions(-) delete mode 100644 web/reNgine/llm.py create mode 100644 web/reNgine/llm/config.py create mode 100644 web/reNgine/llm/llm.py create mode 100644 web/reNgine/llm/utils.py create mode 100644 web/reNgine/llm/validators.py create mode 100644 web/tests/test_llm.py diff --git a/web/api/tests/test_project.py b/web/api/tests/test_project.py index d05a96840..276a022dc 100644 --- a/web/api/tests/test_project.py +++ b/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', - 'TestLLMAttackSuggestion' ] class TestCreateProjectApi(BaseTestCase): @@ -108,24 +106,3 @@ def test_list_todo_notes(self): self.data_generator.todo_note.scan_history.id, ) -class TestLLMAttackSuggestion(BaseTestCase): - """Tests for the LLM Attack Suggestion API.""" - - def setUp(self): - super().setUp() - self.data_generator.create_project_base() - - @patch("reNgine.llm.LLMAttackSuggestionGenerator.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:llm_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") diff --git a/web/api/urls.py b/web/api/urls.py index 703120538..b3503c09c 100644 --- a/web/api/urls.py +++ b/web/api/urls.py @@ -166,6 +166,10 @@ '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(), diff --git a/web/api/views.py b/web/api/views.py index 1ee6c3298..bba54c0cd 100644 --- a/web/api/views.py +++ b/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,16 +31,20 @@ 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 @@ -57,7 +62,7 @@ run_wafw00f, send_hackerone_report ) -from reNgine.llm import LLMAttackSuggestionGenerator +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,7 +174,6 @@ def put(self, request): defaults={ 'selected_model': model_name, 'use_ollama': use_ollama, - 'selected': True } ) return Response({'status': True}) @@ -179,7 +183,7 @@ def put(self, request): 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({ @@ -2892,3 +2896,55 @@ def special_lookup(self, search_value): 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) diff --git a/web/reNgine/definitions.py b/web/reNgine/definitions.py index ee0fee5a6..d37f7cbab 100644 --- a/web/reNgine/definitions.py +++ b/web/reNgine/definitions.py @@ -437,95 +437,5 @@ # 404 page url FOUR_OH_FOUR_URL = '/404/' - -############################################################################### -# OLLAMA DEFINITIONS -############################################################################### -OLLAMA_INSTANCE = 'http://ollama:11434' - -DEFAULT_GPT_MODELS = [ - { - 'name': 'gpt-3', - 'model': 'gpt-3', - 'modified_at': '', - 'details': { - 'family': 'GPT', - 'parameter_size': '~175B', - } - }, - { - 'name': 'gpt-3.5-turbo', - 'model': 'gpt-3.5-turbo', - 'modified_at': '', - 'details': { - 'family': 'GPT', - 'parameter_size': '~7B', - } - }, - { - 'name': 'gpt-4', - 'model': 'gpt-4', - 'modified_at': '', - 'details': { - 'family': 'GPT', - 'parameter_size': '~1.7T', - } - }, - { - 'name': 'gpt-4-turbo', - 'model': 'gpt-4', - 'modified_at': '', - 'details': { - 'family': 'GPT', - 'parameter_size': '~1.7T', - } - } -] - - - -# LLM Vulnerability Report Generator -VULNERABILITY_DESCRIPTION_SYSTEM_MESSAGE = """ - You are a highly skilled penetration tester who has recently completed a penetration testing. - You will be given with a - - Vulnerability title - - Vulnerable URL - - and some description about the vulnerability. - Your job is to write a detailed technical penetration testing report based on the given Vulnerability details. - The purpose of this report is to provide an in-depth analysis of the vulnerabilities discovered during the penetration testing engagement. - - The penetration testing report must contain all separated by \n\n - - - Vulnerability description - Include a detailed vulnerability description, include any known CVE IDs, any known existing vulnerabilities. - - Impact - Include what this vulnerability can impact for web applications. - - Remediation - Include steps to remediate this vulnerability. Separate each new remediations by - and a new line \n - - References - Include any references URL about this vulnerability, any existing CVE ID, or news articles etc. Separate each new references by - and a new line \n. Only include http urls - - Do not write 'Penetration Testing Report:' on the title. -""" - - -ATTACK_SUGGESTION_LLM_SYSTEM_PROMPT = """ - You are a highly skilled penetration tester who has recently completed a reconnaissance on a target. - As a penetration tester, I've conducted a thorough reconnaissance on a specific subdomain. - Based on my reconnaissance you will be given with a - - Subdomain Name - - Subdomain Page Title - - Open Ports if any detected - - HTTP Status - - Technologies Detected - - Content Type - - Web Server - - Page Content Length - I'm seeking insights into potential technical web application attacks that could be executed on this subdomain, along with explanations for why these attacks are feasible given the discovered information. - Please provide a detailed list of these attack types and their underlying technical rationales on every attacks you suggested. - Also suggest if any CVE ID, known exploits, existing vulnerabilities, any news articles URL related to the information provided to you. -""" - - # OSINT GooFuzz Path GOFUZZ_EXEC_PATH = 'GooFuzz' diff --git a/web/reNgine/llm.py b/web/reNgine/llm.py deleted file mode 100644 index d6e40074d..000000000 --- a/web/reNgine/llm.py +++ /dev/null @@ -1,134 +0,0 @@ -import openai -import re -from reNgine.common_func import get_open_ai_key, extract_between -from reNgine.definitions import VULNERABILITY_DESCRIPTION_SYSTEM_MESSAGE, ATTACK_SUGGESTION_LLM_SYSTEM_PROMPT, OLLAMA_INSTANCE -from langchain_community.llms import Ollama - -from dashboard.models import OllamaSettings -import logging - -logger = logging.getLogger(__name__) - -class LLMVulnerabilityReportGenerator: - - def __init__(self): - selected_model = OllamaSettings.objects.first() - self.model_name = selected_model.selected_model if selected_model else 'gpt-3.5-turbo' - self.use_ollama = selected_model.use_ollama if selected_model else False - self.openai_api_key = None - self.ollama = None - - def get_vulnerability_description(self, description): - """Generate Vulnerability Description using LLM. - - Args: - description (str): Vulnerability Description message to pass to LLM. - - Returns: - (dict) of { - 'description': (str) - 'impact': (str), - 'remediation': (str), - 'references': (list) of urls - } - """ - print(f"Generating Vulnerability Description for: {description}") - if self.use_ollama: - prompt = VULNERABILITY_DESCRIPTION_SYSTEM_MESSAGE + "\nUser: " + description - self.ollama = Ollama( - base_url=OLLAMA_INSTANCE, - model=self.model_name - ) - response_content = self.ollama(prompt) - else: - openai_api_key = get_open_ai_key() - if not openai_api_key: - return { - 'status': False, - 'error': 'OpenAI API Key not set' - } - try: - openai.api_key = openai_api_key - llm_response = openai.ChatCompletion.create( - model=self.model_name, - messages=[ - {'role': 'system', 'content': VULNERABILITY_DESCRIPTION_SYSTEM_MESSAGE}, - {'role': 'user', 'content': description} - ] - ) - - response_content = llm_response['choices'][0]['message']['content'] - except Exception as e: - return { - 'status': False, - 'error': str(e) - } - vuln_description_pattern = re.compile( - r"[Vv]ulnerability [Dd]escription:(.*?)(?:\n\n[Ii]mpact:|$)", - re.DOTALL - ) - impact_pattern = re.compile( - r"[Ii]mpact:(.*?)(?:\n\n[Rr]emediation:|$)", - re.DOTALL - ) - remediation_pattern = re.compile( - r"[Rr]emediation:(.*?)(?:\n\n[Rr]eferences:|$)", - re.DOTALL - ) - - description_section = extract_between(response_content, vuln_description_pattern) - impact_section = extract_between(response_content, impact_pattern) - remediation_section = extract_between(response_content, remediation_pattern) - references_start_index = response_content.find("References:") - references_section = response_content[references_start_index + len("References:"):].strip() - - url_pattern = re.compile(r'https://\S+') - urls = url_pattern.findall(references_section) - - return { - 'status': True, - 'description': description_section, - 'impact': impact_section, - 'remediation': remediation_section, - 'references': urls, - } - -class LLMAttackSuggestionGenerator: - - def __init__(self): - self.api_key = get_open_ai_key() - self.model_name = 'gpt-3.5-turbo' - if not self.api_key: - self.ollama = Ollama(base_url='http://ollama:11434', model="llama2-uncensored") - - def get_attack_suggestion(self, input): - ''' - input (str): input for llm - ''' - try: - if not self.api_key: - prompt = ATTACK_SUGGESTION_LLM_SYSTEM_PROMPT + "\nUser: " + input - response_content = self.ollama(prompt) - else: - openai.api_key = self.api_key - llm_response = openai.ChatCompletion.create( - model=self.model_name, - messages=[ - {'role': 'system', 'content': ATTACK_SUGGESTION_LLM_SYSTEM_PROMPT}, - {'role': 'user', 'content': input} - ] - ) - response_content = llm_response['choices'][0]['message']['content'] - - return { - 'status': True, - 'description': response_content, - 'input': input - } - except ValueError as e: - logger.error("Error in get_attack_suggestion: %s", str(e), exc_info=True) - return { - 'status': False, - 'error': "An error occurred while processing your request.", - 'input': input - } diff --git a/web/reNgine/llm/config.py b/web/reNgine/llm/config.py new file mode 100644 index 000000000..df18d4af3 --- /dev/null +++ b/web/reNgine/llm/config.py @@ -0,0 +1,297 @@ +from typing import Dict, Any + +############################################################################### +# OLLAMA DEFINITIONS +############################################################################### + +OLLAMA_INSTANCE = 'http://ollama:11434' + +############################################################################### +# LLM SYSTEM PROMPTS +############################################################################### + +VULNERABILITY_DESCRIPTION_SYSTEM_MESSAGE = """ +You are an expert penetration tester specializing in web application security assessments. Your task is to analyze the following vulnerability information: + - Vulnerability title + - Vulnerable URL + - Vulnerability description + +Required report sections (separate each with \n\n): + +1. TECHNICAL DESCRIPTION + - Detailed technical explanation of the vulnerability + - Associated CVE IDs and CVSS scores if applicable + - Attack vectors and exploitation methods + - Any prerequisites or conditions required for exploitation + +2. BUSINESS IMPACT + - Direct security implications + - Potential business consequences + - Data exposure risks + - Compliance implications + +3. REMEDIATION STEPS + - Provide specific, actionable remediation steps + - Include code examples where relevant + - List configuration changes if needed + - Suggest security controls to prevent similar issues + Format: Each step prefixed with "- " on a new line + +4. REFERENCES + - Only include validated HTTP/HTTPS URLs + - Focus on official documentation, security advisories, and research papers + - Include relevant CVE details and exploit databases + Format: Each reference prefixed with "- " on a new line + +Keep the tone technical and professional. Focus on actionable insights. Avoid generic statements. +""" + +ATTACK_SUGGESTION_LLM_SYSTEM_PROMPT = """ +You are an advanced penetration tester specializing in web application security. Based on the reconnaissance data provided: + - Subdomain Name + - Page Title + - Open Ports + - HTTP Status + - Technologies Stack + - Content Type + - Web Server + - Content Length + +Provide a structured analysis in the following format: + +1. ATTACK SURFACE ANALYSIS + - Enumerate potential entry points + - Identify technology-specific vulnerabilities + - List version-specific known vulnerabilities + - Map attack surface to MITRE ATT&CK framework where applicable + +2. PRIORITIZED ATTACK VECTORS + For each suggested attack: + - Attack name and classification + - Technical rationale based on observed data + - Specific exploitation methodology + - Success probability assessment + - Potential impact rating + +3. RELEVANT SECURITY CONTEXT + - CVE IDs with CVSS scores + - Existing proof-of-concept exploits + - Recent security advisories + - Relevant threat intelligence + Only include verified HTTP/HTTPS URLs + +Focus on actionable, evidence-based suggestions. Prioritize attacks based on feasibility and impact. +Avoid theoretical attacks without supporting evidence from the reconnaissance data. +""" + +############################################################################### +# LLM CONFIGURATION +############################################################################### + +LLM_CONFIG: Dict[str, Any] = { + 'providers': { + 'openai': { + 'default_model': 'gpt-4-turbo-preview', + 'models': [ + 'gpt-4-turbo-preview', + 'gpt-4', + 'gpt-3.5-turbo' + ], + 'api_version': '2024-02-15', + 'max_tokens': 2000, + 'temperature': 0.7, + }, + 'ollama': { + 'default_model': 'llama2', + 'models': [ + 'llama2', + 'mistral', + 'codellama', + 'gemma' + ], + 'timeout': 30, + 'max_retries': 3, + } + }, + 'ollama_url': OLLAMA_INSTANCE, + 'timeout': 30, + 'max_retries': 3, + 'prompts': { + 'vulnerability': VULNERABILITY_DESCRIPTION_SYSTEM_MESSAGE, + 'attack': ATTACK_SUGGESTION_LLM_SYSTEM_PROMPT + } +} + +############################################################################### +# DEFAULT GPT MODELS +############################################################################### + +DEFAULT_GPT_MODELS = [ + { + 'name': 'gpt-3', + 'model': 'gpt-3', + 'modified_at': '', + 'details': { + 'family': 'GPT', + 'parameter_size': '~175B', + } + }, + { + 'name': 'gpt-3.5-turbo', + 'model': 'gpt-3.5-turbo', + 'modified_at': '', + 'details': { + 'family': 'GPT', + 'parameter_size': '~7B', + } + }, + { + 'name': 'gpt-4', + 'model': 'gpt-4', + 'modified_at': '', + 'details': { + 'family': 'GPT', + 'parameter_size': '~1.7T', + } + }, + { + 'name': 'gpt-4-turbo', + 'model': 'gpt-4', + 'modified_at': '', + 'details': { + 'family': 'GPT', + 'parameter_size': '~1.7T', + } + } +] + +############################################################################### +# MODEL CAPABILITIES +############################################################################### + +MODEL_REQUIREMENTS = { + # OpenAI Models + 'gpt-3': { + 'min_tokens': 64, + 'max_tokens': 2048, + 'supports_functions': True, + 'best_for': ['basic_analysis', 'general_purpose'], + 'provider': 'openai' + }, + 'gpt-3.5-turbo': { + 'min_tokens': 64, + 'max_tokens': 4096, + 'supports_functions': True, + 'best_for': ['quick_analysis', 'basic_suggestions', 'cost_effective'], + 'provider': 'openai' + }, + 'gpt-4': { + 'min_tokens': 128, + 'max_tokens': 8192, + 'supports_functions': True, + 'best_for': ['deep_analysis', 'complex_reasoning', 'advanced_security'], + 'provider': 'openai' + }, + 'gpt-4-turbo': { + 'min_tokens': 128, + 'max_tokens': 128000, + 'supports_functions': True, + 'best_for': ['complex_analysis', 'technical_details', 'latest_capabilities'], + 'provider': 'openai' + }, + + # Llama Family Models + 'llama2': { + 'min_tokens': 32, + 'max_tokens': 4096, + 'supports_functions': False, + 'best_for': ['local_processing', 'privacy_focused', 'balanced_performance'], + 'provider': 'ollama' + }, + 'llama2-uncensored': { + 'min_tokens': 32, + 'max_tokens': 4096, + 'supports_functions': False, + 'best_for': ['unfiltered_analysis', 'security_research', 'red_teaming'], + 'provider': 'ollama' + }, + 'llama3': { + 'min_tokens': 64, + 'max_tokens': 8192, + 'supports_functions': False, + 'best_for': ['advanced_reasoning', 'improved_context', 'technical_analysis'], + 'provider': 'ollama' + }, + 'llama3.1': { + 'min_tokens': 64, + 'max_tokens': 8192, + 'supports_functions': False, + 'best_for': ['enhanced_comprehension', 'security_assessment', 'detailed_analysis'], + 'provider': 'ollama' + }, + 'llama3.2': { + 'min_tokens': 64, + 'max_tokens': 16384, + 'supports_functions': False, + 'best_for': ['long_context', 'complex_security_analysis', 'advanced_reasoning'], + 'provider': 'ollama' + }, + + # Other Specialized Models + 'mistral': { + 'min_tokens': 32, + 'max_tokens': 8192, + 'supports_functions': False, + 'best_for': ['efficient_processing', 'technical_analysis', 'good_performance_ratio'], + 'provider': 'ollama' + }, + 'mistral-medium': { + 'min_tokens': 32, + 'max_tokens': 8192, + 'supports_functions': False, + 'best_for': ['balanced_analysis', 'improved_accuracy', 'technical_tasks'], + 'provider': 'ollama' + }, + 'mistral-large': { + 'min_tokens': 64, + 'max_tokens': 16384, + 'supports_functions': False, + 'best_for': ['deep_technical_analysis', 'complex_reasoning', 'high_accuracy'], + 'provider': 'ollama' + }, + 'codellama': { + 'min_tokens': 32, + 'max_tokens': 4096, + 'supports_functions': False, + 'best_for': ['code_analysis', 'vulnerability_assessment', 'technical_details'], + 'provider': 'ollama' + }, + 'qwen2.5': { + 'min_tokens': 64, + 'max_tokens': 8192, + 'supports_functions': False, + 'best_for': ['multilingual_analysis', 'efficient_processing', 'technical_understanding'], + 'provider': 'ollama' + }, + 'gemma': { + 'min_tokens': 32, + 'max_tokens': 4096, + 'supports_functions': False, + 'best_for': ['lightweight_analysis', 'quick_assessment', 'general_tasks'], + 'provider': 'ollama' + }, + 'solar': { + 'min_tokens': 64, + 'max_tokens': 8192, + 'supports_functions': False, + 'best_for': ['creative_analysis', 'unique_perspectives', 'alternative_approaches'], + 'provider': 'ollama' + }, + 'yi': { + 'min_tokens': 64, + 'max_tokens': 8192, + 'supports_functions': False, + 'best_for': ['comprehensive_analysis', 'detailed_explanations', 'technical_depth'], + 'provider': 'ollama' + } +} \ No newline at end of file diff --git a/web/reNgine/llm/llm.py b/web/reNgine/llm/llm.py new file mode 100644 index 000000000..c8bbef81e --- /dev/null +++ b/web/reNgine/llm/llm.py @@ -0,0 +1,213 @@ +from typing import Optional, Dict, Any +import logging +from abc import ABC, abstractmethod +import openai +from langchain_community.llms import Ollama + +from reNgine.llm.config import LLM_CONFIG +from reNgine.llm.utils import RegexPatterns, get_default_llm_model +from reNgine.llm.validators import LLMProvider, LLMInputData, LLMResponse +from reNgine.common_func import get_open_ai_key, extract_between + +logger = logging.getLogger(__name__) + +class BaseLLMGenerator(ABC): + """Base class for LLM generators with common functionality""" + + def __init__(self, provider: Optional[LLMProvider] = None): + """Initialize the LLM generator with optional provider""" + self.api_key = get_open_ai_key() + self.config = LLM_CONFIG + self.model_name = self._get_model_name() + self.provider = provider or self._get_default_provider() + self.ollama = None + + if self.provider == LLMProvider.OLLAMA: + self._setup_ollama() + + @abstractmethod + def _get_model_name(self) -> str: + """Get the model name to use""" + pass + + @abstractmethod + def _get_default_provider(self) -> LLMProvider: + """Get the default provider based on configuration""" + pass + + def _setup_ollama(self) -> None: + """Setup Ollama client with configuration""" + ollama_config = self.config['providers']['ollama'] + self.ollama = Ollama( + base_url=self.config['ollama_url'], + model=self.model_name, + timeout=ollama_config['timeout'] + ) + + def _validate_input(self, input_data: str) -> LLMInputData: + """Validate input data using Pydantic model""" + return LLMInputData( + description=input_data, + llm_model=self.model_name, + provider=self.provider + ) + +class LLMVulnerabilityReportGenerator(BaseLLMGenerator): + """Generator for vulnerability reports using LLM""" + + def _get_model_name(self) -> str: + """Get model name from database or default""" + return get_default_llm_model() + + def _get_default_provider(self) -> LLMProvider: + """Get default provider based on OllamaSettings""" + selected_model = OllamaSettings.objects.first() + return LLMProvider.OLLAMA if selected_model.use_ollama else LLMProvider.OPENAI + + def get_vulnerability_description(self, description: str) -> LLMResponse: + """ + Generate vulnerability description using LLM + + Args: + description: Raw vulnerability description + + Returns: + LLMResponse object containing the structured response + """ + try: + # Validate input + validated_input = self._validate_input(description) + + # Get response from appropriate provider + if self.provider == LLMProvider.OLLAMA: + response_content = self._get_ollama_response(validated_input.description) + else: + response_content = self._get_openai_response(validated_input.description) + + # Extract sections from response + sections = self._extract_sections(response_content) + return LLMResponse( + status=True, + **sections + ) + + except Exception as e: + logger.error(f"Error in get_vulnerability_description: {str(e)}", exc_info=True) + return LLMResponse( + status=False, + error=str(e) + ) + + def _get_ollama_response(self, description: str) -> str: + """Get response from Ollama""" + prompt = f"{self.config['prompts']['vulnerability']}\nUser: {description}" + return self.ollama(prompt) + + def _get_openai_response(self, description: str) -> str: + """Get response from OpenAI""" + if not self.api_key: + raise ValueError("OpenAI API Key not set") + + openai.api_key = self.api_key + + response = openai.ChatCompletion.create( + model=self.model_name, + messages=[ + {'role': 'system', 'content': self.config['prompts']['vulnerability']}, + {'role': 'user', 'content': description} + ], + **self.config['providers']['openai'] + ) + return response['choices'][0]['message']['content'] + + def _extract_sections(self, content: str) -> dict: + """Extract sections from LLM response""" + description = extract_between(content, RegexPatterns.VULN_DESCRIPTION) + impact = extract_between(content, RegexPatterns.IMPACT) + remediation = extract_between(content, RegexPatterns.REMEDIATION) + + references_start = content.find("References:") + references_section = content[references_start + len("References:"):].strip() + urls = RegexPatterns.URL.findall(references_section) + + return { + "description": description, + "impact": impact, + "remediation": remediation, + "references": urls + } + +class LLMAttackSuggestionGenerator(BaseLLMGenerator): + """Generator for attack suggestions using LLM""" + + def _get_model_name(self) -> str: + """Get model name from database or default""" + return get_default_llm_model() + + def _get_default_provider(self) -> LLMProvider: + """Get default provider based on model requirements""" + model_name = self._get_model_name() + if model_name in self.config['providers']['openai']['models']: + return LLMProvider.OPENAI + return LLMProvider.OLLAMA + + def _get_provider_config(self) -> Dict[str, Any]: + """Get provider specific configuration""" + provider_key = self.provider.value + return self.config['providers'][provider_key] + + def get_attack_suggestion(self, input_data: str) -> dict: + """ + Generate attack suggestions using LLM + + Args: + input_data: Reconnaissance data + + Returns: + dict: Response containing status and description + """ + try: + # Validate input + validated_input = self._validate_input(input_data) + + # Get response from appropriate provider + if self.provider == LLMProvider.OLLAMA: + response_content = self._get_ollama_response(validated_input.description) + else: + response_content = self._get_openai_response(validated_input.description) + + return { + 'status': True, + 'description': response_content, + 'input': input_data + } + + except Exception as e: + logger.error(f"Error in get_attack_suggestion: {str(e)}", exc_info=True) + return { + 'status': False, + 'error': str(e), + 'input': input_data + } + + def _get_ollama_response(self, description: str) -> str: + """Get response from Ollama""" + prompt = f"{self.config['prompts']['attack']}\nUser: {description}" + return self.ollama(prompt) + + def _get_openai_response(self, description: str) -> str: + """Get response from OpenAI""" + if not self.api_key: + raise ValueError("OpenAI API Key not set") + + openai.api_key = self.api_key + + response = openai.ChatCompletion.create( + model=self.model_name, + messages=[ + {'role': 'system', 'content': self.config['prompts']['attack']}, + {'role': 'user', 'content': description} + ], + **self._get_provider_config() + ) + return response['choices'][0]['message']['content'] diff --git a/web/reNgine/llm/utils.py b/web/reNgine/llm/utils.py new file mode 100644 index 000000000..4b60db915 --- /dev/null +++ b/web/reNgine/llm/utils.py @@ -0,0 +1,69 @@ +from django.contrib import messages +from dashboard.models import OllamaSettings +from reNgine.llm.config import LLM_CONFIG +import logging +import re + +logger = logging.getLogger(__name__) + +def get_default_llm_model(): + """ + Get the default LLM model from database or fallback to default + Returns the model name as string + """ + try: + ollama_settings = OllamaSettings.objects.first() + if ollama_settings and ollama_settings.selected_model: + return ollama_settings.selected_model + except Exception as e: + logger.error(f"Error while retrieving default LLM model: {e}") + + # Fallback to default model from config based on provider + try: + if ollama_settings and ollama_settings.use_ollama: + return LLM_CONFIG['providers']['ollama']['default_model'] + return LLM_CONFIG['providers']['openai']['default_model'] + except Exception as e: + logger.error(f"Error while getting default model from config: {e}") + return 'gpt-3.5-turbo' # Ultimate fallback + +def validate_llm_model(request, model_name): + """Check if LLM model exists and is available""" + try: + # Check if model exists in LLMToolkit + if not LLMToolkit.is_model_available(model_name): + messages.info( + request, + f"Model {model_name} is not available. " + f'Configure your LLM models here.', + extra_tags='safe' + ) + return False + return True + except Exception as e: + logger.error(f"Error while validating LLM model: {e}") + return False + +class RegexPatterns: + """Regular expression patterns for parsing LLM responses""" + + VULN_DESCRIPTION = re.compile( + r"[Vv]ulnerability [Dd]escription:(.*?)(?:\n\n[Ii]mpact:|$)", + re.DOTALL + ) + + IMPACT = re.compile( + r"[Ii]mpact:(.*?)(?:\n\n[Rr]emediation:|$)", + re.DOTALL + ) + + REMEDIATION = re.compile( + r"[Rr]emediation:(.*?)(?:\n\n[Rr]eferences:|$)", + re.DOTALL + ) + + URL = re.compile(r'https?://\S+') + + # Add other patterns if needed + CVE = re.compile(r'CVE-\d{4}-\d{4,7}') + CVSS = re.compile(r'CVSS:3\.\d/AV:[NALP]/AC:[LH]/PR:[NLH]/UI:[NR]/S:[UC]/C:[NLH]/I:[NLH]/A:[NLH]') diff --git a/web/reNgine/llm/validators.py b/web/reNgine/llm/validators.py new file mode 100644 index 000000000..5eb7aaa3c --- /dev/null +++ b/web/reNgine/llm/validators.py @@ -0,0 +1,51 @@ +from typing import Optional, List, Dict +from pydantic import BaseModel, Field, validator +from enum import Enum + +class LLMProvider(str, Enum): + OPENAI = "openai" + OLLAMA = "ollama" + +class ModelCapabilities(BaseModel): + min_tokens: int + max_tokens: int + supports_functions: bool + best_for: List[str] + provider: str + +class LLMInputData(BaseModel): + description: str + llm_model: Optional[str] = Field(default=None) + provider: Optional[LLMProvider] = Field(default=None) + capabilities: Optional[ModelCapabilities] = Field(default=None) + + @validator('description') + def validate_description(cls, v): + if not v or len(v.strip()) < 10: + raise ValueError("Description must be at least 10 characters long") + return v.strip() + + class Config: + json_schema_extra = { + "example": { + "description": "SQL Injection vulnerability found in login form", + "llm_model": "gpt-3.5-turbo", + "provider": "openai", + "capabilities": { + "min_tokens": 64, + "max_tokens": 2048, + "supports_functions": True, + "best_for": ["quick_analysis"], + "provider": "openai" + } + } + } + +class LLMResponse(BaseModel): + status: bool + description: Optional[str] = None + impact: Optional[str] = None + remediation: Optional[str] = None + references: Optional[List[str]] = None + error: Optional[str] = None + input: Optional[str] = None \ No newline at end of file diff --git a/web/reNgine/tasks.py b/web/reNgine/tasks.py index dd1d93dc9..f6f4fcee6 100644 --- a/web/reNgine/tasks.py +++ b/web/reNgine/tasks.py @@ -28,12 +28,11 @@ from metafinder.extractor import extract_metadata_from_google_search from reNgine.celery import app -from reNgine.llm import LLMVulnerabilityReportGenerator +from reNgine.llm.llm import LLMVulnerabilityReportGenerator from reNgine.celery_custom_task import RengineTask from reNgine.common_func import * from reNgine.definitions import * from reNgine.settings import * -from reNgine.llm import * from reNgine.utilities import * from scanEngine.models import (EngineType, InstalledExternalTool, Notification, Proxy) from startScan.models import * diff --git a/web/scanEngine/views.py b/web/scanEngine/views.py index 4518db361..72d500c36 100644 --- a/web/scanEngine/views.py +++ b/web/scanEngine/views.py @@ -12,7 +12,7 @@ from rolepermissions.decorators import has_permission_decorator from reNgine.common_func import get_open_ai_key -from reNgine.definitions import OLLAMA_INSTANCE, DEFAULT_GPT_MODELS +from reNgine.llm.config import OLLAMA_INSTANCE, DEFAULT_GPT_MODELS from reNgine.tasks import run_command, send_discord_message, send_slack_message, send_lark_message, send_telegram_message, run_gf_list from scanEngine.forms import AddEngineForm, UpdateEngineForm, AddWordlistForm, ExternalToolForm, InterestingLookupForm, NotificationForm, ProxyForm, HackeroneForm, ReportForm from scanEngine.models import EngineType, Wordlist, InstalledExternalTool, InterestingLookupModel, Notification, Hackerone, Proxy, VulnerabilityReportSetting diff --git a/web/static/custom/custom.js b/web/static/custom/custom.js index 8b62bbfb9..e4419b3fd 100644 --- a/web/static/custom/custom.js +++ b/web/static/custom/custom.js @@ -3270,39 +3270,190 @@ async function send_llm__attack_surface_api_request(endpoint_url, subdomain_id){ } -async function show_attack_surface_modal(endpoint_url, id){ - var loader_title = "Loading..."; - var text = 'Please wait while the LLM is generating attack surface.' - try { - showSwalLoader(loader_title, text); - const data = await send_llm__attack_surface_api_request(endpoint_url,id); - Swal.close(); - if (data.status) { - $('#modal_dialog .modal-title').html(`Attack Surface Suggestion for ${data.subdomain_name} (BETA)`); - $('#modal_dialog .modal-text').empty(); - $('#modal_dialog .modal-text').append(data.description.replace(new RegExp('\r?\n','g'), '
')); - $('#modal_dialog').modal('show'); - } - else{ - Swal.close(); - Swal.fire({ - icon: 'error', - title: 'Oops...', - text: data.error, - }); - } - } catch (error) { - console.error(error); - Swal.close(); - Swal.fire({ - icon: 'error', - title: 'Oops...', - text: 'Something went wrong!', - }); - } +async function show_attack_surface_modal(endpoint_url, id) { + // Define generateAttackSurface in the global scope of the function + window.generateAttackSurface = async () => { + const selectedModel = $('input[name="llm_model"]:checked').val(); + if (!selectedModel) { + Swal.fire({ + title: 'Error', + text: 'Please select a model', + icon: 'error' + }); + return; + } + + try { + // Update selected model in database first + const updateResponse = await fetch('/api/tool/ollama/', { + method: 'PUT', + headers: { + 'Content-Type': 'application/json', + 'X-CSRFToken': getCookie('csrftoken') + }, + body: JSON.stringify({ model: selectedModel }) + }); + + const updateData = await updateResponse.json(); + if (!updateData.status) { + throw new Error('Failed to update selected model'); + } + + // Then proceed with attack surface analysis + var loader_title = "Loading..."; + var text = 'Please wait while the LLM is generating attack surface.' + showSwalLoader(loader_title, text); + const data = await send_llm__attack_surface_api_request(endpoint_url, id); + Swal.close(); + + if (data.status) { + $('#modal_dialog .modal-title').html(`Attack Surface Suggestion for ${data.subdomain_name} (BETA)`); + $('#modal_dialog .modal-text').empty(); + $('#modal_dialog .modal-text').append(data.description.replace(new RegExp('\r?\n','g'), '
')); + $('#modal_dialog').modal('show'); + } else { + Swal.fire({ + icon: 'error', + title: 'Oops...', + text: data.error, + }); + } + } catch (error) { + console.error(error); + Swal.close(); + Swal.fire({ + icon: 'error', + title: 'Oops...', + text: 'Something went wrong!', + }); + } + }; + + try { + // Fetch models from the unified endpoint that combines GPT and Ollama models + const response = await fetch('/api/tools/llm_models'); + const data = await response.json(); + + if (!data.status) { + throw new Error(data.error || 'Failed to fetch models'); + } + + const allModels = data.models; + const selectedModel = data.selected_model; + + let modelOptions = ''; + allModels.forEach(model => { + const modelName = model.name; + const capabilities = model.capabilities || {}; + const isLocal = model.is_local || false; + + modelOptions += ` +
+
+
+
+ +
+ ${modelName} + ${modelName === selectedModel ? 'Selected' : ''} + +
+

+ + + ${isLocal ? 'Local model' : 'OpenAI'} + + ${model.details ? ` + + + ${model.details.family} + + ` : ''} +
+ + ${capabilities.best_for ? capabilities.best_for.join(', ') : 'General analysis'} + + ${!isLocal ? '
API Key Required' : ''} +

+
+
+
+
`; + }); + + $('#modal_dialog .modal-title').html('Select LLM Model for Attack Surface Analysis'); + $('#modal_dialog .modal-text').empty(); + $('#modal_dialog .modal-text').append(` +
+

Select the LLM model to use:

+ ${modelOptions} +
+
+ +
+ `); + + $('#modal_dialog').modal('show'); + } catch (error) { + console.error(error); + Swal.fire({ + icon: 'error', + title: 'Error', + text: 'Unable to fetch LLM models. Please check configuration.', + footer: 'Configure LLM models' + }); + } +} + +async function continueWithSelectedModel(endpoint_url, id, callback) { + const selectedModel = $('input[name="llm_model"]:checked').val(); + if (!selectedModel) { + Swal.fire({ + title: 'Error', + text: 'Please select a model', + icon: 'error' + }); + return; + } + + try { + // Update selected model in database + const response = await fetch('/api/tool/ollama/', { + method: 'PUT', + headers: { + 'Content-Type': 'application/json', + 'X-CSRFToken': getCookie('csrftoken') + }, + body: JSON.stringify({ model: selectedModel }) + }); + + const data = await response.json(); + if (data.status) { + $('#modal_dialog').modal('hide'); + // Continue with attack surface analysis using the callback + if (callback) await callback(); + } else { + Swal.fire({ + title: 'Error', + text: 'Unable to set selected model', + icon: 'error' + }); + } + } catch (error) { + console.error(error); + Swal.fire({ + title: 'Error', + text: 'Something went wrong while setting the model', + icon: 'error' + }); + } } - function convertToCamelCase(inputString) { // Converts camel case string to title // Split the input string by underscores @@ -3327,3 +3478,79 @@ function handleHashInUrl(){ } } } + +function showLLMModelSelectionModal(callback) { + $('#modal_dialog .modal-title').html('Select LLM Model'); + $('#modal_dialog .modal-text').empty(); + + // Get available models + fetch('/api/tool/ollama/') + .then(response => response.json()) + .then(data => { + const models = data.models; + const selectedModel = data.selected_model; + + let modelOptions = ''; + models.forEach(model => { + modelOptions += ` +
+ + +
`; + }); + + $('#modal_dialog .modal-text').append(` +
+

Select the LLM model to use for vulnerability analysis:

+ ${modelOptions} +
+
+ +
+ `); + + $('#modal_dialog').modal('show'); + }); +} + +function selectLLMModel() { + const selectedModel = $('input[name="llm_model"]:checked').val(); + if (!selectedModel) { + Swal.fire({ + title: 'Error', + text: 'Please select a model', + icon: 'error' + }); + return; + } + + // Update selected model in database + fetch('/api/tool/ollama/', { + method: 'PUT', + headers: { + 'Content-Type': 'application/json', + 'X-CSRFToken': getCookie('csrftoken') + }, + body: JSON.stringify({ model: selectedModel }) + }) + .then(response => response.json()) + .then(data => { + if (data.status) { + $('#modal_dialog').modal('hide'); + // Continue with scan + startScan(); + } else { + Swal.fire({ + title: 'Error', + text: 'Unable to set selected model', + icon: 'error' + }); + } + }); +} diff --git a/web/tests/test_llm.py b/web/tests/test_llm.py new file mode 100644 index 000000000..e7a819256 --- /dev/null +++ b/web/tests/test_llm.py @@ -0,0 +1,129 @@ +from unittest.mock import patch, MagicMock +from django.test import TestCase +from rest_framework import status + +from reNgine.llm.config import MODEL_REQUIREMENTS +from reNgine.llm.llm import LLMVulnerabilityReportGenerator, LLMAttackSuggestionGenerator +from reNgine.llm.validators import LLMProvider, LLMInputData +from utils.test_base import BaseTestCase + + +class TestLLMBase(BaseTestCase): + """Base test class for LLM functionality.""" + + def setUp(self): + super().setUp() + self.data_generator.create_project_base() + self.mock_llm_response = { + "status": True, + "description": "Test vulnerability description", + "impact": "Test impact description", + "remediation": "Test remediation steps", + "references": ["https://test.com/ref1", "https://test.com/ref2"] + } + + +class TestLLMVulnerabilityReport(TestLLMBase): + """Test cases for LLM Vulnerability Report Generator.""" + + def setUp(self): + super().setUp() + self.generator = LLMVulnerabilityReportGenerator() + + @patch('reNgine.llm.LLMVulnerabilityReportGenerator._get_openai_response') + def test_get_vulnerability_description_success(self, mock_get_response): + """Test successful vulnerability description generation.""" + mock_get_response.return_value = "Test vulnerability description" + + response = self.generator.get_vulnerability_description("Test input") + self.assertTrue(response.status) + self.assertIsNotNone(response.description) + + def test_validate_input_success(self): + """Test input validation success.""" + input_data = "Detailed vulnerability description for testing" + validated = self.generator._validate_input(input_data) + self.assertIsInstance(validated, LLMInputData) + self.assertEqual(validated.description, input_data) + + @patch('reNgine.llm.LLMVulnerabilityReportGenerator._get_openai_response') + def test_get_vulnerability_description_failure(self, mock_get_response): + """Test vulnerability description generation failure.""" + mock_get_response.side_effect = Exception("API Error") + + response = self.generator.get_vulnerability_description("Test input") + self.assertFalse(response.status) + self.assertIsNotNone(response.error) + + +class TestLLMAttackSuggestion(TestLLMBase): + """Test cases for LLM Attack Suggestion Generator.""" + + def setUp(self): + super().setUp() + self.generator = LLMAttackSuggestionGenerator() + + @patch('reNgine.llm.LLMAttackSuggestionGenerator.get_attack_suggestion') + def test_get_attack_suggestion_success(self, mock_get_suggestion): + """Test successful attack suggestion generation.""" + mock_get_suggestion.return_value = { + "status": True, + "description": "Test attack suggestion" + } + + api_url = reverse("api:llm_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") + + def test_validate_input_success(self): + """Test input validation success.""" + input_data = "Detailed reconnaissance data for testing" + validated = self.generator._validate_input(input_data) + self.assertIsInstance(validated, LLMInputData) + self.assertEqual(validated.description, input_data) + + @patch('reNgine.llm.LLMAttackSuggestionGenerator._get_openai_response') + def test_get_attack_suggestion_failure(self, mock_get_response): + """Test attack suggestion generation failure.""" + mock_get_response.side_effect = Exception("API Error") + + response = self.generator.get_attack_suggestion("Test input") + self.assertFalse(response.status) + self.assertIsNotNone(response.error) + + def test_get_provider_config(self): + """Test provider configuration retrieval""" + generator = LLMAttackSuggestionGenerator(provider=LLMProvider.OLLAMA) + config = generator._get_provider_config() + self.assertIn('default_model', config) + self.assertIn('models', config) + self.assertIn('timeout', config) + + def test_model_capabilities(self): + """Test model capabilities access""" + generator = LLMAttackSuggestionGenerator() + model_name = generator._get_model_name() + self.assertIn(model_name, MODEL_REQUIREMENTS) + self.assertIn('provider', MODEL_REQUIREMENTS[model_name]) + + +class TestLLMProviders(TestCase): + """Test cases for LLM providers configuration.""" + + def test_openai_provider_config(self): + """Test OpenAI provider configuration.""" + generator = LLMVulnerabilityReportGenerator(provider=LLMProvider.OPENAI) + self.assertEqual(generator.provider, LLMProvider.OPENAI) + self.assertIsNone(generator.ollama) + + def test_ollama_provider_config(self): + """Test Ollama provider configuration.""" + generator = LLMVulnerabilityReportGenerator(provider=LLMProvider.OLLAMA) + self.assertEqual(generator.provider, LLMProvider.OLLAMA) + self.assertIsNotNone(generator.ollama) \ No newline at end of file From 82fccad0e5528f9cefa027fcfacbe5253ba7d664 Mon Sep 17 00:00:00 2001 From: psyray Date: Mon, 11 Nov 2024 14:31:25 +0100 Subject: [PATCH 03/21] refactor(logging): improve vulnerability logging details - Enhanced logging in tasks.py to include both the title and path of vulnerabilities when exceptions occur. - Updated variable naming from llm to vuln for clarity in vulnerability handling. - Minor text punctuation corrections in custom.js to ensure consistency in user messages. - Remove unused import --- web/reNgine/llm/validators.py | 2 +- web/reNgine/tasks.py | 12 ++++++------ web/static/custom/custom.js | 4 ++-- web/tests/test_llm.py | 2 +- 4 files changed, 10 insertions(+), 10 deletions(-) diff --git a/web/reNgine/llm/validators.py b/web/reNgine/llm/validators.py index 5eb7aaa3c..113e1f46d 100644 --- a/web/reNgine/llm/validators.py +++ b/web/reNgine/llm/validators.py @@ -1,4 +1,4 @@ -from typing import Optional, List, Dict +from typing import Optional, List from pydantic import BaseModel, Field, validator from enum import Enum diff --git a/web/reNgine/tasks.py b/web/reNgine/tasks.py index f6f4fcee6..9da24d0e5 100644 --- a/web/reNgine/tasks.py +++ b/web/reNgine/tasks.py @@ -2290,11 +2290,11 @@ def nuclei_individual_severity_module(self, cmd, severity, enable_http_crawl, sh # Wait for all tasks to complete for future in concurrent.futures.as_completed(future_to_llm): - llm = future_to_llm[future] + vuln = future_to_llm[future] try: future.result() except Exception as e: - logger.error(f"Exception for Vulnerability {vuln}: {e}") + logger.error(f"Exception for Vulnerability {vuln[0]} - {vuln[1]}: {e}") # Display title and path return None @@ -2613,11 +2613,11 @@ def dalfox_xss_scan(self, urls=[], ctx={}, description=None): # Wait for all tasks to complete for future in concurrent.futures.as_completed(future_to_llm): - llm = future_to_llm[future] + vuln = future_to_llm[future] try: future.result() except Exception as e: - logger.error(f"Exception for Vulnerability {vuln}: {e}") + logger.error(f"Exception for Vulnerability {vuln[0]} - {vuln[1]}: {e}") # Display title and path return results @@ -2739,11 +2739,11 @@ def crlfuzz_scan(self, urls=[], ctx={}, description=None): # Wait for all tasks to complete for future in concurrent.futures.as_completed(future_to_llm): - llm = future_to_llm[future] + vuln = future_to_llm[future] try: future.result() except Exception as e: - logger.error(f"Exception for Vulnerability {vuln}: {e}") + logger.error(f"Exception for Vulnerability {vuln[0]} - {vuln[1]}: {e}") # Display title and path return results diff --git a/web/static/custom/custom.js b/web/static/custom/custom.js index e4419b3fd..d0dada426 100644 --- a/web/static/custom/custom.js +++ b/web/static/custom/custom.js @@ -3165,7 +3165,7 @@ async function send_llm_api_request(endpoint_url, vuln_id){ async function fetch_llm_vuln_details(endpoint_url, id, title) { var loader_title = "Loading..."; - var text = 'Please wait while the LLM is generating vulnerability description.' + var text = 'Please wait while the LLM is generating vulnerability description.'; try { showSwalLoader(loader_title, text); const data = await send_llm_api_request(endpoint_url, id); @@ -3301,7 +3301,7 @@ async function show_attack_surface_modal(endpoint_url, id) { // Then proceed with attack surface analysis var loader_title = "Loading..."; - var text = 'Please wait while the LLM is generating attack surface.' + var text = 'Please wait while the LLM is generating attack surface.'; showSwalLoader(loader_title, text); const data = await send_llm__attack_surface_api_request(endpoint_url, id); Swal.close(); diff --git a/web/tests/test_llm.py b/web/tests/test_llm.py index e7a819256..728453b58 100644 --- a/web/tests/test_llm.py +++ b/web/tests/test_llm.py @@ -1,4 +1,4 @@ -from unittest.mock import patch, MagicMock +from unittest.mock import patch from django.test import TestCase from rest_framework import status From b170a31bac3997e453abef2d55768c348e88c4d9 Mon Sep 17 00:00:00 2001 From: psyray Date: Mon, 11 Nov 2024 16:02:39 +0100 Subject: [PATCH 04/21] feat(ui): enhance LLM toolkit UI and refactor model management - Improved the UI of the LLM toolkit by updating the layout and adding badges for model status and capabilities. - Refactored the model management logic to use a centralized API call for fetching model data, improving error handling and reducing code duplication. - Updated the model requirements configuration to enhance readability and consistency in the description of model capabilities. - Adjusted the modal size for displaying model options to provide a better user experience. --- web/reNgine/llm/config.py | 34 ++++---- .../scanEngine/settings/llm_toolkit.html | 78 ++++++++++++------- web/scanEngine/views.py | 43 ++++------ web/static/custom/custom.js | 29 +++++-- 4 files changed, 105 insertions(+), 79 deletions(-) diff --git a/web/reNgine/llm/config.py b/web/reNgine/llm/config.py index df18d4af3..b4e32f352 100644 --- a/web/reNgine/llm/config.py +++ b/web/reNgine/llm/config.py @@ -175,28 +175,28 @@ 'min_tokens': 64, 'max_tokens': 2048, 'supports_functions': True, - 'best_for': ['basic_analysis', 'general_purpose'], + 'best_for': ['Basic analysis', 'General purpose tasks'], 'provider': 'openai' }, 'gpt-3.5-turbo': { 'min_tokens': 64, 'max_tokens': 4096, 'supports_functions': True, - 'best_for': ['quick_analysis', 'basic_suggestions', 'cost_effective'], + 'best_for': ['Quick analysis', 'Basic suggestions', 'Cost effective solutions'], 'provider': 'openai' }, 'gpt-4': { 'min_tokens': 128, 'max_tokens': 8192, 'supports_functions': True, - 'best_for': ['deep_analysis', 'complex_reasoning', 'advanced_security'], + 'best_for': ['Deep security analysis', 'Complex reasoning', 'Advanced security tasks'], 'provider': 'openai' }, 'gpt-4-turbo': { 'min_tokens': 128, 'max_tokens': 128000, 'supports_functions': True, - 'best_for': ['complex_analysis', 'technical_details', 'latest_capabilities'], + 'best_for': ['Complex analysis', 'Technical details', 'Latest AI capabilities'], 'provider': 'openai' }, @@ -205,35 +205,35 @@ 'min_tokens': 32, 'max_tokens': 4096, 'supports_functions': False, - 'best_for': ['local_processing', 'privacy_focused', 'balanced_performance'], + 'best_for': ['Local processing', 'Privacy focused tasks', 'Balanced performance'], 'provider': 'ollama' }, 'llama2-uncensored': { 'min_tokens': 32, 'max_tokens': 4096, 'supports_functions': False, - 'best_for': ['unfiltered_analysis', 'security_research', 'red_teaming'], + 'best_for': ['Unfiltered analysis', 'Security research', 'Red team operations'], 'provider': 'ollama' }, 'llama3': { 'min_tokens': 64, 'max_tokens': 8192, 'supports_functions': False, - 'best_for': ['advanced_reasoning', 'improved_context', 'technical_analysis'], + 'best_for': ['Advanced reasoning', 'Improved context', 'Technical analysis'], 'provider': 'ollama' }, 'llama3.1': { 'min_tokens': 64, 'max_tokens': 8192, 'supports_functions': False, - 'best_for': ['enhanced_comprehension', 'security_assessment', 'detailed_analysis'], + 'best_for': ['Enhanced comprehension', 'Security assessment', 'Detailed analysis'], 'provider': 'ollama' }, 'llama3.2': { 'min_tokens': 64, 'max_tokens': 16384, 'supports_functions': False, - 'best_for': ['long_context', 'complex_security_analysis', 'advanced_reasoning'], + 'best_for': ['Long context', 'Complex security analysis', 'Advanced reasoning'], 'provider': 'ollama' }, @@ -242,56 +242,56 @@ 'min_tokens': 32, 'max_tokens': 8192, 'supports_functions': False, - 'best_for': ['efficient_processing', 'technical_analysis', 'good_performance_ratio'], + 'best_for': ['Efficient processing', 'Technical analysis', 'Performance optimization'], 'provider': 'ollama' }, 'mistral-medium': { 'min_tokens': 32, 'max_tokens': 8192, 'supports_functions': False, - 'best_for': ['balanced_analysis', 'improved_accuracy', 'technical_tasks'], + 'best_for': ['Balanced analysis', 'Improved accuracy', 'Technical tasks'], 'provider': 'ollama' }, 'mistral-large': { 'min_tokens': 64, 'max_tokens': 16384, 'supports_functions': False, - 'best_for': ['deep_technical_analysis', 'complex_reasoning', 'high_accuracy'], + 'best_for': ['Deep technical analysis', 'Complex reasoning', 'High accuracy'], 'provider': 'ollama' }, 'codellama': { 'min_tokens': 32, 'max_tokens': 4096, 'supports_functions': False, - 'best_for': ['code_analysis', 'vulnerability_assessment', 'technical_details'], + 'best_for': ['Code analysis', 'Vulnerability assessment', 'Technical documentation'], 'provider': 'ollama' }, 'qwen2.5': { 'min_tokens': 64, 'max_tokens': 8192, 'supports_functions': False, - 'best_for': ['multilingual_analysis', 'efficient_processing', 'technical_understanding'], + 'best_for': ['Multilingual analysis', 'Efficient processing', 'Technical understanding'], 'provider': 'ollama' }, 'gemma': { 'min_tokens': 32, 'max_tokens': 4096, 'supports_functions': False, - 'best_for': ['lightweight_analysis', 'quick_assessment', 'general_tasks'], + 'best_for': ['Lightweight analysis', 'Quick assessment', 'General tasks'], 'provider': 'ollama' }, 'solar': { 'min_tokens': 64, 'max_tokens': 8192, 'supports_functions': False, - 'best_for': ['creative_analysis', 'unique_perspectives', 'alternative_approaches'], + 'best_for': ['Creative analysis', 'Unique perspectives', 'Alternative approaches'], 'provider': 'ollama' }, 'yi': { 'min_tokens': 64, 'max_tokens': 8192, 'supports_functions': False, - 'best_for': ['comprehensive_analysis', 'detailed_explanations', 'technical_depth'], + 'best_for': ['Comprehensive analysis', 'Detailed explanations', 'Technical depth'], 'provider': 'ollama' } } \ No newline at end of file diff --git a/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html b/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html index cf5434590..f03d6cd21 100644 --- a/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html +++ b/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html @@ -44,43 +44,52 @@
{{installed_models|length}} available Models
Warning: GPT model is currently selected and requires API key to be set. Please set the API key in the API Vault.

{% endif %} -
+
{% for model in installed_models %} -
-
-
-
+

Initializing...

-
`, showConfirmButton: false, @@ -446,19 +443,6 @@
}); } }; - - // Handle cancel button - document.getElementById('cancel-download').addEventListener('click', () => { - if (socket) { - socket.close(); - Swal.close(); - Swal.fire({ - icon: 'info', - title: 'Download Cancelled', - text: 'The download was cancelled by user.' - }); - } - }); } } catch (error) { console.error('Download error:', error); From d9d7af42d765509acd7a985755b9d7329322bbc4 Mon Sep 17 00:00:00 2001 From: psyray Date: Mon, 18 Nov 2024 00:05:37 +0100 Subject: [PATCH 19/21] feat: add model name to progress bar popup Updated the UI text to include the model name in the download progress popup. --- web/api/tests/test_tools.py | 2 +- web/scanEngine/templates/scanEngine/settings/llm_toolkit.html | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/web/api/tests/test_tools.py b/web/api/tests/test_tools.py index 941e8385e..5c688e99a 100644 --- a/web/api/tests/test_tools.py +++ b/web/api/tests/test_tools.py @@ -2,7 +2,7 @@ This file contains the test cases for the API views. """ -from unittest.mock import patch, AsyncMock, Mock +from unittest.mock import patch from django.urls import reverse from rest_framework import status from startScan.models import SubScan diff --git a/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html b/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html index c5ae5f9e5..063e5d068 100644 --- a/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html +++ b/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html @@ -361,7 +361,7 @@
if (result.isConfirmed) { // Show progress popup first Swal.fire({ - title: 'Downloading Model', + title: `Downloading Model ${modelName}`, html: `

Starting download...

From 5d489f723c8bc60bdc320d555a5d5e01948bd7b2 Mon Sep 17 00:00:00 2001 From: psyray Date: Mon, 18 Nov 2024 00:18:21 +0100 Subject: [PATCH 20/21] fix: update model selection API endpoint Modified API endpoint usage in JavaScript to include the selected model in the URL path. --- web/static/custom/custom.js | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/web/static/custom/custom.js b/web/static/custom/custom.js index 877e47251..f1a65cd91 100644 --- a/web/static/custom/custom.js +++ b/web/static/custom/custom.js @@ -3355,7 +3355,7 @@ async function showModelSelectionDialog(endpoint_url, id, force_regenerate = fal try { // Update selected model in database first - const updateResponse = await fetch('/api/tool/ollama/', { + const updateResponse = await fetch(`/api/tool/ollama/${selectedModel}`, { method: 'PUT', headers: { 'Content-Type': 'application/json', @@ -3629,7 +3629,7 @@ function selectLLMModel() { } // Update selected model in database - fetch('/api/tool/ollama/', { + fetch(`/api/tool/ollama/${selectedModel}`, { method: 'PUT', headers: { 'Content-Type': 'application/json', From 380974197b895a8a1e6e20c94f64cea13b6387f9 Mon Sep 17 00:00:00 2001 From: psyray Date: Mon, 18 Nov 2024 02:44:35 +0100 Subject: [PATCH 21/21] feat: enhance model URL handling - Improved URL handling by encoding model names in API requests across multiple JavaScript files. - Modified the URL pattern in the Django application to support path parameters for model names. --- web/api/urls.py | 2 +- .../templates/scanEngine/settings/llm_toolkit.html | 6 ++++-- web/static/custom/custom.js | 6 ++++-- 3 files changed, 9 insertions(+), 5 deletions(-) diff --git a/web/api/urls.py b/web/api/urls.py index a04d90ac9..d5ecb164b 100644 --- a/web/api/urls.py +++ b/web/api/urls.py @@ -195,7 +195,7 @@ OllamaManager.as_view(), name='ollama_manager'), path( - 'tool/ollama//', + 'tool/ollama//', OllamaDetailManager.as_view(), name='ollama_detail_manager'), path( diff --git a/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html b/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html index 063e5d068..6b40f913c 100644 --- a/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html +++ b/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html @@ -161,7 +161,8 @@ {% block page_level_script %}