1 files changed, 189 insertions, 0 deletions
diff --git a/src/wii/u8archive.py b/src/wii/u8archive.py
new file mode 100644
index 0000000..8295d52
--- /dev/null
+++ b/src/wii/u8archive.py
@@ -0,0 +1,189 @@
+from common import WiiStringTableBuilder, alignUp, alignDown
+from filesystem import *
+import struct
+import os
+import cStringIO
+
+class WiiArchiveU8:
+	class ReadInfo:
+		# startPos, stringTable, currentNode
+		pass
+	
+	class WriteInfo:
+		# startPos, stringTableBuilder, currentRecursionLevel,
+		# currentNode, nextDataOffset
+		pass
+
+	def __init__(self, handle=None):
+		self.root = WiiDirectory()
+
+		if handle:
+			if not hasattr(handle, 'read'):
+				handle = cStringIO.StringIO(handle)
+
+			info = WiiArchiveU8.ReadInfo()
+			info.startPos = handle.tell()
+
+			magic = handle.read(4)
+			if magic != "U\xAA8\x2D":
+				print("WiiArchiveU8: tried to load an archive without the U8 magic")
+
+			fstStart, fstSize, dataOffset = struct.unpack('>III', handle.read(12))
+
+			# read the FST
+			handle.seek(info.startPos + fstStart)
+
+			# read the root node
+			handle.seek(8, os.SEEK_CUR) # ignore type, nameOffset and dataOffset
+			rootNodeLastChild = struct.unpack('>I', handle.read(4))[0]
+
+			# but before we do this, read the string table
+			savePos = handle.tell()
+			handle.seek(savePos + ((rootNodeLastChild - 1) * 12))
+
+			stringTableLength = fstSize - (rootNodeLastChild * 12)
+			info.stringTable = handle.read(stringTableLength)
+			
+			# now read the root node
+			handle.seek(savePos)
+
+			info.currentNode = 1
+			self._readDir(handle, self.root, rootNodeLastChild, info)
+	
+	def _readDir(self, handle, directory, lastChild, info):
+		# read every node in this directory
+		while info.currentNode < lastChild:
+			info.currentNode += 1
+
+			value, dataOffset, size = struct.unpack('>III', handle.read(12))
+
+			nameOffset = value & 0xFFFFFF
+			objType = value >> 24
+
+			if objType == 0:
+				newObj = WiiFile()
+			elif objType == 1:
+				newObj = WiiDirectory()
+			else:
+				raise "oh crap! unknown FS obj type %d" % objType
+
+			nameEnd = info.stringTable.find("\0", nameOffset)
+			newObj.name = info.stringTable[nameOffset:nameEnd]
+
+			if newObj.isFile():
+				savePos = handle.tell()
+				handle.seek(info.startPos + dataOffset)
+				newObj.data = handle.read(size)
+				handle.seek(savePos)
+
+			elif newObj.isDirectory():
+				self._readDir(handle, newObj, size, info)
+
+			directory.addChild(newObj)
+
+	def pack(self, handle=None):
+		returnData = False
+		if handle is None:
+			handle = cStringIO.StringIO()
+			returnData = True
+
+		info = WiiArchiveU8.WriteInfo()
+
+		# first off, before we do anything else, create the string table
+		info.stringTableBuilder = WiiStringTableBuilder()
+		self._addNodeToStringTable(self.root, info.stringTableBuilder)
+
+		stringTable = info.stringTableBuilder.data
+
+		# calculate various fun offsets/sizes
+		nodeCount = self._countNode(self.root)
+		info.startPos = handle.tell()
+
+		fstStart = 0x20
+		nodeDataSize = nodeCount * 12
+		stringTableSize = len(stringTable)
+		fstSize = nodeDataSize + stringTableSize
+		dataOffset = alignUp(fstStart + fstSize, 0x20)
+
+		# now write the header
+		handle.write("U\xAA8\x2D")
+		handle.write(struct.pack('>III', fstStart, fstSize, dataOffset))
+		handle.write("\0"*16)
+
+		# write root node
+		info.currentNode = 1
+		info.currentRecursionLevel = 0
+		info.nextDataOffset = dataOffset
+
+		handle.write(struct.pack('>III',
+			0x01000000 | info.stringTableBuilder.add(''),
+			0, nodeCount))
+
+		self._writeDir(handle, self.root, info)
+
+		# write string table
+		handle.write(stringTable)
+
+		# write data (after padding)
+		handle.write("\0" * (dataOffset - fstSize - fstStart))
+		self._writeNodeData(handle, self.root)
+
+		# looks like we are finally done
+		if returnData:
+			data = handle.getvalue()
+			handle.close()
+			return data
+
+	def _addNodeToStringTable(self, node, table):
+		table.add(node.name)
+
+		if node.isDirectory():
+			for obj in node.children:
+				self._addNodeToStringTable(obj, table)
+
+	def _countNode(self, node):
+		if node.isDirectory():
+			return 1 + sum(map(self._countNode, node.children))
+		else:
+			return 1
+
+	def _writeDir(self, handle, directory, info):
+		for obj in directory.children:
+			info.currentNode += 1
+
+			if obj.isDirectory():
+				handle.write(struct.pack('>III',
+					0x01000000 | info.stringTableBuilder.add(obj.name),
+					info.currentRecursionLevel, 0))
+
+				lastChildFieldPos = handle.tell() - 4
+
+				info.currentRecursionLevel += 1
+				self._writeDir(handle, obj, info)
+				info.currentRecursionLevel -= 1
+
+				# write lastChild field
+				dirEndPos = handle.tell()
+
+				handle.seek(lastChildFieldPos)
+				handle.write(struct.pack('>I', info.currentNode))
+				handle.seek(dirEndPos)
+
+			elif obj.isFile():
+				handle.write(struct.pack('>III',
+					info.stringTableBuilder.add(obj.name),
+					info.nextDataOffset, len(obj.data)))
+
+				info.nextDataOffset = alignUp(info.nextDataOffset + len(obj.data), 0x20)
+
+
+	def _writeNodeData(self, handle, node):
+		if node.isDirectory():
+			for obj in node.children:
+				self._writeNodeData(handle, obj)
+			
+		elif node.isFile():
+			size = len(node.data)
+			handle.write(node.data)
+			handle.write("\0" * (alignUp(size, 0x20) - size))
+