diff --git a/kaitaistruct.py b/kaitaistruct.py index 2684fa6..0fcfd44 100644 --- a/kaitaistruct.py +++ b/kaitaistruct.py @@ -15,6 +15,19 @@ __version__ = '0.8' +if PY2: + range = xrange + def integers2bytes(ints): + return bytes(bytearray(ints)) + def bytes2integers(data): + return bytearray(data) +else: + def integers2bytes(ints): + return bytes(ints) + def bytes2integers(data): + return data + + class KaitaiStruct(object): def __init__(self, stream): self._io = stream @@ -335,30 +348,57 @@ def bytes_terminate(data, term, include_term): @staticmethod def process_xor_one(data, key): - if PY2: - return bytes(bytearray(v ^ key for v in bytearray(data))) - else: - return bytes(v ^ key for v in data) + if key == 0: + return data + + return integers2bytes(v ^ key for v in bytes2integers(data)) @staticmethod def process_xor_many(data, key): - if PY2: - return bytes(bytearray(a ^ b for a, b in zip(bytearray(data), itertools.cycle(bytearray(key))))) - else: - return bytes(a ^ b for a, b in zip(data, itertools.cycle(key))) + if len(key) == 1: + return KaitaiStream.process_xor_one(data, ord(key)) + if len(key) <= 64 and key == b'\x00' * len(key): + return data - @staticmethod - def process_rotate_left(data, amount, group_size): - if group_size != 1: - raise Exception( - "unable to rotate group of %d bytes yet" % - (group_size,) - ) + return integers2bytes(a ^ b for a, b in zip(bytes2integers(data), itertools.cycle(bytes2integers(key)))) - mask = group_size * 8 - 1 - anti_amount = -amount & mask + # formula taken from: http://stackoverflow.com/a/812039 + precomputed_single_rotations = {amount: [(i << amount) & 0xff | (i >> (8-amount)) for i in range(256)] for amount in range(1,8)} - r = bytearray(data) - for i in range(len(r)): - r[i] = (r[i] << amount) & 0xff | (r[i] >> anti_amount) - return bytes(r) + @staticmethod + def process_rotate_left(data, amount, group_size): + if group_size < 1: + raise Exception("group size must be at least 1 to be valid") + + amount = amount % (group_size * 8) + if amount == 0: + return data + + amount_bytes = amount // 8 + data_ints = bytes2integers(data) + + if group_size == 1: + translate = KaitaiStream.precomputed_single_rotations[amount] + return integers2bytes(translate[a] for a in data_ints) + + if len(data) % group_size != 0: + raise Exception("data length must be a multiple of group size") + + if amount % 8 == 0: + indices = [(i + amount_bytes) % group_size for i in range(group_size)] + return integers2bytes(data_ints[i+k] for i in range(0,len(data),group_size) for k in indices) + + amount1 = amount % 8 + amount2 = 8 - amount1 + indices_pairs = [ ((i+amount_bytes) % group_size, (i+1+amount_bytes) % group_size) for i in range(group_size)] + return integers2bytes((data_ints[i+k1] << amount1) & 0xff | (data_ints[i+k2] >> amount2) for i in range(0,len(data),group_size) for k1,k2 in indices_pairs) + + # NOTE: unused implementation, left for reference + # + # cap = (1 << 8 * group_size) - 1 + # anti_amount = -amount & (8 * group_size - 1) + # for i in range(0,len(data),group_size): + # group = bytes2combinedinteger(data[i:i+group_size]) + # group = (group << amount) & cap | (group >> anti_amount) + # r.append(combinedinteger2bytes(group, group_size)) + # return b''.join(r)