1 files changed, 29 insertions, 12 deletions
diff --git a/src/add_python/lilliput/ae_mode_1.py b/src/add_python/lilliput/ae_mode_1.py
index b07adf6..1a3c39e 100644
--- a/src/add_python/lilliput/ae_mode_1.py
+++ b/src/add_python/lilliput/ae_mode_1.py
@@ -20,11 +20,12 @@ using Lilliput-AE's nonce-respecting mode based on ΘCB3.
 
 from enum import Enum
 
-from .constants import BLOCK_BYTES, NONCE_BYTES
+from .constants import BLOCK_BYTES, NONCE_BITS
 from .ae_common import (
     bytes_to_block_matrix,
     block_matrix_to_bytes,
     build_auth,
+    integer_to_byte_array,
     pad10,
     TagValidationError,
     xor
@@ -43,19 +44,33 @@ class _MessageTweak(Enum):
     FINAL = 0b0101
 
 
+def _upper_nibble(i):
+    return i >> 4
+
+
+def _lower_nibble(i):
+    return i & 0b00001111
+
+
+def _byte_from_nibbles(lower, upper):
+    return upper<<4 | lower
+
+
 def _tweak_message(N, j, padding):
-    tweak = [0 for byte in range(0, TWEAK_BYTES)]
-    for byte in range(NONCE_BYTES-1, -1, -1):
-        tweak[byte + (TWEAK_BYTES-NONCE_BYTES)] |= (N[byte] & 0xf0) >> 4
-        tweak[byte + (TWEAK_BYTES-NONCE_BYTES-1)] |= (N[byte] & 0x0f) << 4
+    j = integer_to_byte_array(j, (TWEAK_BITS-NONCE_BITS-4)//8+1)
+
+    middle_byte = _byte_from_nibbles(
+        _lower_nibble(j[-1]), _lower_nibble(N[0])
+    )
 
-    tweak[TWEAK_BYTES-NONCE_BYTES-1] |= ((j >> 64) & 0xf)
-    for byte in range(TWEAK_BYTES-NONCE_BYTES-2, -1, -1):
-        tweak[byte] = (j >> (8 * byte)) & 0xff
+    shifted_N = [
+        _byte_from_nibbles(_upper_nibble(N[i-1]), _lower_nibble(N[i]))
+        for i in range(1, NONCE_BITS//8)
+    ]
 
-    tweak[-1] |= padding.value<<4
+    last_byte = _byte_from_nibbles(_upper_nibble(N[-1]), padding.value)
 
-    return tweak
+    return j[:-1] + [middle_byte] + shifted_N + [last_byte]
 
 
 def _treat_message_enc(M, N, key):
@@ -124,9 +139,10 @@ def _treat_message_dec(C, N, key):
 
 def encrypt(A, M, N, key):
     K = list(key)
+    N = list(N)
 
     Auth = build_auth(TWEAK_BITS, A, K)
-    (Final, C) = _treat_message_enc(M, N, K)
+    Final, C = _treat_message_enc(M, N, K)
     tag = xor(Auth, Final)
 
     return block_matrix_to_bytes(C), bytes(tag)
@@ -134,10 +150,11 @@ def encrypt(A, M, N, key):
 
 def decrypt(A, C, N, tag, key):
     K = list(key)
+    N = list(N)
     tag = list(tag)
 
     Auth = build_auth(TWEAK_BITS, A, K)
-    (Final, M) = _treat_message_dec(C, N, K)
+    Final, M = _treat_message_dec(C, N, K)
     tag2 = xor(Auth, Final)
 
     if tag != tag2: