diff --git a/forms/mainwindow.py b/forms/mainwindow.py new file mode 100644 --- /dev/null +++ b/forms/mainwindow.py @@ -0,0 +1,218 @@ +from forms.about import AboutWidget +from ui.mainwindow import Ui_MainWindow +from PyQt5 import QtWidgets, QtCore, QtGui +from devices import get_storage_devices +from mount.udisks import MountManager +import tempfile +import os +import shutil +import stat +import floppyemu +from .format_dlgs import FormatManyDialog, FormatOneDialog + + +class CustomListModel(QtCore.QStringListModel): + def __init__(self, model=()): + self.model = model + super().__init__([m[1] for m in self.model]) + + def setModel(self, model): + self.model = model + self.setStringList([m[1] for m in self.model]) + + def getItem(self, index): + return self.model[index] + + +class CustomTableModel(QtCore.QAbstractTableModel): + headers = ['Label', 'Size'] + + def __init__(self, placeholder): + super().__init__() + self.placeholder = placeholder + self.model = [] + + def setModel(self, model): + self.beginResetModel() + self.model = model + self.endResetModel() + + def data(self, index, role=QtCore.Qt.DisplayRole): + model = self.model[index.row()] + if role == QtCore.Qt.DisplayRole: + if QtCore.QVariant(model[index.column()]): + if index.column() == 1: + return QtCore.QVariant(str(round(model[index.column()] / 1024, 2)) + 'KB') + else: + return QtCore.QVariant(model[index.column()]) + else: + return QtCore.QVariant(self.placeholder) + + def headerData(self, section, orientation, role=QtCore.Qt.DisplayRole): + if orientation == QtCore.Qt.Horizontal and role == QtCore.Qt.DisplayRole: + return QtCore.QVariant(self.headers[section]) + return super().headerData(section, orientation, role) + + def rowCount(self, parent=None): + return len(self.model) + + def columnCount(self, parent=None): + return 2 + + +class MainWindow(Ui_MainWindow, QtWidgets.QMainWindow): + def __init__(self, root=None): + super().__init__() + self.setupUi(self) + + self.root = root if root else tempfile.mktemp(prefix='kusbff_') + os.makedirs(self.root, exist_ok=True) + + self.devices_model = CustomListModel() + self.devices_combobox.setModel(self.devices_model) + self.refresh_disk_drives() + + self.table_model = CustomTableModel('(Not set)') + self.tableView.setModel(self.table_model) + self.refresh_image_list() + + self.action_refresh_devices.triggered.connect(self.refresh_disk_drives) + self.devices_combobox.currentIndexChanged.connect(self.on_device_change) + self.action_open_floppy.triggered.connect(self.open_directory) + self.action_write_floppy.triggered.connect(self.save_directory) + self.action_read_image.triggered.connect(self.read_image) + self.action_write_image.triggered.connect(self.write_image) + self.action_format_disk.triggered.connect(self.format_one) + self.action_format_usb.triggered.connect(self.format_many) + self.action_about.triggered.connect(self.show_about_dialog) + + def show_about_dialog(self): + AboutWidget().exec_() + + def format_one(self): + if self.devices_combobox.currentIndex() != -1: + indexes = self.tableView.selectedIndexes() + device = '/dev/{}'.format(self.devices_model.getItem(self.devices_combobox.currentIndex())[0]) + FormatOneDialog.show_dialog(indexes[0].row(), device, os.path.join(self.root, 'temp')) + self.refresh_image_list() + + def format_many(self): + if self.devices_combobox.currentIndex() != -1: + device = '/dev/{}'.format(self.devices_model.getItem(self.devices_combobox.currentIndex())[0]) + FormatManyDialog.show_dialog(device, os.path.join(self.root, 'temp')) + self.refresh_image_list() + + def refresh_disk_drives(self): + self.devices_model.setModel(get_storage_devices(True, 'disk')) + self.refresh_image_list() + + def clean(self): + MountManager.get().dispose_all() + if os.path.exists(self.root): + shutil.rmtree(self.root) + + def on_device_change(self): + # Unmount and remove all old images + MountManager.get().dispose_all() + if os.path.exists(os.path.join(self.root, 'temp')): + for f in os.listdir(os.path.join(self.root, 'temp')): + if os.path.isfile(os.path.join(self.root, 'temp', f)): + os.remove(os.path.join(self.root, 'temp', f)) + else: + shutil.rmtree(os.path.join(self.root, 'temp', f)) + + os.makedirs(os.path.join(self.root, 'temp'), exist_ok=True) + self.refresh_image_list() + + def refresh_image_list(self): + try: + if self.devices_combobox.currentIndex() != -1: + device = '/dev/{}'.format(self.devices_model.getItem(self.devices_combobox.currentIndex())[0]) + m = floppyemu.list_images(device) + self.table_model.setModel(m) + except PermissionError: + device = '/dev/{}'.format(self.devices_model.getItem(self.devices_combobox.currentIndex())[0]) + QtWidgets.QMessageBox.critical( + self, 'Permission error', + 'You have no access to the device %s. Please restart program as user with proper rights' % device + ) + + def open_directory(self): + try: + if self.devices_combobox.currentIndex() != -1: + device = '/dev/{}'.format(self.devices_model.getItem(self.devices_combobox.currentIndex())[0]) + indexes = self.tableView.selectedIndexes() + for index in indexes: + image = os.path.join(self.root, 'temp/image{}.img'.format(index.row())) + if MountManager.get().is_mounted(image): + MountManager.get().unmount(image) + if os.path.exists(image): + os.remove(image) + floppyemu.read_image(device, image, index.row()) + os.makedirs(os.path.join(self.root, 'folders/{}'.format(index.row())), exist_ok=True) + os.chmod(os.path.join(self.root, 'folders/{}'.format(index.row())), + stat.S_IRWXU | stat.S_IRWXG | stat.S_IRWXO) + folder_path = MountManager.get().mount(image) + QtGui.QDesktopServices.openUrl( + QtCore.QUrl.fromLocalFile(os.path.join(self.root, folder_path))) + except: + import traceback + traceback.print_exc() + + def save_directory(self): + if self.devices_combobox.currentIndex() != -1: + device = '/dev/{}'.format(self.devices_model.getItem(self.devices_combobox.currentIndex())[0]) + indexes = self.tableView.selectedIndexes() + for index in indexes: + image = os.path.join(self.root, 'temp/image{}.img'.format(index.row())) + if MountManager.get().is_mounted(image): + MountManager.get().unmount(image) + floppyemu.write_image(device, image, index.row()) + shutil.rmtree(os.path.join(self.root, 'folders/{}'.format(index.row()))) + else: + QtWidgets.QMessageBox.critical( + self, 'Error', + 'Directory does not exist' + ) + + self.refresh_image_list() + + def write_image(self): + if self.devices_combobox.currentIndex() != -1: + device = '/dev/{}'.format(self.devices_model.getItem(self.devices_combobox.currentIndex())[0]) + indexes = self.tableView.selectedIndexes() + for index in indexes: + fd = QtWidgets.QFileDialog(self) + fd.setAcceptMode(QtWidgets.QFileDialog.AcceptOpen) + fd.setWindowTitle('Select file for image {}'.format(index.row())) + if fd.exec() == QtWidgets.QDialog.Accepted: + image = fd.selectedFiles()[0] + floppyemu.write_image(device, image, index.row()) + else: + return + + self.refresh_image_list() + + def read_image(self): + if self.devices_combobox.currentIndex() != -1: + device = '/dev/{}'.format(self.devices_model.getItem(self.devices_combobox.currentIndex())[0]) + indexes = self.tableView.selectedIndexes() + for index in indexes: + fd = QtWidgets.QFileDialog(self) + fd.setAcceptMode(QtWidgets.QFileDialog.AcceptSave) + fd.setWindowTitle('Select file for image {}'.format(index.row())) + if fd.exec() == QtWidgets.QDialog.Accepted: + image = fd.selectedFiles()[0] + floppyemu.read_image(device, image, index.row()) + else: + return + + def closeEvent(self, event: QtGui.QCloseEvent): + try: + self.clean() + except Exception as e: + event.ignore() + QtWidgets.QMessageBox.critical( + self, 'Error', + 'Unable to clean up. Some resources are in use: {}'.format(e) + ) \ No newline at end of file