Skip to content

Commit bbf212d

Browse files
committed
added 2 helper functions: integers2bytes and bytes2integers
1 parent 836d7a3 commit bbf212d

File tree

1 file changed

+15
-12
lines changed

1 file changed

+15
-12
lines changed

kaitaistruct.py

Lines changed: 15 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,18 @@
1515
__version__ = '0.8'
1616

1717

18+
if PY2:
19+
def integers2bytes(ints):
20+
return bytes(bytearray(ints))
21+
def bytes2integers(data):
22+
return bytearray(data)
23+
else:
24+
def integers2bytes(ints):
25+
return bytes(ints)
26+
def bytes2integers(data):
27+
return data
28+
29+
1830
class KaitaiStruct(object):
1931
def __init__(self, stream):
2032
self._io = stream
@@ -338,10 +350,7 @@ def process_xor_one(data, key):
338350
if key == 0:
339351
return data
340352

341-
if PY2:
342-
return bytes(bytearray(v ^ key for v in bytearray(data)))
343-
else:
344-
return bytes(v ^ key for v in data)
353+
return integers2bytes(v ^ key for v in bytes2integers(data))
345354

346355
@staticmethod
347356
def process_xor_many(data, key):
@@ -350,10 +359,7 @@ def process_xor_many(data, key):
350359
if len(key) <= 64 and key == b'\x00' * len(key):
351360
return data
352361

353-
if PY2:
354-
return bytes(bytearray(a ^ b for a, b in zip(bytearray(data), itertools.cycle(bytearray(key)))))
355-
else:
356-
return bytes(a ^ b for a, b in zip(data, itertools.cycle(key)))
362+
return integers2bytes(a ^ b for a, b in zip(bytes2integers(data), itertools.cycle(bytes2integers(key))))
357363

358364
# formula taken from: http://stackoverflow.com/a/812039
359365
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):
367373
return data
368374

369375
translate = KaitaiStream.precomputed_rotations[amount]
370-
if PY2:
371-
return bytes(bytearray(translate[a] for a in bytearray(data)))
372-
else:
373-
return bytes(translate[a] for a in data)
376+
return integers2bytes(translate[a] for a in bytes2integers(data))

0 commit comments

Comments
 (0)