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

This commit is contained in:
Daniel S. 2021-05-23 03:03:37 +02:00
parent 88737f29e4
commit 9f69934a32
1 changed files with 667 additions and 0 deletions

667
tools/scrapper_gui.py Normal file
View File

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