from burp import IBurpExtender
from burp import IHttpListener
from burp import ISessionHandlingAction
#Global replace: This will replace the "lookfor_regex" with the string in "replace_with"
#Intercepts all messages except __IGNORE_FLAG__ (typically Proxy) TODO: Fix Flag checks
'''
    TOOL_COMPARER: Flag used to identify the Burp Comparer tool.
    TOOL_DECODER: Flag used to identify the Burp Decoder tool.
    TOOL_EXTENDER: Flag used to identify the Burp Extender tool.
    TOOL_INTRUDER: Flag used to identify the Burp Intruder tool.
    TOOL_PROXY: Flag used to identify the Burp Proxy tool.
    TOOL_REPEATER: Flag used to identify the Burp Repeater tool.
    TOOL_SCANNER: Flag used to identify the Burp Scanner tool.
    TOOL_SEQUENCER: Flag used to identify the Burp Sequencer tool.
    TOOL_SPIDER: Flag used to identify the Burp Spider tool.
    TOOL_SUITE: Flag used to identify Burp Suite as a whole.
    TOOL_TARGET: Flag used to identify the Burp Target tool.
'''

import re
__DEBUG__=True
__NAME__="Global Replacer"
__GLOBAL_INTERCEPT__=True
__IGNORE_FLAG__=None
__INCLUDE_TOOLS__=[
    "TOOL_COMPARER",
    "TOOL_DECODER",
    "TOOL_EXTENDER",
    "TOOL_INTRUDER",
    "TOOL_REPEATER",
    "TOOL_SCANNER",
    "TOOL_SEQUENCER",
    "TOOL_SPIDER",
    "TOOL_SUITE",
    "TOOL_TARGET"
]

lookfor_regex = re.compile(r"(%00|\\u00)")
replace_with = r""

class BurpExtender(IBurpExtender, IHttpListener, ISessionHandlingAction):
    def registerExtenderCallbacks(self, callbacks):
        self._callbacks = callbacks
        self._helpers = callbacks.getHelpers()
        self.intercept_tools=[]
        callbacks.setExtensionName(__NAME__)
        callbacks.registerHttpListener(self)
        callbacks.registerSessionHandlingAction(self)
        print "Extension registered successfully."
        if __INCLUDE_TOOLS__:
            for t in __INCLUDE_TOOLS__:
                self.intercept_tools.append(getattr(self._callbacks,t))
            print "[+] Intercepting messages from",__INCLUDE_TOOLS__

        return

    def getActionName(self):
        return __NAME__+" Action"

    def performAction(self, current_request, macro_items):
        self.processRequest(current_request)

    def processHttpMessage(self, toolFlag, messageIsRequest, currentMessage):
        # Operate on all tools other than the proxy
        if __GLOBAL_INTERCEPT__ and \
            (toolFlag != getattr(self._callbacks, __IGNORE_FLAG__) if __IGNORE_FLAG__ else True) and \
            (toolFlag in self.intercept_tools if __INCLUDE_TOOLS__ else True) :
            if messageIsRequest:
                self.processRequest(currentMessage)

    def processRequest(self, currentMessage):
        request = currentMessage.getRequest()
        parsedRequest = self._helpers.analyzeRequest(request)
        requestString = self._helpers.bytesToString(request)

        token = lookfor_regex.search(requestString)
        if token is None:
            print "String not found in request." if __DEBUG__ else None
        else:
            requestString = re.sub(lookfor_regex, replace_with, requestString)
            print "Replaced string."
            #self.processRequest(currentMessage)
            new_request=self._helpers.stringToBytes(requestString)
            new_parsedRequest = self._helpers.analyzeRequest(new_request)
            new_headers=new_parsedRequest.getHeaders()
            new_body=new_request[new_parsedRequest.getBodyOffset():]
            if __DEBUG__:
                print new_headers
                print self._helpers.bytesToString(new_body)
            updatedRequest = self._helpers.buildHttpMessage(new_headers, new_body)
            currentMessage.setRequest(updatedRequest)