Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

update minimal example #27

Merged
merged 1 commit into from
Apr 23, 2024
Merged

update minimal example #27

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
33 changes: 14 additions & 19 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -100,29 +100,35 @@ This work is partially based upon our empiric research into [how hackers work](h
The following would create a new (minimal) linux privilege-escalation agent. Through using our infrastructure, this already uses configurable LLM-connections (e.g., for testing OpenAI or locally run LLMs), logs trace data to a local sqlite database for each run, implements a round limit (after which the agent will stop if root has not been achieved until then) and is able to connect to a linux target over SSH for fully-autonomous command execution (as well as password guessing).

~~~ python
template_dir = pathlib.Path(__file__).parent / "templates"
template_next_cmd = Template(filename=str(template_dir / "query_next_command.txt"))
template_dir = pathlib.Path(__file__).parent
template_next_cmd = Template(filename=str(template_dir / "next_cmd.txt"))

@use_case("linux_privesc", "Linux Privilege Escalation")
@use_case("minimal_linux_privesc", "Showcase Minimal Linux Priv-Escalation")
@dataclass
class LinuxPrivesc(RoundBasedUseCase, UseCase, abc.ABC):
class MinimalLinuxPrivesc(RoundBasedUseCase, UseCase, abc.ABC):

system: str = 'linux'
conn: SSHConnection = None

_sliding_history: SlidingCliHistory = None
_capabilities: Dict[str, Capability] = field(default_factory=dict)

def init(self):
super().init()
# provide a shell history of limited size to the LLM
self._sliding_history = SlidingCliHistory(self.llm)
self._capabilities["run_command"] = SSHRunCommand(conn=self.conn)
self._capabilities["test_credential"] = SSHTestCredential(conn=self.conn)
self._template_size = self.llm.count_tokens(template_next_cmd.source)

def perform_round(self, turn):
got_root : bool = False

with self.console.status("[bold green]Asking LLM for a new command..."):
answer = self.get_next_command()
cmd = answer.result
# get as much history as fits into the target context size
history = self._sliding_history.get_history(self.llm.context_size - llm_util.SAFETY_MARGIN - self._template_size)

# get the next command from the LLM
answer = self.llm.get_response(template_next_cmd, _capabilities=self._capabilities, history=history, conn=self.conn)
cmd = llm_util.cmd_output_fixer(cmd.result)

with self.console.status("[bold green]Executing that command..."):
if answer.result.startswith("test_credential"):
Expand All @@ -136,19 +142,8 @@ class LinuxPrivesc(RoundBasedUseCase, UseCase, abc.ABC):
self._sliding_history.add_command(cmd, result)
self.console.print(Panel(result, title=f"[bold cyan]{cmd}"))

# Output Round Data..
self.console.print(ui.get_history_table(False, False, self._run_id, self.log_db, turn))

# if we got root, we can stop the loop
return got_root

def get_next_command(self):
template_size = self.llm.count_tokens(template_next_cmd.source)
history = self._sliding_history.get_history(self.llm.context_size - llm_util.SAFETY_MARGIN - template_size)

cmd = self.llm.get_response(template_next_cmd, _capabilities=self._capabilities, history=history, conn=self.conn, system=self.system, target_user="root")
cmd.result = llm_util.cmd_output_fixer(cmd.result)
return cmd
~~~

## Setup and Usage
Expand Down