Skip to content

Commit 64164cd

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

File tree

1 file changed

+19
-12
lines changed

1 file changed

+19
-12
lines changed

kaitaistruct.py

Lines changed: 19 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,22 @@
1616
__version__ = '0.8'
1717

1818

19+
if PY2:
20+
# function(integer iterable) -> bytes
21+
def integers2bytes(ints):
22+
return bytes(bytearray(ints))
23+
# iterate(function(bytes)) -> integer iterable
24+
def bytes2integers(data):
25+
return bytearray(data)
26+
else:
27+
# function(integer iterable) -> bytes
28+
def integers2bytes(ints):
29+
return bytes(ints)
30+
# iterate(function(bytes)) -> integer iterable
31+
def bytes2integers(data):
32+
return data
33+
34+
1935
class KaitaiStruct(object):
2036
def __init__(self, stream):
2137
self._io = stream
@@ -339,10 +355,7 @@ def process_xor_one(data, key):
339355
if key == 0:
340356
return data
341357

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)
358+
return integers2bytes(v ^ key for v in bytes2integers(data))
346359

347360
@staticmethod
348361
def process_xor_many(data, key):
@@ -351,10 +364,7 @@ def process_xor_many(data, key):
351364
if len(key) <= 64 and key == b'\x00' * len(key):
352365
return data
353366

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)))
367+
return integers2bytes(a ^ b for a, b in zip(bytes2integers(data), itertools.cycle(bytes2integers(key))))
358368

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

370380
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)
381+
return integers2bytes(translate[a] for a in bytes2integers(data))

0 commit comments

Comments
 (0)