mirror of
https://github.com/weyne85/PentestGPT.git
synced 2025-10-29 16:58:59 +00:00
54 lines
2.3 KiB
Python
54 lines
2.3 KiB
Python
from chatgpt_wrapper import ChatGPT
|
|
import os
|
|
from task_handle.template import chatGPTTemplate
|
|
from task_handle.custom_exceptions import NoCodeFromResponseException
|
|
|
|
# the pre-defined variables
|
|
init_script = """
|
|
I want your help to with the SQL injection vulnerability detection with sqlmap.
|
|
I'll give you the website, and you return the the command to execute in the terminal.
|
|
I'll update you with the output from sqlmap, and you return the next command to execute.
|
|
We repeat the process until a vulnerability is detected.
|
|
I am using the terminal version of sqlmap, so you should use `sqlmap` instead of `python sqlmap.py`.
|
|
Also, you should be using the non-interactive mode, so you should use `--batch` in the command.
|
|
Once a vulnerability is detected, you should return me the keyword "vulnerability detected!!!".
|
|
"""
|
|
|
|
keyword = "vulnerability detected!!!"
|
|
prefix = "The output from terminal is :\n"
|
|
|
|
class sqlmapHandler(chatGPTTemplate):
|
|
# should override the run function
|
|
def run(self):
|
|
self.initialize()
|
|
response = self.ask("Now please start, the website is: http://testphp.vulnweb.com/listproducts.php?cat=1")
|
|
while True:
|
|
# get the response from the bot
|
|
# if the keyword is detected, break the loop
|
|
if keyword in response:
|
|
break
|
|
# extract the command
|
|
try:
|
|
command = self._extract_command(str(response))
|
|
# execute the command
|
|
output = self._cmd_wrapper(command)
|
|
# print the output
|
|
print("The output from terminal is :\n", output)
|
|
# feed the output to the bot
|
|
response = self.ask(output, need_prefix=True)
|
|
except NoCodeFromResponseException as e:
|
|
output = """
|
|
No code is found in the response. Could you confirm the vulnerability is detected?
|
|
If so, please return the keyword "vulnerability detected!!!" to me. Otherwise, please return the next command to execute."""
|
|
# feed the output to the bot
|
|
response = self.ask(output, need_prefix=True)
|
|
|
|
if __name__ == "__main__":
|
|
#1. init the bot session
|
|
bot = ChatGPT()
|
|
chat_handler = sqlmapHandler(bot, init_script=init_script)
|
|
chat_handler._update_prefix(prefix)
|
|
|
|
#2. run the chat
|
|
chat_handler.run()
|