Added "Extract Selected" option to pkviewer.

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("<Return>", 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()

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