rework app, ability to load files

app.py

@@ -5,8 +5,8 @@ import time
 import mapper
 import extract
 from PyQt5 import uic
-from PyQt5.QtGui import QTextCursor
-from PyQt5.QtWidgets import QMessageBox, QMainWindow, QApplication, QFileDialog
+from PyQt5.QtGui import QTextCursor, QStandardItemModel, QStandardItem
+from PyQt5.QtWidgets import QMessageBox, QMainWindow, QApplication, QFileDialog, QHeaderView
 from PyQt5.QtCore import QObject, pyqtSignal, pyqtSlot, QThread, QMetaType
 
 QMetaType.type('QTextCursor')
@@ -30,31 +30,26 @@ class TextEditStream(QObject):
 		self.text_edit.insertPlainText(text)
 		self.text_edit.moveCursor(QTextCursor.End)
 
-class ExtractionWorker(QObject):
-	finished = pyqtSignal()
+class BackgroundWorker(QObject):
+	finished = pyqtSignal(dict)
 	progress = pyqtSignal(list)
 
-	def __init__(self, folders, _map, _format):
+	def __init__(self, action, folders, _map, _format):
 		super().__init__()
+		self.action = action
 		self.folders = folders
 		self.map = _map
 		self.format = _format
 
 	def run(self):
-		_map = None
-		if self.map:
-			_map = mapper.Mapper(os.path.join(os.getcwd(), f"maps/{self.map}"))
-			print("\n==========\n")
-		self.progress.emit(["total", 5])
-		extracter = extract.WwiseExtract(_map, self.format, *self.folders.values(), progress=self.progress.emit)
-		extracter.extract()
-		self.finished.emit()
+		fileStructure = extract.WwiseExtract(self.map, "mp3", *self.folders, progress="").load_folder()
+
+		self.finished.emit(fileStructure)
 
 class AnimeWwise(QMainWindow):
 	def __init__(self):
 		super(AnimeWwise, self).__init__()
 		uic.loadUi("gui.ui", self)
-		# self.setupUi(self)
 		self.maps = self.getMaps()
 		self.folders = {
 			"input": "",
@@ -67,6 +62,8 @@ class AnimeWwise(QMainWindow):
 		# utils
 		self.selectFolder = lambda: QFileDialog.getExistingDirectory(self, "Select Folder")
 
+		# self.updateTreeView()
+
 	def setFolder(self, elem, folder):
 		path = self.selectFolder()
 		self.folders[folder] = path
@@ -108,9 +105,10 @@ class AnimeWwise(QMainWindow):
 			_map = None
 
 		self.extractThread = QThread()
-		self.extractWorker = ExtractionWorker(self.folders, _map, self.outputFormat.currentText())
+		self.extractWorker = BackgroundWorker("load", self.folders, _map, self.outputFormat.currentText())
 		self.extractWorker.moveToThread(self.extractThread)
 		self.extractThread.started.connect(self.extractWorker.run)
+		self.extractWorker.finished.connect(self.handleFinished)
 		self.extractWorker.finished.connect(self.extractThread.quit)
 		self.extractWorker.finished.connect(self.extractWorker.deleteLater)
 		self.extractThread.finished.connect(self.extractThread.deleteLater)
@@ -118,6 +116,12 @@ class AnimeWwise(QMainWindow):
 		self.extractWorker.progress.connect(self.progressBarSlot)
 		self.extractThread.start()
 
+	@pyqtSlot(dict)
+	def handleFinished(self, data):
+		self.fileStructure = data
+		self.updateTreeView()
+		self.tabs.setCurrentIndex(1)
+
 	def _appendText(self, text):
 		cursor = self.console.textCursor()
 		cursor.movePosition(cursor.End)
@@ -125,6 +129,29 @@ class AnimeWwise(QMainWindow):
 		self.console.setTextCursor(cursor)
 		self.console.ensureCursorVisible()
 
+	def updateTreeView(self):
+		model = QStandardItemModel()
+		model.setHorizontalHeaderLabels(["Name", "Offset", "Size"])
+
+		root_item = model.invisibleRootItem()
+		self.addItems(root_item, self.fileStructure)
+
+		self.treeView.setModel(model)
+		self.treeView.expandAll()
+
+		self.treeView.header().setSectionResizeMode(0, QHeaderView.Stretch)
+		self.treeView.header().setSectionResizeMode(1, QHeaderView.ResizeToContents)
+		self.treeView.header().setSectionResizeMode(2, QHeaderView.ResizeToContents)
+
+	def addItems(self, parent, element):
+		for folder_name in sorted(element.get("folders", {}).keys()):
+			folder_content = element["folders"][folder_name]
+			folder_item = QStandardItem(folder_name)
+			parent.appendRow([folder_item, QStandardItem(""), QStandardItem("")])
+			self.addItems(folder_item, folder_content)
+
+		for file in sorted(element.get("files", [])):
+			parent.appendRow([QStandardItem(str(e)) for e in file])
 
 if __name__ == "__main__":
 	app = QApplication(sys.argv)

extract.py

@@ -1,11 +1,14 @@
 import os
+import io
 import sys
 import time
+from mapper import Mapper
 import shutil
 import filecmp
 import tempfile
 import wavescan
 import subprocess
+from filereader import FileReader
 
 cwd = os.getcwd()
 path = lambda path: os.path.join(cwd, path)
@@ -37,8 +40,6 @@ class WwiseExtract:
 
 		self.progress = progress
 
-	# TODO: add skip / select mapping option
-
 	def path(self, base, path):
 		base_path = self.paths[base]
 		if base == "temp":
@@ -302,3 +303,64 @@ class WwiseExtract:
 
 		print("-"*30)
 		print("Done extracting everything !")
+
+
+	### new content ###
+
+	def load_folder(self):
+		self.mapper = None
+		if self.map is not None:
+			self.mapper = Mapper(os.path.join(os.getcwd(), f"maps/{self.map}"))
+		self.file_structure = {"folders": {}, "files": []}
+
+		files = [f for f in os.listdir(self.paths["input"]) if f.endswith(".pck")]
+
+		for file in files:
+			self.load_file(os.path.join(self.paths["input"], file))
+
+		return self.file_structure
+
+	def load_file(self, _input):
+		with open(_input, "rb") as f:
+			data = f.read()
+			f.close()
+		self.get_wems(data, os.path.basename(_input))
+
+	def get_wems(self, data, filename):
+		reader = FileReader(io.BytesIO(data), "little")
+		files = wavescan.get_data(reader)
+		self.map_names(files, filename)
+	
+	def map_names(self, files, filename):
+		mapper = self.mapper
+		base = self.file_structure
+
+		for file in files:
+			if mapper is not None:
+				key = mapper.get_key(file[0].split(".")[0])
+			else:
+				key = None
+			if key is not None:
+				self.add_to_structure(f"{filename}\\{key[0]}.wem".split("\\"), [file[1], file[2]])
+			else:
+				temp = base["folders"]
+				if filename not in temp:
+					temp[filename] = {"folders": {}, "files": []}
+				temp = temp[filename]["folders"]
+				if "unmapped" not in temp:
+					temp["unmapped"] = {"folders": {}, "files": []}
+				temp["unmapped"]["files"].append(file)
+
+		self.file_structure = base
+
+	def add_to_structure(self, parts, meta):
+		current_level = self.file_structure
+		for part in parts[:-1]:
+			if "folders" not in current_level:
+				current_level["folders"] = {}
+			if part not in current_level["folders"]:
+				current_level["folders"][part] = {"folders": {}, "files": []}
+			current_level = current_level["folders"][part]
+		if "files" not in current_level:
+			current_level["files"] = []
+		current_level["files"].append([parts[-1], meta[0], meta[1]])

gui.ui

@@ -35,7 +35,7 @@
       <x>4</x>
       <y>-1</y>
       <width>1091</width>
-      <height>791</height>
+      <height>641</height>
      </rect>
     </property>
     <property name="currentIndex">
@@ -69,7 +69,7 @@
         <x>9</x>
         <y>9</y>
         <width>1071</width>
-        <height>741</height>
+        <height>601</height>
        </rect>
       </property>
       <layout class="QVBoxLayout" name="mainVLayout">
@@ -162,13 +162,6 @@
        <item>
         <layout class="QGridLayout" name="settingsGrid">
          <item row="0" column="0">
-          <widget class="QLabel" name="outputFormatLabel">
-           <property name="text">
-            <string>Output format</string>
-           </property>
-          </widget>
-         </item>
-         <item row="1" column="0">
           <widget class="QLabel" name="assetMapLabel">
            <property name="text">
             <string>Asset map (optional)</string>
@@ -176,95 +169,14 @@
           </widget>
          </item>
          <item row="0" column="1">
-          <widget class="QComboBox" name="outputFormat">
-           <property name="currentText">
-            <string/>
-           </property>
-          </widget>
-         </item>
-         <item row="1" column="1">
           <widget class="QComboBox" name="assetMap"/>
          </item>
-         <item row="0" column="2">
-          <widget class="QLabel" name="spacingLabel">
-           <property name="text">
-            <string/>
-           </property>
-          </widget>
-         </item>
         </layout>
        </item>
        <item>
         <widget class="QPushButton" name="startButton">
          <property name="text">
-          <string>Start</string>
-         </property>
-        </widget>
-       </item>
-       <item>
-        <widget class="Line" name="separatorB">
-         <property name="orientation">
-          <enum>Qt::Horizontal</enum>
-         </property>
-        </widget>
-       </item>
-       <item>
-        <layout class="QVBoxLayout" name="progressWrapperLayout">
-         <item>
-          <layout class="QHBoxLayout" name="progressLayout">
-           <item>
-            <widget class="QLabel" name="progressLabel">
-             <property name="text">
-              <string>Progress</string>
-             </property>
-            </widget>
-           </item>
-           <item>
-            <widget class="QProgressBar" name="progress">
-             <property name="value">
-              <number>0</number>
-             </property>
-            </widget>
-           </item>
-          </layout>
-         </item>
-         <item>
-          <layout class="QHBoxLayout" name="taskProgressLayout">
-           <item>
-            <widget class="QLabel" name="taskProgressLabel">
-             <property name="text">
-              <string>Task Progress</string>
-             </property>
-            </widget>
-           </item>
-           <item>
-            <widget class="QProgressBar" name="taskProgress">
-             <property name="value">
-              <number>0</number>
-             </property>
-            </widget>
-           </item>
-          </layout>
-         </item>
-        </layout>
-       </item>
-       <item>
-        <widget class="Line" name="line">
-         <property name="orientation">
-          <enum>Qt::Horizontal</enum>
-         </property>
-        </widget>
-       </item>
-       <item>
-        <widget class="QTextEdit" name="console">
-         <property name="maximumSize">
-          <size>
-           <width>16777215</width>
-           <height>220</height>
-          </size>
-         </property>
-         <property name="readOnly">
-          <bool>true</bool>
+          <string>Load file(s)</string>
          </property>
         </widget>
        </item>
@@ -275,9 +187,169 @@
      <attribute name="title">
       <string>Browse</string>
      </attribute>
+     <widget class="QTreeView" name="treeView">
+      <property name="geometry">
+       <rect>
+        <x>0</x>
+        <y>0</y>
+        <width>1081</width>
+        <height>611</height>
+       </rect>
+      </property>
+     </widget>
+    </widget>
+    <widget class="QWidget" name="tab_3">
+     <attribute name="title">
+      <string>Page</string>
+     </attribute>
+     <widget class="QWidget" name="layoutWidget">
+      <property name="geometry">
+       <rect>
+        <x>153</x>
+        <y>110</y>
+        <width>521</width>
+        <height>201</height>
+       </rect>
+      </property>
+      <layout class="QVBoxLayout" name="progressWrapperLayout_2">
+       <item>
+        <layout class="QHBoxLayout" name="progressLayout_2">
+         <item>
+          <widget class="QLabel" name="progressLabel_2">
+           <property name="text">
+            <string>Progress</string>
+           </property>
+          </widget>
+         </item>
+         <item>
+          <widget class="QProgressBar" name="progress_2">
+           <property name="value">
+            <number>0</number>
+           </property>
+          </widget>
+         </item>
+        </layout>
+       </item>
+       <item>
+        <layout class="QHBoxLayout" name="taskProgressLayout_2">
+         <item>
+          <widget class="QLabel" name="taskProgressLabel_2">
+           <property name="text">
+            <string>Task Progress</string>
+           </property>
+          </widget>
+         </item>
+         <item>
+          <widget class="QProgressBar" name="taskProgress_2">
+           <property name="value">
+            <number>0</number>
+           </property>
+          </widget>
+         </item>
+        </layout>
+       </item>
+      </layout>
+     </widget>
+     <widget class="QComboBox" name="outputFormat">
+      <property name="geometry">
+       <rect>
+        <x>529</x>
+        <y>400</y>
+        <width>261</width>
+        <height>22</height>
+       </rect>
+      </property>
+      <property name="currentText">
+       <string/>
+      </property>
+     </widget>
+     <widget class="QLabel" name="outputFormatLabel">
+      <property name="geometry">
+       <rect>
+        <x>260</x>
+        <y>400</y>
+        <width>262</width>
+        <height>22</height>
+       </rect>
+      </property>
+      <property name="text">
+       <string>Output format</string>
+      </property>
+     </widget>
     </widget>
    </widget>
+   <widget class="QTextEdit" name="console">
+    <property name="geometry">
+     <rect>
+      <x>10</x>
+      <y>640</y>
+      <width>1081</width>
+      <height>131</height>
+     </rect>
+    </property>
+    <property name="maximumSize">
+     <size>
+      <width>16777215</width>
+      <height>220</height>
+     </size>
+    </property>
+    <property name="acceptDrops">
+     <bool>false</bool>
+    </property>
+    <property name="toolTipDuration">
+     <number>0</number>
+    </property>
+    <property name="readOnly">
+     <bool>true</bool>
+    </property>
+   </widget>
   </widget>
+  <widget class="QMenuBar" name="menubar">
+   <property name="geometry">
+    <rect>
+     <x>0</x>
+     <y>0</y>
+     <width>1100</width>
+     <height>26</height>
+    </rect>
+   </property>
+   <widget class="QMenu" name="menuFile">
+    <property name="title">
+     <string>File</string>
+    </property>
+    <addaction name="actionnot_working_here_yet"/>
+    <addaction name="actionclear_table"/>
+   </widget>
+   <widget class="QMenu" name="menuExtract">
+    <property name="title">
+     <string>Extract</string>
+    </property>
+    <addaction name="actionall"/>
+    <addaction name="actionselected"/>
+   </widget>
+   <addaction name="menuFile"/>
+   <addaction name="menuExtract"/>
+  </widget>
+  <action name="actionnot_working_here_yet">
+   <property name="text">
+    <string>not working here yet</string>
+   </property>
+  </action>
+  <action name="actionall">
+   <property name="text">
+    <string>all</string>
+   </property>
+  </action>
+  <action name="actionselected">
+   <property name="text">
+    <string>selected</string>
+   </property>
+  </action>
+  <action name="actionclear_table">
+   <property name="text">
+    <string>clear table</string>
+   </property>
+  </action>
  </widget>
  <tabstops>
   <tabstop>inputPath</tabstop>
@@ -286,8 +358,6 @@
   <tabstop>changeAltInput</tabstop>
   <tabstop>outputPath</tabstop>
   <tabstop>changeOutput</tabstop>
-  <tabstop>outputFormat</tabstop>
-  <tabstop>assetMap</tabstop>
   <tabstop>startButton</tabstop>
   <tabstop>tabs</tabstop>
  </tabstops>

mapper.py

@@ -108,7 +108,10 @@ class Mapper:
 			
 			name_length = reader.ReadUInt8()
 			prefix = reader.ReadBytes(1)
-			prefix = prefixes[prefix]
+			if prefix != b"\x00":
+				prefix = prefixes[prefix]
+			else:
+				prefix = ""
 			name = raw(name_length)
 			
 			name = f"{prefix}{name}"

wavescan.py

@@ -6,18 +6,20 @@ import os
 
 reader = None
 bank_version = 0
+wwise_data = []
 
 
-def extract(input_file, output_folder):
+def get_data(_reader):
+	global wwise_data
 	global bank_version
 	global reader
 
-	file = open(input_file, "rb")
-	reader = FileReader(file, "little") # defaults to little endian
+	wwise_data = []
+	reader = _reader
 
 	# check file
 	if reader.ReadBytes(4) != b"AKPK":
-		file.close()
+		# file.close()
 		raise Exception("not a valid audio file")
 
 	# check endianness
@@ -29,7 +31,6 @@ def extract(input_file, output_folder):
 	elif endian_check == 0x1000000:
 		endianness = 1 # big
 	else:
-		file.close()
 		raise Exception("couldn't detect endianness")
 
 	# retrieve sectors in header
@@ -52,7 +53,6 @@ def extract(input_file, output_folder):
 	try:
 		lang_array = get_langs(languages_sector_size)
 	except Exception as e:
-		file.close()
 		raise Exception(f"failed to read languages, {e}, {traceback.format_exc()}")
 
 	# extract each sector
@@ -60,18 +60,16 @@ def extract(input_file, output_folder):
 	try:
 		for sector in sectors:
 			curr_sector = sector
-			extract_sector(*sector[1:], endianness, lang_array, bank_version, output_folder)
+			extract_sector(*sector[1:], endianness, lang_array, bank_version)
 
 			if sector[0] and bank_version == 0:
 				if externals_sector_size == 0:
 					print("can't detect bank version")
 				bank_version = 62
 	except Exception as e:
-		file.close()
 		raise Exception(f"failed to extract sector {curr_sector}, {e}, {traceback.format_exc()}")
 
-	# close
-	file.close()
+	return wwise_data
 
 def get_langs(langs_sector_size):
 	string_offset = reader.GetBufferPos()
@@ -125,7 +123,9 @@ def detect_bank_version(offset):
 
 	reader.SetBufferPos(current)
 
-def extract_sector(section_size, is_sounds, is_externals, ext, endianness, lang_array, bank_version, output_folder, filter_bnk_only=0, filter_wem_only=0, include_name=False):
+def extract_sector(section_size, is_sounds, is_externals, ext, endianness, lang_array, bank_version, filter_bnk_only=0, filter_wem_only=0, include_name=False):
+	global wwise_data
+
 	# check sector validity
 	if section_size == 0:
 		return
@@ -210,23 +210,4 @@ def extract_sector(section_size, is_sounds, is_externals, ext, endianness, lang_
 			continue
 
 		# file infos
-		# print(f"NAME - {name} | OFFSET - {offset} | SIZE - {size}")
-
-		# save file into disk
-		current = reader.GetBufferPos()
-		reader.SetBufferPos(offset)
-		file_data = reader.ReadBytes(size)
-
-		if include_name:
-			file_path = os.path.join(output_folder, os.path.dirname(name))
-		else:
-			file_path = output_folder
-		name = os.path.basename(name)
-
-		os.makedirs(file_path, exist_ok=True)
-
-		with open(os.path.join(file_path, name), "wb+") as f:
-			f.write(file_data)
-			f.close()
-
-		reader.SetBufferPos(current)
+		wwise_data.append([os.path.basename(name), offset, size])