Merge pull request #137 from ipa-lab/restructuring

Clean-Up dependencies

Author: andreashappe

pyproject.toml

@@ -38,21 +38,13 @@ dependencies = [
     'pydantic == 2.8.2',
     'openai == 1.65.2',
     'BeautifulSoup4',
-    'nltk',
     'fastapi == 0.114.0',
     'fastapi-utils == 0.7.0',
     'uvicorn[standard] == 0.30.6',
     'dataclasses_json == 0.6.7',
     'websockets == 13.1',
-    'pandas',
     'faker',
     'fpdf',
-    'langchain_core',
-    'langchain_community',
-    'langchain_chroma',
-    'langchain_openai',
-    'markdown',
-    'chromadb',
 ]
 
 [project.urls]
@@ -71,17 +63,26 @@ where = ["src"]
 [tool.pytest.ini_options]
 pythonpath = "src"
 addopts = ["--import-mode=importlib"]
+
 [project.optional-dependencies]
-testing = ['pytest', 'pytest-mock', 'pandas', 'faker',     'langchain_core']
+testing = [
+	'pytest',
+	'pytest-mock',
+	'faker',
+	'langchain_core'
+]
+
 dev = [
 	'ruff',
 ]
-rag-usecase = [
-    'langchain-community',
+
+rag = [
+	'langchain_core',
+	'langchain-community',
+	'langchain-chroma',
 	'langchain-openai',
 	'markdown',
 	'chromadb',
-	'langchain-chroma',
 ]
 
 [project.scripts]

src/hackingBuddyGPT/usecases/linux_privesc.py

@@ -9,10 +9,13 @@
 from hackingBuddyGPT.usecases.usecase import use_case
 from hackingBuddyGPT.utils import llm_util
 from hackingBuddyGPT.utils.logging import log_conversation
-from hackingBuddyGPT.utils.rag import RagBackground
+from hackingBuddyGPT.utils.rag import has_langchain
 from hackingBuddyGPT.utils.connectors.ssh_connection import SSHConnection
 from hackingBuddyGPT.utils.shell_root_detection import got_root
 
+if has_langchain():
+    from hackingBuddyGPT.utils.rag import RagBackground
+
 template_analyze = Template("""Your task is to analyze the result of an executed command to determina 
 a way to escalate your privileges into a root shell. Describe your findings including all needed
 information while being as concise as possible.
@@ -132,6 +135,10 @@ def init(self):
         guidance = []
 
         if self.rag_path != '':
+            if not has_langchain():
+                self.log.console.print("[red]RAG path provided but langchain is not installed. Please install langchain to use RAG functionality, e.g., through `pip install -e .\[rag]`.[/red]")
+                raise ImportError("langchain is not installed")
+
             self._enable_rag = True
             self._rag_data = RagBackground(self.rag_path, self.llm)
 
@@ -254,4 +261,4 @@ def check_success(self, cmd:str, result:str) -> bool:
         ansi_escape = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
         last_line = result.split("\n")[-1] if result else ""
         last_line = ansi_escape.sub("", last_line)
-        return got_root(self.conn.hostname, last_line)
+        return got_root(self.conn.hostname, last_line)

src/hackingBuddyGPT/utils/prompt_generation/information/pentesting_information.py

@@ -1,14 +1,13 @@
 import base64
 import copy
+import csv
 import glob
 import os
 import random
 import re
 import secrets
 from typing import List
 
-import pandas
-
 from hackingBuddyGPT.utils.openapi.openapi_parser import OpenAPISpecificationParser
 from hackingBuddyGPT.utils.prompt_generation.information.prompt_information import (
     PromptPurpose,
@@ -43,10 +42,15 @@ def __init__(self, openapi_spec_parser: OpenAPISpecificationParser, config) -> N
         self.available_numbers = []
         self.config = config
         file = self.get_file(self.config.get("csv_file"))
-        if file == "Not found":
-            self.df = pandas.DataFrame()
-        else:
-            self.df = pandas.read_csv(file[0], names=["username", "password"])
+        self._credentials = []
+        if file != "Not found":
+            with open(file[0]) as csvfile:
+                reader = csv.reader(csvfile, delimiter=',')
+                for row in reader:
+                    self._credentials.append({
+                        'username': row[0],
+                        'password': row[1]
+                })
 
         # Parse endpoints and their categorization from the given parser instance
         categorized_endpoints = openapi_spec_parser.classify_endpoints(self.config.get("name"))
@@ -1486,14 +1490,11 @@ def mechanic_report(self, endpoint, account, prompts):
         return prompts
 
     def random_common_users(self, endpoint, login_path, login_schema, prompts):
-        if len(self.df) >= 10:
-            random_entries = self.df.sample(n=10,
-                                            random_state=42)  # Adjust random_state for different samples
-        else:
-            # Either raise an error, sample fewer, or handle gracefully
-            random_entries = self.df.sample(n=len(self.df)) if len(self.df) > 0 else pandas.DataFrame()
+        random_entries = random.sample(self._credentials)
+        if len(random_entries) >= 10:
+            random_entries = random_entries[:10]
 
-        for index, random_entry in random_entries.iterrows():
+        for index, random_entry in random_entries:
             username = random_entry['username']
             password = random_entry['password']
             # Now you can print or use username and password as needed

src/hackingBuddyGPT/utils/rag.py

@@ -1,7 +1,15 @@
-from langchain_community.document_loaders import DirectoryLoader, TextLoader
-from langchain_chroma import Chroma
-from langchain_openai import OpenAIEmbeddings
-from langchain_text_splitters import MarkdownTextSplitter
+try:
+    from langchain_community.document_loaders import DirectoryLoader, TextLoader
+    from langchain_chroma import Chroma
+    from langchain_openai import OpenAIEmbeddings
+    from langchain_text_splitters import MarkdownTextSplitter
+except ImportError:
+    _has_langchain = False
+else:
+    _has_langchain = True
+
+def has_langchain():
+    return _has_langchain
 
 class RagBackground: