#!/usr/bin/env python
import sys
from time import sleep
from re import sub, findall

from PySide.QtCore import *
from PySide.QtGui import *

from ui.exploit import Ui_MainWindow


class ComboBox(QComboBox):

    def __init__(self, dict_, key, items, default=None, tooltip=None, parent=None):
        super(ComboBox, self).__init__(parent)

        self.setObjectName('%sComboBox' % key.replace(' ', ''))

        # properties
        self.dict_ = dict_
        self.key = key
        self.setItems(items)

        # UI stuff
        if tooltip is not None:
            self.setToolTip(tooltip)

        self.currentIndexChanged.connect(self.updateDataSource)

        self.changeCurrentValue(default)


    def changeCurrentValue(self, value):
        index = self.findData(value)
        if index == -1:
            index = 0
        self.setCurrentIndex(index)
        self.currentIndexChanged.emit(index)


    def setItems(self, value):
        self.clear()
        if isinstance(value, dict):
            for k, v in value.iteritems():
                self.addItem(v, k)
        else:
            for i in value:
                self.addItem(i, i)

    def updateDataSource(self, index):
        self.dict_[self.key] = self.itemData(self.currentIndex())


class SpinBox(QSpinBox):

    def __init__(self, dict_, key, default=0, range_=(-2147483648, 2147483647), tooltip=None, parent=None):
        super(SpinBox, self).__init__(parent)

        self.setObjectName('%sSpinBox' % key.replace(' ', ''))
        self.setRange(*range_)

        # properties
        self.dict_ = dict_
        self.key = key

        # UI stuff
        if tooltip is not None:
            self.setToolTip(tooltip)

        self.valueChanged.connect(self.updateDataSource)

        self.setValue(default)

    def updateDataSource(self, value):
        self.dict_[self.key] = value


class CheckBox(QCheckBox):

    def __init__(self, dict_, key, default=False, tooltip=None, parent=None):
        super(CheckBox, self).__init__(parent)

        self.setObjectName('%sCheckBox' % key.replace(' ', ''))

        # properties
        self.dict_ = dict_
        self.key = key

        # UI stuff
        if tooltip is not None:
            self.setToolTip(tooltip)

        self.stateChanged.connect(self.updateDataSource)

        self.changeState(default)

    def changeState(self, state):
        state = Qt.CheckState.Checked if state else Qt.CheckState.Unchecked
        self.setCheckState(state)
        self.stateChanged.emit(state)

    def updateDataSource(self, value):
        self.dict_[self.key] = bool(value)


class LineEdit(QLineEdit):

    def __init__(self, dict_, key, default='', tooltip=None, parent=None):
        super(LineEdit, self).__init__(parent)

        self.setObjectName('%sLineEdit' % key.replace(' ', ''))

        # properties
        self.dict_ = dict_
        self.key = key

        # UI stuff
        if tooltip is not None:
            self.setToolTip(tooltip)

        self.editingFinished.connect(self.updateDataSource)

        self.changeText(default)

    def changeText(self, text):
        if not isinstance(text, basestring):
            text = str(text)
        self.setText(text)
        self.editingFinished.emit()

    def updateDataSource(self):
        self.dict_[self.key] = self.text()


class Label(QLabel):

    def __init__(self, name, label, tooltip=None, bold=False, parent=None):
        super(Label, self).__init__(parent)

        self.setObjectName('%sLabel' % name.replace(' ', ''))

        self.setText(label)

        if tooltip is not None:
            self.setToolTip(tooltip)

        if bold:
            font = QFont()
            font.setBold(True)
            font.setWeight(75)
            self.setFont(font)


class ExploitWindow(QMainWindow, Ui_MainWindow):

    def __init__(self, msfrpc, defaults, filter_=None, parent=None):
        QMainWindow.__init__(self, parent)
        self.defaults = defaults
        self.sessionid = -1
        self.exploits = []
        self.setupUi(self)
        self.exploitCommandLinkButton.clicked.connect(self.exploit)
        self._msfInit(msfrpc, filter_)

    def _msfInit(self, msfrpc, filter_):
        self.rpc = msfrpc
        self.m = self.rpc.modules

        if filter_ is not None:
            c = self.rpc.consoles.console()
            c.read() # get banner out of the way
            c.write('search %s' % filter_)
            while True:
                d = c.read()
                if not d['busy']:
                    self.exploits.extend(findall('exploit/([^\s]+)', d['data']))
                    break
                sleep(1)
        if not self.exploits:
            self.exploits.extend(self.rpc.modules.exploits)

        for i, e in enumerate(self.exploits):
            self.exploitComboBox.addItem(e)

        self.exploitComboBox.currentIndexChanged.connect(self._exploitChanged)
        self.exploitComboBox.setCurrentIndex(0)
        self.exploitComboBox.currentIndexChanged.emit(0)

    def _exploitChanged(self, index):
        self.e = self.m.use('exploit', self.exploits[index])
        self.setWindowTitle(self.e.name)
        self._setDescription()
        self._populatePages()
        if hasattr(self, 'targetComboBox'):
            del self.targetComboBox
            del self.payloadComboBox
        self._initTargetComboBox()

    def _setDescription(self):
        d = sub(r'\t+', '', self.e.description)
        d = sub(r'([^\n])\n([^\n])', r'\1 \2', d).strip(' ').strip('\n')
        self.descriptionTextBrowser.setText(d)

    def _populatePages(self):

        # Clear all layouts
        for layout in [self.requiredFormLayout, self.evasionFormLayout, self.advancedFormLayout]:
            self.clearLayout(layout)

        # Create dynamic forms
        for i in self.e.options:
            o = self.e.optioninfo(i)

            # set any defaults
            if i in self.defaults:
                self.e.optioninfo(i)['default'] = self.defaults[i]

            # default is required page
            page = self.requiredOptionsPage
            layout = self.requiredFormLayout

            # is it evasion option?
            if o['evasion']:
                page = self.evasionOptionsPage
                layout = self.evasionFormLayout
            # is it advanced option?
            elif o['advanced']:
                page = self.advancedOptionsPage
                layout = self.advancedFormLayout
                # add to layout
            self._populate(page, layout, i, o, self.e)



    def exploit(self):
        j = None
        try:
            j = self.e.execute(payload=self.p)
        except TypeError, e:
            qmb = QMessageBox()
            qmb.setWindowTitle('Error!')
            qmb.setText(str(e))
            qmb.exec_()
            return
        self.hide()
        if j['job_id'] is not None:
            while j['job_id'] in self.rpc.jobs.list:
                sleep(1)
            timeout = self.e.runoptions.get('ConnectTimeout', 10)
            for i in range(timeout):
                sessions = self.rpc.sessions.list
                for k in sessions:
                    if sessions[k]['exploit_uuid'] == j['uuid']:
                        self.sessionid = k
                        break
                sleep(1)
        self.close()


    def _initTargetComboBox(self):
        self.targetComboBox = self._populate(
            self.requiredOptionsPage,
            self.requiredFormLayout,
            'TARGET',
            {
                'default': 0,
                'type': 'targets',
                'enums': self.e.targets,
                'required': True
            },
            self
        )[1]

    def __setitem__(self, key, value):
        if key == 'TARGET' and value is not None:
            self.e.target = value
            self._targetChanged(value)
        elif key == 'PAYLOAD' and value is not None:
            self._payloadChanged(value)

    def _targetChanged(self, target):
        if not hasattr(self, 'payloadComboBox'):
            self._initPayloadComboBox()
        self.payloadComboBox.setItems(self.e.targetpayloads(target))

    def _initPayloadComboBox(self):
        payloads = self.e.targetpayloads(0)

        self.payloadComboBox = self._populate(
            self.requiredOptionsPage,
            self.requiredFormLayout,
            'PAYLOAD',
            {
                'default': payloads[0],
                'type': 'enum',
                'enums': payloads,
                'required' : True
            },
            self
        )[1]

    def clearLayout(self, layout):
        while layout.count():
            c = layout.takeAt(0)
            c.widget().deleteLater()

    def _payloadChanged(self, payload):

        self.clearLayout(self.payloadFormLayout)

        self.p = self.m.use('payload', payload)

        for i in self.p.options:
            if i not in ['RHOST', 'WORKSPACE']:
                self._populate(
                    self.payloadGroupBox,
                    self.payloadFormLayout,
                    i,
                    self.p.optioninfo(i),
                    self.p
                )



    def _populate(self, page, layout, name, optioninfo, dataset):
        type_ = optioninfo['type']
        default = optioninfo.get('default')
        description = optioninfo['desc'] if 'desc' in optioninfo else None

        if default is None:
            if type_ == 'bool':
                default = False
            elif type_ in ['integer', 'port']:
                default = 0
            else:
                default = ''

        l = Label(
            '%sLabel' % name.replace(' ', ''),
            '%s:' % name,
            description,
            optioninfo['required'],
            page
        )

        f = None

        if type_ == 'bool':
            f = CheckBox(dataset, name, default, description, page)
        elif type_ == 'integer':
            f = SpinBox(dataset, name, default, tooltip=description, parent=page)
        elif type_ == 'port':
            f = SpinBox(dataset, name, default, (0, 65535), description, page)
        elif type_ == 'enum':
            f = ComboBox(dataset, name, optioninfo['enums'], default, description, page)
        elif type_ == 'targets':
            f = ComboBox(dataset, name, optioninfo['enums'], default, description, page)
        else:
            f = LineEdit(dataset, name, default, description, page)

        layout.addRow(l, f)
        return l, f

def launch(msfrpc, defaults, filter_=None):
    app = QApplication(sys.argv)
    w = ExploitWindow(msfrpc, defaults, filter_)
    w.show()
    if sys.platform == 'darwin':
        from subprocess import Popen
        Popen(['osascript', '-e', 'tell application "Python" to activate'])
    app.exec_()
    return w.sessionid
