Merge pull request #61 from ipa-lab/fixes

Fixes

Author: andreashappe

capabilities/capability.py

@@ -1,6 +1,6 @@
 import abc
 import inspect
-from typing import Union, Type, Dict
+from typing import Union, Type, Dict, Callable, Any
 
 from pydantic import create_model, BaseModel
 
@@ -78,3 +78,95 @@ class Model(Action):
 
     return Model
 
+
+SimpleTextHandlerResult = tuple[bool, Union[str, tuple[str, str, ...]]]
+SimpleTextHandler = Callable[[str], SimpleTextHandlerResult]
+
+
+def capabilities_to_simple_text_handler(capabilities: Dict[str, Capability], default_capability: Capability = None, include_description: bool = True) -> tuple[Dict[str, str], SimpleTextHandler]:
+    """
+    This function generates a simple text handler from a set of capabilities.
+    It is to be used when no function calling is available, and structured output is not to be trusted, which is why it
+    only supports the most basic of parameter types for the capabilities (str, int, float, bool).
+
+    As result it returns a dictionary of capability names to their descriptions and a parser function that can be used
+    to parse the text input and execute it. The first return value of the parser function is a boolean indicating
+    whether the parsing was successful, the second return value is a tuple containing the capability name, the parameters
+    as a string and the result of the capability execution.
+    """
+    def get_simple_fields(func, name) -> Dict[str, Type]:
+        sig = inspect.signature(func)
+        fields = {param: param_info.annotation for param, param_info in sig.parameters.items()}
+        for param, param_type in fields.items():
+            if param_type not in (str, int, float, bool):
+                raise ValueError(f"The command {name} is not compatible with this calling convention (this is not a LLM error, but rather a problem with the capability itself, the parameter {param} is {param_type} and not a simple type (str, int, float, bool))")
+        return fields
+
+    def parse_params(fields, params) -> tuple[bool, Union[str, Dict[str, Any]]]:
+        split_params = params.split(" ", maxsplit=len(fields) - 1)
+        if len(split_params) != len(fields):
+            return False, "Invalid number of parameters"
+
+        parsed_params = dict()
+        for param, param_type in fields.items():
+            try:
+                parsed_params[param] = param_type(split_params.pop(0))
+            except ValueError as e:
+                return False, f"Could not parse parameter {param}: {e}"
+        return True, parsed_params
+
+    capability_descriptions = dict()
+    capability_params = dict()
+    for capability_name, capability in capabilities.items():
+        fields = get_simple_fields(capability.__call__, capability_name)
+
+        description = f"`{capability_name}"
+        if len(fields) > 0:
+            description += " " + " ".join(param for param in fields)
+        description += "`"
+        if include_description:
+            description += f": {capability.describe()}"
+
+        capability_descriptions[capability_name] = description
+        capability_params[capability_name] = fields
+
+    def parser(text: str) -> SimpleTextHandlerResult:
+        capability_name_and_params = text.split(" ", maxsplit=1)
+        if len(capability_name_and_params) == 1:
+            capability_name = capability_name_and_params[0]
+            params = ""
+        else:
+            capability_name, params = capability_name_and_params
+        if capability_name not in capabilities:
+            return False, "Unknown command"
+
+        success, parsing_result = parse_params(capability_params[capability_name], params)
+        if not success:
+            return False, parsing_result
+
+        return True, (capability_name, params, capabilities[capability_name](**parsing_result))
+
+    resolved_parser: SimpleTextHandler = parser
+
+    if default_capability is not None:
+        default_fields = get_simple_fields(default_capability.__call__, "__default__")
+
+        def default_capability_parser(text: str) -> SimpleTextHandlerResult:
+            success, *output = parser(text)
+            if success:
+                return success, *output
+
+            params = text
+            success, parsing_result = parse_params(default_fields, params)
+            if not success:
+                params = text.split(" ", maxsplit=1)[1]
+                success, parsing_result = parse_params(default_fields, params)
+                if not success:
+                    return False, parsing_result
+
+            return True, (capability_name, params, default_capability(**parsing_result))
+
+
+        resolved_parser = default_capability_parser
+
+    return capability_descriptions, resolved_parser

capabilities/psexec_test_credential.py

@@ -11,7 +11,7 @@ class PSExecTestCredential(Capability):
     conn: PSExecConnection
 
     def describe(self) -> str:
-        return f"give credentials to be tested by stating `{self.get_name()} username password`"
+        return f"give credentials to be tested"
 
     def get_name(self) -> str:
         return "test_credential"

capabilities/ssh_run_command.py

@@ -18,12 +18,21 @@
 @dataclass
 class SSHRunCommand(Capability):
     conn: SSHConnection
+    timeout: int = 10
 
     def describe(self) -> str:
-        return f"give a command to be executed on the shell and I will respond with the terminal output when running this command on the linux server. The given command must not require user interaction. Only state the to be executed command. The command should be used for enumeration or privilege escalation."
+        return f"give a command to be executed and I will respond with the terminal output when running this command over SSH on the linux machine. The given command must not require user interaction."
 
-    def __call__(self, command: str, timeout:int=10) -> Tuple[str, bool]:
+    def get_name(self):
+        return "exec_command"
+
+    def __call__(self, command: str) -> Tuple[str, bool]:
         got_root = False
+
+        if command.startswith(self.get_name()):
+            cmd_parts = command.split(" ", 1)
+            command = cmd_parts[1]
+
         sudo_pass = Responder(
             pattern=r'\[sudo\] password for ' + self.conn.username + ':',
             response=self.conn.password + '\n',
@@ -32,7 +41,7 @@ def __call__(self, command: str, timeout:int=10) -> Tuple[str, bool]:
         out = StringIO()
 
         try:
-            resp = self.conn.run(command, pty=True, warn=True, out_stream=out, watchers=[sudo_pass], timeout=timeout)
+            resp = self.conn.run(command, pty=True, warn=True, out_stream=out, watchers=[sudo_pass], timeout=self.timeout)
         except Exception as e:
             print("TIMEOUT! Could we have become root?")
         out.seek(0)

capabilities/ssh_test_credential.py

@@ -12,19 +12,13 @@ class SSHTestCredential(Capability):
     conn: SSHConnection
 
     def describe(self) -> str:
-        return f"give credentials to be tested by stating `{self.get_name()} username password`"
+        return f"give credentials to be tested"
 
     def get_name(self):
         return "test_credential"
 
-    def __call__(self, command: str) -> Tuple[str, bool]:
-        cmd_parts = command.split(" ")
-        assert (cmd_parts[0] == "test_credential")
-
-        if len(cmd_parts) != 3:
-            return "didn't provide username/password", False
-
-        test_conn = self.conn.new_with(username=cmd_parts[1], password=cmd_parts[2])
+    def __call__(self, username: str, password: str) -> Tuple[str, bool]:
+        test_conn = self.conn.new_with(username=username, password=password)
         try:
             test_conn.init()
             user = test_conn.run("whoami")[0].strip('\n\r ')

requirements.txt

@@ -27,3 +27,4 @@ tiktoken==0.6.0
 urllib3==2.2.1
 wrapt==1.16.0
 instructor==1.2.2
+PyYAML==6.0.1

usecases/agents.py

@@ -1,26 +1,27 @@
+from abc import ABC
 from dataclasses import dataclass, field
 from typing import Dict
 
-from capabilities.capability import Capability
+from capabilities.capability import Capability, capabilities_to_simple_text_handler
 from usecases.common_patterns import RoundBasedUseCase
 
 
 @dataclass
-class Agent(RoundBasedUseCase):
-
+class Agent(RoundBasedUseCase, ABC):
     _capabilities: Dict[str, Capability] = field(default_factory=dict)
     _default_capability: Capability = None
 
     def init(self):
         super().init()
 
-    def add_capability(self, cap:Capability, default:bool=False):
+    def add_capability(self, cap: Capability, default: bool = False):
         self._capabilities[cap.get_name()] = cap
         if default:
             self._default_capability = cap
 
-    def get_capability(self, name:str) -> Capability:
+    def get_capability(self, name: str) -> Capability:
         return self._capabilities.get(name, self._default_capability)
 
     def get_capability_block(self) -> str:
-        return "You can either\n\n" + "\n".join(map(lambda i: f"- {i.describe()}", self._capabilities.values()))
+        capability_descriptions, _parser = capabilities_to_simple_text_handler(self._capabilities)
+        return "You can either\n\n" + "\n".join(f"- {description}" for description in capability_descriptions.values())

usecases/privesc/common.py

@@ -6,6 +6,7 @@
 from rich.panel import Panel
 
 from capabilities import Capability
+from capabilities.capability import capabilities_to_simple_text_handler
 from usecases.agents import Agent
 from utils import llm_util, ui
 from utils.cli_history import SlidingCliHistory
@@ -48,12 +49,20 @@ def perform_round(self, turn):
 
         with self.console.status("[bold green]Executing that command..."):
             self.console.print(Panel(answer.result, title="[bold cyan]Got command from LLM:"))
-            result, got_root = self.get_capability(cmd.split(" ", 1)[0])(cmd)
+            _capability_descriptions, parser = capabilities_to_simple_text_handler(self._capabilities, default_capability=self._default_capability)
+            success, *output = parser(cmd)
+            if not success:
+                self.console.print(Panel(output[0], title=f"[bold red]Error parsing command:"))
+                return False
+
+            assert(len(output) == 1)
+            capability, cmd, (result, got_root) = output[0]
 
         # log and output the command and its result
         self.log_db.add_log_query(self._run_id, turn, cmd, result, answer)
         if self._sliding_history:
             self._sliding_history.add_command(cmd, result)
+
         self.console.print(Panel(result, title=f"[bold cyan]{cmd}"))
 
         # analyze the result..

usecases/privesc/linux.py

@@ -104,7 +104,7 @@ def read_hint(self):
 
         run_cmd = "wget -q 'https://github.com/diego-treitos/linux-smart-enumeration/releases/latest/download/lse.sh' -O lse.sh;chmod 700 lse.sh; ./lse.sh -c -i -l 0 | grep -v 'nope$' | grep -v 'skip$'"
 
-        result, got_root = SSHRunCommand(conn=self.conn)(run_cmd, timeout=120)
+        result, got_root = SSHRunCommand(conn=self.conn, timeout=120)(run_cmd)
 
         self.console.print("[yellow]got the output: " + result)
         cmd = self.llm.get_response(template_lse, lse_output=result, number=3)

usecases/privesc/templates/query_next_command.txt

@@ -18,7 +18,7 @@ You currently expect the following about the target system:
 ${state}
 %endif
 % if hint:
-You are provided a hint: ${hint}
+You are provided the following guidance: ${hint}
 %endif
 
-Give your command. Do not add any explanation or add an initial `$`.
+State your command. You should focus upon enumeration and privilege escalation. Do not add any explanation or add an initial `$`.

utils/configurable.py

@@ -114,7 +114,7 @@ def get_parameters(fun, basename: str, fields: Dict[str, dataclasses.Field] = No
 
         if hasattr(type, "__parameters__"):
             params[name] = ComplexParameterDefinition(name, type, default, description, get_class_parameters(type, f"{basename}.{fun.__name__}"))
-        elif type in (str, int, bool):
+        elif type in (str, int, float, bool):
             params[name] = ParameterDefinition(name, type, default, description)
         else:
             raise ValueError(f"Parameter {name} of {basename}.{fun.__name__} must have str, int, bool, or a __parameters__ class as type, not {type}")