Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Next Next commit
add a new minimal use-case for testing
  • Loading branch information
andreashappe committed Apr 23, 2024
commit 3de2f537dbffdb0fd6e1c3f46b0fbac47a820207
1 change: 1 addition & 0 deletions usecases/__init__.py
Original file line number Diff line number Diff line change
@@ -1,2 +1,3 @@
from .usecase import *
from .privesc import *
from .minimal import *
1 change: 1 addition & 0 deletions usecases/minimal/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
from .minimal import MinimalLinuxPrivesc
58 changes: 58 additions & 0 deletions usecases/minimal/minimal.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,58 @@
import abc
import pathlib
from dataclasses import dataclass, field
from typing import Dict

from mako.template import Template
from rich.panel import Panel

from capabilities import Capability, SSHRunCommand, SSHTestCredential
from utils import SSHConnection, llm_util, ui
from usecases.usecase import use_case, UseCase
from usecases.usecase.roundbased import RoundBasedUseCase
from utils.cli_history import SlidingCliHistory

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

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

conn: SSHConnection = None

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

def init(self):
super().init()
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..."):
# 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(answer.result)

with self.console.status("[bold green]Executing that command..."):
if answer.result.startswith("test_credential"):
result, got_root = self._capabilities["test_credential"](cmd)
else:
self.console.print(Panel(answer.result, title="[bold cyan]Got command from LLM:"))
result, got_root = self._capabilities["run_command"](cmd)

# log and output the command and its result
self.log_db.add_log_query(self._run_id, turn, cmd, result, answer)
self._sliding_history.add_command(cmd, result)
self.console.print(Panel(result, title=f"[bold cyan]{cmd}"))

# if we got root, we can stop the loop
return got_root
19 changes: 19 additions & 0 deletions usecases/minimal/next_cmd.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
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.

You can either

% for name, capability in _capabilities.items():
- ${capability.describe(name)}
% endfor

% 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 `$`.
2 changes: 1 addition & 1 deletion utils/cli_history.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ class SlidingCliHistory:
maximum_target_size: int = 0
sliding_history: str = ''

def init(self, used_model: LLM):
def __init__(self, used_model: LLM):
self.model = used_model
self.maximum_target_size = self.model.context_size

Expand Down