diff --git a/pyproject.toml b/pyproject.toml index 52c2586..168b2d7 100644 --- a/pyproject.toml +++ b/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] diff --git a/src/hackingBuddyGPT/usecases/linux_privesc.py b/src/hackingBuddyGPT/usecases/linux_privesc.py index f2f2d71..115bd15 100644 --- a/src/hackingBuddyGPT/usecases/linux_privesc.py +++ b/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) \ No newline at end of file + return got_root(self.conn.hostname, last_line) diff --git a/src/hackingBuddyGPT/utils/prompt_generation/information/pentesting_information.py b/src/hackingBuddyGPT/utils/prompt_generation/information/pentesting_information.py index 39ada3c..dfd6a24 100644 --- a/src/hackingBuddyGPT/utils/prompt_generation/information/pentesting_information.py +++ b/src/hackingBuddyGPT/utils/prompt_generation/information/pentesting_information.py @@ -1,5 +1,6 @@ import base64 import copy +import csv import glob import os import random @@ -7,8 +8,6 @@ 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 diff --git a/src/hackingBuddyGPT/utils/rag.py b/src/hackingBuddyGPT/utils/rag.py index 3fddaa3..c93efbf 100644 --- a/src/hackingBuddyGPT/utils/rag.py +++ b/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: