U8 archive stuff ported from LayoutStudio which may or may not work

3 files changed, 338 insertions, 0 deletions

diff --git a/src/wii/common.py b/src/wii/common.py
new file mode 100644
index 0000000..40c3954
--- /dev/null
+++ b/src/wii/common.py
@@ -0,0 +1,25 @@
+class WiiStringTableBuilder(object):
+	def __init__(self):
+		self.nextOffset = 0
+		self.data = ''
+		self.lookup = {}
+	
+	def add(self, string):
+		if string in self.lookup:
+			return self.lookup[string]
+
+		offset = self.nextOffset
+		self.lookup[string] = offset
+		
+		self.data = "%s%s\0" % (self.data, string.encode('Shift-JIS'))
+		self.nextOffset = len(self.data)
+
+		return offset
+
+
+def alignUp(value, alignTo):
+	return (value + alignTo - 1) & ~(alignTo - 1)
+
+def alignDown(value, alignTo):
+	return value & ~(alignTo - 1)
+
diff --git a/src/wii/filesystem.py b/src/wii/filesystem.py
new file mode 100644
index 0000000..b6b8103
--- /dev/null
+++ b/src/wii/filesystem.py
@@ -0,0 +1,124 @@
+import weakref
+
+class WiiFSObject(object):
+	def __init__(self):
+		self._parent_ref = None
+		self._name = ''
+		self._low_name = ''
+	
+	@property
+	def parent(self):
+		ref = self._parent_ref
+		return (ref() if ref else None)
+
+	@parent.setter
+	def parent(self, value):
+		self._parent_ref = weakref.ref(value)
+	
+	@parent.deleter
+	def parent(self):
+		self._parent_ref = None
+	
+
+	@property
+	def name(self):
+		return self._name
+
+	@name.setter
+	def name(self, value):
+		self._name = value
+		self._low_name = value.lower()
+
+	
+	def unlinkFromParent(self):
+		if self.parent == None:
+			return
+
+		if self.parent.isDirectory():
+			self.parent.children.remove(self)
+
+		del self.parent
+
+	def isFile(self): return False
+	def isDirectory(self): return False
+
+
+class WiiFile(WiiFSObject):
+	def __init__(self):
+		WiiFSObject.__init__(self)
+		self.data = None
+	
+	def isFile(self): return True
+
+
+class WiiDirectory(WiiFSObject):
+	def __init__(self):
+		WiiFSObject.__init__(self)
+		self.children = []
+	
+	def isDirectory(self): return True
+
+	def findByName(self, name, recursive):
+		return self._findByName(name.lower(), recursive)
+
+	def _findByName(self, lowercase_name, recursive):
+		for obj in self.children:
+			if obj._low_name == lowercase_name:
+				return obj
+
+			if recursive and obj.isDirectory():
+				tryThis = obj._findByName(lowercase_name, recursive)
+				if tryThis:
+					return tryThis
+
+		return None
+
+	def resolvePath(self, path):
+		components = path.split('/')
+		currentDir = self
+
+		# special case: handle absolute paths
+		if components[0] == '':
+			while currentDir.parent:
+				currentDir = currentDir.parent
+
+			del components[0]
+
+		while len(components) > 0:
+			bit = components.pop(0)
+			nextObj = None
+
+			if bit == '.':
+				nextObj = currentDir
+			elif bit == '..':
+				nextObj = currentDir.parent
+			else:
+				nextObj = currentDir.findByName(bit, False)
+
+			if nextObj is None:
+				print("failed to resolve path %s: missing component %s" % (path, bit))
+				return None
+
+			if len(components) == 0:
+				return nextObj
+
+			if not nextObj.isDirectory():
+				print("failed to resolve path %s: component %s is not a directory" % (path, bit))
+				return None
+
+			# has to be a directory so just continue
+			currentDir = nextObj
+
+		print("wtf? %s" % path)
+		return None
+
+	def addChild(self, obj):
+		if self.findByName(obj.name, False):
+			return False
+
+		obj.unlinkFromParent()
+		self.children.append(obj)
+		obj.parent = self
+
+		return True
+
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))
+