From 88e41acb0980ad36254acba96a6f5f4207b72135 Mon Sep 17 00:00:00 2001 From: Arkadiusz Bulski Date: Sun, 1 Apr 2018 17:28:28 +0200 Subject: [PATCH 1/4] process xor: skips when key is zero, recognizes 1-byte case --- kaitaistruct.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/kaitaistruct.py b/kaitaistruct.py index 2684fa6..4e141d1 100644 --- a/kaitaistruct.py +++ b/kaitaistruct.py @@ -335,6 +335,9 @@ def bytes_terminate(data, term, include_term): @staticmethod def process_xor_one(data, key): + if key == 0: + return data + if PY2: return bytes(bytearray(v ^ key for v in bytearray(data))) else: @@ -342,6 +345,11 @@ def process_xor_one(data, key): @staticmethod def process_xor_many(data, 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 + if PY2: return bytes(bytearray(a ^ b for a, b in zip(bytearray(data), itertools.cycle(bytearray(key))))) else: From 836d7a324348373cc6ca3e68c8e23dd35d41a74b Mon Sep 17 00:00:00 2001 From: Arkadiusz Bulski Date: Sun, 1 Apr 2018 17:23:40 +0200 Subject: [PATCH 2/4] process rotate: precomputed translations, comprehension instead of a loop --- kaitaistruct.py | 23 ++++++++++++----------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/kaitaistruct.py b/kaitaistruct.py index 4e141d1..856003f 100644 --- a/kaitaistruct.py +++ b/kaitaistruct.py @@ -355,18 +355,19 @@ def process_xor_many(data, key): else: return bytes(a ^ b for a, b in zip(data, itertools.cycle(key))) + # formula taken from: http://stackoverflow.com/a/812039 + precomputed_rotations = {amount:[(i << amount) & 0xff | (i >> (-amount & 7)) for i in range(256)] for amount in range(8)} + @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,) - ) - - mask = group_size * 8 - 1 - anti_amount = -amount & mask + raise Exception("unable to rotate groups other than 1 byte") + amount = amount % 8 + if amount == 0: + return data - r = bytearray(data) - for i in range(len(r)): - r[i] = (r[i] << amount) & 0xff | (r[i] >> anti_amount) - return bytes(r) + translate = KaitaiStream.precomputed_rotations[amount] + if PY2: + return bytes(bytearray(translate[a] for a in bytearray(data))) + else: + return bytes(translate[a] for a in data) From bbf212d7c84d43571d9b6f559e06095d9a959725 Mon Sep 17 00:00:00 2001 From: Arkadiusz Bulski Date: Wed, 4 Apr 2018 12:51:18 +0200 Subject: [PATCH 3/4] added 2 helper functions: integers2bytes and bytes2integers --- kaitaistruct.py | 27 +++++++++++++++------------ 1 file changed, 15 insertions(+), 12 deletions(-) diff --git a/kaitaistruct.py b/kaitaistruct.py index 856003f..14ff940 100644 --- a/kaitaistruct.py +++ b/kaitaistruct.py @@ -15,6 +15,18 @@ __version__ = '0.8' +if PY2: + 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 @@ -338,10 +350,7 @@ def process_xor_one(data, key): if key == 0: return data - if PY2: - return bytes(bytearray(v ^ key for v in bytearray(data))) - else: - return bytes(v ^ key for v in data) + return integers2bytes(v ^ key for v in bytes2integers(data)) @staticmethod def process_xor_many(data, key): @@ -350,10 +359,7 @@ def process_xor_many(data, key): if len(key) <= 64 and key == b'\x00' * len(key): return data - 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))) + return integers2bytes(a ^ b for a, b in zip(bytes2integers(data), itertools.cycle(bytes2integers(key)))) # formula taken from: http://stackoverflow.com/a/812039 precomputed_rotations = {amount:[(i << amount) & 0xff | (i >> (-amount & 7)) for i in range(256)] for amount in range(8)} @@ -367,7 +373,4 @@ def process_rotate_left(data, amount, group_size): return data translate = KaitaiStream.precomputed_rotations[amount] - if PY2: - return bytes(bytearray(translate[a] for a in bytearray(data))) - else: - return bytes(translate[a] for a in data) + return integers2bytes(translate[a] for a in bytes2integers(data)) From 3511a6a08384e998c4a7a11b11db1e2d111eaf81 Mon Sep 17 00:00:00 2001 From: Arkadiusz Bulski Date: Wed, 4 Apr 2018 13:23:14 +0200 Subject: [PATCH 4/4] rotations of larger groups --- kaitaistruct.py | 40 ++++++++++++++++++++++++++++++++++------ 1 file changed, 34 insertions(+), 6 deletions(-) diff --git a/kaitaistruct.py b/kaitaistruct.py index 14ff940..0fcfd44 100644 --- a/kaitaistruct.py +++ b/kaitaistruct.py @@ -16,6 +16,7 @@ if PY2: + range = xrange def integers2bytes(ints): return bytes(bytearray(ints)) def bytes2integers(data): @@ -362,15 +363,42 @@ def process_xor_many(data, key): return integers2bytes(a ^ b for a, b in zip(bytes2integers(data), itertools.cycle(bytes2integers(key)))) # formula taken from: http://stackoverflow.com/a/812039 - precomputed_rotations = {amount:[(i << amount) & 0xff | (i >> (-amount & 7)) for i in range(256)] for amount in range(8)} + precomputed_single_rotations = {amount: [(i << amount) & 0xff | (i >> (8-amount)) for i in range(256)] for amount in range(1,8)} @staticmethod def process_rotate_left(data, amount, group_size): - if group_size != 1: - raise Exception("unable to rotate groups other than 1 byte") - amount = amount % 8 + 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 - translate = KaitaiStream.precomputed_rotations[amount] - return integers2bytes(translate[a] for a in bytes2integers(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)