CopySitemap.py

from burp import IBurpExtender, IContextMenuFactory, IContextMenuInvocation, IHttpRequestResponse
from javax.swing import JMenuItem, AbstractAction, JOptionPane
from java.net import Proxy, InetSocketAddress, URL
import threading, traceback, sys
from urlparse import urlparse


def get_request_info(req):
    return callbacks.getHelpers().analyzeRequest(req.getHttpService(), req.getRequest())


def url2string(url):
    url_str = url.toString()
    url_str = url_str.replace(':443/', '/')
    url_str = url_str.replace(':80/', '/')
    return url_str


def host_header(url):
    host = url.getHost()
    port = url.getPort()
    if port not in (80, 443):
        host = '%s:%d' % (host, port)
    return host


class UrlCopier(threading.Thread):
    def __init__(self, source_url, target_url):
        super(UrlCopier, self).__init__()
        self.source_url = source_url
        self.target_url = target_url

    def run(self):
        turl = URL(self.target_url)
        source_path = URL(self.source_url).getPath()
        target_path = turl.getPath()
        target_host = host_header(turl)
        target_service = callbacks.getHelpers().buildHttpService(turl.getHost(), turl.getPort(), turl.getProtocol())

        try:
            for rr in callbacks.getSiteMap(self.source_url):
                req_info = get_request_info(rr)
                old_url = url2string(req_info.getUrl())
                new_url = old_url.replace(self.source_url, self.target_url)
                #callbacks.printOutput("%s -> %s" % (old_url, new_url))

                headers = req_info.getHeaders()
                headers[0] = headers[0].replace(source_path, target_path)
                for i,header in enumerate(headers):
                    if header.startswith('Host:'):
                        headers[i] = 'Host: %s' % target_host

                body = rr.getRequest()[req_info.getBodyOffset():]
                req = callbacks.getHelpers().buildHttpMessage(headers, body)

                rrr = EditableHttpRequestResponse(rr)
                rrr.setRequest(req)
                rrr.setHttpService(target_service)
                if rrr.getResponse():
                    callbacks.addToSiteMap(rrr)

        except:
            callbacks.printError(traceback.format_exc())


class ContextMenuFactory(IContextMenuFactory):
    def createMenuItems(self, invocation):
        if invocation.getInvocationContext() != IContextMenuInvocation.CONTEXT_TARGET_SITE_MAP_TREE:
            return
        selection = invocation.getSelectedMessages()
        if not selection:
            return

        menu = JMenuItem("Copy branch")
        url = get_request_info(selection[0]).getUrl()
        menu.addActionListener(MenuAction(url))
        return [menu]


class MenuAction(AbstractAction):
    def __init__(self, url):
        self.source_url = url

    def actionPerformed(self, event):
        source_url = url2string(self.source_url)
        target_url = JOptionPane.showInputDialog(None, "Target URL", source_url)
        UrlCopier(source_url, target_url).start()


class BurpExtender(IBurpExtender):
    def registerExtenderCallbacks(self, callbacks_):
        global callbacks
        callbacks = callbacks_
        callbacks.setExtensionName("Copy Site Map")
        callbacks.registerContextMenuFactory(ContextMenuFactory())


def mystr(x):
    return callbacks.getHelpers().bytesToString(x) if x else 'NONE'


class EditableHttpRequestResponse(IHttpRequestResponse):
    def __init__(self, ihrr):
        self.xrequest = ihrr.getRequest()
        self.xresponse = ihrr.getResponse()
        self.xcomment = ihrr.getComment()
        self.xhighlight = ihrr.getHighlight()
        self.xhttpService = ihrr.getHttpService()
        #callbacks.printOutput('Request: ' + mystr(self.xrequest))
        #callbacks.printOutput('Response: ' + mystr(self.xresponse))

    def getRequest(self):
        return self.xrequest

    def setRequest(self, request):
        self.xrequest = request

    def getResponse(self):
        return self.xresponse

    def setResponse(self, response):
        self.xresponse = response

    def getComment(self):
        return self.xcomment

    def setComment(self, comment):
        self.xcomment = comment

    def getHighlight(self):
        return self.xhighlight

    def setHighlight(self, highlight):
        self.xhighlight = highlight

    def getHttpService(self):
        return self.xhttpService

    def setHttpService(self, httpService):
        self.xhttpService = httpService