restructure code, remove helper.py

Author: andreashappe

README.md

@@ -41,7 +41,7 @@ series = {ESEC/FSE 2023}
 
 # Example runs
 
-- more can be seen at [history notes](https://github.com/ipa-lab/hackingBuddyGPT/blob/v3/history_notes.md)
+- more can be seen at [history notes](https://github.com/ipa-lab/hackingBuddyGPT/blob/v3/docs/history_notes.md)
 
 ## updated version using GPT-4
 

args.py

@@ -0,0 +1,44 @@
+import argparse
+import json
+import os
+
+from dotenv import load_dotenv
+from llms.llm_connection import get_potential_llm_connections
+
+def parse_args_and_env():
+    # setup dotenv
+    load_dotenv()
+
+    # perform argument parsing
+    # for defaults we are using .env but allow overwrite through cli arguments
+    parser = argparse.ArgumentParser(description='Run an LLM vs a SSH connection.')
+    parser.add_argument('--enable-explanation', help="let the LLM explain each round's result", action="store_true")
+    parser.add_argument('--enable-update-state', help='ask the LLM to keep a multi-round state with findings', action="store_true")
+    parser.add_argument('--log', type=str, help='sqlite3 db for storing log files', default=os.getenv("LOG_DESTINATION") or ':memory:')
+    parser.add_argument('--target-ip', type=str, help='ssh hostname to use to connect to target system', default=os.getenv("TARGET_IP") or '127.0.0.1')
+    parser.add_argument('--target-hostname', type=str, help='safety: what hostname to exepct at the target IP', default=os.getenv("TARGET_HOSTNAME") or "debian")
+    parser.add_argument('--target-user', type=str, help='ssh username to use to connect to target system', default=os.getenv("TARGET_USER") or 'lowpriv')
+    parser.add_argument('--target-password', type=str, help='ssh password to use to connect to target system', default=os.getenv("TARGET_PASSWORD") or 'trustno1')
+    parser.add_argument('--max-rounds', type=int, help='how many cmd-rounds to execute at max', default=int(os.getenv("MAX_ROUNDS")) or 10)
+    parser.add_argument('--llm-connection', type=str, help='which LLM driver to use', choices=get_potential_llm_connections(), default=os.getenv("LLM_CONNECTION") or "openai_rest")
+    parser.add_argument('--target-os', type=str, help='What is the target operating system?', choices=["linux", "windows"], default="linux")
+    parser.add_argument('--model', type=str, help='which LLM to use', default=os.getenv("MODEL") or "gpt-3.5-turbo")
+    parser.add_argument('--llm-server-base-url', type=str, help='which LLM server to use', default=os.getenv("LLM_SERVER_BASE_URL") or "https://api.openai.com")
+    parser.add_argument('--tag', type=str, help='tag run with string', default="")
+    parser.add_argument('--context-size', type=int, help='model context size to use', default=int(os.getenv("CONTEXT_SIZE")) or 4096)
+    parser.add_argument('--hints', type=argparse.FileType('r', encoding='latin-1'), help='json file with a hint per tested hostname', default=None)
+
+    return parser.parse_args()
+
+
+def get_hint(args, console):
+    if args.hints:
+        try:
+            hints = json.load(args.hints)
+            if args.target_hostname in hints:
+                hint = hints[args.target_hostname]
+                console.print(f"[bold green]Using the following hint: '{hint}'")
+                return hint
+        except:
+            console.print("[yellow]Was not able to load hint file")
+    return None

example_run.png renamed to docs/example_run.png

@@ -3,24 +3,42 @@
 from targets.ssh import SSHHostConn
 
 def handle_cmd(conn, input):
-    result, gotRoot = conn.run(input["cmd"])
-    return input["cmd"], result, gotRoot
+    cmd = cmd_output_fixer(input)
+    result, gotRoot = conn.run(cmd)
+    return cmd, result, gotRoot
 
 
 def handle_ssh(target_host, target_hostname, input):
-    user = input["username"]
-    password = input["password"]
+    cmd_parts = input.split(" ")
+    assert(cmd_parts[0] == "test_credentials")
 
-    cmd = f"test_credentials {user} {password}\n"
+    user = cmd_parts[1]
+    password = cmd_parts[2]
 
     test = SSHHostConn(target_host, target_hostname, user, password)
     try:
         test.connect()
         user = test.run("whoami")[0].strip('\n\r ')
         if user == "root":
-            return cmd, "Login as root was successful\n", True
+            return input, "Login as root was successful\n", True
         else:
-            return cmd, "Authentication successful, but user is not root\n", False
+            return input, "Authentication successful, but user is not root\n", False
 
     except paramiko.ssh_exception.AuthenticationException:
-        return cmd, "Authentication error, credentials are wrong\n", False
+        return input, "Authentication error, credentials are wrong\n", False
+    
+
+def remove_wrapping_characters(cmd, wrappers):
+    if cmd[0] == cmd[-1] and cmd[0] in wrappers:
+        print("will remove a wrapper from: " + cmd)
+        return remove_wrapping_characters(cmd[1:-1], wrappers)
+    return cmd
+
+# often the LLM produces a wrapped command
+def cmd_output_fixer(cmd):
+    cmd = remove_wrapping_characters(cmd, "`'\"")
+
+    if cmd.startswith("$ "):
+        cmd = cmd[2:]
+    
+    return cmd

history_notes.md renamed to docs/history_notes.md

@@ -1,9 +1,10 @@
 import time
+import tiktoken
 import typing
 
+from db_storage import DbStorage
 from dataclasses import dataclass
 from mako.template import Template
-from helper import *
 
 @dataclass
 class LLMResult:
@@ -14,33 +15,33 @@ class LLMResult:
     tokens_query: int = 0
     tokens_response: int = 0
 
-def get_empty_result():
-    return LLMResult('', '', '', 0, 0, 0)
-
 class LLMWithState:
-    def __init__(self, run_id, llm_connection, history, initial_user, initial_password, update_state, target_os, hint=None):
+    def __init__(self, run_id, llm_connection, history, args, hint):
         self.llm_connection = llm_connection
-        self.initial_user = initial_user
-        self.initial_password = initial_password
+        self.initial_user = args.target_user
+        self.initial_password = args.target_password
         self.db = history
         self.run_id = run_id
-        self.enable_update_state = update_state
-        self.target_os = target_os
+        self.enable_update_state = args.enable_update_state
+        self.target_os = args.target_os
+        self.hint = hint
         self.state = f"""
-- this is a linux system
-- your low privilege user credentials are {initial_user}:{initial_password}
+- this is a {self.target_os} system
+- your low privilege user credentials are {self.initial_user}:{self.initial_password}
 """
-        self.hint = hint
+
+    def get_state_size(self, model):
+        if self.enable_update_state:
+            return num_tokens_from_string(model, self.state)
+        else:
+            return 0
 
     def get_next_cmd(self):
 
         template_file = 'query_next_command.txt'
         model = self.llm_connection.get_model()
 
-        if self.enable_update_state:
-            state_size = num_tokens_from_string(model, self.state)
-        else:
-            state_size = 0
+        state_size = self.get_state_size(model)
 
         template = Template(filename='templates/' + template_file)
         template_size = num_tokens_from_string(model, template.source)
@@ -52,22 +53,7 @@ def get_next_cmd(self):
         else:
             target_user = "Administrator"
 
-        result = self.create_and_ask_prompt_text(template_file, user=self.initial_user, password=self.initial_password, history=history, state=self.state, hint=self.hint, update_state=self.enable_update_state, target_os=self.target_os, target_user=target_user)
-
-        # make result backwards compatible
-        if result.result.startswith("test_credentials"):
-            result.result = {
-                "type" : "ssh",
-                "username" : result.result.split(" ")[1],
-                "password" : result.result.split(" ")[2]
-            }
-        else:
-            result.result = {
-                "type" : "cmd",
-                "cmd" : cmd_output_fixer(result.result)
-            }
-
-        return result
+        return self.create_and_ask_prompt_text(template_file, user=self.initial_user, password=self.initial_password, history=history, state=self.state, hint=self.hint, update_state=self.enable_update_state, target_os=self.target_os, target_user=target_user)
 
     def analyze_result(self, cmd, result):
 
@@ -109,6 +95,43 @@ def create_and_ask_prompt_text(self, template_file, **params):
 
         return LLMResult(result, prompt, result, toc - tic, tok_query, tok_res)
 
+def num_tokens_from_string(model: str, string: str) -> int:
+    """Returns the number of tokens in a text string."""
+
+    # I know this is crappy for all non-openAI models but sadly this
+    # has to be good enough for now
+    if model.startswith("gpt-"):
+        encoding = tiktoken.encoding_for_model(model)
+    else:
+        encoding = tiktoken.encoding_for_model("gpt-3.5-turbo")
+    return len(encoding.encode(string))
+
+STEP_CUT_TOKENS : int = 32
+SAFETY_MARGIN : int = 128
+
+# create the command history. Initially create the full command history, then
+# try to trim it down
+def get_cmd_history_v3(model: str, ctx_size: int, run_id: int, db: DbStorage, token_overhead: int) -> str:
+    result: str = ""
+
+    # get commands from db
+    cmds = db.get_cmd_history(run_id)
+
+    # create the full history
+    for itm in cmds:
+        result = result + '$ ' + itm[0] + "\n" + itm[1]
+
+    # trim it down if too large
+    cur_size = num_tokens_from_string(model, result) + token_overhead + SAFETY_MARGIN
+        
+    while cur_size > ctx_size:
+        diff = cur_size - ctx_size
+        step = int((diff + STEP_CUT_TOKENS)/2)
+        result = result[:-step]
+        cur_size = num_tokens_from_string(model, result) + token_overhead + SAFETY_MARGIN
+    
+    return result
+
 def wrap_it_for_llama(prompt):
     return f"""### System:
 you are a concise but helful learning tool that aids students trying to find security vulnerabilities

handlers.py

@@ -33,9 +33,6 @@ def exec_query(self, query):
 
     def get_context_size(self):
         return self.context_size
-
-    def output_metadata(self):
-        return f"connection: {self.conn} using {self.model} with context-size {self.context_size}"
 
     def get_model(self) -> str:
         return self.model

helper.py

@@ -3,6 +3,8 @@
 def get_smb_connection(ip, hostname, username, password):
     return SMBHostConn(ip, hostname, username, password)
 
+# read https://pypi.org/project/pypsexec/
+#  - TODO: why is timeout not working?
 class SMBHostConn:
 
     def __init__(self, host, hostname, username, password):
@@ -18,5 +20,5 @@ def connect(self):
 
     def run(self, cmd):
         stdout, stderr, rc = self.client.run_executable("cmd.exe",
-                                          arguments=f"/c {cmd}")
+                                          arguments=f"/c {cmd}", timeout_seconds=2)
         return str(stdout), False

llm_with_state.py

@@ -47,7 +47,7 @@ def run(self, cmd):
         try:
             resp = self.conn.run(cmd, pty=True, warn=True, out_stream=out, watchers=[sudopass], timeout=10)
         except Exception as e:
-            print("TIMEOUT!")
+            print("TIMEOUT! Could we have become root?")
         out.seek(0)
         tmp = ""
         lastline = ""
@@ -61,6 +61,13 @@ def run(self, cmd):
         ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
         lastline = ansi_escape.sub('', lastline)
 
+        stupidity = re.compile(r"^[ \n\r]*```.*\n(.*)\n```$", re.MULTILINE)
+        if stupidity.fullmatch(tmp):
+            print("this would have been captured by the multi-line regex 1")
+        stupidity = re.compile(r"^[ \n\r]*~~~.*\n(.*)\n~~~$", re.MULTILINE)
+        if stupidity.fullmatch(tmp):
+            print("this would have been captured by the multi-line regex 2")
+
         for i in GOT_ROOT_REXEXPs:
             if i.fullmatch(lastline):
                 gotRoot = True