[implem-python] Réécriture de certains range() dans ae_common.py

Et réutilisation de fonctions Python natives.

1 files changed, 8 insertions, 14 deletions

diff --git a/src/add_python/lilliput/ae_common.py b/src/add_python/lilliput/ae_common.py
index f212353..033b5b0 100644
--- a/src/add_python/lilliput/ae_common.py
+++ b/src/add_python/lilliput/ae_common.py
@@ -35,32 +35,26 @@ def pad10(X):
 
 
 def _tweak_associated_data(t, i, padded):
-    t_bytes = t//8
-    tweak = [0]*(t_bytes)
+    tweak = list(i.to_bytes(t//8, 'little'))
 
-    mask = 0xff
-    for byte in range(t_bytes-1):
-        tweak[byte] = (i & mask) >> (byte * 8)
-        mask = mask << 8
+    prefix = 0b0110 if padded else 0b0010
 
-    mask = (0xf << (8 * t_bytes-1))
-    tweak[-1] = (i & mask) >> ((t_bytes-1)*8)
-    if not padded:
-        tweak[-1] |= 0x20
-    else:
-        tweak[-1] |= 0x60
+    # Clear upper 4 bits and set them to prefix.
+    tweak[-1] &= 0b00001111
+    tweak[-1] = prefix << 4
 
     return tweak
 
 
 def build_auth(t, A, key):
-    Auth = [0 for byte in range(0, BLOCK_BYTES)]
+    Auth = [0]*BLOCK_BYTES
+
     l_a = len(A)//BLOCK_BYTES
     need_padding = len(A)%BLOCK_BYTES > 0
 
     A = bytes_to_block_matrix(A)
 
-    for i in range(0, l_a):
+    for i in range(l_a):
         tweak = _tweak_associated_data(t, i, padded=False)
         enc = tbc.encrypt(tweak, key, A[i])
         Auth = xor(Auth, enc)