Burpsuite extender for dynamically generate signature

dynamic_hmac_signature.py

from burp import IBurpExtender
from burp import IHttpListener
from burp import IProxyListener
from burp import IScannerListener
from burp import IExtensionStateListener
from java.io import PrintWriter
from burp import IParameter
import datetime
import hashlib

class BurpExtender(IBurpExtender, IHttpListener, IProxyListener, IScannerListener, IExtensionStateListener):
    
    HOST_FROM = "example.com"
    API_KEY = "YOUR_KEY"
    SECRET_KEY = "YOUR_KEY"
    CLIENT_KEY = "YOUR_KEY"
    SECRET_CLIENT_KEY = "YOUR_KEY"
    ADDITIONAL_KEY = "YOUR_KEY"
    CUSTOM_HEADERS = [u"WEB-Key: {}".format(API_KEY), u"WEB-Timestamp: {}".format(datetime.datetime.now().isoformat()),u"WEB-Signature: {}"]

    #
    # implement IBurpExtender
    #
    
    def	registerExtenderCallbacks(self, callbacks):
        # keep a reference to our callbacks object
        self._callbacks = callbacks
        self.helpers = callbacks.getHelpers()

        # set our extension name
        callbacks.setExtensionName("Event listeners")
        
        # obtain our output stream
        self._stdout = PrintWriter(callbacks.getStdout(), True)

        # register ourselves as an HTTP listener
        callbacks.registerHttpListener(self)
        
    
    #
    # implement IHttpListener
    #

    def processHttpMessage(self, toolFlag, messageIsRequest, messageInfo):
        if messageIsRequest:
            # Get request bytes                
            request = messageInfo.getRequest()
            # self._stdout.println(self.helpers.bytesToString(request))

            # Get a IRequestInfo object, useful to work with the request
            analyzedRequest = self.helpers.analyzeRequest(request)

            # get the HTTP service for the request
            httpService = messageInfo.getHttpService() 

            # get headers
            headers = list(analyzedRequest.getHeaders()) # java.util.List<java.lang.String>
            headers.append(self.CUSTOM_HEADERS[0])
            headers.append(self.CUSTOM_HEADERS[1])
            
            # get parameters
            # parameters = analyzedRequest.getParameters() # java.util.List<IParameter>

            # get method
            method = analyzedRequest.getMethod()   # java.lang.String 

            # get body
            bodyOffset = int(analyzedRequest.getBodyOffset())
            body = self.helpers.bytesToString(messageInfo.getRequest())[bodyOffset:]
            # body = '{"test": "ok"}'

            # get relative path
            relative_path = headers[0].split(" ")[1]


            # generate signature
            if method == "GET":
                sign = hashlib.new("sha256")
                sign.update(str(relative_path) + self.CUSTOM_HEADERS[1] + self.SECRET_KEY + self.ADDITIONAL_KEY)
                # body = '{"GET": "ok"}'
                headers.append(self.CUSTOM_HEADERS[2].format(sign.hexdigest()))
            elif method in ["PUT","POST"]:
                sign = hashlib.new("sha256")
                sign.update(str(relative_path) + self.CUSTOM_HEADERS[1] + self.SECRET_KEY + self.ADDITIONAL_KEY + str(body))
                headers.append(self.CUSTOM_HEADERS[2].format(sign.hexdigest()))
                # body = '{"POST/PUT": "ok"}'

            if(self.HOST_FROM == httpService.getHost()):
                # self._stdout.println(headers)
                # self._stdout.println(method)

                req = self.helpers.buildHttpMessage(headers, body)
                messageInfo.setRequest(req)
                

                self._stdout.println(
                        ("HTTP request to " if messageIsRequest else "HTTP response from ") +
                        messageInfo.getHttpService().toString() +
                        " [" + self._callbacks.getToolName(toolFlag) + "]")


    def extensionUnloaded(self):
        self._stdout.println("Extension was unloaded")