diff --git a/kbinxml/bytebuffer.py b/kbinxml/bytebuffer.py index 20620bb..500d4b4 100644 --- a/kbinxml/bytebuffer.py +++ b/kbinxml/bytebuffer.py @@ -20,6 +20,11 @@ class ByteBuffer(): else: return self.endian + str(count) + type + def get_bytes(self, count): + start = self.offset + self.offset += count + return self.data[start:self.offset] + def get(self, type, count = None): ret = self.peek(type, count) size = calcsize(type) @@ -33,6 +38,10 @@ class ByteBuffer(): ret = unpack_from(fmt, self.data, self.offset) return ret[0] if count is None else ret + def append_bytes(self, data): + self.data.extend(data) + self.offset += len(data) + def append(self, data, type, count = None): fmt = self._format_type(type, count) self.offset += calcsize(fmt) diff --git a/kbinxml/kbinxml.py b/kbinxml/kbinxml.py index 6464364..c33c488 100644 --- a/kbinxml/kbinxml.py +++ b/kbinxml/kbinxml.py @@ -124,13 +124,13 @@ class KBinXML(): def data_grab_auto(self): size = self.dataBuf.get_s32() - ret = self.dataBuf.get('B', size) + ret = self.dataBuf.get_bytes(size) self.dataBuf.realign_reads() return ret def data_append_auto(self, data): self.dataBuf.append_s32(len(data)) - self.dataBuf.append(data, 'B', len(data)) + self.dataBuf.append_bytes(data) self.dataBuf.realign_writes() def data_grab_string(self): @@ -309,7 +309,7 @@ class KBinXML(): name = unpack_sixbit(self.nodeBuf) else: length = (self.nodeBuf.get_u8() & ~64) + 1 - name = self.nodeBuf.get('B', length) + name = self.nodeBuf.get_bytes(length) name = bytes(name).decode(self.encoding) debug_print(name) diff --git a/kbinxml/sixbit.py b/kbinxml/sixbit.py index 8bd22b4..dbaf1a8 100644 --- a/kbinxml/sixbit.py +++ b/kbinxml/sixbit.py @@ -2,55 +2,27 @@ from builtins import bytes from bitarray import bitarray +charmap = '0123456789:ABCDEFGHIJKLMNOPQRSTUVWXYZ_abcdefghijklmnopqrstuvwxyz' +bytemap = {charmap[i] : bytes(chr(i).encode()) for i in range(len(charmap))} + +@profile def pack_sixbit(string, byteBuf): - chars = str_to_sixbit(string) + chars = [bytemap[x] for x in string] bits = bitarray(endian='big') for c in chars: - bits.frombytes(c.encode()) + bits.frombytes(c) # leave only the 6 bits we care for del bits[-8:-6] data = bytes(bits.tobytes()) - byteBuf.append_u8(len(string)) - byteBuf.append(data, 'B', len(data)) + byteBuf.append_bytes((len(string),)) + byteBuf.append_bytes(data) def unpack_sixbit(byteBuf): length = byteBuf.get_u8() - length_bytes = (length * 6 + 7) // 8 + length_bits = length * 6 + length_bytes = (length_bits + 7) // 8 bitBuf = bitarray(endian='big') - bitBuf.frombytes(bytes(byteBuf.get('B', length_bytes))) - result = [] - offset = 0 - for i in range(length): - result.append(ord(bitBuf[offset:offset+6].tobytes()) >> (8 - 6)) - offset += 6 - return sixbit_to_str(result) - -# 0-9 for numbers, 10 is ':', 11 to 36 for capitals, 37 for underscore, 38-63 for lowercase -def sixbit_to_str(decompressed): - string = '' - for d in decompressed: - if d <= 10: - d += ord('0') - elif d < 37: - d += 54 - elif d == 37: - d += 58 - else: - d += 59 - string += chr(d) - return string - -def str_to_sixbit(string): - compress = [] - for c in string: - if c >= '0' and c <= ':': - compress.append(ord(c) - ord('0')) - elif c >= 'A' and c <= 'Z': - compress.append(ord(c) - 54) - elif c == '_': - compress.append(ord(c) - 58) - elif c >= 'a' and c <= 'z': - compress.append(ord(c) - 59) - else: - raise ValueError('Node or attribute name can only contain alphanumeric + underscore') - return ''.join(map(chr, compress)) + bitBuf.frombytes(bytes(byteBuf.get_bytes(length_bytes))) + result = [bitBuf[offset:offset+6].tobytes()[0] >> 2 + for offset in range(0, length_bits, 6)] + return ''.join([charmap[x] for x in result])