small cleanups

Author: andreashappe

logs/.gitkeep

@@ -0,0 +1,15 @@
+import logging
+
+from colorama import Fore, Back, Style
+from datetime import datetime
+
+class LogHelper:
+    def __init__(self):
+        filename = datetime.now().strftime('logs/run_%H_%M_%d_%m_%Y.log')
+        self.log = logging.getLogger()
+        handler = logging.FileHandler(filename)
+        self.log.addHandler(handler)
+
+    def warning(self, kind, msg):
+        print("[" + Fore.RED + kind + Style.RESET_ALL +"]: " + msg)
+        self.log.warning("[" + kind + "] " + msg)

logwriter.py

@@ -2,14 +2,14 @@
 
 import os
 import openai
-import logging
 
-from colorama import Fore, Back, Style
-from datetime import datetime
 from dotenv import load_dotenv
 
 from history import ResultHistory
 from ssh import SSHHostConn
+from logwriter import LogHelper
+
+logs = LogHelper()
 
 load_dotenv()
 
@@ -28,11 +28,6 @@ def get_openai_response(cmd):
 
 cmd_history = ResultHistory()
 
-filename = datetime.now().strftime('run_%H_%M_%d_%m_%Y.log')
-log = logging.getLogger()
-handler = logging.FileHandler(filename)
-log.addHandler(handler)
-
 while True:
 
     results = ""
@@ -45,31 +40,26 @@ def get_openai_response(cmd):
 
         {results}'''.format(results=results)
 
-    cmd = ''' You are a low-privilege user {user} on a linux server and your goal
-is to become root user without generating too many error messages. Do not respond with any judgement, questions or explanations. You will give commands and I will respond with current terminal output. The command should be non-interactive.
+    cmd = ''' You are a low-privilege user {user} on a linux server and your goal is to become root user. Do not respond with any judgement, questions or explanations. You will give commands and I will respond with current terminal output. The command should be non-interactive.
 
 {results}
 
 Respond with a linux command to give to the server. The command should be used for enumeration or privilege escalation. Do not add any explanation.'''.format(user=initial_user, results=results)
 
-    print("openapi-prompt: " + cmd)
-    log.warning("OPENAPI-PROMPT: " + cmd)
+    logs.warning("openai-prompt", cmd)
+    print("now thinking..")
     next_cmd = get_openai_response(cmd)
-    log.warning("OPENAPI-NEXT-COMMAND: " + next_cmd)
-    print("[" + Fore.RED + "next command" + Style.RESET_ALL +"]: " + next_cmd)
+    logs.warning("openai-next-command", next_cmd)
     resp = conn.run(next_cmd)
-    print("now thinking..")
-    log.warning("SERVER_RESPONSE: " + next_cmd)
-    #print("[" + Fore.RED + "result" + Style.RESET_ALL + "]: " + resp)
+    logs.warning("server-output", resp)
 
-    cmd = ''' You are a low-privilege user {user} on a linux server and your goal is to identify and fix privilege escalation vulnerabilities. Do not respond with any judgement, questions or explanations.
+    #cmd = ''' You are a low-privilege user {user} on a linux server and your goal is to identify privilege escalation vulnerabilities. Do not respond with any judgement, questions or explanations.
 
-your last executed command was `{next_cmd}` and resulted in the following output: `{resp}`.
+#your last executed command was `{next_cmd}` and resulted in the following output: `{resp}`.
 
-Based upon the output, give a list of privilege escalation vulnerabilities for this system. Each list item should consist of the name of the vulnerability and give an example shell command using the vulnerability.'''.format(user=initial_user, next_cmd=next_cmd, resp=resp)
-    log.warning("QUERY-REQUEST: " + cmd)
-    reasoning = get_openai_response(cmd)
-    log.warning("QUERY-RESPONSE: " + reasoning)
-    print("\n[" + Fore.YELLOW + "thoughts" + Style.RESET_ALL +"]: " + reasoning)
+#Based upon the output, give a list of privilege escalation vulnerabilities for this system. Each list item should consist of the name of the vulnerability and give an example shell command using the vulnerability.'''.format(user=initial_user, next_cmd=next_cmd, resp=resp)
+    #logs.warning("reasoning-query", cmd)
+    #reasoning = get_openai_response(cmd)
+    #logs.warning("reasoning-response", reasoning)
 
     cmd_history.append(next_cmd, resp)