#!/usr/bin/python3
from PySide6.QtWidgets import (
    QApplication,
    QListWidgetItem,
    QFileDialog,
    QInputDialog,
    QProgressDialog,
    QMessageBox,
    QFileSystemModel,
    QHeaderView,
)
from PySide6.QtGui import QDesktopServices, QTextOption
from PySide6.QtCore import Qt, QUrl, Signal, QThread
from PySide6.QtUiTools import QUiLoader

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

__import__("sys").path.append(dirname(realpath(__file__)))

from utils.common import (
    extractors,
    transports,
    BASE_PATH,
    assert_installed,
    extractor_save,
    insert_endpoint,
    load_proto_msgs,
)
from views.fuzzer import ProtobufItem, ProtocolDataItem, MyFrame
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/"

        loader = QUiLoader()
        loader.registerCustomWidget(MyFrame)

        self.welcome = loader.load(views + "welcome.ui")
        self.choose_extractor = loader.load(views + "choose_extractor.ui")
        self.choose_proto = loader.load(views + "choose_proto.ui")
        self.create_endpoint = loader.load(views + "create_endpoint.ui")
        self.choose_endpoint = loader.load(views + "choose_endpoint.ui")
        self.fuzzer = loader.load(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.getAdd.clicked.connect(self.add_tab_data)

        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.itemExpanded.connect(
                lambda item: item.expanded() if hasattr(item, "expanded") else None
            )
            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.start()

    def extraction_progress(self, info, progress):
        self.view.setLabelText(info)

        if progress is not None:
            self.view.setRange(0, 100)
            self.view.setValue(int(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 not getattr(self, "only_resp_combo", False):
                self.create_endpoint.pbRequestCombo.clear()
            self.create_endpoint.pbRespCombo.clear()

            has_msgs = False
            for name, _ 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 hasattr(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)

    """
        Step 3: Fuzz and test endpoints live
    """

    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, object_pairs_hook=OrderedDict):
                        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:
                        path_and_qs = (
                            "/" + endpoint["request"]["url"].split("/", 3).pop()
                        )
                        item = QListWidgetItem(
                            " " * 8 + path_and_qs, self.choose_endpoint.endpoints
                        )
                        item.setData(Qt.UserRole, endpoint)

        self.set_view(self.choose_endpoint)

    def launch_fuzzer(self, item):
        if isinstance(item, int):
            data, sample_id = self.fuzzer.comboBox.itemData(item)
        else:
            data, sample_id = item.data(Qt.UserRole), 0

        if data:
            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"]]()

            self.pb_resp = None

            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 = {}

            for ds in self.pb_request.DESCRIPTOR.fields:
                ProtobufItem(self.fuzzer.pbTree, ds, self, [ds.full_name])
            self.parse_fields(self.pb_request)

            # Do the same for transport-specific data
            self.fuzzer.getTree.clear()
            self.fuzzer.tabs.setTabText(1, self.transport_meta.get("ui_tab", ""))
            if self.get_params:
                for key, val in self.get_params.items():
                    ProtocolDataItem(self.fuzzer.getTree, key, val, self)

            # Fill the request samples combo box if we're loading a new
            # endpoint.
            if not isinstance(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.
    """

    def parse_fields(self, msg, base_path=None):
        if base_path is None:
            base_path = []
        for ds, val in msg.ListFields():
            path = [*base_path + ds.full_name]

            if ds.label == ds.LABEL_REPEATED:
                for val_index, val_value in enumerate(val):
                    if ds.cpp_type == ds.CPPTYPE_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)

                    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.cpp_type == ds.CPPTYPE_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)

                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, object_pairs_hook=OrderedDict)
            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()

    def add_tab_data(self):
        text, good = QInputDialog.getText(self.view, " ", "Field name:")
        if text:
            ProtocolDataItem(self.fuzzer.getTree, text, "", self).edit()

    """
        Utility methods follow
    """

    def set_view(self, view):
        if hasattr(self, "view"):
            self.view.hide()
        view.show()
        self.view = view

        screen = self.primaryScreen()
        resolution = screen.geometry()
        view.move(
            int((resolution.width() / 2) - (view.frameSize().width() / 2)),
            int((resolution.height() / 2) - (view.frameSize().height() / 2)),
        )


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


class Worker(QThread):
    finished = Signal(object)
    progress = Signal(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)
                else:
                    output[folder].append((name, contents))

        self.finished.emit(output)

def main():
    PBTKGUI()

if __name__ == "__main__":
    main()
