Merge branch 'fix-concatenation'

1 files changed, 17 insertions, 15 deletions

diff --git a/src/add_python/lilliput/ae_mode_2.py b/src/add_python/lilliput/ae_mode_2.py
index 79d1bcd..a55ecb8 100644
--- a/src/add_python/lilliput/ae_mode_2.py
+++ b/src/add_python/lilliput/ae_mode_2.py
@@ -18,6 +18,8 @@ This module provides the functions for authenticated encryption and decryption
 using Lilliput-AE's nonce-misuse-resistant mode based on SCT-2.
 """
 
+from enum import Enum
+
 from .constants import BLOCK_BYTES
 from .ae_common import (
     bytes_to_block_matrix,
@@ -35,22 +37,24 @@ TWEAK_BITS = 128
 TWEAK_BYTES = TWEAK_BITS//8
 
 
-def _tweak_tag(j, padded):
-    tweak = integer_to_byte_array(j, TWEAK_BYTES)
+class _TagTweak(Enum):
+    BLOCK = 0b0000
+    PAD = 0b0100
 
-    prefix = 0b0100 if padded else 0b0000
+
+def _tweak_tag(j, prefix):
+    tweak = integer_to_byte_array(j, TWEAK_BYTES)
 
     # Clear upper 4 bits and set them to prefix.
-    tweak[-1] &= 0b00001111
-    tweak[-1] = prefix << 4
+    tweak[0] &= 0b00001111
+    tweak[0] |= prefix.value << 4
 
     return tweak
 
 
 def _add_tag_j(tag, j):
-    array_j = integer_to_byte_array(j, TWEAK_BYTES)
-    tweak = xor(tag, array_j)
-    tweak[-1] |= 0b10000000
+    tweak = xor(tag, integer_to_byte_array(j, TWEAK_BYTES))
+    tweak[0] |= 0b10000000
 
     return tweak
 
@@ -63,18 +67,16 @@ def _message_auth_tag(M, N, Auth, key):
     M = bytes_to_block_matrix(M)
 
     for j in range(0, l):
-        tweak = _tweak_tag(j, False)
+        tweak = _tweak_tag(j, _TagTweak.BLOCK)
         encryption = tbc.encrypt(tweak, key, M[j])
         tag = xor(tag, encryption)
 
     if need_padding:
-        tweak = _tweak_tag(l, True)
+        tweak = _tweak_tag(l, _TagTweak.PAD)
         encryption = tbc.encrypt(tweak, key, pad10(M[l]))
         tag = xor(tag, encryption)
 
-    tweak = N + [0b00010000]
-    encryption = tbc.encrypt(tweak, key, tag)
-    tag = encryption
+    tag = tbc.encrypt([0b00010000]+N, key, tag)
 
     return tag
 
@@ -88,12 +90,12 @@ def _message_encryption(M, N, tag, key):
 
     for j in range(0, l):
         tweak = _add_tag_j(tag, j)
-        encryption = tbc.encrypt(tweak, key, N+[0b00000000])
+        encryption = tbc.encrypt(tweak, key, [0b00000000]+N)
         C.append(xor(M[j], encryption))
 
     if need_padding:
         tweak = _add_tag_j(tag, l)
-        encryption = tbc.encrypt(tweak, key, N+[0b00000000])
+        encryption = tbc.encrypt(tweak, key, [0b00000000]+N)
         C.append(xor(M[l], encryption))
 
     return C