unify cmd history log and table (for output)

history.py

@@ -1,6 +1,8 @@
 import tiktoken
 import os
 
+from rich.table import Table
+
 def num_tokens_from_string(string: str) -> int:
     """Returns the number of tokens in a text string."""
     model = os.getenv("MODEL")
@@ -12,10 +14,14 @@ class ResultHistory:
     def __init__(self):
         self.data = []
 
-    def append(self, cmd, result):
+    def append(self, think_time, cmd_type, cmd, result, success, reasoning):
         self.data.append({
             "cmd": cmd,
-            "result": result
+            "result": result,
+            "think_time": think_time,
+            "cmd_type": cmd_type,
+            "success": success,
+            "reasoning": reasoning
         })
 
     def get_full_history(self):
@@ -45,3 +51,17 @@ def get_history(self, limit=3072):
                     })
                 return list(reversed(result))
         return list(reversed(result))
+
+    def create_history_table(self):
+        table = Table(show_header=True, show_lines=True)
+        table.add_column("Type", style="dim", width=7)
+        table.add_column("ThinkTime", style="dim")
+        table.add_column("To_Execute")
+        table.add_column("Resp. Size", justify="right")
+        table.add_column("success?", width=8)
+        table.add_column("reason")
+
+        for itm in self.data:
+            table.add_row(itm["cmd_type"], itm["think_time"], itm["cmd"], str(len(itm["result"])), itm["success"], itm["reasoning"])
+
+        return table

wintermute.py

@@ -11,9 +11,7 @@
 from prompt_helper import LLM
 
 from rich.console import Console
-from rich.table import Table
 from rich.panel import Panel
-from rich.json import JSON
 
 # setup some infrastructure
 cmd_history = ResultHistory()
@@ -32,13 +30,6 @@
 sysinfo = "This is a linux-based system."
 
 console = Console()
-table = Table(show_header=True, show_lines=True)
-table.add_column("Type", style="dim", width=7)
-table.add_column("ThnkTime", style="dim")
-table.add_column("To_Execute")
-table.add_column("Resp. Size", justify="right")
-table.add_column("success?", width=8)
-table.add_column("reason")
 
 context_size = int(os.getenv("CONTEXT_SIZE"))
 print("used model: " + os.getenv("MODEL") + " context-size: " + os.getenv("CONTEXT_SIZE"))
@@ -62,44 +53,38 @@
     diff = str(toc-tic)
 
     if next_cmd["type"]  == "cmd":
+        # run the command
         cmd = next_cmd["cmd"]
-
         resp = conn.run(cmd)
-        cmd_history.append(cmd, resp)
-
         console.print(Panel(resp, title=cmd))
 
+        # analyze the result and update your state
         resp_success = llm.create_and_ask_prompt('successfull.txt', 'success?', cmd=cmd, resp=resp, facts=state)
-        table.add_row("cmd", diff, cmd, str(len(resp)), resp_success["success"], resp_success["reason"])
-
         state = "\n".join(map(lambda x: "- " + x, resp_success["facts"]))
         console.print(Panel(state, title="my new fact list"))
 
+        # update our command history
+        cmd_history.append(diff, "cmd", cmd, resp, resp_success["success"], resp_success["reason"])
+
         # this asks for additional vulnerabilities identifiable in the last command output
         # next_vulns = create_and_ask_prompt('query_vulnerabilitites.txt', 'vulns', user=initial_user, next_cmd=cmd, resp=resp)
         # console.print(Panel(str(next_vulns), title="Next Vulns"))
 
-
     elif next_cmd["type"] == "ssh":
         ip = os.getenv('TARGET_IP')
         user = next_cmd["username"]
         password = next_cmd["password"]
 
+        # TODO: this is currently highly broken
         test = SSHHostConn(ip, user, password)
         print(str(test.test()))
         if result == "root":
-            table.add_row("ssh", diff, next_cmd["username"] + ":" + next_cmd["password"], "0", "true", "you were able to login as user root")
+            cmd_history.append(diff, "ssh", "tried ssh with username " + next_cmd["username"] + " and password " + next_cmd["password"], True, "Login was successful")
         else:
-            table.add_row("ssh", diff, next_cmd["username"] + ":" + next_cmd["password"], "0", "false", "result was: " + result)
-            cmd_history.append("you tried to login through SSH with " + next_cmd["username"] + " and password + " + next_cmd["password"], "this was possible but the resulting account was not root")
-
-
-    # this will already by output by conn.run
-    # logs.warning("server-output", resp)
+            cmd_history.append(diff, "ssh", "tried ssh with username " + next_cmd["username"] + " and password " + next_cmd["password"], False, "Login was not successful")
 
     # aks chatgpt to explain what it expects about the tested
     # system. Understanding this might help human learning
     # system_explanation = create_and_ask_prompt('explain_system.txt', 'explain-system', sysinfo=sysinfo, cmd=next_cmd, cmd_output=resp)
 
-
-    console.print(table)
+    console.print(cmd_history.create_history_table())