#
# SiteLogger - Log sitemap and findings to database
#
# Copyright (c) 2017 Andrea Brancaleoni.
#

from java.awt import Component

from burp import IBurpExtenderCallbacks, IExtensionHelpers, IHttpRequestResponse
from burp import IScanIssue, ITab, IBurpExtender
from com.mongodb import BasicDBObject, DB, DBCollection, MongoClient

from java.io import PrintWriter
from java.net import MalformedURLException, URL, UnknownHostException

from javax.swing import (BoxLayout, ImageIcon, JButton, JFrame, JPanel,
        JPasswordField, JLabel, JTextArea, JTextField, JScrollPane,
        SwingConstants, WindowConstants, GroupLayout)
import javax
from java.lang import Short, Integer


# Original code from src/burp/BurpExtender.java class
class BurpExtender(IBurpExtender):

    def registerExtenderCallbacks(self, callbacks):
        self.callbacks = callbacks
        helpers = callbacks.getHelpers()
        callbacks.setExtensionName("SiteLogger")
        callbacks.addSuiteTab(SiteLoggerTab(callbacks, helpers))


# Original code from src/com/doyensec/SiteLoggerTab.java class
class SiteLoggerTab(ITab):

    def __init__(self, callbacks, helpers):
        self.callbacks = callbacks
        self.helpers = helpers

    def getTabCaption(self):
        return "SiteLogger"

    def getUiComponent(self):
        panel = SiteLoggerPanel(self.callbacks, self.helpers)
        self.callbacks.customizeUiComponent(panel.this)
        return panel.this


# Original code from src/com/doyensec/SiteLoggerPanel.java class
# XXX: inheriting from Java classes is very tricky. It is preferable to use
#      the decorator pattern instead.
class SiteLoggerPanel:

    def __init__(self, callbacks, helpers):
        self.callbacks = callbacks
        self.helpers = helpers
        self.this = JPanel()
        self.initComponents()

    # This method is called from within the constructor to initialize the form.
    # WARNING: Do NOT modify this code. The content of this method is always
    # regenerated by the Form Editor.
    #
    # <editor-fold defaultstate="collapsed" desc="Generated Code">//GEN-BEGIN:initComponents
    def initComponents(self):
        self.jLabel1 = javax.swing.JLabel()
        self.website = javax.swing.JTextField()
        self.jLabel2 = javax.swing.JLabel()
        self.mongohost = javax.swing.JTextField()
        self.mongoport = javax.swing.JTextField()
        self.logButton = javax.swing.JButton()
        self.jLabel3 = javax.swing.JLabel()
        self.jLabel1.setText("Website:")
        self.jLabel2.setText("MongoDB Host: ")

        self.mongohost.setText("127.0.0.1")

        self.mongoport.setText("27017")

        self.logButton.setText("Log to Database")
        self.logButton.addActionListener(
            lambda evt: self.logButtonActionPerformed(evt))

        self.jLabel3.setText("MongoDB Port: ")
        layout = javax.swing.GroupLayout(self.this)
        self.this.setLayout(layout)

        layout.setHorizontalGroup(
            layout.createParallelGroup(javax.swing.GroupLayout.Alignment.LEADING)
            .addGroup(layout.createSequentialGroup()
                .addGroup(layout.createParallelGroup(javax.swing.GroupLayout.Alignment.LEADING)
                    .addGroup(layout.createSequentialGroup()
                        .addGap(48, 48, 48)
                        .addGroup(layout.createParallelGroup(javax.swing.GroupLayout.Alignment.LEADING)
                            .addGroup(layout.createSequentialGroup()
                                .addComponent(self.jLabel3)
                                .addPreferredGap(javax.swing.LayoutStyle.ComponentPlacement.UNRELATED)
                                .addComponent(self.mongoport, javax.swing.GroupLayout.PREFERRED_SIZE, 226, javax.swing.GroupLayout.PREFERRED_SIZE))
                            .addGroup(layout.createParallelGroup(javax.swing.GroupLayout.Alignment.TRAILING, False)
                                .addGroup(javax.swing.GroupLayout.Alignment.LEADING, layout.createSequentialGroup()
                                    .addComponent(self.jLabel2)
                                    .addPreferredGap(javax.swing.LayoutStyle.ComponentPlacement.RELATED)
                                    .addComponent(self.mongohost, javax.swing.GroupLayout.PREFERRED_SIZE, 227, javax.swing.GroupLayout.PREFERRED_SIZE))
                                .addGroup(javax.swing.GroupLayout.Alignment.LEADING, layout.createSequentialGroup()
                                    .addComponent(self.jLabel1)
                                    .addPreferredGap(javax.swing.LayoutStyle.ComponentPlacement.RELATED)
                                    .addComponent(self.website)))))
                    .addGroup(layout.createSequentialGroup()
                        .addGap(99, 99, 99)
                        .addComponent(self.logButton, javax.swing.GroupLayout.PREFERRED_SIZE, 169, javax.swing.GroupLayout.PREFERRED_SIZE)))
                .addContainerGap(818, Short.MAX_VALUE))
        )
        layout.setVerticalGroup(
            layout.createParallelGroup(javax.swing.GroupLayout.Alignment.LEADING)
            .addGroup(layout.createSequentialGroup()
                .addGap(34, 34, 34)
                .addGroup(layout.createParallelGroup(javax.swing.GroupLayout.Alignment.BASELINE)
                    .addComponent(self.jLabel1)
                    .addComponent(self.website, javax.swing.GroupLayout.PREFERRED_SIZE, javax.swing.GroupLayout.DEFAULT_SIZE, javax.swing.GroupLayout.PREFERRED_SIZE))
                .addGap(18, 18, 18)
                .addGroup(layout.createParallelGroup(javax.swing.GroupLayout.Alignment.BASELINE)
                    .addComponent(self.jLabel2)
                    .addComponent(self.mongohost, javax.swing.GroupLayout.PREFERRED_SIZE, javax.swing.GroupLayout.DEFAULT_SIZE, javax.swing.GroupLayout.PREFERRED_SIZE))
                .addGap(18, 18, 18)
                .addGroup(layout.createParallelGroup(javax.swing.GroupLayout.Alignment.BASELINE)
                    .addComponent(self.jLabel3)
                    .addComponent(self.mongoport, javax.swing.GroupLayout.PREFERRED_SIZE, javax.swing.GroupLayout.DEFAULT_SIZE, javax.swing.GroupLayout.PREFERRED_SIZE))
                .addGap(32, 32, 32)
                .addComponent(self.logButton, javax.swing.GroupLayout.PREFERRED_SIZE, 51, javax.swing.GroupLayout.PREFERRED_SIZE)
                .addContainerGap(115, Short.MAX_VALUE))
        )
# </editor-fold>//GEN-END:initComponents

    def logButtonActionPerformed(self, evt): #GEN-FIRST:event_logButtonActionPerformed
        stdout = PrintWriter(self.callbacks.getStdout(), True)
        stderr = PrintWriter(self.callbacks.getStderr(), True)

        try:
            #Connect to the database and create the collections
            mongo = MongoClient(self.mongohost.getText(), Integer.parseInt(self.mongoport.getText()))
            db = mongo.getDB("sitelogger")
            siteUrl = URL(self.website.getText())
            tableSite = db.getCollection(siteUrl.getHost().replace("\\.", "_") + "_site")
            tableVuln = db.getCollection(siteUrl.getHost().replace("\\.", "_") + "_vuln")

            #Retrieve SiteMap HTTP Requests and Responses and save to the database
            allReqRes = self.callbacks.getSiteMap(self.website.getText())
            for rc in xrange(0, len(allReqRes)):
                document = BasicDBObject()
                document.put("host", allReqRes[rc].getHost())
                document.put("port", allReqRes[rc].getPort())
                document.put("protocol", allReqRes[rc].getProtocol())
                document.put("URL", allReqRes[rc].getUrl().toString())
                document.put("status_code", allReqRes[rc].getStatusCode())
                if (allReqRes[rc].getRequest() != None):
                    document.put("request", self.helpers.base64Encode(allReqRes[rc].getRequest()))

                if (allReqRes[rc].getResponse() != None):
                    document.put("response", self.helpers.base64Encode(allReqRes[rc].getResponse()))

                tableSite.insert(document)


            #Retrieve Scan findings and save to the database
            allVulns = self.callbacks.getScanIssues(self.website.getText())
            for vc in xrange(0, len(allVulns)):
                document = BasicDBObject()
                document.put("type", allVulns[vc].getIssueType())
                document.put("name", allVulns[vc].getIssueName())
                document.put("detail", allVulns[vc].getIssueDetail())
                document.put("severity", allVulns[vc].getSeverity())
                document.put("confidence", allVulns[vc].getConfidence())
                document.put("host", allVulns[vc].getHost())
                document.put("port", allVulns[vc].getPort())
                document.put("protocol", allVulns[vc].getProtocol())
                document.put("URL", allVulns[vc].getUrl().toString())
                if (len(allVulns[vc].getHttpMessages()) > 1):
                    if (allVulns[vc].getHttpMessages()[0].getRequest() != None):
                        document.put("request", self.helpers.base64Encode(allVulns[vc].getHttpMessages()[0].getRequest()))

                    if (allVulns[vc].getHttpMessages()[0].getResponse() != None):
                        document.put("response", self.helpers.base64Encode(allVulns[vc].getHttpMessages()[0].getResponse()))


                tableVuln.insert(document)


            self.callbacks.issueAlert("Data Saved!")

        except UnknownHostException as ex:
            stderr.println("Mongo DB Connection Error:" + ex.toString())
        except MalformedURLException as ex:
            stderr.println("Malformed URL:" + ex.toString())
