mtkclient/mtk_gui

#!/usr/bin/env python3
# MTK Flash Client (c) G.Kreileman, B.Kerler 2021.
# Licensed under GPLv3 License
import sys
import time
import mock
import traceback
import math
import logging
import ctypes
from PySide2.QtCore import Qt, QVariantAnimation, Signal, QObject, QSize, QTranslator, QLocale, QLibraryInfo
from PySide2.QtGui import QTextOption, QPixmap, QTransform, QIcon
from PySide2.QtWidgets import QMainWindow, QApplication, QAction, QWidget

from mtkclient.Library.mtk import Mtk
from mtkclient.Library.mtk_da_cmd import DA_handler
from mtkclient.Library.gpt import gpt_settings
from mtkclient.Library.mtk_main import Main, Mtk_Config

from mtkclient.gui.readFlashPartitions import ReadFlashWindow
from mtkclient.gui.readFull import ReadFullFlashWindow
from mtkclient.gui.toolsMenu import generateKeysMenu
from mtkclient.gui.toolkit import asyncThread, trap_exc_during_debug
from mtkclient.config.payloads import pathconfig
from mtkclient.gui.main_gui import Ui_MainWindow
import os
os.environ['QT_MAC_WANTS_LAYER'] = '1' #This fixes a bug in pyside2 on MacOS Big Sur
# TO do Move all GUI modifications to signals!
# install exception hook: without this, uncaught exception would cause application to exit
sys.excepthook = trap_exc_during_debug

# Initiate MTK classes
variables = mock.Mock()
variables.cmd = "stage"
variables.debugmode = True
path=pathconfig()
# if sys.platform.startswith('darwin'):
#    config.ptype = "kamakiri" #Temp for Mac testing
MtkTool = Main(variables)

guiState = "welcome"
phoneInfo = {"chipset": "", "bootMode": "", "daInit": False, "cdcInit": False}

class DeviceHandler(QObject):
    sendToLogSignal = Signal(str)
    mtkClass = None
    da_handler = None

    def __init__(self, parent, preloader: str = None, loglevel=logging.INFO, *args, **kwargs):
        super().__init__(parent, *args, **kwargs)
        config = Mtk_Config(loglevel=logging.INFO, gui=self.sendToLogSignal)
        config.gpt_settings = gpt_settings(gpt_num_part_entries='0', gpt_part_entry_size='0',
                                           gpt_part_entry_start_lba='0')  # This actually sets the right GPT settings..
        self.mtkClass = Mtk(config=config, loglevel=logging.INFO)
        self.loglevel = logging.DEBUG
        self.da_handler = DA_handler(self.mtkClass, loglevel)

def getDevInfo(self, parameters):
    loglevel = parameters[0]
    phoneInfo = parameters[1]
    devhandler = parameters[2]

    mtkClass = devhandler.mtkClass
    da_handler = devhandler.da_handler
    try:
        if mtkClass.port.cdc.connect() == False:
            mtkClass.preloader.init()
        else:
            phoneInfo['cdcInit'] = True
    except:
        phoneInfo['cantConnect'] = True
    phoneInfo['chipset'] = str(mtkClass.config.chipconfig.name) + \
                           " (" + str(mtkClass.config.chipconfig.description) + ")"
    self.sendUpdateSignal.emit()
    mtkClass = da_handler.configure_da(mtkClass, preloader=None)
    if mtkClass:
        phoneInfo['daInit'] = True
        phoneInfo['chipset'] = str(mtkClass.config.chipconfig.name) + \
                               " (" + str(mtkClass.config.chipconfig.description) + ")"
        if mtkClass.config.is_brom:
            phoneInfo['bootMode'] = "Bootrom mode"
        elif mtkClass.config.chipconfig.damode:
            phoneInfo['bootMode'] = "DA mode"
        else:
            phoneInfo['bootMode'] = "Preloader mode"
        self.sendUpdateSignal.emit()
    else:
        phoneInfo['cantConnect'] = True
        self.sendUpdateSignal.emit()
    # try:
    #    print(mtkClass.daloader.get_partition_data(parttype="user"))
    # except Exception:
    #    print(traceback.format_exc())
    # MtkTool.cmd_stage(mtkClass, None, None, None, False)


def sendToLog(info):
    t = time.localtime()
    mainwindow.logBox.appendPlainText(time.strftime("[%H:%M:%S", t) + "]: " + info)
    mainwindow.logBox.verticalScrollBar().setValue(mainwindow.logBox.verticalScrollBar().maximum())


def updateGui():
    global phoneInfo
    if (phoneInfo['cdcInit'] and phoneInfo['bootMode'] == "" ):
        mainwindow.phoneInfoTextbox.setText("Phone detected:\nReading model info...")
    else:
        mainwindow.phoneInfoTextbox.setText("Phone detected:\n" + phoneInfo['chipset'] + "\n" + phoneInfo['bootMode'])
    mainwindow.status.setText("Device detected, please wait.\nThis can take a while...")
    if phoneInfo['daInit']:
        mainwindow.status.setText("Device connected :)")
        mainwindow.menubar.setEnabled(True)
        pixmap = QPixmap(path.get_images_path("phone_connected.png"))
        mainwindow.pic.setPixmap(pixmap)
        spinnerAnim.stop()
        mainwindow.spinner_pic.setHidden(True)
    else:
        if 'cantConnect' in phoneInfo:
            mainwindow.status.setText("Error initialising. Did you install the drivers?")
        spinnerAnim.start()
        mainwindow.spinner_pic.setHidden(False)


def openReadflashWindow(q):
    # status.setText("OH YEAH"+str(q.text()))
    curmenutext = q.text()
    if curmenutext == "Read partition(s)":
        readFlashWindowVal = ReadFlashWindow(w, devhandler.mtkClass, devhandler.da_handler, sendToLog)
    elif curmenutext == "Read full flash":
        readFlashWindowVal = ReadFullFlashWindow(w, devhandler.mtkClass, devhandler.da_handler, sendToLog)
    # w.close()


def openToolsMenu(q):
    # status.setText("OH YEAH"+str(q.text()))
    curmenutext=q.text()
    if curmenutext == "Generate RPMB keys":
        readFlashWindowVal = generateKeysMenu(w, devhandler.mtkClass, devhandler.da_handler, sendToLog)
    # w.close()


def load_translations(application):
    """
    Load application translations and the QT base translations for the current locale
    :param application: QApplication object
    """
    locale = QLocale.system()
    translator = QTranslator(application)
    directory = os.path.dirname(__file__)
    lang = 'mtkclient/gui/i18n/' + locale.name()
    if translator.load(lang, directory):
        application.installTranslator(translator)

    translations_path = QLibraryInfo.location(QLibraryInfo.TranslationsPath)
    base_translator = QTranslator(application)
    if base_translator.load(locale, "qtbase", "_", translations_path):
        application.installTranslator(base_translator)

if __name__ == '__main__':
    # Enable nice 4K Scaling
    QApplication.setAttribute(Qt.AA_EnableHighDpiScaling, True)
    os.environ["QT_AUTO_SCREEN_SCALE_FACTOR"] = "1"

    # Init the app window
    app = QApplication(sys.argv)
    win = QMainWindow()
    w = QWidget(win)
    load_translations(app)
    mainwindow=Ui_MainWindow()
    mainwindow.setupUi(win)
    icon = QIcon()
    icon.addFile(path.get_images_path('logo_32.png'), QSize(32, 32))
    icon.addFile(path.get_images_path('logo_64.png'), QSize(64, 64))
    icon.addFile(path.get_images_path('logo_256.png'), QSize(256, 256))
    icon.addFile(path.get_images_path('logo_512.png'), QSize(512, 512))
    app.setWindowIcon(icon)
    win.setWindowIcon(icon)
    if sys.platform.startswith('win'):
        ctypes.windll.shell32.SetCurrentProcessExplicitAppUserModelID('MTKTools.Gui')
    dpiMultiplier = win.logicalDpiX()
    if dpiMultiplier == 72:
        dpiMultiplier = 2
    else:
        dpiMultiplier = 1
    addTopMargin = 20
    if sys.platform.startswith('darwin'):  # MacOS has the toolbar in the top bar insted of in the app...
        addTopMargin = 0
    win.setWindowTitle("MTKClient - Version 2.0 beta")
    # lay = QVBoxLayout(self)

    # Menubar
    mainwindow.readFlashMenu.triggered[QAction].connect(openReadflashWindow)
    mainwindow.toolsFlashMenu.triggered[QAction].connect(openToolsMenu)

    # phone spinner
    pixmap = QPixmap(path.get_images_path("phone_loading.png")).scaled(96, 96, Qt.KeepAspectRatio, Qt.SmoothTransformation)
    pixmap.setDevicePixelRatio(2)
    mainwindow.spinner_pic.setPixmap(pixmap)
    mainwindow.spinner_pic.show()

    nfpixmap = QPixmap(path.get_images_path("phone_notfound.png"))
    mainwindow.pic.setPixmap(nfpixmap)

    logo = QPixmap(path.get_images_path("logo_256.png"))
    mainwindow.logoPic.setPixmap(logo)
    mainwindow.logBox.setHidden(True)

    def spinnerAnimRot(angle):
        # print(angle)
        trans = QTransform()
        # trans.translate(pixmap.width()+angle / 2, pixmap.height() / 2)
        dimension = pixmap.width() / math.sqrt(2)
        #trans.translate(100 , mainwindow.spinner_pic.height()/2.0)
        #trans.rotate(angle)
        #newPixmap = pixmap.transformed(trans, Qt.SmoothTransformation)
        newPixmap = pixmap.transformed(QTransform().rotate(angle), Qt.SmoothTransformation)
        xoffset = (newPixmap.width() - pixmap.width())/2
        yoffset = (newPixmap.height() - pixmap.height())/2
        rotated = newPixmap.copy(xoffset, yoffset, pixmap.width(), pixmap.height())
        mainwindow.spinner_pic.setPixmap(rotated)


    spinnerAnim = QVariantAnimation()
    spinnerAnim.setDuration(3000)
    spinnerAnim.setStartValue(0)
    spinnerAnim.setEndValue(360)
    spinnerAnim.setLoopCount(-1)
    spinnerAnim.valueChanged.connect(spinnerAnimRot)

    mainwindow.spinner_pic.setHidden(True)

    def showDebugInfo():
        if mainwindow.logBox.isHidden():
            mainwindow.logBox.setHidden(False)
            win.setFixedSize(746, 600 + addTopMargin)
            mainwindow.debugBtn.setText("Hide debug info")
        else:
            mainwindow.logBox.setHidden(True)
            win.setFixedSize(746, 400 + addTopMargin)
            mainwindow.debugBtn.setText("Show debug info")

    mainwindow.debugBtn.clicked.connect(showDebugInfo)
    mainwindow.logBox.setWordWrapMode(QTextOption.NoWrap)
    mainwindow.menubar.setEnabled(False)
    win.show()
    win.setFixedSize(746, 400 + addTopMargin)

    # Device setup
    loglevel = logging.INFO
    devhandler = DeviceHandler(parent=app,preloader=None,loglevel=loglevel)
    devhandler.sendToLogSignal.connect(sendToLog)

    # Get the device info
    thread = asyncThread(parent=app, n=0, function=getDevInfo, parameters=[loglevel,phoneInfo,devhandler])
    thread.sendToLogSignal.connect(sendToLog)
    thread.sendUpdateSignal.connect(updateGui)
    thread.start()

    # Run loop the app
    app.exec_()
    # Prevent thread from not being closed and call error end codes
    thread.terminate()
    thread.wait()