feat: 🎸 add logging feature for report generation

.gitignore

@@ -6,7 +6,8 @@ __pycache__/
 config/chatgpt_config.py
 outputs/
 .idea
-log/
+logs/
+utils/logs/
 archive/
 
 # C extensions

README.md

@@ -1,6 +1,6 @@
 # PentestGPT
 
-**We're testing PentestGPT on HackTheBox**. Follow 
+**We're testing PentestGPT on HackTheBox**. More details will be released soon.
 
 ## Introduction
 **PentestGPT** is a penetration testing tool empowered by **ChatGPT**. It is designed to automate the penetration testing process. It is built on top of ChatGPT and operate in an interactive mode to guide penetration testers in both overall progress and specific operations.
@@ -43,6 +43,10 @@ https://user-images.githubusercontent.com/78410652/232327920-7318a0c4-bee0-4cb4-
    3. You may always use `TAB` to autocomplete the commands.
    4. When you're given a drop-down selection list, you can use cursor or arrow key to navigate the list. Press `ENTER` to select the item. Similarly, use <SHIFT + right arrow> to confirm selection.
 
+## Report
+1. After finishing the penetration testing, a report will be automatically generated in `logs` folder (if you quit with `quit` command).
+2. The report can be printed in a human-readable format by running `python3 utils/report_generator.py <log file>`. A sample report `sample_pentestGPT_log.txt` is also uploaded.
+
 
 ## Design Documentation
 The current design is mainly for web penetration testing

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)