Merge branch 'master' of tsn:Koopatlas

1 files changed, 364 insertions, 0 deletions

diff --git a/Puzzle+/common.py b/Puzzle+/common.py
new file mode 100755
index 0000000..792acec
--- /dev/null
+++ b/Puzzle+/common.py
@@ -0,0 +1,364 @@
+import os.path, struct, sys

+

+

+class StructType(tuple):

+	def __getitem__(self, value):

+		return [self] * value

+	def __call__(self, value, endian='<'):

+		if isinstance(value, str):

+			return struct.unpack(endian + tuple.__getitem__(self, 0), value[:tuple.__getitem__(self, 1)])[0]

+		else:

+			return struct.pack(endian + tuple.__getitem__(self, 0), value)

+

+class StructException(Exception):

+	pass

+

+class Struct(object):

+	__slots__ = ('__attrs__', '__baked__', '__defs__', '__endian__', '__next__', '__sizes__', '__values__')

+	int8 = StructType(('b', 1))

+	uint8 = StructType(('B', 1))

+	

+	int16 = StructType(('h', 2))

+	uint16 = StructType(('H', 2))

+	

+	int32 = StructType(('l', 4))

+	uint32 = StructType(('L', 4))

+	

+	int64 = StructType(('q', 8))

+	uint64 = StructType(('Q', 8))

+	

+	float = StructType(('f', 4))

+

+	def string(cls, len, offset=0, encoding=None, stripNulls=False, value=''):

+		return StructType(('string', (len, offset, encoding, stripNulls, value)))

+	string = classmethod(string)

+	

+	LE = '<'

+	BE = '>'

+	__endian__ = '<'

+	

+	def __init__(self, func=None, unpack=None, **kwargs):

+		self.__defs__ = []

+		self.__sizes__ = []

+		self.__attrs__ = []

+		self.__values__ = {}

+		self.__next__ = True

+		self.__baked__ = False

+		

+		if func == None:

+			self.__format__()

+		else:

+			sys.settrace(self.__trace__)

+			func()

+			for name in func.func_code.co_varnames:

+				value = self.__frame__.f_locals[name]

+				self.__setattr__(name, value)

+		

+		self.__baked__ = True

+		

+		if unpack != None:

+			if isinstance(unpack, tuple):

+				self.unpack(*unpack)

+			else:

+				self.unpack(unpack)

+		

+		if len(kwargs):

+			for name in kwargs:

+				self.__values__[name] = kwargs[name]

+	

+	def __trace__(self, frame, event, arg):

+		self.__frame__ = frame

+		sys.settrace(None)

+	

+	def __setattr__(self, name, value):

+		if name in self.__slots__:

+			return object.__setattr__(self, name, value)

+		

+		if self.__baked__ == False:

+			if not isinstance(value, list):

+				value = [value]

+				attrname = name

+			else:

+				attrname = '*' + name

+			

+			self.__values__[name] = None

+			

+			for sub in value:

+				if isinstance(sub, Struct):

+					sub = sub.__class__

+				try:

+					if issubclass(sub, Struct):

+						sub = ('struct', sub)

+				except TypeError:

+					pass

+				type_, size = tuple(sub)

+				if type_ == 'string':

+					self.__defs__.append(Struct.string)

+					self.__sizes__.append(size)

+					self.__attrs__.append(attrname)

+					self.__next__ = True

+					

+					if attrname[0] != '*':

+						self.__values__[name] = size[3]

+					elif self.__values__[name] == None:

+						self.__values__[name] = [size[3] for val in value]

+				elif type_ == 'struct':

+					self.__defs__.append(Struct)

+					self.__sizes__.append(size)

+					self.__attrs__.append(attrname)

+					self.__next__ = True

+					

+					if attrname[0] != '*':

+						self.__values__[name] = size()

+					elif self.__values__[name] == None:

+						self.__values__[name] = [size() for val in value]

+				else:

+					if self.__next__:

+						self.__defs__.append('')

+						self.__sizes__.append(0)

+						self.__attrs__.append([])

+						self.__next__ = False

+					

+					self.__defs__[-1] += type_

+					self.__sizes__[-1] += size

+					self.__attrs__[-1].append(attrname)

+					

+					if attrname[0] != '*':

+						self.__values__[name] = 0

+					elif self.__values__[name] == None:

+						self.__values__[name] = [0 for val in value]

+		else:

+			try:

+				self.__values__[name] = value

+			except KeyError:

+				raise AttributeError(name)

+	

+	def __getattr__(self, name):

+		if self.__baked__ == False:

+			return name

+		else:

+			try:

+				return self.__values__[name]

+			except KeyError:

+				raise AttributeError(name)

+	

+	def __len__(self):

+		ret = 0

+		arraypos, arrayname = None, None

+		

+		for i in range(len(self.__defs__)):

+			sdef, size, attrs = self.__defs__[i], self.__sizes__[i], self.__attrs__[i]

+			

+			if sdef == Struct.string:

+				size, offset, encoding, stripNulls, value = size

+				if isinstance(size, str):

+					size = self.__values__[size] + offset

+			elif sdef == Struct:

+				if attrs[0] == '*':

+					if arrayname != attrs:

+						arrayname = attrs

+						arraypos = 0

+					size = len(self.__values__[attrs[1:]][arraypos])

+				size = len(self.__values__[attrs])

+			

+			ret += size

+		

+		return ret

+	

+	def unpack(self, data, pos=0):

+		for name in self.__values__:

+			if not isinstance(self.__values__[name], Struct):

+				self.__values__[name] = None

+			elif self.__values__[name].__class__ == list and len(self.__values__[name]) != 0:

+				if not isinstance(self.__values__[name][0], Struct):

+					self.__values__[name] = None

+		

+		arraypos, arrayname = None, None

+		

+		for i in range(len(self.__defs__)):

+			sdef, size, attrs = self.__defs__[i], self.__sizes__[i], self.__attrs__[i]

+			

+			if sdef == Struct.string:

+				size, offset, encoding, stripNulls, value = size

+				if isinstance(size, str):

+					size = self.__values__[size] + offset

+				

+				temp = data[pos:pos+size]

+				if len(temp) != size:

+					raise StructException('Expected %i byte string, got %i' % (size, len(temp)))

+				

+				if encoding != None:

+					temp = temp.decode(encoding)

+				

+				if stripNulls:

+					temp = temp.rstrip('\0')

+				

+				if attrs[0] == '*':

+					name = attrs[1:]

+					if self.__values__[name] == None:

+						self.__values__[name] = []

+					self.__values__[name].append(temp)

+				else:

+					self.__values__[attrs] = temp

+				pos += size

+			elif sdef == Struct:

+				if attrs[0] == '*':

+					if arrayname != attrs:

+						arrayname = attrs

+						arraypos = 0

+					name = attrs[1:]

+					self.__values__[attrs][arraypos].unpack(data, pos)

+					pos += len(self.__values__[attrs][arraypos])

+					arraypos += 1

+				else:

+					self.__values__[attrs].unpack(data, pos)

+					pos += len(self.__values__[attrs])

+			else:

+				values = struct.unpack(self.__endian__+sdef, data[pos:pos+size])

+				pos += size

+				j = 0

+				for name in attrs:

+					if name[0] == '*':

+						name = name[1:]

+						if self.__values__[name] == None:

+							self.__values__[name] = []

+						self.__values__[name].append(values[j])

+					else:

+						self.__values__[name] = values[j]

+					j += 1

+		

+		return self

+	

+	def pack(self):

+		arraypos, arrayname = None, None

+		

+		ret = ''

+		for i in range(len(self.__defs__)):

+			sdef, size, attrs = self.__defs__[i], self.__sizes__[i], self.__attrs__[i]

+			

+			if sdef == Struct.string:

+				size, offset, encoding, stripNulls, value = size

+				if isinstance(size, str):

+					size = self.__values__[size]+offset

+				

+				if attrs[0] == '*':

+					if arrayname != attrs:

+						arraypos = 0

+						arrayname = attrs

+					temp = self.__values__[attrs[1:]][arraypos]

+					arraypos += 1

+				else:

+					temp = self.__values__[attrs]

+				

+				if encoding != None:

+					temp = temp.encode(encoding)

+				

+				temp = temp[:size]

+				ret += temp + ('\0' * (size - len(temp)))

+			elif sdef == Struct:

+				if attrs[0] == '*':

+					if arrayname != attrs:

+						arraypos = 0

+						arrayname = attrs

+					ret += self.__values__[attrs[1:]][arraypos].pack()

+					arraypos += 1

+				else:

+					ret += self.__values__[attrs].pack()

+			else:

+				values = []

+				for name in attrs:

+					if name[0] == '*':

+						if arrayname != name:

+							arraypos = 0

+							arrayname = name

+						values.append(self.__values__[name[1:]][arraypos])

+						arraypos += 1

+					else:

+						values.append(self.__values__[name])

+				

+				ret += struct.pack(self.__endian__+sdef, *values)

+		return ret

+	

+	def __getitem__(self, value):

+		return [('struct', self.__class__)] * value

+

+

+class WiiObject(object):

+	def load(cls, data, *args, **kwargs):

+		self = cls()

+		self._load(data, *args, **kwargs)

+		return self

+	load = classmethod(load)

+

+	def loadFile(cls, filename, *args, **kwargs):

+		return cls.load(open(filename, "rb").read(), *args, **kwargs)

+	loadFile = classmethod(loadFile)

+

+	def dump(self, *args, **kwargs):

+		return self._dump(*args, **kwargs)

+	def dumpFile(self, filename, *args, **kwargs):

+		open(filename, "wb").write(self.dump(*args, **kwargs))

+		return filename

+

+

+class WiiArchive(WiiObject):

+	def loadDir(cls, dirname):

+		self = cls()

+		self._loadDir(dirname)

+		return self

+	loadDir = classmethod(loadDir)

+		

+	def dumpDir(self, dirname):

+		if(not os.path.isdir(dirname)):

+			os.mkdir(dirname)

+		self._dumpDir(dirname)

+		return dirname

+

+

+class WiiHeader(object):

+	def __init__(self, data):

+		self.data = data

+	def addFile(self, filename):

+		open(filename, "wb").write(self.add())

+	def removeFile(self, filename):

+		open(filename, "wb").write(self.remove())

+	def loadFile(cls, filename, *args, **kwargs):

+		return cls(open(filename, "rb").read(), *args, **kwargs)

+	loadFile = classmethod(loadFile)

+

+

+

+def align(x, boundary):

+	while x % boundary != 0:

+		x += 1

+	return x

+	

+def clamp(var, min, max):

+	if var < min: var = min

+	if var > max: var = max

+	return var

+

+def abs(var):

+	if var < 0:

+		var = var + (2 * var)

+	return var

+

+def hexdump(s, sep=" "): # just dumps hex values

+	return sep.join(map(lambda x: "%02x" % ord(x), s))

+		

+def hexdump2(src, length = 16): # dumps to a "hex editor" style output

+	result = []

+	for i in xrange(0, len(src), length):

+		s = src[i:i + length]

+		if(len(s) % 4 == 0):

+			mod = 0 

+		else: 

+			mod = 1 

+		hexa = ''

+		for j in range((len(s) / 4) + mod):

+			hexa += ' '.join(["%02X" % ord(x) for x in s[j * 4:j * 4 + 4]])

+			if(j != ((len(s) / 4) + mod) - 1):

+				hexa += '  '

+		printable = s.translate(''.join([(len(repr(chr(x))) == 3) and chr(x) or '.' for x in range(256)]))

+		result.append("0x%04X   %-*s   %s\n" % (i, (length * 3) + 2, hexa, printable))

+	return ''.join(result)