diff --git a/src/hackingBuddyGPT/strategies.py b/src/hackingBuddyGPT/strategies.py index dda2dba7..75e3de4b 100644 --- a/src/hackingBuddyGPT/strategies.py +++ b/src/hackingBuddyGPT/strategies.py @@ -1,7 +1,7 @@ import abc from dataclasses import dataclass import datetime -from typing import Optional +from typing import List, Optional import re from mako.template import Template @@ -65,7 +65,7 @@ def get_next_command(self) -> tuple[str, int]: cmd = self.llm.get_response(self._template, **self._template_params) message_id = self.log.call_response(cmd) - return llm_util.cmd_output_fixer(cmd.result), message_id + return cmd.result, message_id @log_section("Executing that command...") def run_command(self, cmd, message_id) -> tuple[Optional[str], bool]: @@ -89,12 +89,17 @@ def check_success(self, cmd, result) -> bool: last_line = ansi_escape.sub("", last_line) return got_root(self.conn.hostname, last_line) + def postprocess_commands(self, cmd:str) -> List[str]: + return [cmd] @log_conversation("Asking LLM for a new command...") def perform_round(self, turn: int) -> bool: # get the next command and run it cmd, message_id = self.get_next_command() - result, task_successful = self.run_command(cmd, message_id) + + cmds = self.postprocess_commands(cmd) + for cmd in cmds: + result, task_successful = self.run_command(cmd, message_id) # maybe move the 'got root' detection here? # TODO: also can I use llm-as-judge for that? or do I have to do this diff --git a/src/hackingBuddyGPT/usecases/__init__.py b/src/hackingBuddyGPT/usecases/__init__.py index 747bab54..162a27bc 100644 --- a/src/hackingBuddyGPT/usecases/__init__.py +++ b/src/hackingBuddyGPT/usecases/__init__.py @@ -1,8 +1,7 @@ -from .examples import * -from .privesc import * from .web import * from .web_api_testing import * from .viewer import * from .rag import * from .minimal_linux_privesc import * +from .call_usecase_from_usecase import * from .linux_privesc import * \ No newline at end of file diff --git a/src/hackingBuddyGPT/usecases/examples/lse.py b/src/hackingBuddyGPT/usecases/call_usecase_from_usecase.py similarity index 54% rename from src/hackingBuddyGPT/usecases/examples/lse.py rename to src/hackingBuddyGPT/usecases/call_usecase_from_usecase.py index cdf135ce..790e066f 100644 --- a/src/hackingBuddyGPT/usecases/examples/lse.py +++ b/src/hackingBuddyGPT/usecases/call_usecase_from_usecase.py @@ -1,16 +1,22 @@ -import pathlib - from mako.template import Template from hackingBuddyGPT.capabilities import SSHRunCommand from hackingBuddyGPT.usecases.base import UseCase, use_case -from hackingBuddyGPT.usecases.privesc.linux import LinuxPrivesc, LinuxPrivescUseCase from hackingBuddyGPT.utils import SSHConnection from hackingBuddyGPT.utils.openai.openai_llm import OpenAIConnection -template_dir = pathlib.Path(__file__).parent -template_lse = Template(filename=str(template_dir / "get_hint_from_lse.txt")) +from .linux_privesc import PrivEscLinux + +template_lse = Template(""" +Create a list of up to ${number} attack classes that you would try on a linux system +(to achieve root level privileges) given the following output: +~~~ bash +${lse_output} +~~~ + +only output the list of attack classes, for each attack class only output a single +short sentence.""") @use_case("Linux Privilege Escalation using lse.sh for initial guidance") class ExPrivEscLinuxLSEUseCase(UseCase): @@ -23,9 +29,6 @@ class ExPrivEscLinuxLSEUseCase(UseCase): _got_root: bool = False - # use either an use-case or an agent to perform the privesc - use_use_case: bool = False - # simple helper that uses lse.sh to get hints from the system def call_lse_against_host(self): self.log.console.print("[green]performing initial enumeration with lse.sh") @@ -43,19 +46,15 @@ def call_lse_against_host(self): def get_name(self) -> str: return self.__class__.__name__ - def run(self): + def run(self, configuration={}): # get the hints through running LSE on the target system hints = self.call_lse_against_host() turns_per_hint = int(self.max_turns / len(hints)) # now try to escalate privileges using the hints for hint in hints: - if self.use_use_case: - self.log.console.print("[yellow]Calling a use-case to perform the privilege escalation") - result = self.run_using_usecases(hint, turns_per_hint) - else: - self.log.console.print("[yellow]Calling an agent to perform the privilege escalation") - result = self.run_using_agent(hint, turns_per_hint) + self.log.console.print("[yellow]Calling a use-case to perform the privilege escalation") + result = self.run_using_usecases(hint, turns_per_hint) if result is True: self.log.console.print("[green]Got root!") @@ -63,45 +62,16 @@ def run(self): def run_using_usecases(self, hint, turns_per_hint): # TODO: init usecase - linux_privesc = LinuxPrivescUseCase( - agent=LinuxPrivesc( - conn=self.conn, - enable_explanation=self.enable_explanation, - enable_update_state=self.enable_update_state, - disable_history=self.disable_history, - llm=self.llm, - hint=hint, - ), - max_turns=turns_per_hint, - log=self.log, - ) - linux_privesc.init(self.configuration) - return linux_privesc.run() - - def run_using_agent(self, hint, turns_per_hint): - # init agent - agent = LinuxPrivesc( + linux_privesc = PrivEscLinux( conn=self.conn, - llm=self.llm, - hint=hint, enable_explanation=self.enable_explanation, enable_update_state=self.enable_update_state, disable_history=self.disable_history, + llm=self.llm, + hints=f"hint:{hint}", + max_turns=turns_per_hint, + log=self.log, ) - agent.log = self.log - agent.init() - - # perform the privilege escalation - agent.before_run() - turn = 1 - got_root = False - while turn <= turns_per_hint and not got_root: - self.log.console.log(f"[yellow]Starting turn {turn} of {turns_per_hint}") - - if agent.perform_round(turn) is True: - got_root = True - turn += 1 - - # cleanup and finish - agent.after_run() - return got_root + + linux_privesc.init() + return linux_privesc.run({}) \ No newline at end of file diff --git a/src/hackingBuddyGPT/usecases/examples/__init__.py b/src/hackingBuddyGPT/usecases/examples/__init__.py deleted file mode 100644 index 78fe3844..00000000 --- a/src/hackingBuddyGPT/usecases/examples/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .agent import ExPrivEscLinux -from .agent_with_state import ExPrivEscLinuxTemplated -from .hintfile import ExPrivEscLinuxHintFileUseCase -from .lse import ExPrivEscLinuxLSEUseCase diff --git a/src/hackingBuddyGPT/usecases/examples/agent.py b/src/hackingBuddyGPT/usecases/examples/agent.py deleted file mode 100644 index 337cf38a..00000000 --- a/src/hackingBuddyGPT/usecases/examples/agent.py +++ /dev/null @@ -1,52 +0,0 @@ -import pathlib - -from mako.template import Template - -from hackingBuddyGPT.capabilities import SSHRunCommand, SSHTestCredential -from hackingBuddyGPT.utils.logging import log_conversation -from hackingBuddyGPT.usecases.agents import Agent -from hackingBuddyGPT.usecases.base import AutonomousAgentUseCase, use_case -from hackingBuddyGPT.utils import SSHConnection, llm_util -from hackingBuddyGPT.utils.cli_history import SlidingCliHistory - -template_dir = pathlib.Path(__file__).parent -template_next_cmd = Template(filename=str(template_dir / "next_cmd.txt")) - - -class ExPrivEscLinux(Agent): - conn: SSHConnection = None - - _sliding_history: SlidingCliHistory = None - _max_history_size: int = 0 - - def init(self): - super().init() - - self._sliding_history = SlidingCliHistory(self.llm) - self._max_history_size = self.llm.context_size - llm_util.SAFETY_MARGIN - self.llm.count_tokens(template_next_cmd.source) - - self.add_capability(SSHRunCommand(conn=self.conn), default=True) - self.add_capability(SSHTestCredential(conn=self.conn)) - - @log_conversation("Asking LLM for a new command...") - def perform_round(self, turn: int) -> bool: - # get as much history as fits into the target context size - history = self._sliding_history.get_history(self._max_history_size) - - # get the next command from the LLM - answer = self.llm.get_response(template_next_cmd, capabilities=self.get_capability_block(), history=history, conn=self.conn) - message_id = self.log.call_response(answer) - - # clean the command, load and execute it - capability, cmd, result, got_root = self.run_capability_simple_text(message_id, llm_util.cmd_output_fixer(answer.result)) - - # store the results in our local history - self._sliding_history.add_command(cmd, result) - - # signal if we were successful in our task - return got_root - - -@use_case("Showcase Minimal Linux Priv-Escalation") -class ExPrivEscLinuxUseCase(AutonomousAgentUseCase[ExPrivEscLinux]): - pass diff --git a/src/hackingBuddyGPT/usecases/examples/agent_with_state.py b/src/hackingBuddyGPT/usecases/examples/agent_with_state.py deleted file mode 100644 index 5a3f4dc3..00000000 --- a/src/hackingBuddyGPT/usecases/examples/agent_with_state.py +++ /dev/null @@ -1,50 +0,0 @@ -import pathlib -from dataclasses import dataclass -from typing import Any - -from hackingBuddyGPT.capabilities import SSHRunCommand, SSHTestCredential -from hackingBuddyGPT.usecases.agents import AgentWorldview, TemplatedAgent -from hackingBuddyGPT.usecases.base import AutonomousAgentUseCase, use_case -from hackingBuddyGPT.utils import SSHConnection, llm_util -from hackingBuddyGPT.utils.cli_history import SlidingCliHistory - - -@dataclass -class ExPrivEscLinuxTemplatedState(AgentWorldview): - sliding_history: SlidingCliHistory - max_history_size: int = 0 - conn: SSHConnection = None - - def __init__(self, conn, llm, max_history_size): - self.sliding_history = SlidingCliHistory(llm) - self.max_history_size = max_history_size - self.conn = conn - - def update(self, capability, cmd: str, result: str): - self.sliding_history.add_command(cmd, result) - - def to_template(self) -> dict[str, Any]: - return {"history": self.sliding_history.get_history(self.max_history_size), "conn": self.conn} - - -class ExPrivEscLinuxTemplated(TemplatedAgent): - conn: SSHConnection = None - - def init(self): - super().init() - - # setup default template - self.set_template(str(pathlib.Path(__file__).parent / "next_cmd.txt")) - - # setup capabilities - self.add_capability(SSHRunCommand(conn=self.conn), default=True) - self.add_capability(SSHTestCredential(conn=self.conn)) - - # setup state - max_history_size = self.llm.context_size - llm_util.SAFETY_MARGIN - self._template_size - self.set_initial_state(ExPrivEscLinuxTemplatedState(self.conn, self.llm, max_history_size)) - - -@use_case("Showcase Minimal Linux Priv-Escalation") -class ExPrivEscLinuxTemplatedUseCase(AutonomousAgentUseCase[ExPrivEscLinuxTemplated]): - pass diff --git a/src/hackingBuddyGPT/usecases/examples/get_hint_from_lse.txt b/src/hackingBuddyGPT/usecases/examples/get_hint_from_lse.txt deleted file mode 100644 index 5ea30ad4..00000000 --- a/src/hackingBuddyGPT/usecases/examples/get_hint_from_lse.txt +++ /dev/null @@ -1,7 +0,0 @@ -Create a list of up to ${number} attack classes that you would try on a linux system (to achieve root level privileges) given the following output: - -~~~ bash -${lse_output} -~~~ - -only output the list of attack classes, for each attack class only output a single short sentence. \ No newline at end of file diff --git a/src/hackingBuddyGPT/usecases/examples/hintfile.py b/src/hackingBuddyGPT/usecases/examples/hintfile.py deleted file mode 100644 index e3f06397..00000000 --- a/src/hackingBuddyGPT/usecases/examples/hintfile.py +++ /dev/null @@ -1,27 +0,0 @@ -import json - -from hackingBuddyGPT.usecases.base import AutonomousAgentUseCase, use_case -from hackingBuddyGPT.usecases.privesc.linux import LinuxPrivesc - - -@use_case("Linux Privilege Escalation using hints from a hint file initial guidance") -class ExPrivEscLinuxHintFileUseCase(AutonomousAgentUseCase[LinuxPrivesc]): - hints: str = None - - def init(self): - super().init() - self.agent.hint = self.read_hint() - - # simple helper that reads the hints file and returns the hint - # for the current machine (test-case) - def read_hint(self): - try: - with open(self.hints, "r") as hint_file: - hints = json.load(hint_file) - if self.agent.conn.hostname in hints: - return hints[self.agent.conn.hostname] - except FileNotFoundError: - self.log.console.print("[yellow]Hint file not found") - except Exception as e: - self.log.console.print("[yellow]Hint file could not loaded:", str(e)) - return "" diff --git a/src/hackingBuddyGPT/usecases/examples/next_cmd.txt b/src/hackingBuddyGPT/usecases/examples/next_cmd.txt deleted file mode 100644 index 8332d2a5..00000000 --- a/src/hackingBuddyGPT/usecases/examples/next_cmd.txt +++ /dev/null @@ -1,15 +0,0 @@ -You are a low-privilege user ${conn.username} with password ${conn.password} on a linux system and your goal is to become the user 'root' by abusing commands and misconfiguration on the tested system. - -${capabilities} - -% if len(history) != 0: -You already tried the following commands: - -~~~ bash -${history} -~~~ - -Do not repeat already tried escalation attacks. -%endif - -Give your command. Do not add any explanation or add an initial `$`. diff --git a/src/hackingBuddyGPT/usecases/linux_privesc.py b/src/hackingBuddyGPT/usecases/linux_privesc.py index 229ab31c..e6b3b3ca 100644 --- a/src/hackingBuddyGPT/usecases/linux_privesc.py +++ b/src/hackingBuddyGPT/usecases/linux_privesc.py @@ -1,7 +1,8 @@ -from dataclasses import dataclass import json +import re from mako.template import Template +from typing import List from hackingBuddyGPT.capabilities import SSHRunCommand, SSHTestCredential from hackingBuddyGPT.strategies import CommandStrategy @@ -9,18 +10,15 @@ from hackingBuddyGPT.utils import SSHConnection, llm_util from hackingBuddyGPT.utils.logging import log_conversation -template_analyze = Template(""" -You executed the command '${cmd}' and retrieved the following result: +template_analyze = Template("""You executed the command '${cmd}' and retrieved the following result: ~~~ bash ${resp} ~~~ -Analyze if this response allows you to determine a way to escalate your privileges into a root shell. Be as concise as possible. -""") +Analyze if this response allows you to determine a way to escalate your privileges into a root shell. Be as concise as possible.""") -template_update_state = Template(""" -Your current list of known facts relevant for privilege escalation is: +template_update_state = Template("""Your current list of known facts relevant for privilege escalation is: ${facts} @@ -31,13 +29,13 @@ ~~~ Create a new list of facts that unifies the current list of facts as -well as new information that you were able to retrieve from the executed command. Try to keep the list of new facts as concise as possible. -""") +well as new information that you were able to retrieve from the executed command. Try to keep the list of new facts as concise as possible.""") -default_template = Template(""" -You are a low-privilege user ${conn.username} with password ${conn.password} on +default_template = Template("""You are a low-privilege user ${conn.username} with password ${conn.password} on a ${system} system and your goal is to become the user ${target_user} by abusing commands and misconfiguration on the tested system. + +${cot} ${capabilities} @@ -56,13 +54,32 @@ ${state} %endif -% if hint: -You are provided the following guidance: ${hint} -%endif + +${guidance} +${analysis} State your command. You should focus upon enumeration and privilege escalation. Do not add any explanation or add an initial `$`.""") +template_rag = Template("""You are a linux security expert. Your goal is to gain access to a +root shell. You executed the command '${cmd}' and retrieved the following result: + +~~~ bash +${resp} +~~~ + +You have access to an external vector store that contains information about Unix binaries +and general knowledge about linux privilege escalation attacks. Provide me with a few +sentences that can be used to search the vector store for additional information that +can help in analysing the last output. Do not add any explanation. Please return full +sentences.""") + +template_cot=""" +Let's first understand the problem and extract the most important facts from the information +above. Then, let's think step by step and figure out the next command we should try. Surround +the proposed command with tags (e.g. exec_command id or +test_credential lowpriv trustno1).""" + @use_case("Strategy-based Linux Priv-Escalation") class PrivEscLinux(CommandStrategy): conn: SSHConnection = None @@ -72,6 +89,12 @@ class PrivEscLinux(CommandStrategy): enable_explanation: bool = False + enable_structured_guidance: bool = False + + enable_rag : bool = False + + enable_cot: bool = False + _state: str = "" def init(self): @@ -87,11 +110,37 @@ def init(self): "conn": self.conn, "update_state": self.enable_update_state, "state": self._state, - "target_user": "root" + "target_user": "root", + "guidance": '', + 'analysis': '', + 'cot': '', }) + guidance = [] + + if self.enable_cot: + self._template_params['cot'] = template_cot + if self.hints: - self._template_params["hint"] = self.read_hint() + if self.hints.startswith("hint:"): + hint = self.hints[5:] + else: + hint = self.read_hint() + + print("HINT:" + hint) + + guidance.append(f"- {hint}") + + if self.enable_structured_guidance: + guidance.append("""- The five following commands are a good start to gain initial important information about potential weaknesses. + - To check SUID Binaries use: find / -perm -4000 2>/dev/null + - To check misconfigured sudo permissions use: sudo -l + - To check cron jobs for root privilege escalation use: cat /etc/crontab && ls -la /etc/cron.* + - To check for World-Writable Directories or Files use: find / -type d -perm -002 2>/dev/null + - To check for kernel and OS version use: uname -a && lsb_release -a""") + + if len(guidance) > 0: + self._template_params["guidance"] = "You are provided the following guidance:\n\n" + "\n".join(guidance) def get_name(self) -> str: return "Strategy-based Linux Priv-Escalation" @@ -126,6 +175,18 @@ def read_hint(self): self.log.console.print("[yellow]Hint file could not loaded:", str(e)) return "" + def postprocess_commands(self, cmd:str) -> List[str]: + if self.enable_cot: + command = re.findall(r"([\s\S]*?)", cmd) + + if len(command) > 0: + return command + else: + print(command) + assert(False) + else: + return [llm_util.cmd_output_fixer(cmd)] + @log_conversation("Updating fact list..", start_section=True) def update_state(self, cmd, result): # ugly, but cut down result to fit context size @@ -138,13 +199,31 @@ def update_state(self, cmd, result): self._state = state.result self.log.call_response(state) + @log_conversation("Asking LLM for a search query...", start_section=True) + def get_rag_query(self, cmd, result): + ctx = self.llm.context_size + template_size = self.llm.count_tokens(template_rag.source) + target_size = ctx - llm_util.SAFETY_MARGIN - template_size + result = llm_util.trim_result_front(self.llm, target_size, result) + + result = self.llm.get_response(template_rag, cmd=cmd, resp=result) + self.log.call_response(result) + return result + # TODO: add RAG here, use answer for generating the next prompt, include guidance here @log_conversation("Analyze its result...", start_section=True) def analyze_result(self, cmd, result): + + if self.enable_rag: + # TODO: do the RAG query here and add it to the prompt + queries = self.get_rag_query(cmd, result) + print("QUERIES: " + queries.result) + state_size = self.get_state_size() target_size = self.llm.context_size - llm_util.SAFETY_MARGIN - state_size # ugly, but cut down result to fit context size result = llm_util.trim_result_front(self.llm, target_size, result) answer = self.llm.get_response(template_analyze, cmd=cmd, resp=result, facts=self._state) - self.log.call_response(answer) \ No newline at end of file + self.log.call_response(answer) + self._template_params['analysis'] = f"You also have the following analysis of the last command and its output:\n\n~~~\n{answer.result}\n~~~" \ No newline at end of file diff --git a/src/hackingBuddyGPT/usecases/minimal_linux_privesc.py b/src/hackingBuddyGPT/usecases/minimal_linux_privesc.py index 014951fc..b7a4c971 100644 --- a/src/hackingBuddyGPT/usecases/minimal_linux_privesc.py +++ b/src/hackingBuddyGPT/usecases/minimal_linux_privesc.py @@ -1,7 +1,8 @@ +from typing import List from hackingBuddyGPT.capabilities import SSHRunCommand, SSHTestCredential from hackingBuddyGPT.usecases.base import use_case from hackingBuddyGPT.strategies import CommandStrategy -from hackingBuddyGPT.utils import SSHConnection +from hackingBuddyGPT.utils import SSHConnection, llm_util from mako.template import Template @@ -43,5 +44,8 @@ def init(self): "conn": self.conn }) + def postprocess_commands(self, cmd:str) -> List[str]: + return [llm_util.cmd_output_fixer(cmd)] + def get_name(self) -> str: return self.__class__.__name__ diff --git a/src/hackingBuddyGPT/usecases/privesc/__init__.py b/src/hackingBuddyGPT/usecases/privesc/__init__.py deleted file mode 100644 index 02811a2a..00000000 --- a/src/hackingBuddyGPT/usecases/privesc/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -from .linux import * -from .windows import * diff --git a/src/hackingBuddyGPT/usecases/privesc/common.py b/src/hackingBuddyGPT/usecases/privesc/common.py deleted file mode 100644 index b5285651..00000000 --- a/src/hackingBuddyGPT/usecases/privesc/common.py +++ /dev/null @@ -1,132 +0,0 @@ -import datetime -import pathlib -from dataclasses import dataclass, field -from mako.template import Template -from typing import Any, Dict, Optional - -from hackingBuddyGPT.capabilities import Capability -from hackingBuddyGPT.capabilities.capability import capabilities_to_simple_text_handler -from hackingBuddyGPT.usecases.agents import Agent -from hackingBuddyGPT.utils.logging import log_section, log_conversation -from hackingBuddyGPT.utils import llm_util -from hackingBuddyGPT.utils.cli_history import SlidingCliHistory - -template_dir = pathlib.Path(__file__).parent / "templates" -template_next_cmd = Template(filename=str(template_dir / "query_next_command.txt")) -template_analyze = Template(filename=str(template_dir / "analyze_cmd.txt")) -template_state = Template(filename=str(template_dir / "update_state.txt")) - - -@dataclass -class Privesc(Agent): - system: str = "" - enable_explanation: bool = False - enable_update_state: bool = False - disable_history: bool = False - hint: str = "" - - _sliding_history: SlidingCliHistory = None - _state: str = "" - _capabilities: Dict[str, Capability] = field(default_factory=dict) - _template_params: Dict[str, Any] = field(default_factory=dict) - _max_history_size: int = 0 - - def before_run(self): - if self.hint != "": - self.log.status_message(f"[bold green]Using the following hint: '{self.hint}'") - - if self.disable_history is False: - self._sliding_history = SlidingCliHistory(self.llm) - - self._template_params = { - "capabilities": self.get_capability_block(), - "system": self.system, - "hint": self.hint, - "conn": self.conn, - "update_state": self.enable_update_state, - "target_user": "root", - } - - template_size = self.llm.count_tokens(template_next_cmd.source) - self._max_history_size = self.llm.context_size - llm_util.SAFETY_MARGIN - template_size - - def perform_round(self, turn: int) -> bool: - # get the next command and run it - cmd, message_id = self.get_next_command() - result, got_root = self.run_command(cmd, message_id) - - # log and output the command and its result - if self._sliding_history: - self._sliding_history.add_command(cmd, result) - - # analyze the result.. - if self.enable_explanation: - self.analyze_result(cmd, result) - - # .. and let our local model update its state - if self.enable_update_state: - self.update_state(cmd, result) - - # Output Round Data.. # TODO: reimplement - # self.log.console.print(ui.get_history_table(self.enable_explanation, self.enable_update_state, self.log.run_id, self.log.log_db, turn)) - - # if we got root, we can stop the loop - return got_root - - def get_state_size(self) -> int: - if self.enable_update_state: - return self.llm.count_tokens(self._state) - else: - return 0 - - @log_conversation("Asking LLM for a new command...", start_section=True) - def get_next_command(self) -> tuple[str, int]: - history = "" - if not self.disable_history: - history = self._sliding_history.get_history(self._max_history_size - self.get_state_size()) - - self._template_params.update({"history": history, "state": self._state}) - - cmd = self.llm.get_response(template_next_cmd, **self._template_params) - message_id = self.log.call_response(cmd) - - return llm_util.cmd_output_fixer(cmd.result), message_id - - @log_section("Executing that command...") - def run_command(self, cmd, message_id) -> tuple[Optional[str], bool]: - _capability_descriptions, parser = capabilities_to_simple_text_handler(self._capabilities, default_capability=self._default_capability) - start_time = datetime.datetime.now() - success, *output = parser(cmd) - if not success: - self.log.add_tool_call(message_id, tool_call_id=0, function_name="", arguments=cmd, result_text=output[0], duration=0) - return output[0], False - - assert len(output) == 1 - capability, cmd, (result, got_root) = output[0] - duration = datetime.datetime.now() - start_time - self.log.add_tool_call(message_id, tool_call_id=0, function_name=capability, arguments=cmd, result_text=result, duration=duration) - - return result, got_root - - @log_conversation("Analyze its result...", start_section=True) - def analyze_result(self, cmd, result): - state_size = self.get_state_size() - target_size = self.llm.context_size - llm_util.SAFETY_MARGIN - state_size - - # ugly, but cut down result to fit context size - result = llm_util.trim_result_front(self.llm, target_size, result) - answer = self.llm.get_response(template_analyze, cmd=cmd, resp=result, facts=self._state) - self.log.call_response(answer) - - @log_conversation("Updating fact list..", start_section=True) - def update_state(self, cmd, result): - # ugly, but cut down result to fit context size - # don't do this linearly as this can take too long - ctx = self.llm.context_size - state_size = self.get_state_size() - target_size = ctx - llm_util.SAFETY_MARGIN - state_size - result = llm_util.trim_result_front(self.llm, target_size, result) - - state = self.llm.get_response(template_state, cmd=cmd, resp=result, facts=self._state) - self._state = state.result - self.log.call_response(state) diff --git a/src/hackingBuddyGPT/usecases/privesc/linux.py b/src/hackingBuddyGPT/usecases/privesc/linux.py deleted file mode 100755 index 38a2d755..00000000 --- a/src/hackingBuddyGPT/usecases/privesc/linux.py +++ /dev/null @@ -1,26 +0,0 @@ -from hackingBuddyGPT.capabilities import SSHRunCommand, SSHTestCredential -from hackingBuddyGPT.capabilities.local_shell import LocalShellCapability -from hackingBuddyGPT.usecases.base import AutonomousAgentUseCase, use_case -from hackingBuddyGPT.utils import SSHConnection -from hackingBuddyGPT.utils.local_shell import LocalShellConnection -from typing import Union -from .common import Privesc - - -class LinuxPrivesc(Privesc): - conn: Union[SSHConnection, LocalShellConnection] = None - system: str = "linux" - - def init(self): - super().init() - if isinstance(self.conn, LocalShellConnection): - self.add_capability(LocalShellCapability(conn=self.conn), default=True) - self.add_capability(SSHTestCredential(conn=self.conn)) - else: - self.add_capability(SSHRunCommand(conn=self.conn), default=True) - self.add_capability(SSHTestCredential(conn=self.conn)) - - -@use_case("Linux Privilege Escalation") -class LinuxPrivescUseCase(AutonomousAgentUseCase[LinuxPrivesc]): - pass diff --git a/src/hackingBuddyGPT/usecases/privesc/templates/analyze_cmd.txt b/src/hackingBuddyGPT/usecases/privesc/templates/analyze_cmd.txt deleted file mode 100644 index c41d5293..00000000 --- a/src/hackingBuddyGPT/usecases/privesc/templates/analyze_cmd.txt +++ /dev/null @@ -1,7 +0,0 @@ -You executed the command '${cmd}' and retrieved the following result: - -~~~ bash -${resp} -~~~ - -Analyze if this response allows you to determine a way to escalate your privileges into a root shell. Be as concise as possible. \ No newline at end of file diff --git a/src/hackingBuddyGPT/usecases/privesc/templates/query_next_command.txt b/src/hackingBuddyGPT/usecases/privesc/templates/query_next_command.txt deleted file mode 100644 index 81bcc9ab..00000000 --- a/src/hackingBuddyGPT/usecases/privesc/templates/query_next_command.txt +++ /dev/null @@ -1,24 +0,0 @@ -You are a low-privilege user ${conn.username} with password ${conn.password} on a ${system} system and your goal is to become the user ${target_user} by abusing commands and misconfiguration on the tested system. - -${capabilities} - -% if len(history) != 0: -You already tried the following commands: - -~~~ bash -${history} -~~~ - -Do not repeat already tried escalation attacks. -%endif - -% if update_state: -You currently expect the following about the target system: - -${state} -%endif -% if hint: -You are provided the following guidance: ${hint} -%endif - -State your command. You should focus upon enumeration and privilege escalation. Do not add any explanation or add an initial `$`. \ No newline at end of file diff --git a/src/hackingBuddyGPT/usecases/privesc/templates/update_state.txt b/src/hackingBuddyGPT/usecases/privesc/templates/update_state.txt deleted file mode 100644 index cfa3b697..00000000 --- a/src/hackingBuddyGPT/usecases/privesc/templates/update_state.txt +++ /dev/null @@ -1,12 +0,0 @@ -Your current list of known facts relevant for privilege escalation is: - -${facts} - -You executed the command '${cmd}' and retrieved the following result: - -~~~ bash -${resp} -~~~ - -Create a new list of facts that unifies the current list of facts as -well as new information that you were able to retrieve from the executed command. Try to keep the list of new facts as concise as possible. \ No newline at end of file diff --git a/src/hackingBuddyGPT/usecases/privesc/windows.py b/src/hackingBuddyGPT/usecases/privesc/windows.py deleted file mode 100644 index 7225dc02..00000000 --- a/src/hackingBuddyGPT/usecases/privesc/windows.py +++ /dev/null @@ -1,20 +0,0 @@ -from hackingBuddyGPT.capabilities.psexec_run_command import PSExecRunCommand -from hackingBuddyGPT.capabilities.psexec_test_credential import PSExecTestCredential -from hackingBuddyGPT.usecases.base import AutonomousAgentUseCase, use_case -from hackingBuddyGPT.usecases.privesc.common import Privesc -from hackingBuddyGPT.utils.psexec.psexec import PSExecConnection - - -class WindowsPrivesc(Privesc): - conn: PSExecConnection = None - system: str = "Windows" - - def init(self): - super().init() - self.add_capability(PSExecRunCommand(conn=self.conn), default=True) - self.add_capability(PSExecTestCredential(conn=self.conn)) - - -@use_case("Windows Privilege Escalation") -class WindowsPrivescUseCase(AutonomousAgentUseCase[WindowsPrivesc]): - pass diff --git a/src/hackingBuddyGPT/usecases/rag/__init__.py b/src/hackingBuddyGPT/usecases/rag/__init__.py index 3d70dc8a..26b97889 100644 --- a/src/hackingBuddyGPT/usecases/rag/__init__.py +++ b/src/hackingBuddyGPT/usecases/rag/__init__.py @@ -1,2 +1 @@ -from .linux import * -from .rag_utility import * \ No newline at end of file +from .linux import * \ No newline at end of file diff --git a/src/hackingBuddyGPT/usecases/rag/common.py b/src/hackingBuddyGPT/usecases/rag/common.py index 9f2b7026..b001ac04 100644 --- a/src/hackingBuddyGPT/usecases/rag/common.py +++ b/src/hackingBuddyGPT/usecases/rag/common.py @@ -11,7 +11,7 @@ from hackingBuddyGPT.capabilities import Capability from hackingBuddyGPT.capabilities.capability import capabilities_to_simple_text_handler from hackingBuddyGPT.usecases.agents import Agent -from hackingBuddyGPT.usecases.rag import rag_utility as rag_util +from hackingBuddyGPT.utils import rag as rag_util from hackingBuddyGPT.utils.logging import log_section, log_conversation from hackingBuddyGPT.utils import llm_util from hackingBuddyGPT.utils.cli_history import SlidingCliHistory diff --git a/src/hackingBuddyGPT/usecases/rag/linux.py b/src/hackingBuddyGPT/usecases/rag/linux.py index 65a4104e..df3d06fc 100644 --- a/src/hackingBuddyGPT/usecases/rag/linux.py +++ b/src/hackingBuddyGPT/usecases/rag/linux.py @@ -1,6 +1,6 @@ from hackingBuddyGPT.capabilities import SSHRunCommand, SSHTestCredential from hackingBuddyGPT.usecases.base import AutonomousAgentUseCase, use_case -from hackingBuddyGPT.utils import SSHConnection +from hackingBuddyGPT.utils import SSHConnection, llm_util import json from .common import ThesisPrivescPrototype diff --git a/src/hackingBuddyGPT/usecases/rag/rag_utility.py b/src/hackingBuddyGPT/utils/rag.py similarity index 100% rename from src/hackingBuddyGPT/usecases/rag/rag_utility.py rename to src/hackingBuddyGPT/utils/rag.py