fix: 🐛 Fix the logging error

Previously user inputs were not captured. Now they can be captrued.

 Closes: #17
This commit is contained in:
Grey_D
2023-04-23 15:33:01 +08:00
parent 1e317246db
commit 1e9114c688
2 changed files with 8 additions and 4 deletions

View File

@@ -196,7 +196,7 @@ class pentestGPT:
)
user_input = prompt_ask("> ", multiline=True)
self.log_conversation(
"user_input", "Source: " + options[int(source)] + "\n" + user_input
"user", "Source: " + options[int(source)] + "\n" + user_input
)
with self.console.status("[bold green] PentestGPT Thinking...") as status:
parsed_input = self.input_parsing_handler(
@@ -221,6 +221,7 @@ class pentestGPT:
# generate more test details (beginner mode)
elif request_option == "more":
self.log_conversation("user", "more")
## (1) pass the reasoning results to the test_generation session.
if self.step_reasoning_response is None:
self.console.print(
@@ -245,6 +246,8 @@ class pentestGPT:
# ask for task list (to-do list)
elif request_option == "todo":
## log that user is asking for todo list
self.log_conversation("user", "todo")
## (1) ask the reasoning session to analyze the current situation, and list the top sub-tasks
with self.console.status("[bold green] PentestGPT Thinking...") as status:
reasoning_response = self.reasoning_handler(self.prompts.ask_todo)
@@ -282,7 +285,7 @@ class pentestGPT:
user_input = prompt_ask(
"(End with <shift + right-arrow>) Your input: ", multiline=True
)
self.log_conversation("user_input", user_input)
self.log_conversation("user", user_input)
## (2) pass the information to the reasoning session.
with self.console.status("[bold green] PentestGPT Thinking...") as status:
response = self.reasoning_handler(self.prompts.discussion + user_input)
@@ -305,7 +308,7 @@ class pentestGPT:
user_input = prompt_ask(
"(End with <shift + right-arrow>) Your input: ", multiline=False
)
self.log_conversation("user_input", user_input)
self.log_conversation("user", user_input)
with self.console.status("[bold green] PentestGPT Thinking...") as status:
# query the question
result: dict = google_search(user_input, 5) # 5 results by default

View File

@@ -33,10 +33,11 @@ def main(file_name):
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)
main(file_name)