diff options
Diffstat (limited to 'Koopuzzle/common.py')
| -rwxr-xr-x | Koopuzzle/common.py | 364 |
1 files changed, 364 insertions, 0 deletions
diff --git a/Koopuzzle/common.py b/Koopuzzle/common.py new file mode 100755 index 0000000..792acec --- /dev/null +++ b/Koopuzzle/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)
|