Skip to content

Commit c340e40

Browse files
committed
added 2 helper functions: integers2bytes and bytes2integers
1 parent ebc0925 commit c340e40

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
@@ -16,6 +16,18 @@
1616
__version__ = '0.8'
1717

1818

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

342-
if PY2:
343-
return bytes(bytearray(v ^ key for v in bytearray(data)))
344-
else:
345-
return bytes(v ^ key for v in data)
354+
return integers2bytes(v ^ key for v in bytes2integers(data))
346355

347356
@staticmethod
348357
def process_xor_many(data, key):
@@ -351,10 +360,7 @@ def process_xor_many(data, key):
351360
if len(key) <= 64 and key == b'\x00' * len(key):
352361
return data
353362

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

359365
# formula taken from: http://stackoverflow.com/a/812039
360366
precomputed_rotations = {amount:[(i << amount) & 0xff | (i >> (-amount & 7)) for i in range(256)] for amount in range(8)}
@@ -368,7 +374,4 @@ def process_rotate_left(data, amount, group_size):
368374
return data
369375

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

0 commit comments

Comments
 (0)