gui.py

#!/usr/bin/python3
#-*- encoding: Utf-8 -*-
from PyQt5.QtWidgets import QApplication, QListWidgetItem, QDesktopWidget, QFileDialog, QInputDialog, QProgressDialog, QMessageBox, QFileSystemModel, QHeaderView
from PyQt5.QtCore import Qt, QUrl, pyqtSignal, QThread
from PyQt5.QtGui import QDesktopServices, QTextOption
from PyQt5.uic import loadUi

from signal import signal, SIGINT, SIG_DFL
from os.path import dirname, realpath
from collections import defaultdict
from urllib.parse import urlparse
from os import listdir, remove
from functools import partial
from json import load, dump
from binascii import crc32
from pathlib import Path
from sys import argv

from utils.common import extractors, transports, BASE_PATH, assert_installed, extractor_save, insert_endpoint, load_proto_msgs
from views.fuzzer import ProtobufItem, ProtocolDataItem
from utils.transports import *
from extractors import *

"""
    This script runs the main code for the PBTK GUI, and essentially
    links loaded QtDesigner *.uis within themselves using signals, and
    with the other parts of code composing PBTK. The order of methods
    more or less follows the action flow of a graphical user.
"""

class PBTKGUI(QApplication):
    def __init__(self):
        super().__init__(argv)
        signal(SIGINT, SIG_DFL)
        
        views = dirname(realpath(__file__)) + '/views/'
        
        self.welcome = loadUi(views + 'welcome.ui')
        self.choose_extractor = loadUi(views + 'choose_extractor.ui')
        self.choose_proto = loadUi(views + 'choose_proto.ui')
        self.create_endpoint = loadUi(views + 'create_endpoint.ui')
        self.choose_endpoint = loadUi(views + 'choose_endpoint.ui')
        self.fuzzer = loadUi(views + 'fuzzer.ui')

        self.welcome.step1.clicked.connect(self.load_extractors)
        self.choose_extractor.rejected.connect(partial(self.set_view, self.welcome))
        self.choose_extractor.extractors.itemClicked.connect(self.prompt_extractor)
        
        self.welcome.step2.clicked.connect(self.load_protos)
        self.proto_fs = QFileSystemModel()
        self.choose_proto.protos.setModel(self.proto_fs)
        self.proto_fs.directoryLoaded.connect(self.choose_proto.protos.expandAll)
        
        for i in range(1, self.proto_fs.columnCount()):
            self.choose_proto.protos.hideColumn(i)
        self.choose_proto.protos.setRootIndex(self.proto_fs.index(str(BASE_PATH / 'protos')))
        self.choose_proto.rejected.connect(partial(self.set_view, self.welcome))
        self.choose_proto.protos.clicked.connect(self.new_endpoint)
        
        self.create_endpoint.transports.itemClicked.connect(self.pick_transport)
        self.create_endpoint.loadRespPbBtn.clicked.connect(self.load_another_pb)
        self.create_endpoint.rejected.connect(partial(self.set_view, self.choose_proto))
        self.create_endpoint.buttonBox.accepted.connect(self.write_endpoint)
                
        self.welcome.step3.clicked.connect(self.load_endpoints)
        self.choose_endpoint.rejected.connect(partial(self.set_view, self.welcome))
        self.choose_endpoint.endpoints.itemClicked.connect(self.launch_fuzzer)
        
        self.fuzzer.rejected.connect(partial(self.set_view, self.choose_endpoint))
        self.fuzzer.fuzzFields.clicked.connect(self.fuzz_endpoint)
        self.fuzzer.deleteThis.clicked.connect(self.delete_endpoint)
        self.fuzzer.comboBox.activated.connect(self.launch_fuzzer)

        self.fuzzer.urlField.setWordWrapMode(QTextOption.WrapAnywhere)
        
        for tree in (self.fuzzer.pbTree, self.fuzzer.getTree):
            tree.itemEntered.connect(lambda item, _: item.edit() if hasattr(item, 'edit') else None)
            tree.itemClicked.connect(lambda item, col: item.update_check(col=col))
            tree.header().setSectionResizeMode(QHeaderView.ResizeToContents)
        
        self.welcome.mydirLabel.setText(self.welcome.mydirLabel.text() % BASE_PATH)
        self.welcome.mydirBtn.clicked.connect(partial(QDesktopServices.openUrl, QUrl.fromLocalFile(str(BASE_PATH))))
        
        self.set_view(self.welcome)
        self.exec_()
    
    """
        Step 1 - Extract .proto structures from apps
    """
    
    def load_extractors(self):
        self.choose_extractor.extractors.clear()
        
        for name, meta in extractors.items():
            item = QListWidgetItem(meta['desc'], self.choose_extractor.extractors)
            item.setData(Qt.UserRole, name)
        
        self.set_view(self.choose_extractor)
    
    def prompt_extractor(self, item):
        extractor = extractors[item.data(Qt.UserRole)]
        inputs = []
        if not assert_installed(self.view, **extractor.get('depends', {})):
            return
        
        if not extractor.get('pick_url', False):
            files, mime = QFileDialog.getOpenFileNames()
            for path in files:
                inputs.append((path, Path(path).stem))
        else:
            text, good = QInputDialog.getText(self.view, ' ', 'Input an URL:')
            if text:
                url = urlparse(text)
                inputs.append((url.geturl(), url.netloc))
        
        if inputs:
            wait = QProgressDialog('Extracting .proto structures...', None, 0, 0)
            wait.setWindowTitle(' ')
            self.set_view(wait)
            
            self.worker = Worker(inputs, extractor)
            self.worker.progress.connect(self.extraction_progress)
            self.worker.finished.connect(self.extraction_done)
            self.worker.signal_proxy.connect(self.signal_proxy)
            self.worker.start()
    
    def extraction_progress(self, info, progress):
        self.view.setLabelText(info)
        
        if progress is not None:
            self.view.setRange(0, 100)
            self.view.setValue(progress * 100)
        else:
            self.view.setRange(0, 0)
    
    def extraction_done(self, outputs):
        nb_written_all, wrote_endpoints = 0, False
        
        for folder, output in outputs.items():
            nb_written, wrote_endpoints = extractor_save(BASE_PATH, folder, output)
            nb_written_all += nb_written
        
        if wrote_endpoints:
            self.set_view(self.welcome)
            QMessageBox.information(self.view, ' ', '%d endpoints and their <i>.proto</i> structures have been extracted! You can now reuse the <i>.proto</i>s or fuzz the endpoints.' % nb_written_all)
        
        elif nb_written_all:
            self.set_view(self.welcome)
            QMessageBox.information(self.view, ' ', '%d <i>.proto</i> structures have been extracted! You can now reuse the <i>.protos</i> or define endpoints for them to fuzz.' % nb_written_all)
        
        else:
            self.set_view(self.choose_extractor)
            QMessageBox.warning(self.view, ' ', 'This extractor did not find Protobuf structures in the corresponding format for specified files.')
    
    """
        Step 2 - Link .protos to endpoints
    """
    
    # Don't load .protos from the filesystem until asked to, in order
    # not to slow down startup.
    
    def load_protos(self):
        self.proto_fs.setRootPath(str(BASE_PATH / 'protos'))
        self.set_view(self.choose_proto)
    
    def new_endpoint(self, path):
        if not self.proto_fs.isDir(path):
            path = self.proto_fs.filePath(path)
            
            if assert_installed(self.choose_proto, binaries=['protoc']):
                if not getattr(self, 'only_resp_combo', False):
                    self.create_endpoint.pbRequestCombo.clear()
                self.create_endpoint.pbRespCombo.clear()
                
                has_msgs = False
                for name, cls in load_proto_msgs(path):
                    has_msgs = True
                    if not getattr(self, 'only_resp_combo', False):
                        self.create_endpoint.pbRequestCombo.addItem(name, (path, name))
                    self.create_endpoint.pbRespCombo.addItem(name, (path, name))
                if not has_msgs:
                    QMessageBox.warning(self.view, ' ', 'There is no message defined in this .proto.')
                    return
                
                self.create_endpoint.reqDataSubform.hide()

                if not getattr(self, 'only_resp_combo', False):
                    self.create_endpoint.endpointUrl.clear()
                    self.create_endpoint.transports.clear()
                    self.create_endpoint.sampleData.clear()
                    self.create_endpoint.pbParamKey.clear()
                    self.create_endpoint.parsePbCheckbox.setChecked(False)
                    
                    for name, meta in transports.items():
                        item = QListWidgetItem(meta['desc'], self.create_endpoint.transports)
                        item.setData(Qt.UserRole, (name, meta.get('ui_data_form')))
                
                elif getattr(self, 'saved_transport_choice'):
                    self.create_endpoint.transports.setCurrentItem(self.saved_transport_choice)
                    self.pick_transport(self.saved_transport_choice)
                    self.saved_transport_choice = None
                
                self.only_resp_combo = False
                self.set_view(self.create_endpoint)
    
    def pick_transport(self, item):
        name, desc = item.data(Qt.UserRole)
        self.has_pb_param = desc and 'regular' in desc
        self.create_endpoint.reqDataSubform.show()
        if self.has_pb_param:
            self.create_endpoint.pbParamSubform.show()
        else:
            self.create_endpoint.pbParamSubform.hide()
        self.create_endpoint.sampleDataLabel.setText('Sample request data, one per line (in the form of %s):' % desc)
    
    def load_another_pb(self):
        self.only_resp_combo = True
        self.saved_transport_choice = self.create_endpoint.transports.currentItem()
        self.set_view(self.choose_proto)
    
    def write_endpoint(self):
        request_pb = self.create_endpoint.pbRequestCombo.itemData(self.create_endpoint.pbRequestCombo.currentIndex())
        url = self.create_endpoint.endpointUrl.text()
        transport = self.create_endpoint.transports.currentItem()
        sample_data = self.create_endpoint.sampleData.toPlainText()
        pb_param = self.create_endpoint.pbParamKey.text()
        has_resp_pb = self.create_endpoint.parsePbCheckbox.isChecked()
        resp_pb = self.create_endpoint.pbRespCombo.itemData(self.create_endpoint.pbRespCombo.currentIndex())
        
        if not (request_pb and urlparse(url).netloc and transport and (not self.has_pb_param or pb_param) and (not has_resp_pb or resp_pb)):
            QMessageBox.warning(self.view, ' ', 'Please fill all relevant information fields.')
        
        else:
            json = {
                'request': {
                    'transport': transport.data(Qt.UserRole)[0],
                    'proto_path': request_pb[0].replace(str(BASE_PATH / 'protos'), '').strip('/\\'),
                    'proto_msg': request_pb[1],
                    'url': url
                }
            }
            if self.has_pb_param:
                json['request']['pb_param'] = pb_param
            
            sample_data = list(filter(None, sample_data.split('\n')))
            if sample_data:
                transport_obj = transports[transport.data(Qt.UserRole)[0]]
                transport_obj = transport_obj['func'](pb_param, url)
                
                for sample_id, sample in enumerate(sample_data):
                    try:
                        sample = transport_obj.serialize_sample(sample)
                    except Exception:
                        return QMessageBox.warning(self.view, ' ', 'Some of your sample data is not in the specified format.')
                    if not sample:
                        return QMessageBox.warning(self.view, ' ', "Some of your sample data didn't contain the Protobuf parameter key you specified.")
                    sample_data[sample_id] = sample
                
                json['request']['samples'] = sample_data
            
            if has_resp_pb:
                json['response'] = {
                    'format': 'raw_pb',
                    'proto_path': resp_pb[0].replace(str(BASE_PATH / 'protos'), '').strip('/\\'),
                    'proto_msg': resp_pb[1]
                }
            insert_endpoint(BASE_PATH / 'endpoints', json)
            
            QMessageBox.information(self.view, ' ', 'Endpoint created successfully.')
            self.set_view(self.welcome)
    
    def load_endpoints(self):
        self.choose_endpoint.endpoints.clear()
        
        for name in listdir(str(BASE_PATH / 'endpoints')):
            if name.endswith('.json'):
                item = QListWidgetItem(name.split('.json')[0], self.choose_endpoint.endpoints)
                item.setFlags(item.flags() & ~Qt.ItemIsEnabled)
                
                pb_msg_to_endpoints = defaultdict(list)
                with open(str(BASE_PATH / 'endpoints' / name)) as fd:
                    for endpoint in load(fd):
                        pb_msg_to_endpoints[endpoint['request']['proto_msg'].split('.')[-1]].append(endpoint)
                
                for pb_msg, endpoints in pb_msg_to_endpoints.items():
                    item = QListWidgetItem(' ' * 4 + pb_msg, self.choose_endpoint.endpoints)
                    item.setFlags(item.flags() & ~Qt.ItemIsEnabled)
                    
                    for endpoint in endpoints:
                        item = QListWidgetItem(' ' * 8 + (urlparse(endpoint['request']['url']).path or '/'), self.choose_endpoint.endpoints)
                        item.setData(Qt.UserRole, endpoint)
        
        self.set_view(self.choose_endpoint)
    
    """
        Step 3: Fuzz and test endpoints live
    """
    
    def launch_fuzzer(self, item):
        if type(item) == int:
            data, sample_id = self.fuzzer.comboBox.itemData(item)
        else:
            data, sample_id = item.data(Qt.UserRole), 0
        
        if data and assert_installed(self.view, binaries=['protoc']):
            self.current_req_proto = BASE_PATH / 'protos' / data['request']['proto_path']
            
            self.pb_request = load_proto_msgs(self.current_req_proto)
            self.pb_request = dict(self.pb_request)[data['request']['proto_msg']]()
            
            if data.get('response') and data['response']['format'] == 'raw_pb':
                self.pb_resp = load_proto_msgs(BASE_PATH / 'protos' / data['response']['proto_path'])
                self.pb_resp = dict(self.pb_resp)[data['response']['proto_msg']]

            self.pb_param = data['request'].get('pb_param')
            self.base_url = data['request']['url']
            self.endpoint = data
            
            self.transport_meta = transports[data['request']['transport']]
            self.transport = self.transport_meta['func'](self.pb_param, self.base_url)
            
            sample = ''
            if data['request'].get('samples'):
                sample = data['request']['samples'][sample_id]
            self.get_params = self.transport.load_sample(sample, self.pb_request)
            
            # Get initial data into the Protobuf tree view
            self.fuzzer.pbTree.clear()
            
            self.ds_items = defaultdict(dict)
            self.ds_full_names = {}
            
            self.parse_desc(self.pb_request.DESCRIPTOR, self.fuzzer.pbTree)
            self.parse_fields(self.pb_request)
            
            # Do the same for transport-specific data
            self.fuzzer.getTree.clear()
            if self.transport_meta.get('ui_tab'):
                self.fuzzer.tabs.setTabText(1, self.transport_meta['ui_tab'])
                if self.get_params:
                    for key, val in self.get_params.items():
                        ProtocolDataItem(self.fuzzer.getTree, key, val, self)
            else:
                self.fuzzer.tabs.setTabText(1, '(disabled)')
                # how to hide it ?
            
            # Fill the request samples combo box if we're loading a new
            # endpoint.
            if type(item) != int:
                if len(data['request'].get('samples', [])) > 1:
                    self.fuzzer.comboBox.clear()
                    for sample_id, sample in enumerate(data['request']['samples']):
                        self.fuzzer.comboBox.addItem(sample[self.pb_param] if self.pb_param else str(sample), (data, sample_id))
                    self.fuzzer.comboBoxLabel.show()
                    self.fuzzer.comboBox.show()
                else:
                    self.fuzzer.comboBoxLabel.hide()
                    self.fuzzer.comboBox.hide()
                
                self.set_view(self.fuzzer)
            
            self.fuzzer.frame.setUrl(QUrl("about:blank"))
            self.update_fuzzer()

    """
        Parsing and rendering the Protobuf message to a tree view:

        Every Protobuf field is fed to ProtobufItem (a class inheriting
        QTreeWidgetItem), and the created object is saved in the ds_items
        entry for the corresponding descriptor.
    """
    
    # First, parse the descriptor (structure) of the Protobuf message.
    
    def parse_desc(self, msg, item, path=[]):
        for ds in msg.fields:
            new_item = ProtobufItem(item, ds, self, path)
            if ds.type == ds.TYPE_MESSAGE and ds.full_name not in path:
                self.parse_desc(ds.message_type, new_item, path + [ds.full_name])

    # Then, parse the fields (contents) of the Protobuf message.
    
    def parse_fields(self, msg, path=[]):
        for ds, val in msg.ListFields():
            if ds.label == ds.LABEL_REPEATED:
                for val_index, val_value in enumerate(val):
                    if ds.type == ds.TYPE_MESSAGE:
                        self.ds_items[id(ds)][tuple(path)].setExpanded(True)
                        self.ds_items[id(ds)][tuple(path)].setDefault(parent=msg, msg=val, index=val_index)
                        self.parse_fields(val_value, path + [ds.full_name])
                    
                    else:
                        self.ds_items[id(ds)][tuple(path)].setDefault(val_value, parent=msg, msg=val, index=val_index)
                    
                    self.ds_items[id(ds)][tuple(path)].duplicate(True)
            
            else:
                if ds.type == ds.TYPE_MESSAGE:
                    self.ds_items[id(ds)][tuple(path)].setExpanded(True)
                    self.ds_items[id(ds)][tuple(path)].setDefault(parent=msg, msg=val)
                    self.parse_fields(val, path + [ds.full_name])
                
                else:
                    self.ds_items[id(ds)][tuple(path)].setDefault(val, parent=msg, msg=val)
    
    def update_fuzzer(self):
        resp = self.transport.perform_request(self.pb_request, self.get_params)
        
        data, text, url, mime = resp.content, resp.text, resp.url, resp.headers['Content-Type'].split(';')[0]
        
        meta = '%s %d %08x\n%s' % (mime, len(data), crc32(data) & 0xffffffff, resp.url)
        self.fuzzer.urlField.setText(meta)
        
        self.fuzzer.frame.update_frame(data, text, url, mime, getattr(self, 'pb_resp', None))
    
    def fuzz_endpoint(self):
        QMessageBox.information(self.view, ' ', 'Automatic fuzzing is not implemented yet.')
    
    def delete_endpoint(self):
        if QMessageBox.question(self.view, ' ', 'Delete this endpoint?') == QMessageBox.Yes:
            path = str(BASE_PATH / 'endpoints' / (urlparse(self.base_url).netloc + '.json'))
            
            with open(path) as fd:
                json = load(fd)
            json.remove(self.endpoint)
            
            with open(path, 'w') as fd:
                dump(json, fd, ensure_ascii=False, indent=4)
            if not json:
                remove(path)
            
            self.load_endpoints()
    
    """
        Utility methods follow
    """
    
    def set_view(self, view):
        if hasattr(self, 'view'):
            self.view.hide()
        view.show()
        self.view = view
        
        resolution = QDesktopWidget().screenGeometry()
        view.move((resolution.width() / 2) - (view.frameSize().width() / 2),
                  (resolution.height() / 2) - (view.frameSize().height() / 2))
    
    """
        signal() can't be called from inside a thread, and some
        extractors need it in order not to have their GUI child process
        interrupt signal catched by our main thread, so here is an ugly
        way to reach signal() through a slot.
    """
    def signal_proxy(self, *args):
        signal(*args)

"""
    Simple wrapper for running extractors in background.
"""

class Worker(QThread):
    finished = pyqtSignal(object)
    progress = pyqtSignal(object, object)
    signal_proxy = pyqtSignal(object, object)

    def __init__(self, inputs, extractor):
        super().__init__()
        self.inputs = inputs
        self.extractor = extractor
    
    def run(self):
        output = defaultdict(list)
        for input_, folder in self.inputs:
            # Extractor is runned here
            for name, contents in self.extractor['func'](input_):
                if name == '_progress':
                    self.progress.emit(*contents)
                elif name == '_signal':
                    self.signal_proxy.emit(*contents)
                else:
                    output[folder].append((name, contents))
        
        self.finished.emit(output)

PBTKGUI()