Add GUI version of scrapper.py, needs PyQt5 and construct

tools/scrapper_gui.py

@@ -0,0 +1,667 @@
+import sys
+from PyQt5.QtWidgets import *
+from PyQt5.QtCore import *
+from PyQt5.QtGui import *
+from PyQt5 import QtGui, QtCore, QtWidgets
+from threading import Thread
+from queue import Queue
+from construct import *
+from glob import glob
+import os
+
+ScrapFile = Struct(
+    'path'/PrefixedArray(Int32ul, Byte),
+    'size'/Int32ul,
+    'offset'/Int32ul,
+)
+
+PackedFile = Struct(
+    Const(b'BFPK'),
+    Const(b'\0\0\0\0'),
+    'files'/PrefixedArray(Int32ul, ScrapFile),
+)
+
+
+def get_path(file):
+    return str(bytes(file.path), "utf-8",
+               "backslashreplace")
+
+
+def fsize(n):
+    idx = 0
+    l = ["", "K", "M", "G", "T", "P", "E"]
+    while n > 1024:
+        idx += 1
+        n /= 1024
+    return "{:.02f} {}B".format(n, l[idx])
+
+
+class Tree(dict):
+    def __init__(self, *args, **kwargs):
+        self._size = 0
+        self.data = None
+        self.path = None
+        super(type(self), self).__init__(*args, **kwargs)
+
+    def __missing__(self, k):
+        self[k] = type(self)()
+        return self[k]
+
+    @property
+    def size(self):
+        if self.data:
+            ret = self.data.size
+        else:
+            ret = 0
+        for v in self.values():
+            ret += v.size
+        return ret
+
+    @size.setter
+    def size(self, s):
+        self._size = s
+
+
+def load_data(path):
+    ret = Tree()
+    fn = os.path.split(path)[-1]
+    ret[fn] = Tree()
+    with open(path, "rb") as fh:
+        data = PackedFile.parse_stream(fh)
+        for entry in data.files:
+            path = get_path(entry).split("/")
+            root = ret[fn]
+            for elem in path:
+                root = root[elem]
+            root.data = entry
+    return ret
+
+
+class Patcher(object):
+    def __init__(self, ed):
+        self.editor = ed
+
+    def debug(self, *args):
+        "Enable Scengraph Debugging Console"
+        print("DBG")
+
+    def test(self, *args):
+        "TEST!"
+        print(self.editor.selected())
+        pass
+
+
+class PackerThread(QThread):
+    signal = pyqtSignal('PyQt_PyObject')
+
+    def __init__(self, node, path):
+        super().__init__()
+
+        self.node = node
+        self.path = path
+
+    def func(self, *args, **kwargs):
+        self.signal.emit(*args, **kwargs)
+
+    def run(self):
+        try:
+            self.__run()
+        except:
+            (type, value, traceback) = sys.exc_info()
+            print(type, value, traceback)
+            sys.excepthook(type, value, traceback)
+
+    def __run(self):
+        node = self.node
+        path = self.path
+        PAK = node.pak
+        file_list = []
+
+        self.func(("wtl", "Preprocessing"))
+        self.func(("range", 0,  0))
+        for item in walk(node):
+            if item.struct:
+                file_list.append(item)
+                self.func(("upd", get_loc(item)[1], 1))
+
+        header = PackedFile.build(dict(files=[f.struct for f in file_list]))
+        total_size = len(header)
+        self.func(("wtl", "Updating offsets"))
+        self.func(("range", 0,  len(file_list)))
+        for idx, file in enumerate(file_list):
+            file.struct.offset = total_size
+            total_size += file.struct.size
+            self.func(("upd", get_loc(file)[1], idx))
+
+        header = PackedFile.build(dict(files=[f.struct for f in file_list]))
+        cnt = 0
+        if path:
+            self.func(("wtl", "Writing output"))
+            self.func(("range", 0, total_size))
+            with open(path, "wb") as of:
+                of.write(header)
+                written = len(header)
+                self.func(("upd", "Header", written))
+                for file in file_list:
+                    loc = get_loc(file)
+                    data = get_data(file)[0]
+                    written += len(data)
+                    self.func(("upd", get_loc(file)[1], written))
+                    of.write(data)
+            print("XXXXXXXXXXXXXXXXXXXXX")
+        self.func(("close",))
+        return
+
+
+class PackerWindow(QProgressDialog):
+
+    def __init__(self, node, path, *args, **kwargs):
+        self.node = node
+        self.path = path
+        self.mtx = QMutex()
+        super().__init__()
+        self.initUI()
+        self.setAutoClose(True)
+        self.setWindowTitle("Packing {}".format(node.pak))
+        self.setCancelButtonText("Stop")
+        self.run()
+
+    def initUI(self):
+        self.show()
+        self.ensurePolished()
+        self.setFixedSize(300, self.height())
+
+    def closeEvent(self, event):
+        print("BYE")
+        if self.th:
+            self.th.wait()
+        event.accept()
+
+    def sig(self, val):
+        lock = QMutexLocker(self.mtx)
+        #print("SIG:", val)
+        typ, *args = val
+        if typ == "wtl":
+            self.setWindowTitle(*args)
+        elif typ == "close":
+            self.th.wait()
+            self.th = None
+            self.close()
+        elif typ == "range":
+            self.setRange(*args)
+        elif typ == "upd":
+            self.setLabelText(args[0])
+            self.setValue(args[1])
+        else:
+            raise ValueError("Unknown signal: {}".format(val))
+
+    def run(self):
+        self.queue = Queue()
+        self.th = PackerThread(self.node, self.path)
+        self.th.signal.connect(self.sig)
+        self.th.start()
+
+
+class MainWindow(QMainWindow):
+    def __init__(self, parent=None):
+        super(MainWindow, self).__init__(parent)
+        self.setWindowTitle("SMT - Scrap Mod Tool")
+        self.resize(800, 600)
+        self.center()
+
+        menu_close_all = QAction("Close &All", self)
+        menu_close_all.setShortcut("Ctrl+Shift+C")
+        menu_close_all.triggered.connect(self.menu_close_all)
+
+        menu_create = QAction("&Create new .packed file", self)
+        menu_create.setShortcut("Ctrl+N")
+        menu_create.triggered.connect(self.menu_new)
+
+        menu_load = QAction("&Open .packed file", self)
+        menu_load.setShortcut("Ctrl+O")
+        menu_load.triggered.connect(self.menu_load)
+
+        menu_exit = QAction("&Exit", self)
+        menu_exit.setShortcut("Ctrl+Q")
+        menu_exit.triggered.connect(qApp.quit)
+
+        menu = self.menuBar()
+        fileMenu = menu.addMenu('&File')
+        fileMenu.addAction(menu_load)
+        fileMenu.addAction(menu_create)
+        fileMenu.addSeparator()
+        fileMenu.addAction(menu_close_all)
+        fileMenu.addSeparator()
+        fileMenu.addAction(menu_exit)
+
+        self.editor = PackedEditor(self)
+        self.setCentralWidget(self.editor)
+        self.show()
+        self.menu_load(sys.argv[1:])
+
+    def menu_new(self):
+        data = PackedFile.build({'files': []})
+        loc = QFileDialog.getSaveFileName(
+            self, 'Save file', os.getcwd(), "Scrapland Data Files (*.packed)")[0]
+        with open(loc, "wb") as of:
+            of.write(data)
+        self.editor.clear()
+        self.editor.load(loc)
+
+    def menu_load(self, paths=None):
+        if not paths:
+            paths = QFileDialog.getOpenFileNames(
+                self, 'Open file', os.getcwd(), "Scrapland Data Files (*.packed)")[0]
+        for path in paths:
+            path = os.path.abspath(path)
+            self.editor.load(path)
+
+    def menu_close_all(self):
+        self.editor.clear()
+
+    def center(self):
+        qr = self.frameGeometry()
+        cp = QDesktopWidget().availableGeometry().center()
+        qr.moveCenter(cp)
+        self.move(qr.topLeft())
+
+
+def get_loc(item):
+    node = item
+    path = []
+    while node:
+        path.insert(0, node.text(0))
+        node = node.parent()
+    if path:
+        file = path.pop(0)
+    else:
+        file = None
+    path = "/".join(path)
+    return (file, path)
+
+
+def walk(node):
+    yield node
+    for child_idx in range(node.childCount()):
+        yield from walk(node.child(child_idx))
+
+
+def read_file(path, size=None, offset=None):
+    data = None
+    with open(path, "rb") as fh:
+        if offset:
+            fh.seek(offset)
+        data = fh.read(size)
+    return data
+
+
+def get_data(node, size=None):
+    bin_data = None
+    more = None
+    if node.struct:
+        if size:
+            if node.struct.size > size:
+                more = True
+            size = min(node.struct.size, size)
+        else:
+            size = node.struct.size
+        if node.path:
+            bin_data = read_file(node.path, size)
+        else:
+            bin_data = read_file(node.pak, size, node.struct.offset)
+    return bin_data, more
+
+
+def hexdump(data, addr=0):
+    res = []
+    bin_data = list(data)
+    while bin_data:
+        chunk = bin_data[:16]
+        bin_data = bin_data[16:]
+        res.append("0x{:04x} | {}".format(
+            addr, " ".join("{:02x}".format(b) for b in chunk)))
+        addr += len(chunk)
+    return "\n".join(res)
+
+
+def build_tree(node, url):
+    pak = get_loc(node)[0]
+    subtree = Tree()
+    path = url.toLocalFile().replace("/", "\\")
+    if os.path.isfile(path):
+        full_path = os.path.abspath(path)
+        path, file = os.path.split(full_path)
+        loc = os.path.join(get_loc(node)[1], file)
+        root_node = subtree[pak]
+        for elem in loc.split("/"):
+            root_node = root_node[elem]
+        root_node.data = Container(path=bytes(loc, "utf-8"), size=os.path.getsize(
+            full_path), offset=0)
+        root_node.size = root_node.data.size
+        root_node.path = full_path
+    else:
+        for root, _, files in os.walk(path):
+            root = root.replace("/", os.sep)
+            for file in files:
+                full_path = os.path.join(root, file)
+                loc = full_path.replace(path.replace(
+                    "/", os.sep), "").strip(os.sep)
+                loc = os.path.join(os.path.split(
+                    path)[-1], loc).replace(os.sep, "/")
+                root_node = subtree[pak]
+                for elem in loc.split("/"):
+                    root_node = root_node[elem]
+                root_node.data = Container(path=bytes(loc, "utf-8"), size=os.path.getsize(
+                    full_path), offset=0)
+                root_node.size = root_node.data.size
+                root_node.path = full_path
+    return subtree[pak]
+
+
+class DataTree(QTreeWidget):
+    def dragEnterEvent(self, event):
+        m = event.mimeData()
+        if m.hasUrls():
+            for url in m.urls():
+                if url.isLocalFile():
+                    event.accept()
+                    return
+        event.accept()
+
+    def dragMoveEvent(self, event):
+        if event.mimeData().hasUrls:
+            event.setDropAction(QtCore.Qt.CopyAction)
+        else:
+            event.setDropAction(QtCore.Qt.MoveAction)
+        event.accept()
+
+    def dropEvent(self, event):
+        editor = self.window().editor
+        target = self.itemAt(event.pos())
+        if target.struct:
+            return
+        print("Drop: ", get_loc(target))
+        if event.source():
+            selected = event.source().selectedItems()
+            for item in selected:
+                labels = [item.text(i) for i in range(item.columnCount())]
+                node = QTreeWidgetItem(target, labels)
+                node.struct = item.struct
+                node.struct.offset = 0
+                node.path = item.path
+                node.pak = item.pak
+        else:
+            m = event.mimeData()
+            if m.hasUrls():
+                for url in m.urls():
+                    tree = build_tree(target, url)
+                    editor.make_subtree(target, tree, target.pak)
+        editor.update_tree()
+        event.accept()
+
+
+class DataLoader(QThread):
+    signal = pyqtSignal('PyQt_PyObject')
+
+    def __init__(self, path, subtree=None):
+        super().__init__()
+        self.path = path
+
+    def run(self):
+        self.signal.emit((self.path, load_data(self.path)))
+
+
+class DataExtractor(QThread):
+    signal = pyqtSignal('PyQt_PyObject')
+
+    def __init__(self, nodes, dest):
+        super().__init__()
+        self.dest = dest
+        self.nodes = nodes
+
+    def run(self):
+        for node in self.nodes:
+            for ch in walk(node):
+                if ch.struct:
+                    self.__extract(ch, self.dest)
+        self.signal.emit("Done!")
+        return
+
+    def __extract(self, node, dest):
+        path = get_path(node.struct)
+        self.signal.emit("Extracting {}".format(path))
+        folder, file = os.path.split(path)
+        folder = os.path.join(dest, folder)
+        folder = folder.replace("/", os.sep).replace("\\", os.sep)
+        file = os.path.join(folder, file)
+        os.makedirs(folder, exist_ok=True)
+        with open(node.pak, "rb") as pak:
+            with open(file, "wb") as of:
+                pak.seek(node.struct.offset)
+                of.write(pak.read(node.struct.size))
+
+
+class PackedEditor(QWidget):
+    def __init__(self, parent=None):
+        super(PackedEditor, self).__init__(parent)
+        self.initUI()
+        self.data = {}
+        self.threads = []
+        self.log_mutex = QMutex()
+
+    def initUI(self):
+        self.initLayout()
+        self.show()
+
+    def initLayout(self):
+        grid = QGridLayout()
+        hgrid = QSplitter(Qt.Horizontal)
+        vgrid = QSplitter(Qt.Vertical)
+        patcher_grid = QGridLayout()
+
+        self.info_tab_hex = QTextEdit()
+        self.info_tab_text = QTextEdit()
+        self.info_tab_info = QTextEdit()
+        self.info_tabs = QTabWidget()
+
+        self.info_tab_hex.setReadOnly(True)
+        self.info_tab_text.setReadOnly(True)
+        self.info_tab_info.setReadOnly(True)
+
+        self.info_tabs.addTab(self.info_tab_info, "Info")
+        self.info_tabs.addTab(self.info_tab_text, "Text")
+        self.info_tabs.addTab(self.info_tab_hex, "Hexdump")
+
+        self.util_tabs = QTabWidget()
+
+        self.util_tab_log = QTextEdit()
+        self.util_tab_log.setReadOnly(True)
+
+        self.util_tab_patch = QWidget()
+
+        self.util_tab_patch.setLayout(patcher_grid)
+        self.util_tabs.addTab(self.util_tab_log, "Log")
+        self.util_tabs.addTab(self.util_tab_patch, "Patcher")
+        self.patcher = Patcher(self)
+        i = 0
+        patcher_grid_w = 0
+        for func in dir(self.patcher):
+            if func.startswith("__") or not hasattr(getattr(self.patcher, func), "__call__"):
+                continue
+            patcher_grid_w += 1
+        patcher_grid_w = int(patcher_grid_w**0.5)
+        for func in dir(self.patcher):
+            if func.startswith("__") or not hasattr(getattr(self.patcher, func), "__call__"):
+                continue
+            func = getattr(self.patcher, func)
+            x, y = divmod(i, patcher_grid_w)
+            button = QPushButton(func.__name__.replace("_", " ").title())
+            button.setToolTip(func.__doc__)
+            button.clicked.connect(func)
+            patcher_grid.addWidget(button, x, y)
+            i += 1
+
+        self.tree = DataTree()
+        self.tree.setHeaderLabels(["Path", "Size", "Offset"])
+        self.tree.setSelectionMode(self.tree.ExtendedSelection)
+        self.tree.currentItemChanged.connect(self.tree_changed)
+        self.tree.setDragDropMode(self.tree.DragDrop | self.tree.InternalMove)
+        self.tree.setDragEnabled(True)
+        self.tree.setAcceptDrops(True)
+        self.tree.setDropIndicatorShown(True)
+
+        hgrid.addWidget(self.tree)
+        hgrid.addWidget(self.info_tabs)
+        vgrid.addWidget(hgrid)
+        vgrid.addWidget(self.util_tabs)
+        grid.addWidget(vgrid)
+        self.setLayout(grid)
+
+    def info(self, msg):
+        lock = QMutexLocker(self.log_mutex)
+        self.util_tab_log.insertPlainText(msg + "\n")
+        self.util_tab_log.moveCursor(QTextCursor.End)
+
+    def load(self, path):
+        if path in self.data:
+            return
+        self.data[path] = DataLoader(path)
+        self.data[path].signal.connect(self.done_loading)
+        self.data[path].start()
+        self.info("Loading {}".format(path))
+
+    def done_loading(self, result):
+        path, data = result
+        self.data[path] = data
+        self.make_subtree(self.tree, data, path)
+
+    def clear(self):
+        self.data.clear()
+        self.tree.clear()
+        self.info_tab_hex.setText("")
+        self.info_tab_text.setText("")
+        self.info_tab_info.setText("")
+        self.util_tab_log.setText("")
+
+    def make_tree(self):
+        for pak in self.data:
+            if isinstance(self.data[pak], DataLoader):
+                continue
+            self.make_subtree(self.tree, self.data[pak], pak)
+
+    def update_tree(self, node=None):
+        total_size = 0
+        if node is None:
+            node = self.tree.invisibleRootItem()
+        else:
+            if node.struct:
+                total_size += node.struct.size
+        for child_idx in range(node.childCount()):
+            total_size += self.update_tree(node.child(child_idx))
+        node.setText(1, fsize(total_size))
+        return total_size
+
+    @classmethod
+    def make_subtree(cls, tree, data, pak):
+        total_size = 0
+        for name, children in sorted(data.items()):
+            total_size += children.size
+            if children.data:
+                labels = [name, fsize(children.size),
+                          hex(children.data.offset)]
+            else:
+                labels = [name, fsize(children.size), ""]
+            node = QTreeWidgetItem(tree, labels)
+            node.struct = children.data
+            node.path = children.path
+            node.pak = pak
+            cls.make_subtree(node, children, pak)
+
+    def contextMenuEvent(self, event):
+        cmenu = QMenu(self)
+        actions = {}
+        actions[cmenu.addAction("Extract")] = self.extract_handler
+        actions[cmenu.addAction("Delete")] = self.del_handler
+        if all(node.parent() == None for node in self.selected()):
+            actions[cmenu.addAction("Save")] = self.save_handler
+        try:
+            act_fn = actions.get(
+                cmenu.exec_(self.mapToGlobal(event.pos())))
+        except TypeError:
+            return
+        if not act_fn:
+            return
+        th = act_fn(self.selected())
+        if th:
+            th.signal.connect(self.info)
+            th.start()
+            self.threads.append(th)
+
+    def extract_handler(self, nodes):
+        dest = QFileDialog.getExistingDirectory(
+            self, 'Open file', os.getcwd())
+        if not nodes:
+            return
+        if not dest:
+            return
+        return DataExtractor(nodes, dest)
+
+    def make_packed(self, node, path=None):
+        a = PackerWindow(node, path)
+        a.exec_()
+
+    def save_handler(self, nodes):
+        for node in nodes:
+            path = QFileDialog.getSaveFileName(
+                self, 'Save file', node.pak, "Scrapland Data Files (*.packed)")[0]
+            self.make_packed(node, path)
+
+    def del_handler(self, nodes):
+        if not nodes:
+            return
+        root = self.tree.invisibleRootItem()
+        for node in nodes:
+            (node.parent() or root).removeChild(node)
+        self.update_tree()
+
+    def get_info(self, node):
+        info = []
+        if node.struct:
+            info.append(("Size", node.struct.size))
+            info.append(("Offset", hex(node.struct.offset)))
+            info.append(("Path", get_path(node.struct)))
+        ret = ""
+        for k, v in info:
+            ret += "{}:\t{}\n".format(k, v)
+        return ret.strip()
+
+    def tree_changed(self, new, old):
+        if not new:
+            return
+        self.info_tab_text.setText("")
+        self.info_tab_hex.setText("")
+        self.info_tab_info.setText(self.get_info(new))
+        data, more = get_data(new, 1024)
+        if data:
+            more_text = "\n<...>" if more else ""
+            self.info_tab_hex.setText(hexdump(data)+more_text)
+            try:
+                data = str(data, "cp1252")
+            except UnicodeDecodeError:
+                data = repr(data)
+            self.info_tab_text.setText(data+more_text)
+            self.info_tab_info.setText(self.get_info(new))
+
+    def selected(self):
+        return self.tree.selectedItems()
+
+
+def main():
+    app = QApplication(sys.argv)
+    win = MainWindow()
+    win.show()
+    sys.exit(app.exec_())
+
+
+if __name__ == '__main__':
+    main()