From 3bbfa326f78c51a5c860c8e2e0afe51cbb78063b Mon Sep 17 00:00:00 2001 From: lcdr Date: Sat, 21 Apr 2018 14:48:06 +0200 Subject: [PATCH] Added "Extract Selected" option to pkviewer. --- pkviewer.pyw | 94 ++++++++++++++++++++++++++++++++++++++++------------ viewer.py | 9 +++-- 2 files changed, 79 insertions(+), 24 deletions(-) diff --git a/pkviewer.pyw b/pkviewer.pyw index 0331fd2..d1904fa 100644 --- a/pkviewer.pyw +++ b/pkviewer.pyw @@ -1,12 +1,13 @@ import hashlib import os +import struct import subprocess import sys import tempfile import tkinter.filedialog as filedialog from tkinter import BOTH, END, Menu, RIGHT, X, Y -from tkinter.ttk import Entry, Scrollbar, Treeview +from tkinter.ttk import Entry, Progressbar, Scrollbar, Treeview import decompress_sd0 import viewer @@ -33,6 +34,7 @@ class PKViewer(viewer.Viewer): menubar = Menu() menubar.add_command(label="Open", command=self.askopen) + menubar.add_command(label="Extract Selected", command=self.extract_selected) self.master.config(menu=menubar) columns = "Size", @@ -44,25 +46,37 @@ class PKViewer(viewer.Viewer): self.tree.bind("", self.extract_and_show_selected) def askopen(self): - dir = filedialog.askdirectory() + dir = filedialog.askdirectory(title="Select LU root folder (containing /client/, /versions/)") if dir: self.load(dir) def load(self, dir): self.filenames = {} self.records = {} - self.tree.set_children("") + self.reattach_all() + self.tree.delete(*self.tree.get_children()) + self.progress = Progressbar() + self.progress.pack(fill=X) + for filename in ("trunk.txt", "hotfix.txt"): self.load_filehashes(os.path.join(dir, "versions", filename)) print("Loaded hashes") + pks = [] for dir, _, filenames in os.walk(os.path.join(dir, "client/res/pack")): for filename in filenames: if filename.endswith(".pk"): - self.load_pk(os.path.join(dir, filename)) + pks.append(os.path.join(dir, filename)) + + self.progress.configure(maximum=len(pks)+1) + for pk in pks: + self.load_pk(pk) + self.progress.step() + self.update() print("Loaded records") for filename in sorted(self.records.keys()): self.create_tree(filename, self.records[filename][3]) + self.progress.pack_forget() def create_tree(self, path, values=()): dir, filename = os.path.split(path) @@ -96,8 +110,8 @@ class PKViewer(viewer.Viewer): with open(path, "rb") as file: assert file.read(7) == b"ndpk\x01\xff\x00" file.seek(-8, 2) - number_of_records_address = c_uint.unpack(file.read(c_uint.size))[0] - unknown = c_uint.unpack(file.read(c_uint.size))[0] + number_of_records_address = struct.unpack("I", file.read(4))[0] + unknown = struct.unpack("I", file.read(4))[0] if unknown != 0: print(unknown, path) file.seek(number_of_records_address) @@ -123,26 +137,30 @@ class PKViewer(viewer.Viewer): self.filenames[original_md5] = "unlisted/"+original_md5 self.records[self.filenames[original_md5]] = path, data_position, is_compressed, original_size, original_md5, compressed_size, compressed_md5 + def extract_path(self, path): + pk_path, data_position, is_compressed, original_size, original_md5, compressed_size, compressed_md5 = self.records[path] + + with open(pk_path, "rb") as file: + file.seek(data_position) + if is_compressed: + data = file.read(compressed_size) + else: + data = file.read(original_size) + assert file.read(5) == b"\xff\x00\x00\xdd\x00" + + if is_compressed: + assert hashlib.md5(data).hexdigest() == compressed_md5 + data = decompress_sd0.decompress(data) + + assert hashlib.md5(data).hexdigest() == original_md5 + return data + def extract_and_show_selected(self, _): for path in self.tree.selection(): if self.tree.get_children(path): continue # is directory - pk_path, data_position, is_compressed, original_size, original_md5, compressed_size, compressed_md5 = self.records[path] - - with open(pk_path, "rb") as file: - file.seek(data_position) - if is_compressed: - data = file.read(compressed_size) - else: - data = file.read(original_size) - assert file.read(5) == b"\xff\x00\x00\xdd\x00" - - if is_compressed: - assert hashlib.md5(data).hexdigest() == compressed_md5 - data = decompress_sd0.decompress(data) - - assert hashlib.md5(data).hexdigest() == original_md5 + data = self.extract_path(path) tempfile_path = os.path.join(tempfile.gettempdir(), os.path.basename(path)) with open(tempfile_path, "wb") as file: file.write(data) @@ -153,6 +171,40 @@ class PKViewer(viewer.Viewer): opener = "open" if sys.platform == "darwin" else "xdg-open" subprocess.call([opener, tempfile_path]) + def extract_selected(self): + outdir = filedialog.askdirectory(title="Select output directory") + if not outdir: + return + paths = set() + for path in self.tree.selection(): + paths.update(self.get_leaves(path)) + + self.progress = Progressbar(maximum=len(paths)+1) + self.progress.pack(fill=X) + for path in paths: + self.save_path(outdir, path) + self.progress.pack_forget() + + def save_path(self, outdir, path): + data = self.extract_path(path) + dir, filename = os.path.split(path) + out = os.path.join(outdir, dir) + os.makedirs(out, exist_ok=True) + with open(os.path.join(out, filename), "wb") as file: + file.write(data) + + self.progress.step() + self.update() + + def get_leaves(self, path): + output = set() + if self.tree.get_children(path): + for child in self.tree.get_children(path): + output.update(self.get_leaves(child)) + elif path in self.records: + output.add(path) + return output + if __name__ == "__main__": app = PKViewer() app.mainloop() diff --git a/viewer.py b/viewer.py index 2b0c1a2..b11f231 100644 --- a/viewer.py +++ b/viewer.py @@ -51,16 +51,19 @@ class Viewer(Frame): tags = list(self.tree.item(item, "tags")) tags.remove("match") self.tree.item(item, tags=tags) + self.reattach_all() + if query: + self.filter_items(query) + + def reattach_all(self): for parent, detached_children in self.detached_items.items(): for index, item in detached_children: self.tree.reattach(item, parent, index) self.detached_items.clear() - if query: - self.filter_items(query) def filter_items(self, query, parent=""): all_children = self.tree.get_children(parent) - detached_children = [item for item in all_children if not any(query in i.lower() for i in self.tree.item(item, "values"))] # first, find all children that don't match + detached_children = [item for item in all_children if not any(query in i.lower() for i in self.tree.item(item, "values")) and not query in self.tree.item(item, "text").lower()] # first, find all children that don't match for item in all_children: if item not in detached_children: tags = list(self.tree.item(item, "tags"))