Merge pull request #1259 from yogeshojha/feature/llm

Add LLM Toolkit

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!*
 

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
 
 # --------------------------
 

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:

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:

web/api/urls.py

@@ -178,6 +178,10 @@
         'tool/uninstall/',
         UninstallTool.as_view(),
         name='uninstall_tool'),
+	path(
+        'tool/ollama/',
+        OllamaManager.as_view(),
+        name='ollama_manager'),
     path(
         'rengine/update/',
         RengineUpdateCheck.as_view(),

web/api/views.py

@@ -33,6 +33,82 @@
 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

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)),
+            ],
+        ),
+    ]

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),
+        ),
+    ]

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',
+        ),
+    ]

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),
+        ),
+    ]

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',
+        ),
+    ]

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):

web/reNgine/definitions.py

@@ -436,6 +436,52 @@
 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.