pymuleparser/pyMuleAnalyzerQT.py

#! /usr/bin/env python3
"""
"""

import sys
import os
import traceback
import time
import random

import PyQt5
from PyQt5.QtCore import Qt, QDateTime, QObject, pyqtSignal, pyqtSlot, QRunnable, QThreadPool
from PyQt5.QtWidgets import QApplication, QMainWindow, QFileDialog, QAbstractItemView, \
    QTableWidgetItem, QSplashScreen, QProgressBar
from PyQt5.QtGui import QStandardItemModel, QPixmap, QMovie, QPainter
from MuleAnalyzerUI import Ui_MainWindow

from ED2K import load_knownfiles
from HashDatabase import load_hashsets


class ExceptionHandler(QObject):
    errorSignal = pyqtSignal()

    def __init__(self):
        super(ExceptionHandler, self).__init__()

    def handler(self, exctype, value, traceback):
        self.errorSignal.emit()
        sys._excepthook(exctype, value, traceback)


exceptionHandler = ExceptionHandler()
sys._excepthook = sys.excepthook
sys.excepthook = exceptionHandler.handler


class MovieSplashScreen(QSplashScreen):

    def __init__(self, movie, parent=None):
        movie.jumpToFrame(0)
        pixmap = QPixmap(movie.frameRect().size())
        QSplashScreen.__init__(self, pixmap, Qt.WindowStaysOnTopHint)
        self.movie = movie
        self.movie.frameChanged.connect(self.repaint)

    def showEvent(self, event):
        self.movie.start()

    def hideEvent(self, event):
        self.movie.stop()

    def paintEvent(self, event):
        painter = QPainter(self)
        pixmap = self.movie.currentPixmap()
        self.setMask(pixmap.mask())
        painter.drawPixmap(0, 0, pixmap)

    def sizeHint(self):
        return self.movie.scaledSize()


class WorkerSignals(QObject):
    finished = pyqtSignal()
    error = pyqtSignal(tuple)
    result = pyqtSignal(object)
    progress = pyqtSignal(str)


class Worker(QRunnable):

    def __init__(self, fn, *args, **kwargs):
        QRunnable.__init__(self)
        self.fn = fn
        self.args = args
        self.kwargs = kwargs
        self.signals = WorkerSignals()
        kwargs['progress_callback'] = self.signals.progress

    @pyqtSlot()
    def run(self):
        # Retrieve args/kwargs here; and fire processing using them
        try:
            result = self.fn(*self.args, **self.kwargs)
        except:
            traceback.print_exc()
            exctype, value = sys.exc_info()[:2]
            self.signals.error.emit((exctype, value, traceback.format_exc()))
        else:
            self.signals.result.emit(result)
        finally:
            self.signals.finished.emit()


class MainWindow(QMainWindow, Ui_MainWindow):

    # Todo: Replace TreeHeaders
    Headers = (
        ('ID', None, None, 0, Qt.AlignRight),
        ('Filename', 'FT_FILENAME', str, '', Qt.AlignLeft),
        ('ED2K Hash', 'ED2KHASH', str, '', Qt.AlignLeft),
        ('Hashset', None, None, '', Qt.AlignLeft),
        ('Last Changed', 'LASTCHANGED', QDateTime, 'Never', Qt.AlignLeft),
        ('Filesize', 'FT_FILESIZE', int, 0, Qt.AlignRight),
        ('Transfered', 'FT_ATTRANSFERRED', int, 0, Qt.AlignRight),
        ('Transfered (HI)', 'FT_ATTRANSFERREDHI', int, 0, Qt.AlignRight),
        ('Requested', 'FT_ATREQUESTED', int, 0, Qt.AlignRight),
        ('Accepted', 'FT_ATACCEPTED', int, 0, Qt.AlignRight),
        ('AICHashes', 'FT_AICHHASHSET', str, '', Qt.AlignLeft),
        ('AICHash', 'FT_AICH_HASH', str, '', Qt.AlignLeft),
        ('Part Filename', 'FT_PARTFILENAME', str, '', Qt.AlignLeft),
        ('Number of Parts', 'NUMPARTS', int, 0, Qt.AlignRight),
        ('Parts Hash', 'PARTHASHES', list, '', Qt.AlignLeft),
    )

    def __init__(self, *args, **kwargs):
        super(MainWindow, self).__init__()
        self.setupUi(self)

        self.treeModel = self.createED2KModel(self)
        self.treeView.setModel(self.treeModel)
        self.treeView.setRootIsDecorated(False)
        self.treeView.setAlternatingRowColors(True)
        self.treeView.setSortingEnabled(True)
        self.treeView.setEditTriggers(QAbstractItemView.NoEditTriggers)
        #self.treeView.setSelectionMode(QAbstractItemView.MultiSelection)

        self.tableWidget.setEditTriggers(QAbstractItemView.NoEditTriggers)

        self.actionFile_open.setShortcut('Ctrl+O')
        self.actionFile_open.triggered.connect(self.slot_load_met_file)
        self.actionExit.triggered.connect(self.close)

        self.threadpool = QThreadPool()
        self.hashsets = None
        self.splash_screen = kwargs['splash_screen']
        worker = Worker(self.load_hashsets)
        worker.signals.result.connect(self.cb_got_hashset)
        worker.signals.progress.connect(self.cb_progress)
        self.threadpool.start(worker)

    def cb_progress(self, msg):
        self.splash_screen.showMessage('<h1>%s</h1>' % msg, Qt.AlignTop | Qt.AlignCenter, Qt.black)

    def load_hashsets(self, progress_callback, folder='Hashsets'):
        self.hashsets = load_hashsets(folder, progress_callback)
        return self.hashsets

    def cb_print_msg(self, message):
        print('>>>', message)

    def cb_got_hashset(self, hashes):
        num_hashsets = len(hashes.keys())
        num_hashes = 0
        for key in hashes:
            num_hashes += len(hashes[key])
        self.statusBar.showMessage('Loaded {} hashsets with {:,} hashes.'.format(num_hashsets,
                                                                                 num_hashes))

    def createED2KModel(self, parent):
        model = QStandardItemModel(0, len(self.Headers), parent)
        for i, fields in enumerate(self.Headers):
            model.setHeaderData(i, Qt.Horizontal, fields[0])
        return model

    def addED2K(self, model, metinfo):
        model.insertRow(0)
        for i, field in enumerate(self.Headers):
            if field[2] == None:
                if i == 0:
                    root = self.treeModel.invisibleRootItem()
                    row_count = root.rowCount()
                    model.setData(model.index(0, i), row_count+1)
                if i == 3:
                    found_in = []
                    for hashset, hashes in self.hashsets.items():
                        if metinfo.ED2KHASH.upper() in hashes or metinfo.ED2KHASH.lower() in hashes:
                            found_in.append(hashset)
                            break
                    # TODO Hash lookup
                    model.setData(model.index(0, i), ', '.join(found_in))
            else:
                value = getattr(metinfo, field[1]) or field[3]
                if field[2] is QDateTime:
                    last_changed = QDateTime.fromSecsSinceEpoch(value, 0)
                    model.setData(model.index(0, i), last_changed.toString('dd.MM.yyyy hh:mm:ss'))
                elif field[2] is int:
                    model.setData(model.index(0, i), '{0:,}'.format(value))
                elif field[2] is list:
                    model.setData(model.index(0, i), ''.join(value))
                else:
                    model.setData(model.index(0, i), value)
            model.item(0, i).setTextAlignment(field[4])

    def loadMetFile(self, filename):
        for met_entry in load_knownfiles(filename):
            self.addED2K(self.treeModel, met_entry)

    @pyqtSlot()
    def slot_load_met_file(self):
        options = QFileDialog.Options()
        options |= QFileDialog.DontUseNativeDialog
        files, _ = QFileDialog.getOpenFileNames(self, 'Open .met files...', '',
                                                'MET Files (*.met);;All Files (*)',
                                                options=options)
        for filename in files:
            self.loadMetFile(filename)

        root = self.treeModel.invisibleRootItem()
        row_count = root.rowCount()
        transfered = 0
        for i in range(row_count):
            row = self.treeModel.item(i, column=5)
            transfered += int(row.text().replace(',', ''))
        self.tableWidget.setRowCount(1)
        self.tableWidget.setColumnCount(2)
        self.tableWidget.setItem(0, 0, QTableWidgetItem('Transfered overall:'))
        self.tableWidget.setItem(0, 1, QTableWidgetItem(str('{0:,}'.format(transfered))))


def main(argv):
    app = QApplication(argv)

    resources = './resources/loaders/'
    movies = []
    for filename in os.listdir(resources):
        filepath = os.path.join(resources, filename)
        if os.path.isfile(filepath):
            movies.append(filepath)
    movie = QMovie(random.choice(movies))
    splash_screen = MovieSplashScreen(movie)
    splash_screen.setEnabled(False)
    splash_screen.show()

    window = MainWindow(splash_screen=splash_screen)
    while not window.hashsets:
        app.processEvents()
    window.show()
    splash_screen.finish(window)

    sys.exit(app.exec_())


if __name__ == '__main__':
    main(sys.argv)