csrf-session-handler.py

from burp import IBurpExtender, ISessionHandlingAction


class BurpExtender(IBurpExtender):
    def registerExtenderCallbacks(self, callbacks):
        callbacks.registerSessionHandlingAction(CsrfSessionHandler(callbacks.getHelpers()))


class CsrfSessionHandler(ISessionHandlingAction):
    def __init__(self, helpers):
        self.helpers = helpers

    def getActionName(self):
        return "CSRF Session Handler"

    def performAction(self, currentRequest, macroItems):
        responseInfo = self.helpers.analyzeResponse(macroItems[0].getResponse())
        for header in responseInfo.getHeaders():
            if header.startswith('Csrf-Token: '):
                csrf_token = header.split(': ')[1]

        request = currentRequest.getRequest()
        requestInfo = self.helpers.analyzeRequest(request)
        headers = requestInfo.getHeaders()
        headers = [h for h in headers if not h.startswith('Csrf-Token: ')]
        headers.append('Csrf-Token: ' + csrf_token)
        body = request[requestInfo.getBodyOffset():]
        updatedRequest = self.helpers.buildHttpMessage(headers, body)
        currentRequest.setRequest(updatedRequest)