ProcessXor updated, optimised for neutral keys, 1-length bytes

construct · Apr 8, 2018 · ecdfc27 · ecdfc27
1 parent deb8938
commit ecdfc27
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 4 deletions.
diff --git a/construct/core.py b/construct/core.py
@@ -4949,11 +4949,15 @@ def _parse(self, stream, context, path):
         pad = evaluate(self.padfunc, context)
         if not isinstance(pad, (integertypes, bytestringtype)):
             raise StringError("ProcessXor needs integer or bytes pad")
+        if isinstance(pad, bytestringtype) and len(pad) == 1:
+            pad = byte2int(pad)
         data = stream_read_entire(stream)
         if isinstance(pad, integertypes):
-            data = bytes(bytearray((b ^ pad) for b in iterateints(data)))
+            if not (pad == 0):
+                data = bytes(bytearray( (b ^ pad) for b in iterateints(data) ))
         if isinstance(pad, bytestringtype):
-            data = bytes(bytearray((b ^ p) for b,p in zip(iterateints(data), itertools.cycle(iterateints(pad)))))
+            if not (len(pad) <= 64 and pad == bytes(len(pad))):
+                data = bytes(bytearray( (b ^ p) for b,p in zip(iterateints(data), itertools.cycle(iterateints(pad))) ))
         if self.subcon is GreedyBytes:
             return data
         if type(self.subcon) is GreedyString:
@@ -4964,13 +4968,17 @@ def _build(self, obj, stream, context, path):
         pad = evaluate(self.padfunc, context)
         if not isinstance(pad, (integertypes, bytestringtype)):
             raise StringError("ProcessXor needs integer or bytes pad")
+        if isinstance(pad, bytestringtype) and len(pad) == 1:
+            pad = byte2int(pad)
         stream2 = io.BytesIO()
         buildret = self.subcon._build(obj, stream2, context, path)
         data = stream2.getvalue()
         if isinstance(pad, integertypes):
-            data = bytes(bytearray((b ^ pad) for b in iterateints(data)))
+            if not (pad == 0):
+                data = bytes(bytearray( (b ^ pad) for b in iterateints(data) ))
         if isinstance(pad, bytestringtype):
-            data = bytes(bytearray((b ^ p) for b,p in zip(iterateints(data), itertools.cycle(iterateints(pad)))))
+            if not (len(pad) <= 64 and pad == bytes(len(pad))):
+                data = bytes(bytearray( (b ^ p) for b,p in zip(iterateints(data), itertools.cycle(iterateints(pad))) ))
         stream_write(stream, data)
         return buildret
 

diff --git a/tests/test_core.py b/tests/test_core.py
@@ -1130,8 +1130,12 @@ def test_restreamed_partial_read():
     assert raises(d.parse, b"") == StreamError
 
 def test_processxor():
+    d = ProcessXor(0, Int16ub)
+    common(d, b"\xf0\x0f", 0xf00f, 2)
     d = ProcessXor(0xf0, Int16ub)
     common(d, b"\x00\xff", 0xf00f, 2)
+    d = ProcessXor(bytes(10), Int16ub)
+    common(d, b"\xf0\x0f", 0xf00f, 2)
     d = ProcessXor(b"\xf0\xf0\xf0\xf0\xf0", Int16ub)
     common(d, b"\x00\xff", 0xf00f, 2)