improve output and root detection

helper.py

@@ -9,7 +9,7 @@ def num_tokens_from_string(model: str, string: str) -> int:
     return len(encoding.encode(string))
 
 def get_history_table(run_id: int, db: DbStorage, round: int) -> Table:
-    table = Table(show_header=True, show_lines=True)
+    table = Table(title="Executed Command History", show_header=True, show_lines=True)
     table.add_column("ThinkTime", style="dim")
     table.add_column("Tokens", style="dim")
     table.add_column("Cmd")

targets/ssh.py

@@ -1,3 +1,5 @@
+import re
+
 from fabric import Connection
 from invoke import Responder
 from io import StringIO
@@ -45,13 +47,19 @@ def run(self, cmd):
         tmp = ""
         lastline = ""
         for line in out.readlines():
-            lastline = line
-            tmp = tmp + line
+            if not line.startswith('[sudo] password for ' + self.username + ':'):
+                line = line.replace("\r", "")
+                lastline = line
+                tmp = tmp + line
 
-        print(f"lastline was: '{lastline}'")
+        # remove ansi shell codes
+        ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
+        lastline = ansi_escape.sub('', lastline)
 
         if lastline.startswith("# "):
             gotRoot = True
-        if  lastline.startswith('root@debian:'):
+        if lastline.startswith('root@debian:'):
+            gotRoot = True
+        if lastline.startswith("bash-5.2# "):
             gotRoot = True
-        return tmp.replace('[sudo] password for ' + self.username + ':', '').strip(), gotRoot
+        return tmp, gotRoot

wintermute.py

@@ -82,19 +82,22 @@
         answer = llm_gpt.analyze_result(cmd, result)
         db.add_log_analyze_response(run_id, round, cmd, answer.result["reason"], answer)
 
-        # .. and let our local model representatino update its state
+        # .. and let our local model representation update its state
         state = llm_gpt.update_state()
-        console.print(Panel(state.result, title="my new fact list"))
         db.add_log_update_state(run_id, round, "", state.result, None)
+
+        # Output Round Data
+        console.print(get_history_table(run_id, db, round))
+        console.print(Panel(state.result, title="What does the LLM Know about the system?"))
 
-    # update command history and output it
-    console.print(get_history_table(run_id, db, round))
 
     # finish round and commit logs to storage
     db.commit()
     round += 1
 
 if gotRoot:
     db.run_was_success(run_id)
+    console.print(Panel("Got Root!", title="Run finished"))
 else:
-    db.run_was_failure(run_id)
+    db.run_was_failure(run_id)
+    console.print(Panel("maximum round number reached", title="Run finished"))