Burp plugin (Handling Action) for multi-step requests from within the plugin when Macros are not enough

burp-multistep-request.py

# XXX: Example / Untested code snippet

from burp import IBurpExtender
from burp import ISessionHandlingAction
from burp import ICookie
import re

__DEBUG__ = True

class Cookie(ICookie):
    # Needed in order to Update Burps cookie jar with new cookies
    # AnalyzeResponse cookies not always contain domain name / without domain name cookie jar cannot be updated
    def getDomain(self):
        return self.cookie_domain
    def getPath(self):
        return self.cookie_path
    def getExpiration(self):
        return self.cookie_expiration
    def getName(self):
        return self.cookie_name
    def getValue(self):
        return self.cookie_value

    def __init__(self, cookie_domain=None, cookie_name=None, cookie_value=None, cookie_path=None, cookie_expiration=None):
        self.cookie_domain = cookie_domain
        self.cookie_name = cookie_name
        self.cookie_value = cookie_value
        self.cookie_path = cookie_path
        self.cookie_expiration = cookie_expiration

class BurpExtender(IBurpExtender, ISessionHandlingAction):
    token_regex = re.compile(r"\"Token:(.*?)\"")
    host="this.is.just.an.example.com"
    default_headers=[
        "User-Agent: Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Ubuntu Chromium/73.0.3683.86 Chrome/73.0.3683.86 Safari/537.36",
        "Accept: text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,image/apng,*/*;q=0.8",
        "Accept-Encoding: gzip, deflate",
        "Accept-Language: en-GB,en-US;q=0.9,en;q=0.8"
    ]

    def registerExtenderCallbacks(self, callbacks):
        self._callbacks = callbacks
        self._helpers = callbacks.getHelpers()
        callbacks.setExtensionName("Multistep Requests")
        callbacks.registerHttpListener(self)
        callbacks.registerSessionHandlingAction(self)
        print "Multistep Requests"
        return

    def getActionName(self):
        return "Multistep Requests"

    def do_request(self, host, url_path, headers, body=None, port=443, use_Https=True, cookies=None):
        h=("POST " if body else "GET ") + \
        url_path + \
        " HTTP/1.1"
        host="Host: "+host
        headers=[h]+host+headers
        headers.append(cookies) if cookies else pass
        request=self._helpers.buildHttpMessage(headers,body)
        response=self._callbacks.makeHttpRequest(host, port, use_Https, request)
        print response if __DEBUG__ else None
        return response

    def get_cookies_from_response(self, response):
        return "Cookie: {}".format(
            ";".join('{}={}'.format(
                c.getName, c.GetValue()
                for c in self._helpers.analyzeResponse(response).getCookies()
            )        
        )
   
    def update_burp_cookie_jar(self, response, domain):
        for c in self._helpers.analyzeResponse(response).getCookies():
            new_cookie=Cookie(c.getDomain() if c.getDomain() else domain, c.getName(), c.getValue(), c.getPath(),c.getExpiration())
            self._callbacks.updateCookieJar(new_cookie)

    def performAction(self, current_request, macro_items):
        host=self.host
        headers=self.default_headers
        # Initial get request to get Token (example) and cookies
        response=self.do_request(host,"/initial_step/index.html", default_headers)
        cookies=self.get_cookies_from_response(response)
        token=token_regex.search(self._helpers.bytesToString(response)).group(1)    #TODO: error handling

        # Set up second/final request
        username="Test"
        password="Test"

        # Additional headers for POST request
        headers=self.default_headers + \
        [
            "Connection: close",
            #"Content-Length: 128",     #TODO: Test if needed
            "Cache-Control: max-age=0",
            "Origin: https://"+host+"/",
            "Content-Type: application/x-www-form-urlencoded",
            "Referer: https://"+host+"/"
        ]

        # Set up body
        body="Token="+token+"&username="+username+"&password="+password

        final_response=self.do_request(host,"/second_step/index.html", headers, body,cookies=cookies)

        # Update burp's cookie jar
        cookie_domain=host
        self.update_burp_cookie_jar(final_response, cookie_domain)