Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 13 additions & 12 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand All @@ -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]
Expand Down
11 changes: 9 additions & 2 deletions src/hackingBuddyGPT/usecases/linux_privesc.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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)
Original file line number Diff line number Diff line change
@@ -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,
Expand Down Expand Up @@ -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"))
Expand Down Expand Up @@ -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
Expand Down
16 changes: 12 additions & 4 deletions src/hackingBuddyGPT/utils/rag.py
Original file line number Diff line number Diff line change
@@ -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:

Expand Down
Loading