Initial commit

alan/attack.py

@@ -0,0 +1,122 @@
+#!/usr/bin/env python3
+import numpy as np
+import matplotlib.pyplot as plt
+import aes
+
+D = 6000  # Number of power traces (Number of Samples)
+T = 87  # Number of data points per power trace (Points in time)
+KEY_GUESSES = np.arange(256, dtype=np.uint8)
+
+
+def calculate_models(
+    ciphertexts: np.ndarray[np.ndarray[np.uint8]],
+) -> np.ndarray[np.ndarray[np.ndarray[np.uint8]]]:
+    """
+    Calculates for each sample ciphertext c the state just before the last
+    SubBytes operation of the AddKey step in the last round of AES.
+
+    In order to do this we calculate: sbox^-1 (c ⊕ k_hyp)
+    and return an numpy array with this form
+    [
+        [ # first ciphertext sample c1
+            [c_1_0  ⊕ 0, c_1_0  ⊕ 1, ..., c_1_0  ⊕ 255], # first byte of c1
+            [c_1_1  ⊕ 0, c_1_1  ⊕ 1, ..., c_1_1  ⊕ 255], # second byte of c1
+            ...,
+            [c_1_15 ⊕ 0, c_1_15 ⊕ 1, ..., c_1_15 ⊕ 255], # 16'th byte of c1
+        ],
+        [ # second ciphertext sample c2
+            [c_2_0  ⊕ 0, c_2_0  ⊕ 1, ..., c_2_0  ⊕ 255],
+            [c_2_1  ⊕ 0, c_2_1  ⊕ 1, ..., c_2_1  ⊕ 255],
+            ...,
+            [c_2_15 ⊕ 0, c_2_15 ⊕ 1, ..., c_2_15 ⊕ 255],
+        ],
+        ...,
+        [ # last ciphertext sample cD
+            [c_D_0  ⊕ 0, c_D_0  ⊕ 1, ..., c_D_0  ⊕ 255],
+            [c_D_1  ⊕ 0, c_D_1  ⊕ 1, ..., c_D_1  ⊕ 255],
+            ...,
+            [c_D_15 ⊕ 0, c_D_15 ⊕ 1, ..., c_D_15 ⊕ 255],
+        ],
+    ]
+
+    Args:
+        ciphertexts:
+            An array of all samples (ciphertexts) as an np.ndarray of np.ndarrays
+            for each byte as np.uint8.
+
+    Returns:
+        models:
+            An np.ndarray with each model of the state before the SubBytes
+            operation of the AddKey step in the last round of AES.
+                - axis 0: Samples/Ciphertexts
+                - axis 1: Bytes per sample
+                - axis 2: Byte xor k_hyp
+    """
+
+    # duplicate each ciphertext 256 times to xor with all possible keys.
+    models = np.repeat(ciphertexts, 256, axis=0).reshape((D, 256, 16))
+
+    # create a view of models with the inner axis swaped, so when we xor with
+    # KEY_GUESSES numpy can use broadcast.
+    models_view = np.swapaxes(models, 1, 2)
+
+    # c ⊕ k_hyp
+    np.bitwise_xor(models_view, KEY_GUESSES, out=models_view)
+
+    # apply reverse rbox to all bytes. (rsbox(c ⊕ k_hyp))
+    models = np.vectorize(lambda x: aes.core.rsbox(x))(models)
+    return models
+
+
+def read_msgs(file_name: str) -> np.ndarray[np.ndarray[np.ndarray[np.uint8]]]:
+    msgs = np.empty((D, 3, 16), dtype=np.uint8)
+
+    with open(file_name, 'r') as fd:
+        for idx, (key, plain_text) in enumerate(
+            (line.strip().split(',') for line in fd)
+        ):
+            msgs[idx][0] = np.frombuffer(bytes.fromhex(key), dtype=np.uint8)         # key
+            msgs[idx][1] = np.frombuffer(bytes.fromhex(plain_text), dtype=np.uint8)  # plain text
+            msgs[idx][2] = np.array(                                                 # ciphertext
+                aes.aes(int(key, 16), 128).enc_once(int(plain_text, 16)), dtype=np.uint8
+            )
+
+    return msgs
+
+def read_traces(file_name: str) -> np.ndarray[np.ndarray[np.uint8]]:
+    return np.loadtxt(file_name, delimiter=",", dtype=np.uint8)
+
+
+if __name__ == "__main__":
+    msgs = read_msgs("Task-3-example_traces/test_msgs.csv")
+    traces = read_traces("Task-3-example_traces/test_traces.csv")
+    models = calculate_models(msgs[:, 2])
+
+    np.set_printoptions(formatter={"int": hex})
+    last_round_key = aes.core.key_expansion(msgs[0][0].tolist())[-16:]
+    
+
+    for bit in range(128):
+        # i'th row, and j'th col is the correlation coefficient of key_hyp = i and time sample j
+        r = np.corrcoef(
+            np.bitwise_and(models[:, :, bit//8], np.array([2**(bit % 8)], dtype=np.uint8)),
+            traces,
+            rowvar=False
+        )[:256, -87:]
+
+        guess = np.argmax(np.max(np.abs(r), axis=1))
+        # print(guess)
+        
+    exit(0)
+        # tmp = np.sort(r.flatten())
+        # confidence = max(abs(tmp[0] - tmp[1]), abs(tmp[-2] - tmp[-1]))
+        # if confidence > 0.005:
+        #     fig, axs = plt.subplots(1, 1, layout='constrained')
+        #     axs.set_title(f"Bit {bit%8 + 1} of Byte {bit//8 + 1} (Confidence: {confidence:.6f})")
+        #     axs.plot(r.transpose(), alpha=0.3, color='grey')
+        #     axs.plot(r[last_round_key[bit//8]], color="blue")
+        #     axs.plot(r[guess], color="red")
+        #     axs.set_xlabel("Time Samples")
+        #     axs.set_ylabel("Correlation")
+        #     axs.grid(True)
+        #     plt.show()

alan/multi-processing.py

@@ -0,0 +1,83 @@
+#!/usr/bin/python3
+import numpy as np
+import matplotlib.pyplot as plt
+from multiprocessing import Pool
+import aes
+
+D = 6000  # Number of power traces (Number of Samples)
+T = 87  # Number of data points per power trace (Points in time)
+KEY_GUESSES = np.arange(256, dtype=np.uint8)
+
+
+def calculate_models(
+    ciphertext: np.ndarray[np.ndarray[np.uint8]],
+) -> np.ndarray[np.ndarray[np.ndarray[np.uint8]]]:
+
+    # duplicate each ciphertext 256 times to xor with all possible keys.
+    models = np.repeat(ciphertext, 256).reshape(16,256)
+
+    # create a view of models with the inner axis swaped, so when we xor with
+    # KEY_GUESSES numpy can use broadcast.
+    models_view = np.swapaxes(models, 0, 1)
+
+    # c ⊕ k_hyp
+    np.bitwise_xor(models, KEY_GUESSES, out=models)
+
+    # apply reverse rbox to all bytes. (rsbox(c ⊕ k_hyp))
+    models = np.vectorize(lambda x: aes.core.rsbox(x))(models_view)
+    return models
+
+
+def read_msgs(file_name: str) -> np.ndarray[np.ndarray[np.ndarray[np.uint8]]]:
+    msgs = np.empty((D, 3, 16), dtype=np.uint8)
+
+    with open(file_name, 'r') as fd:
+        for idx, (key, plain_text) in enumerate(
+            (line.strip().split(',') for line in fd)
+        ):
+            msgs[idx][0] = np.frombuffer(bytes.fromhex(key), dtype=np.uint8)         # key
+            msgs[idx][1] = np.frombuffer(bytes.fromhex(plain_text), dtype=np.uint8)  # plain text
+            msgs[idx][2] = np.array(                                                 # ciphertext
+                aes.aes(int(key, 16), 128).enc_once(int(plain_text, 16)), dtype=np.uint8
+            )
+
+    return msgs
+
+def read_traces(file_name: str) -> np.ndarray[np.ndarray[np.uint8]]:
+    return np.loadtxt(file_name, delimiter=",", dtype=np.uint8)
+
+
+if __name__ == "__main__":
+    msgs = read_msgs("Task-3-example_traces/test_msgs.csv")
+    traces = read_traces("Task-3-example_traces/test_traces.csv")
+    with Pool() as pool:
+        models = pool.map(calculate_models, msgs[:, 2])
+
+    models = np.stack(models)
+    
+    # np.set_printoptions(formatter={"int": hex})
+    last_round_key = aes.core.key_expansion(msgs[0][0].tolist())[-16:]
+
+    for bit in range(128):
+        # i'th row, and j'th col is the correlation coefficient of key_hyp i and time sample j
+        model = np.bitwise_and(models[:, :, bit//8], np.array([2**(bit % 8)], dtype=np.uint8))
+        r = np.corrcoef(
+            model,
+            traces,
+            rowvar=False
+        )[:256, -87:]
+
+        guess = np.argmax(np.max(np.abs(r), axis=1))
+        
+        # tmp = np.sort(r.flatten())
+        # confidence = max(abs(tmp[0] - tmp[1]), abs(tmp[-2] - tmp[-1]))
+        # if confidence > 0.005:
+        #     fig, axs = plt.subplots(1, 1, layout='constrained')
+        #     axs.set_title(f"Bit {bit%8 + 1} of Byte {bit//8 + 1} (Confidence: {confidence:.6f})")
+        #     axs.plot(r.transpose(), alpha=0.3, color='grey')
+        #     axs.plot(r[last_round_key[bit//8]], color="blue")
+        #     axs.plot(r[guess], color="red")
+        #     axs.set_xlabel("Time Samples")
+        #     axs.set_ylabel("Correlation")
+        #     axs.grid(True)
+        #     plt.show()