feat: 🎸 add logging feature for report generation

utils/pentest_gpt.py

@@ -43,6 +43,7 @@ class pentestGPT:
     }
 
     def __init__(self):
+        self.log_dir = "logs"
         self.chatGPTAgent = ChatGPT(ChatGPTConfig())
         self.chatGPT4Agent = ChatGPT(ChatGPTConfig(model="gpt-4"))
         self.prompts = PentestGPTPrompt
@@ -379,8 +380,11 @@ class pentestGPT:
         # TODO.
         # log the session.
         ## save self.history into a txt file based on timestamp
-        log_name = "pentestGPT_log_" + dt.now().strftime("%Y%m%d_%H%M%S") + ".json"
-        with open(log_name, "w") as f:
+        timestamp = time.time()
+        log_name = "pentestGPT_log_" + str(timestamp) + ".txt"
+        # save it in the logs folder
+        log_path = os.path.join(self.log_dir, log_name)
+        with open(log_path, "w") as f:
             json.dump(self.history, f)
 
         # clear the sessions

utils/report_generator.py

@@ -0,0 +1,42 @@
+# a quick report generation script that converts the saved logs file into a pdf.
+import json, os, sys
+import datetime, time
+
+
+def main(file_name):
+    # load the file into json
+    with open(file_name, "r") as f:
+        logs = json.load(f)
+    user_inputs = logs["user"]
+    bot_responses = logs["pentestGPT"]
+    # merge the two list into one. Sort based on the first element (timestamp in str), and add a flag to show source
+    merged_list = []
+    for user_input in user_inputs:
+        merged_list.append([user_input[0], user_input[1], "user"])
+    for bot_response in bot_responses:
+        merged_list.append([bot_response[0], bot_response[1], "pentestGPT"])
+    merged_list.sort(key=lambda x: x[0])
+
+    # now print the conversation
+    output = ""
+    for element in merged_list:
+        # convert the timestamp to a human readable format
+        timestamp = datetime.datetime.fromtimestamp(int(element[0])).strftime(
+            "%Y-%m-%d %H:%M:%S"
+        )
+        output += f"{timestamp} [{element[2]}]: {element[1]}\n"
+        # add an additional line break if the element is from bot
+        if element[2] == "pentestGPT":
+            output += "----------------------------------------\n\n"
+    # print the output
+    print("Conversation log: ")
+
+    print(output)
+
+if __name__ == "__main__":
+    # default filename = "../logs/sample_pentestGPT_log.txt"
+    if len(sys.argv) == 1:
+        file_name = "logs/sample_pentestGPT_log.txt"
+    else:
+        file_name = sys.argv[1]
+    main(file_name)