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) )