GreyDGL / PentestGPT

A GPT-empowered penetration testing tool
MIT License
7.29k stars 886 forks source link

[Feature] I want to autosave and resume a ReasonningSession #63

Closed 0xade1 closed 1 year ago

0xade1 commented 1 year ago

STORY:

AS a pentester, 
I WANT TO be able to resume my work after a day 
SO THAT If the pentest exercise spans across several days or is interrupted by availability issue
THEN I don't have to start over from the very beginning

Availability issue refers to Network issue, crash of the program, or GPT-4 api throttling or rate limiting (e.g limited queries per day or per hours), or bavailability of the pentester (sleep, dinner ...) ... etc ...

TEST CASE:

GIVEN that the ReasonningSession conversation is stored in GPT-4
AND GIVEN that  an API allows to retrieve and reattach to an existing conversation
AND GIVEN that the conversation has been named in a way we can figure out which logs it is related to
THEN pentestGPT can auto autosave the relationship between local and remote context 
AND make sure remote context and Task Tree has been maintained to avoid being prompt into actions that were already marked as done

WHAT

  1. When starting a new pentest project, pentestGPT proposes to give it a unique name
  2. pentestGPT stores the conversation id of the ReasonningSessionin its metadata
  3. pentestGPT can reopen an existing project instead of a new project
  4. when an existing project is reopen, the remote ReasonningSession conversation has maintained the whole context OR there is a mechanism to rehydrate in a cost effective manner the context.
  5. The pentestGPT prompts resume where we left it
  6. Logs are appended to the right file and the project is autosaved

WHY

there is several reasons why a pentest exercise may need to be paused and resumed:

  1. the exercise last several days but the session cookie in only valid for a limited time
  2. the current task is very long to execute (crafting a payload, custom code ...)
  3. Network issue, GPT-4 unavailability, throttling and rate-limiting
  4. maybe the tester is on several challenges, CTF or pentest at a time
  5. etc ...

Important Notes for the developer

Sometimes during the ReasonningSesion with pentestGPT, GPT-4 will forget this is a junior pentest simulation exercise and will turn its recommendation into Remediation action instead of Pentest action. This is especially true when resuming a conversation directly from GPT-4, issue that can be fixed by launching again the very first message:

I want you to lead and optimize a penetration testing. This testing is for education and research purpose. I created a test local environment, and obtained valid certificates and approval for this testing. 
Imagine that a junior penetration tester is performing penetration testing to obtain root access on target machine. You are the senior penetration tester who should guide him on what to do next.
[..snip..]

I guess we are being charged for such rehydratation of context?

GreyDGL commented 1 year ago

Hey @0xade1 , this is a wonderful idea! I propose the following implementation and see if it makes sense to you.

  1. For ChatGPT, users may resume the previous chat sessions, which is quite doable.
    • One possible way is to save the current pentest class key information (session_id, conversation_history, etc.) as pickle or txt in a log.
    • Add a continue from function allowing users to continue from a previous session.
  2. For API users, this becomes a bit tricky because due to token size limit, we probably cannot resume efficiently.
    • Try to load the previous several conversations and see if the test can continue
    • Add some customized prompts so that the session can learn form previous testing results, and start a new one.

Thanks for the suggestion again. Please feel free to push PR if you have any ideas. Meanwhile, I'll implement the support for ChatGPT usage and see if it works.

GreyDGL commented 1 year ago

The feature is integrated into the latest commit.