diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index 00cf16f0..297b9193 100644 --- a/CONTRIBUTORS.md +++ b/CONTRIBUTORS.md @@ -8,6 +8,7 @@ Thanks to these individuals for making reNgine awesome by fixing bugs, resolving * [Suprita-25](https://github.com/Suprita-25) * [TheBinitGhimire](https://github.com/TheBinitGhimire) * [Vinay Leo](https://github.com/vinaynm) +* [Erdem Ozgen](https://github.com/ErdemOzgen) *If you have created a Pull request, feel free to add your name here, because we know you are awesome and deserve thanks from the community!* diff --git a/Makefile b/Makefile index f1559dcf..8e23217d 100644 --- a/Makefile +++ b/Makefile @@ -7,7 +7,7 @@ include .env COMPOSE_PREFIX_CMD := COMPOSE_DOCKER_CLI_BUILD=1 COMPOSE_ALL_FILES := -f docker-compose.yml -SERVICES := db web proxy redis celery celery-beat +SERVICES := db web proxy redis celery celery-beat ollama # -------------------------- diff --git a/docker-compose.dev.yml b/docker-compose.dev.yml index ffd733e3..2afaf3dd 100644 --- a/docker-compose.dev.yml +++ b/docker-compose.dev.yml @@ -115,6 +115,17 @@ services: - celery-beat networks: - rengine_network + ollama: + image: ollama/ollama + container_name: ollama + volumes: + - ollama_data:/root/.ollama + # ports: + # - "11434:11434" + networks: + - rengine_network + restart: always + # command: ["ollama", "run", "llama2-uncensored"] networks: rengine_network: @@ -126,3 +137,4 @@ volumes: github_repos: wordlist: scan_results: + ollama_data: diff --git a/docker-compose.yml b/docker-compose.yml index 3d2dc052..b22c6c67 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -143,7 +143,15 @@ services: - scan_results:/usr/src/scan_results networks: - rengine_network - + ollama: + image: ollama/ollama + container_name: ollama + volumes: + - ollama_data:/root/.ollama + ports: + - "11434:11434" + networks: + - rengine_network networks: rengine_network: @@ -157,6 +165,7 @@ volumes: wordlist: scan_results: static_volume: + ollama_data: secrets: proxy.ca: diff --git a/web/api/urls.py b/web/api/urls.py index f3a09b54..f4ee867e 100644 --- a/web/api/urls.py +++ b/web/api/urls.py @@ -178,6 +178,10 @@ urlpatterns = [ 'tool/uninstall/', UninstallTool.as_view(), name='uninstall_tool'), + path( + 'tool/ollama/', + OllamaManager.as_view(), + name='ollama_manager'), path( 'rengine/update/', RengineUpdateCheck.as_view(), diff --git a/web/api/views.py b/web/api/views.py index c81bd149..45fe5df4 100644 --- a/web/api/views.py +++ b/web/api/views.py @@ -33,6 +33,82 @@ from .serializers import * logger = logging.getLogger(__name__) +class OllamaManager(APIView): + def get(self, request): + """ + API to download Ollama Models + sends a POST request to download the model + """ + req = self.request + model_name = req.query_params.get('model') + response = { + 'status': False + } + try: + pull_model_api = f'{OLLAMA_INSTANCE}/api/pull' + _response = requests.post( + pull_model_api, + json={ + 'name': model_name, + 'stream': False + } + ).json() + if _response.get('error'): + response['status'] = False + response['error'] = _response.get('error') + else: + response['status'] = True + except Exception as e: + response['error'] = str(e) + return Response(response) + + def delete(self, request): + req = self.request + model_name = req.query_params.get('model') + delete_model_api = f'{OLLAMA_INSTANCE}/api/delete' + response = { + 'status': False + } + try: + _response = requests.delete( + delete_model_api, + json={ + 'name': model_name + } + ).json() + if _response.get('error'): + response['status'] = False + response['error'] = _response.get('error') + else: + response['status'] = True + except Exception as e: + response['error'] = str(e) + return Response(response) + + def put(self, request): + req = self.request + model_name = req.query_params.get('model') + # check if model_name is in DEFAULT_GPT_MODELS + response = { + 'status': False + } + use_ollama = True + if any(model['name'] == model_name for model in DEFAULT_GPT_MODELS): + use_ollama = False + try: + OllamaSettings.objects.update_or_create( + defaults={ + 'selected_model': model_name, + 'use_ollama': use_ollama + }, + id=1 + ) + response['status'] = True + except Exception as e: + response['error'] = str(e) + return Response(response) + + class GPTAttackSuggestion(APIView): def get(self, request): req = self.request diff --git a/web/dashboard/migrations/0010_ollamasettings.py b/web/dashboard/migrations/0010_ollamasettings.py new file mode 100644 index 00000000..1e47e5d0 --- /dev/null +++ b/web/dashboard/migrations/0010_ollamasettings.py @@ -0,0 +1,20 @@ +# Generated by Django 3.2.4 on 2024-04-21 04:35 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('dashboard', '0009_delete_openaikeys'), + ] + + operations = [ + migrations.CreateModel( + name='OllamaSettings', + fields=[ + ('id', models.AutoField(primary_key=True, serialize=False)), + ('selected_model', models.CharField(max_length=500)), + ], + ), + ] diff --git a/web/dashboard/migrations/0011_ollamasettings_is_ollama.py b/web/dashboard/migrations/0011_ollamasettings_is_ollama.py new file mode 100644 index 00000000..aefdef37 --- /dev/null +++ b/web/dashboard/migrations/0011_ollamasettings_is_ollama.py @@ -0,0 +1,18 @@ +# Generated by Django 3.2.4 on 2024-04-21 05:06 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('dashboard', '0010_ollamasettings'), + ] + + operations = [ + migrations.AddField( + model_name='ollamasettings', + name='is_ollama', + field=models.BooleanField(default=False), + ), + ] diff --git a/web/dashboard/migrations/0012_rename_is_ollama_ollamasettings_is_openai.py b/web/dashboard/migrations/0012_rename_is_ollama_ollamasettings_is_openai.py new file mode 100644 index 00000000..52bdf6ff --- /dev/null +++ b/web/dashboard/migrations/0012_rename_is_ollama_ollamasettings_is_openai.py @@ -0,0 +1,18 @@ +# Generated by Django 3.2.4 on 2024-04-21 05:06 + +from django.db import migrations + + +class Migration(migrations.Migration): + + dependencies = [ + ('dashboard', '0011_ollamasettings_is_ollama'), + ] + + operations = [ + migrations.RenameField( + model_name='ollamasettings', + old_name='is_ollama', + new_name='is_openai', + ), + ] diff --git a/web/dashboard/migrations/0013_auto_20240421_0507.py b/web/dashboard/migrations/0013_auto_20240421_0507.py new file mode 100644 index 00000000..11ab6594 --- /dev/null +++ b/web/dashboard/migrations/0013_auto_20240421_0507.py @@ -0,0 +1,22 @@ +# Generated by Django 3.2.4 on 2024-04-21 05:07 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('dashboard', '0012_rename_is_ollama_ollamasettings_is_openai'), + ] + + operations = [ + migrations.RemoveField( + model_name='ollamasettings', + name='is_openai', + ), + migrations.AddField( + model_name='ollamasettings', + name='is_ollama', + field=models.BooleanField(default=True), + ), + ] diff --git a/web/dashboard/migrations/0014_rename_is_ollama_ollamasettings_use_ollama.py b/web/dashboard/migrations/0014_rename_is_ollama_ollamasettings_use_ollama.py new file mode 100644 index 00000000..a201df3a --- /dev/null +++ b/web/dashboard/migrations/0014_rename_is_ollama_ollamasettings_use_ollama.py @@ -0,0 +1,18 @@ +# Generated by Django 3.2.4 on 2024-04-21 05:08 + +from django.db import migrations + + +class Migration(migrations.Migration): + + dependencies = [ + ('dashboard', '0013_auto_20240421_0507'), + ] + + operations = [ + migrations.RenameField( + model_name='ollamasettings', + old_name='is_ollama', + new_name='use_ollama', + ), + ] diff --git a/web/dashboard/models.py b/web/dashboard/models.py index 3d600174..8ed77dd4 100644 --- a/web/dashboard/models.py +++ b/web/dashboard/models.py @@ -24,6 +24,15 @@ class OpenAiAPIKey(models.Model): def __str__(self): return self.key + + +class OllamaSettings(models.Model): + id = models.AutoField(primary_key=True) + selected_model = models.CharField(max_length=500) + use_ollama = models.BooleanField(default=True) + + def __str__(self): + return self.selected_model class NetlasAPIKey(models.Model): diff --git a/web/reNgine/definitions.py b/web/reNgine/definitions.py index 7034fbca..d4abe300 100644 --- a/web/reNgine/definitions.py +++ b/web/reNgine/definitions.py @@ -436,6 +436,52 @@ PERM_INITATE_SCANS_SUBSCANS = 'initiate_scans_subscans' 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', + } + } +] + + + # GPT Vulnerability Report Generator VULNERABILITY_DESCRIPTION_SYSTEM_MESSAGE = """ You are a highly skilled penetration tester who has recently completed a penetration testing. diff --git a/web/reNgine/gpt.py b/web/reNgine/gpt.py index 4bae2563..389da1fc 100644 --- a/web/reNgine/gpt.py +++ b/web/reNgine/gpt.py @@ -1,14 +1,20 @@ 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 +from reNgine.definitions import VULNERABILITY_DESCRIPTION_SYSTEM_MESSAGE, ATTACK_SUGGESTION_GPT_SYSTEM_PROMPT, OLLAMA_INSTANCE +from langchain_community.llms import Ollama + +from dashboard.models import OllamaSettings class GPTVulnerabilityReportGenerator: def __init__(self): - self.api_key = get_open_ai_key() - self.model_name = 'gpt-3.5-turbo' - + 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 GPT. @@ -23,94 +29,103 @@ class GPTVulnerabilityReportGenerator: 'references': (list) of urls } """ - if not self.api_key: - return { - 'status': False, - 'error': 'No OpenAI keys provided.' - } - openai.api_key = self.api_key - try: - gpt_response = openai.ChatCompletion.create( - model=self.model_name, - messages=[ - {'role': 'system', 'content': VULNERABILITY_DESCRIPTION_SYSTEM_MESSAGE}, - {'role': 'user', 'content': description} - ] + 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 + gpt_response = openai.ChatCompletion.create( + model=self.model_name, + messages=[ + {'role': 'system', 'content': VULNERABILITY_DESCRIPTION_SYSTEM_MESSAGE}, + {'role': 'user', 'content': description} + ] + ) - response_content = gpt_response['choices'][0]['message']['content'] + response_content = gpt_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 + ) - 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() - 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, - } - except Exception as e: - return { - 'status': False, - 'error': str(e) - } + 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 GPTAttackSuggestionGenerator: 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 gpt ''' if not self.api_key: - return { - 'status': False, - 'error': 'No OpenAI keys provided.', - 'input': input - } - openai.api_key = self.api_key - print(input) - try: - gpt_response = openai.ChatCompletion.create( - model=self.model_name, - messages=[ - {'role': 'system', 'content': ATTACK_SUGGESTION_GPT_SYSTEM_PROMPT}, - {'role': 'user', 'content': input} - ] - ) - response_content = gpt_response['choices'][0]['message']['content'] - return { - 'status': True, - 'description': response_content, - 'input': input - } - except Exception as e: - return { - 'status': False, - 'error': str(e), - 'input': input - } + prompt = ATTACK_SUGGESTION_GPT_SYSTEM_PROMPT + "\nUser: " + input + response_content = self.ollama(prompt) + else: + openai.api_key = self.api_key + print(input) + try: + gpt_response = openai.ChatCompletion.create( + model=self.model_name, + messages=[ + {'role': 'system', 'content': ATTACK_SUGGESTION_GPT_SYSTEM_PROMPT}, + {'role': 'user', 'content': input} + ] + ) + response_content = gpt_response['choices'][0]['message']['content'] + except Exception as e: + return { + 'status': False, + 'error': str(e), + 'input': input + } + return { + 'status': True, + 'description': response_content, + 'input': input + } + \ No newline at end of file diff --git a/web/requirements.txt b/web/requirements.txt index 08c2d49c..fd548693 100644 --- a/web/requirements.txt +++ b/web/requirements.txt @@ -17,6 +17,7 @@ drf-yasg==1.21.3 gunicorn==22.0.0 gevent==24.2.1 humanize==4.3.0 +langchain==0.1.0 Markdown==3.3.4 metafinder==1.2 netaddr==0.8.0 diff --git a/web/scanEngine/templates/scanEngine/settings/api.html b/web/scanEngine/templates/scanEngine/settings/api.html index f4a2d6a4..150e912a 100644 --- a/web/scanEngine/templates/scanEngine/settings/api.html +++ b/web/scanEngine/templates/scanEngine/settings/api.html @@ -2,7 +2,7 @@ {% load static %} {% load custom_tags %} {% block title %} -reNgine Settings +API Vault {% endblock title %} {% block custom_js_css_link %} @@ -27,8 +27,8 @@ API Vault
- -

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 GPT.

{% if openai_key %} {% else %} @@ -38,7 +38,7 @@ API Vault
-

Netlas keys will be used to get whois information and other OSINT data.

+

Netlas keys will be used to get whois information and other OSINT related data.

{% if netlas_key %} {% else %} diff --git a/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html b/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html new file mode 100644 index 00000000..fa6d304e --- /dev/null +++ b/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html @@ -0,0 +1,286 @@ +{% extends 'base/base.html' %} +{% load static %} +{% load humanize %} +{% block title %} +LLM Toolkit +{% endblock title %} + +{% block custom_js_css_link %} +{% endblock custom_js_css_link %} + +{% block breadcrumb_title %} +
  • Settings
    • +
  • LLM Toolkit
    • +{% endblock breadcrumb_title %} + +{% block page_title %} +LLM Toolkit (Beta) +{% endblock page_title %} + +{% block main_content %} +
    +
    +
    +
    + LLM Toolkit includes the ability to download new LLMs, view available models, and delete models no longer needed, and also choose between various models. +
    +

    reNgine makes use of various LLMs to enhance reporting process. Using various LLM AI Models penetration testers will be able to to generate detailed, insightful penetration testing reports. +
    + If you are using custom LLM models, it is expected that response time are much slower in CPU. We recommend using GPU for better performance. Models such as llama2, or llama3 requires significant computation and GPU are required. Having only CPU will result in slow response time. +
    + OpenAI GPT models do not run locally, hence the requirement of GPU is not necessary. +

    +
    +
    +
    +
    +
    + Add new model +
    +
    +
    {{installed_models|length}} available Models
    + {% if openai_key_error %} + + {% endif %} +
    + {% for model in installed_models %} +
    +
    +
    +
    + +
    + {% if model.is_local %} + Delete + {% endif %} + {% if not model.selected %} + Use Model + {% endif %} +
    +
    +

    + {{model.name}} {% if model.selected %}Selected Model{% endif %} +

    +

    +

    + + + Modified {% if model.modified_at %}{{model.modified_at|naturaltime}} {% else %} NA{% endif %} + +
    + + + {% if model.is_local %} + Locally installed model + {% else %} + Open AI Model + {% endif %} + +
    + + + {{model.details.parameter_size}} Parameters + + + + {{model.details.family}} Family + +

    +
    +
    +
    + {% empty %} + + {% endfor %} +
    +
    +{% endblock main_content %} + + +{% block page_level_script %} + +{% endblock page_level_script %} diff --git a/web/scanEngine/urls.py b/web/scanEngine/urls.py index ff7d2812..89c5084b 100644 --- a/web/scanEngine/urls.py +++ b/web/scanEngine/urls.py @@ -56,6 +56,10 @@ urlpatterns = [ '/tool_arsenal', views.tool_arsenal_section, name='tool_arsenal'), + path( + '/llm_toolkit', + views.llm_toolkit_section, + name='llm_toolkit'), path( '/rengine_settings', views.rengine_settings, diff --git a/web/scanEngine/views.py b/web/scanEngine/views.py index 561006c0..0404e8c4 100644 --- a/web/scanEngine/views.py +++ b/web/scanEngine/views.py @@ -4,6 +4,7 @@ import re import shutil import subprocess +from datetime import datetime from django import http from django.contrib import messages from django.shortcuts import get_object_or_404, render @@ -456,6 +457,41 @@ def tool_arsenal_section(request, slug): return render(request, 'scanEngine/settings/tool_arsenal.html', context) +@has_permission_decorator(PERM_MODIFY_SYSTEM_CONFIGURATIONS, redirect_url=FOUR_OH_FOUR_URL) +def llm_toolkit_section(request, slug): + context = {} + list_all_models_url = f'{OLLAMA_INSTANCE}/api/tags' + response = requests.get(list_all_models_url) + all_models = [] + selected_model = None + all_models = DEFAULT_GPT_MODELS.copy() + if response.status_code == 200: + models = response.json() + ollama_models = models.get('models') + date_format = "%Y-%m-%dT%H:%M:%S" + for model in ollama_models: + all_models.append({**model, + 'modified_at': datetime.strptime(model['modified_at'].split('.')[0], date_format), + 'is_local': True, + }) + # find selected model name from db + selected_model = OllamaSettings.objects.first() + if selected_model: + selected_model = {'selected_model': selected_model.selected_model} + else: + # use gpt3.5-turbo as default + selected_model = {'selected_model': 'gpt-3.5-turbo'} + for model in all_models: + if model['name'] == selected_model['selected_model']: + model['selected'] = True + context['installed_models'] = all_models + # show error message for openai key, if any gpt is selected + openai_key = get_open_ai_key() + if not openai_key and 'gpt' in selected_model['selected_model']: + context['openai_key_error'] = True + return render(request, 'scanEngine/settings/llm_toolkit.html', context) + + @has_permission_decorator(PERM_MODIFY_SYSTEM_CONFIGURATIONS, redirect_url=FOUR_OH_FOUR_URL) def api_vault(request, slug): context = {} diff --git a/web/templates/base/_items/top_nav.html b/web/templates/base/_items/top_nav.html index dff468eb..9f9c9275 100644 --- a/web/templates/base/_items/top_nav.html +++ b/web/templates/base/_items/top_nav.html @@ -81,6 +81,7 @@ {% endif %} {% if user|can:'modify_system_configurations' %} API Vault + LLM Toolkit Tools Arsenal {% endif %} {% if user|can:'modify_scan_report' %}