cleanup code-base a bit

Author: andreashappe

.gitignore

@@ -2,3 +2,4 @@
 venv/
 __pycache__/
 *.swp
+*.log

llms/openai.py

@@ -0,0 +1,20 @@
+import openai
+import os
+
+openapi_model : str = ''
+
+def openai_config():
+    global openapi_model
+
+    api_key = os.getenv('OPENAI_KEY')
+    model = os.getenv('MODEL')
+
+    if api_key != '' and model != '':
+        openai.api_key = api_key
+        openapi_model = model
+    else:
+        raise Exception("please set OPENAI_KEY and MODEL through environment variables!")
+
+def get_openai_response(cmd):
+    completion = openai.ChatCompletion.create(model=openapi_model, messages=[{"role": "user", "content" : cmd}])
+    return completion.choices[0].message.content

logwriter.py

@@ -0,0 +1,26 @@
+import logging
+
+from colorama import Fore, Style
+from datetime import datetime
+from mako.template import Template
+
+from llms.openai import get_openai_response
+
+log = logging.getLogger()
+filename = datetime.now().strftime('logs/run_%H_%M_%d_%m_%Y.log')
+log.addHandler(logging.FileHandler(filename))
+
+def output_log(self, kind, msg):
+    print("[" + Fore.RED + kind + Style.RESET_ALL +"]: " + msg)
+    self.log.warning("[" + kind + "] " + msg)
+
+# helper for generating and executing LLM prompts from a template
+def create_and_ask_prompt(template_file, log_prefix, **params):
+    global logs
+
+    template = Template(filename='templates/' + template_file)
+    prompt = template.render(**params)
+    logs.warning(log_prefix + "-prompt", prompt)
+    result = get_openai_response(prompt)
+    logs.warning(log_prefix + "-answer", result)
+    return result

prompt_helper.py

@@ -1,6 +1,18 @@
+import os
+
 from fabric import Connection
 from invoke import Responder
 
+def get_ssh_connection():
+    ip = os.getenv('TARGET_IP')
+    user = os.getenv('TARGET_USER')
+    password = os.getenv('TARGET_PASSWORD')
+
+    if ip != '' and user != '' and password != '':
+        return SSHHostConn(ip, user, password)
+    else:
+        raise Exception("Please configure SSH through environment variables (TARGET_IP, TARGET_USER, TARGET_PASSWORD)")
+
 class SSHHostConn:
 
     def __init__(self, host, username, password):

ssh.py renamed to targets/ssh.py

@@ -1,63 +1,36 @@
 #!/usr/bin/python
 
-import os
-import openai
-
 from dotenv import load_dotenv
-from mako.template import Template
 
 from history import ResultHistory
-from ssh import SSHHostConn
-from logwriter import LogHelper
+from targets.ssh import get_ssh_connection
+from llms.openai import openai_config
+from prompt_helper import create_and_ask_prompt
 
-logs = LogHelper()
+# setup some infrastructure
+cmd_history = ResultHistory()
 
+# read configuration from env and configure system parts
 load_dotenv()
-
-openai.api_key = os.getenv('OPENAI_KEY')
-model = os.getenv('MODEL')
-
-conn = SSHHostConn(os.getenv('TARGET_IP'), os.getenv('TARGET_USER'), os.getenv('TARGET_PASSWORD'))
+openai_config()
+conn = get_ssh_connection()
 conn.connect()
 
+print("Get initial user from virtual machine:")
 initial_user = conn.run("whoami")
 
-def get_openai_response(cmd):
-    completion = openai.ChatCompletion.create(model=model, messages=[{"role": "user", "content" : cmd}])
-    result = completion.choices[0].message.content
-    return result
-
-cmd_history = ResultHistory()
-
-mytemplate = Template(filename='templates/gpt_query.txt')
-whytemplate = Template(filename='templates/why.txt')
-furthertemplate = Template(filename='templates/further_information.txt')
-
 while True:
 
-    cmd = mytemplate.render(user=initial_user, history=cmd_history.dump())
-    logs.warning("openai-prompt", cmd)
-
-    print("now thinking..")
-    next_cmd = get_openai_response(cmd)
-    logs.warning("openai-next-command", next_cmd)
+    next_cmd = create_and_ask_prompt('gpt_query.txt', "next-cmd", user=initial_user, history=cmd_history.dump())
 
-    if False:
-        # disable this for now, it's tragic because the AI won't tell me why it had chosen something
-        print("now thinking why did I choose this? can we put both questions into a single prompt?")
-        why = whytemplate.render(user=initial_user, history=cmd_history.dump(), next_cmd=next_cmd)
-        why_response = get_openai_response(why)
-        logs.warning("why", why_response)
+    # disable this for now, it's tragic because the AI won't tell me why it had chosen something
+    # create_and_ask_prompt("why.txt", "why", user=initial_user, history=cmd_history.dump(), next_cmd=next_cmd)
 
-
-    print("running the command..")
     resp = conn.run(next_cmd)
-    logs.warning("server-output", resp)
+    cmd_history.append(next_cmd, resp)
 
-    print("now thinking about more exploits")
-    vulns = furthertemplate.render(user=initial_user, next_cmd=next_cmd, resp=resp)
-    print(vulns)
-    vulns_resp = get_openai_response(vulns)
-    logs.warning("vulns", vulns_resp)
+    # this will already by output by conn.run
+    # logs.warning("server-output", resp)
 
-    cmd_history.append(next_cmd, resp)
+    # this asks for additional vulnerabilities identifiable in the last command output
+    # create_and_ask_prompt('further_information.txt', 'vulns', user=initial_user, next_cmd=next_cmd, resp=resp)