speedup

1 files changed, 30 insertions, 27 deletions

diff --git a/aesgcmanalysis.py b/aesgcmanalysis.py
index 52b4012..4b7e32c 100644
--- a/aesgcmanalysis.py
+++ b/aesgcmanalysis.py
@@ -1,5 +1,7 @@
 import random, struct, hmac, itertools, math
 from Crypto.Cipher import AES
+import cryptography
+from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
 import numpy as np
 
 ## Computation in GF(2^128)/(x^128 + x^7 + x^2 + x^1 + 1)
@@ -609,8 +611,11 @@ def gen_blocks(n, js):
     return blocks
 
 squarer = np.array(Ms())
-matsqlookup = np.load(open('squares.np', 'rb'))
 adlookup = np.load(open('ad.np', 'rb'))
+mcsqlookup = np.load(open('square_basis.np', 'rb'))
+
+def mc_squared(c, j):
+    return sum(mcsqlookup[i, j] for i in range(128) if 1 == (c >> i) & 1) % 2
 
 def gen_ad(blocks):
     matret = np.zeros((128, 128))
@@ -621,12 +626,11 @@ def gen_ad(blocks):
         j = i + 1 # first is taken up by length block
         mat = None
         d = bytes_to_gf128(block)
-        matd = Mc(d)
-        if len(matsqlookup) > j:
-            matsq = matsqlookup[j]
-        else:
+        try:
+            mat = mc_squared(d, j)
+        except IndexError:
             matsq = np.linalg.matrix_power(squarer, j) % 2
-        mat = matd @ matsq
+            mat = Mc(d) @ matsq
         matret += mat
     return matret % 2 # Because the elements of Ad are in GF2 we can mod 2
 
@@ -666,10 +670,11 @@ def chunk(xs, n=16):
     return [xs[i*n:(i+1)*n] for i in range(len(xs)//16)]
 
 def find_b(n, basis, ct, mac, nonce, aad, oracle):
+    orig_base = bytearray(ct).copy()
     base = bytearray(ct)
     idx = 0
     while True:
-        choice = random.sample(basis, random.randint(1, 14))
+        choice = random.sample(basis, random.randint(1, 10))
         b = sum(choice) % 2
         flips = gen_flips(b)
         blocks = gen_blocks(n, flips)
@@ -679,11 +684,8 @@ def find_b(n, basis, ct, mac, nonce, aad, oracle):
         try:
             oracle(base[len(aad):], base[:len(aad)], mac, nonce)
             return b
-        except ValueError as e:
-            assert str(e) == 'MAC check failed'
-            for i, block in enumerate(chunk(blocks)):
-                j = (len(base)//16)-(2**(i+1)-1)
-                base[j*16:(j+1)*16] = xor(base[j*16:(j+1)*16], block)
+        except (cryptography.exceptions.InvalidTag, ValueError):
+            base = orig_base.copy()
         idx += 1
 
 def nonce_truncation_recover_secrets(ct, mac, nonce, mac_bytes, aad, oracle, compute_T_once=False):
@@ -722,7 +724,7 @@ def nonce_truncation_recover_secrets(ct, mac, nonce, mac_bytes, aad, oracle, com
         _, _, basisKerK = kernel(K, rref_mod_2)
         X = np.array(basisKerK).transpose()
     _, _, kerK = kernel(K, rref_mod_2)
-    assert len(kerK) == 1
+    assert len(kerK) == 1, len(kerK)
     h = kerK[0]
 
     zero_tag = gf128_to_vec(bytes_to_gf128(gmac(vec_to_gf128(h), 0, aad, orig_ct)))[:mac_bytes*8]
@@ -774,24 +776,25 @@ def forbidden_attack_demo():
     assert succeeded
 
 def nonce_truncation_demo():
-    # Should work with non-block size multiples.
+    # Doesn't work with non-block size multiples.
     # Need to modify to consider padding, but we can't mess with the bits in the padding,
     # nor can we extend ad/ct unless we also change length block.
-    k = b'YELLOGIOJEWARINE'
-    aad = b'YELLOW_SUBMAFINERELLOWPUBMARINF_'
-    MACBYTES=1
-    pt = b'CELERYPATCHWORKS'*(2**5)
-    nonce = b'JORGELBORGES'
-    ct, mac = gcm_encrypt(k, nonce, aad, pt)
-    # ct, mac = gcm_encrypt(k, nonce, aad, pt, mac_bytes=MACBYTES)
-    mac = mac[:1]
+    k = b'tlonorbistertius'
+    aad = b'yellow_submarine'
+    mac_bytes=2
+    pt = b'celerypatchworks'*(2**9)
+    nonce = b'jorgelborges'
+    ct, mac = gcm_encrypt(k, nonce, aad, pt, mac_bytes=mac_bytes)
     def oracle(base, aad, mac, nonce):
-        cipher = AES.new(k, mode=AES.MODE_GCM, nonce=nonce, mac_len=MACBYTES)
-        cipher.update(aad)
-        pt = cipher.decrypt_and_verify(base, mac)
-    h, s = nonce_truncation_recover_secrets(ct, mac, nonce, MACBYTES, aad, oracle, compute_T_once=True)
+        decryptor = Cipher(
+            algorithms.AES(k),
+            modes.GCM(nonce, mac, min_tag_length=mac_bytes),
+        ).decryptor()
+        decryptor.authenticate_additional_data(aad)
+        decryptor.update(base) + decryptor.finalize()
+
+    h, s = nonce_truncation_recover_secrets(ct, mac, nonce, mac_bytes, aad, oracle, compute_T_once=mac_bytes==1)
     assert h == authentication_key(k)
-    return h, s
 
 if __name__ == "__main__":
     nonce_truncation_demo()