Crack AES with a Neural Network

Train a CNN to recover an AES-128 secret key from electromagnetic traces.

Introduction

30 Second Summary

Every time you tap your credit card or send a message, encryption keeps your data safe. But what if the physical chip running that encryption accidentally broadcasts the secret key through its electromagnetic emissions?

In this project, you will train a convolutional neural network to crack AES-128 encryption by analyzing real electromagnetic traces captured from a microcontroller. You will build a complete attack pipeline in Google Colab that recovers the secret key from physical leakage.

What You'll Build

A deep learning side-channel attack notebook that downloads real electromagnetic traces, trains a CNN classifier, and uses model predictions to recover an AES-128 secret key byte.

By the end of this project, you'll have:

  • A working Google Colab notebook implementing the full profiling attack pipeline from data loading through key recovery.
  • A trained CNN model that classifies 256 possible S-box outputs from electromagnetic trace data, achieving Key Rank 0.
  • Visualizations of EM traces, training curves, and key rank convergence proving your model cracked the encryption.
  • Secret Mission: Attack desynchronized traces that simulate real-world jitter countermeasures and compare your model's robustness.

Are there any prerequisites?

You need basic Python proficiency and familiarity with neural network concepts like layers, training, and loss functions. No prior cryptography or hardware security experience is needed.

Set Up Your Attack Environment

To crack AES encryption with a neural network, you need three things: a GPU to train your model, real electromagnetic traces captured from cryptographic hardware, and the right Python libraries to process them.

In this step, you will set up a Google Colab notebook with GPU acceleration, download the ASCAD benchmark dataset, and verify everything is ready for building your attack pipeline.

In this step, get ready to:
  • Create a Google Colab notebook and connect it to a T4 GPU.
  • Download the ASCAD fixed-key dataset (4.2 GB of electromagnetic traces).
  • Import libraries and verify the dataset structure.
Create a Colab notebook and connect to a GPU

Your CNN will train on 50,000 electromagnetic traces. A GPU makes this feasible in minutes instead of hours.

  • Open Google Colab in your browser.
  • Click New notebook to create a fresh notebook.
  • In the top menu, click Runtime.
  • Click Change runtime type.
  • In the Hardware accelerator dropdown, select T4 GPU.
  • Click Save.
  • Click the Connect button in the top-right corner to provision your GPU instance.

Once connected, verify your GPU is active by running this code in the first cell:

# Verify GPU access and PyTorch version
import torch
print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}")

What should I see?

You should see output confirming PyTorch version 2.10.0 (or similar), CUDA available: True, and your GPU listed as a Tesla T4.

CUDA available shows False?

Make sure you changed the runtime type before clicking Connect. If you connected first, go to Runtime > Change runtime type, select T4 GPU, save, then click Reconnect. Colab will restart with the new hardware.

help me connect my Google Colab notebook to a T4 GPU

Download the ASCAD dataset

The ASCAD dataset contains real electromagnetic traces captured from an ATMega8515 microcontroller running masked AES-128. This is the data your neural network will learn to exploit.

  • Download and unzip the ASCAD fixed-key dataset by running this cell (the download is approximately 4.2 GB, so it takes a few minutes):
# Download ASCAD fixed-key dataset from data.gouv.fr (~4.2 GB)
!wget https://www.data.gouv.fr/api/1/datasets/r/e7ab6f9e-79bf-431f-a5ed-faf0ebe9b08e -O ASCAD_data.zip
# Unzip quietly into ASCAD_data directory
!unzip -q ASCAD_data.zip
# Verify the HDF5 file exists
!ls ASCAD_data/ASCAD_databases/

What does this cell do?

The ! prefix runs shell commands directly from a Colab cell. This downloads the zip file from France's open data portal, extracts it, and lists the contents to confirm ASCAD.h5 is present.

The .h5 file is in HDF5 format, a container that holds the electromagnetic traces, their labels, and metadata about the encryption operations.

  • Check the output of the ls command. You should see ASCAD.h5 listed in the ASCAD_data/ASCAD_databases/ directory.

Download failed or timed out?

The file is 4.2 GB so it may take 3-5 minutes. If the download fails partway through, re-run the cell. The wget command will start fresh.

If you see a 404 error, the dataset URL may have changed. Check the ASCAD GitHub repository for the current download link.

help me download the ASCAD dataset in Google Colab

Import libraries and verify the dataset structure

Now that the data is downloaded, you need to set up the Python environment for the rest of the project. This includes your deep learning framework, numerical libraries, and the cryptographic constants your attack needs.

  • Create a new cell and add all project imports plus the AES S-box lookup table by pasting this code:
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import numpy as np
import h5py
import matplotlib.pyplot as plt

# AES S-box lookup table
SBOX = np.array([
    0x63, 0x7C, 0x77, 0x7B, 0xF2, 0x6B, 0x6F, 0xC5, 0x30, 0x01, 0x67, 0x2B, 0xFE, 0xD7, 0xAB, 0x76,
    0xCA, 0x82, 0xC9, 0x7D, 0xFA, 0x59, 0x47, 0xF0, 0xAD, 0xD4, 0xA2, 0xAF, 0x9C, 0xA4, 0x72, 0xC0,
    0xB7, 0xFD, 0x93, 0x26, 0x36, 0x3F, 0xF7, 0xCC, 0x34, 0xA5, 0xE5, 0xF1, 0x71, 0xD8, 0x31, 0x15,
    0x04, 0xC7, 0x23, 0xC3, 0x18, 0x96, 0x05, 0x9A, 0x07, 0x12, 0x80, 0xE2, 0xEB, 0x27, 0xB2, 0x75,
    0x09, 0x83, 0x2C, 0x1A, 0x1B, 0x6E, 0x5A, 0xA0, 0x52, 0x3B, 0xD6, 0xB3, 0x29, 0xE3, 0x2F, 0x84,
    0x53, 0xD1, 0x00, 0xED, 0x20, 0xFC, 0xB1, 0x5B, 0x6A, 0xCB, 0xBE, 0x39, 0x4A, 0x4C, 0x58, 0xCF,
    0xD0, 0xEF, 0xAA, 0xFB, 0x43, 0x4D, 0x33, 0x85, 0x45, 0xF9, 0x02, 0x7F, 0x50, 0x3C, 0x9F, 0xA8,
    0x51, 0xA3, 0x40, 0x8F, 0x92, 0x9D, 0x38, 0xF5, 0xBC, 0xB6, 0xDA, 0x21, 0x10, 0xFF, 0xF3, 0xD2,
    0xCD, 0x0C, 0x13, 0xEC, 0x5F, 0x97, 0x44, 0x17, 0xC4, 0xA7, 0x7E, 0x3D, 0x64, 0x5D, 0x19, 0x73,
    0x60, 0x81, 0x4F, 0xDC, 0x22, 0x2A, 0x90, 0x88, 0x46, 0xEE, 0xB8, 0x14, 0xDE, 0x5E, 0x0B, 0xDB,
    0xE0, 0x32, 0x3A, 0x0A, 0x49, 0x06, 0x24, 0x5C, 0xC2, 0xD3, 0xAC, 0x62, 0x91, 0x95, 0xE4, 0x79,
    0xE7, 0xC8, 0x37, 0x6D, 0x8D, 0xD5, 0x4E, 0xA9, 0x6C, 0x56, 0xF4, 0xEA, 0x65, 0x7A, 0xAE, 0x08,
    0xBA, 0x78, 0x25, 0x2E, 0x1C, 0xA6, 0xB4, 0xC6, 0xE8, 0xDD, 0x74, 0x1F, 0x4B, 0xBD, 0x8B, 0x8A,
    0x70, 0x3E, 0xB5, 0x66, 0x48, 0x03, 0xF6, 0x0E, 0x61, 0x35, 0x57, 0xB9, 0x86, 0xC1, 0x1D, 0x9E,
    0xE1, 0xF8, 0x98, 0x11, 0x69, 0xD9, 0x8E, 0x94, 0x9B, 0x1E, 0x87, 0xE9, 0xCE, 0x55, 0x28, 0xDF,
    0x8C, 0xA1, 0x89, 0x0D, 0xBF, 0xE6, 0x42, 0x68, 0x41, 0x99, 0x2D, 0x0F, 0xB0, 0x54, 0xBB, 0x16,
], dtype=np.uint8)

TARGET_BYTE = 2

What does this code do?

The imports bring in PyTorch for building and training your CNN, h5py for reading the HDF5 dataset, numpy for array operations, and matplotlib for plotting.

The SBOX array is the AES substitution box. It maps each possible byte (0-255) to a scrambled output byte. This is the nonlinear operation inside AES that creates the electromagnetic leakage your model will exploit.

TARGET_BYTE = 2 means you are attacking the 3rd byte of the encryption key. The ASCAD dataset was designed with this byte as the default target.

  • Run the cell. You should see no errors (no output means everything imported successfully).
  • Now confirm the dataset structure by creating a new cell and running this verification code:
# Verify dataset structure
f = h5py.File("ASCAD_data/ASCAD_databases/ASCAD.h5", "r")
print(f"Profiling traces shape: {np.array(f['Profiling_traces']['traces']).shape}")
print(f"Attack traces shape: {np.array(f['Attack_traces']['traces']).shape}")
print(f"Metadata fields: {list(f['Profiling_traces']['metadata'].dtype.names)}")
f.close()
print(f"\nSBOX length: {len(SBOX)}")
print(f"Target byte index: {TARGET_BYTE}")

What should I see?

The output should show profiling traces shape (50000, 700) and attack traces shape (10000, 700). This means you have 50,000 training traces and 10,000 attack traces, each containing 700 time samples of electromagnetic activity.

The metadata fields should include plaintext and key. These tell you what was being encrypted during each trace capture.

FileNotFoundError for ASCAD.h5?

This means the download or unzip did not complete correctly. Re-run the download cell from the previous substep. Make sure the file path matches ASCAD_data/ASCAD_databases/ASCAD.h5 by running !ls ASCAD_data/ASCAD_databases/ in a new cell.

help me fix FileNotFoundError when loading the ASCAD HDF5 dataset in Colab

✔️ Awesome, I've got everything!

Great. Make sure both cells are saved in your notebook. The imports and SBOX cell is the foundation your entire attack pipeline builds on.

ⓧ I'd like to double check the full code

Here is the complete code you should have in your main setup cell at this point:

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import numpy as np
import h5py
import matplotlib.pyplot as plt

# AES S-box lookup table
SBOX = np.array([
    0x63, 0x7C, 0x77, 0x7B, 0xF2, 0x6B, 0x6F, 0xC5, 0x30, 0x01, 0x67, 0x2B, 0xFE, 0xD7, 0xAB, 0x76,
    0xCA, 0x82, 0xC9, 0x7D, 0xFA, 0x59, 0x47, 0xF0, 0xAD, 0xD4, 0xA2, 0xAF, 0x9C, 0xA4, 0x72, 0xC0,
    0xB7, 0xFD, 0x93, 0x26, 0x36, 0x3F, 0xF7, 0xCC, 0x34, 0xA5, 0xE5, 0xF1, 0x71, 0xD8, 0x31, 0x15,
    0x04, 0xC7, 0x23, 0xC3, 0x18, 0x96, 0x05, 0x9A, 0x07, 0x12, 0x80, 0xE2, 0xEB, 0x27, 0xB2, 0x75,
    0x09, 0x83, 0x2C, 0x1A, 0x1B, 0x6E, 0x5A, 0xA0, 0x52, 0x3B, 0xD6, 0xB3, 0x29, 0xE3, 0x2F, 0x84,
    0x53, 0xD1, 0x00, 0xED, 0x20, 0xFC, 0xB1, 0x5B, 0x6A, 0xCB, 0xBE, 0x39, 0x4A, 0x4C, 0x58, 0xCF,
    0xD0, 0xEF, 0xAA, 0xFB, 0x43, 0x4D, 0x33, 0x85, 0x45, 0xF9, 0x02, 0x7F, 0x50, 0x3C, 0x9F, 0xA8,
    0x51, 0xA3, 0x40, 0x8F, 0x92, 0x9D, 0x38, 0xF5, 0xBC, 0xB6, 0xDA, 0x21, 0x10, 0xFF, 0xF3, 0xD2,
    0xCD, 0x0C, 0x13, 0xEC, 0x5F, 0x97, 0x44, 0x17, 0xC4, 0xA7, 0x7E, 0x3D, 0x64, 0x5D, 0x19, 0x73,
    0x60, 0x81, 0x4F, 0xDC, 0x22, 0x2A, 0x90, 0x88, 0x46, 0xEE, 0xB8, 0x14, 0xDE, 0x5E, 0x0B, 0xDB,
    0xE0, 0x32, 0x3A, 0x0A, 0x49, 0x06, 0x24, 0x5C, 0xC2, 0xD3, 0xAC, 0x62, 0x91, 0x95, 0xE4, 0x79,
    0xE7, 0xC8, 0x37, 0x6D, 0x8D, 0xD5, 0x4E, 0xA9, 0x6C, 0x56, 0xF4, 0xEA, 0x65, 0x7A, 0xAE, 0x08,
    0xBA, 0x78, 0x25, 0x2E, 0x1C, 0xA6, 0xB4, 0xC6, 0xE8, 0xDD, 0x74, 0x1F, 0x4B, 0xBD, 0x8B, 0x8A,
    0x70, 0x3E, 0xB5, 0x66, 0x48, 0x03, 0xF6, 0x0E, 0x61, 0x35, 0x57, 0xB9, 0x86, 0xC1, 0x1D, 0x9E,
    0xE1, 0xF8, 0x98, 0x11, 0x69, 0xD9, 0x8E, 0x94, 0x9B, 0x1E, 0x87, 0xE9, 0xCE, 0x55, 0x28, 0xDF,
    0x8C, 0xA1, 0x89, 0x0D, 0xBF, 0xE6, 0x42, 0x68, 0x41, 0x99, 0x2D, 0x0F, 0xB0, 0x54, 0xBB, 0x16,
], dtype=np.uint8)

TARGET_BYTE = 2

Your attack environment is fully operational. You have GPU acceleration, 60,000 electromagnetic traces, and the cryptographic constants needed to model AES leakage. Next up, you will dig into how AES actually leaks information and prepare this data for training.

Understand AES Leakage and Prepare the Data

Your attack environment is ready with the ASCAD dataset loaded and GPU connected. Now you need to understand what your neural network will actually learn from those 50,000 electromagnetic traces.

Every time the ATMega8515 microcontroller runs AES-128 encryption, the chip emits electromagnetic radiation that correlates with the data it processes. Your CNN will learn to read that radiation and reverse-engineer the secret key byte.

In this step, get ready to:
  • Understand how AES-128's first-round S-box operation leaks through electromagnetic emissions.
  • Build a PyTorch Dataset and DataLoader pipeline for the EM traces.
  • Understand the profiling attack framework that makes key recovery possible.
Understand AES-128 first-round leakage

AES-128 encrypts data in 10 rounds. The very first operation of round one is where the secret key first touches the plaintext.

How does AES leak through EM emissions?

In AES-128's first round, each plaintext byte is XORed with the corresponding key byte, then passed through the S-box (a nonlinear 256-byte lookup table). The S-box output's Hamming Weight (the number of bits set to 1) directly influences how much electromagnetic energy the chip emits at that instant.

This creates a measurable dependency between the secret key and the physical emissions. If you know the plaintext and can classify the S-box output from the EM trace, you can deduce the key byte that was used.

The leakage model for our attack is: label = SBOX[plaintext[2] XOR key[2]]. We target byte index 2 (the 3rd key byte) because it is the default target in the ASCAD dataset.

  • Visualize several EM traces to see the leakage patterns by running this code in a new cell:
# Load a few traces to visualize the electromagnetic leakage
filepath = "ASCAD_data/ASCAD_databases/ASCAD.h5"
f = h5py.File(filepath, "r")
sample_traces = np.array(f["Profiling_traces"]["traces"][:5], dtype=np.float32)
f.close()

# Plot EM traces to see where leakage patterns appear
plt.figure(figsize=(12, 4))
for i in range(5):
    plt.plot(sample_traces[i], alpha=0.6, label=f"Trace {i}")
plt.xlabel("Sample (time)")
plt.ylabel("EM Amplitude")
plt.title("Electromagnetic Traces from ATMega8515 Running AES-128")
plt.legend()
plt.grid(True)
plt.show()

What does this plot show?

Each line is a 700-sample window of electromagnetic emissions captured while the chip encrypts one plaintext block. The amplitude spikes correspond to moments when the processor is computing the S-box lookup.

The patterns vary between traces because different plaintexts produce different intermediate values. Your CNN will learn to map these amplitude patterns to the 256 possible S-box output classes.

You should see five overlapping waveforms with similar overall shape but visible amplitude differences at certain sample positions.

Not seeing a plot?

Make sure you ran the Step 1 cell that downloads and unzips the ASCAD dataset first. The file path ASCAD_data/ASCAD_databases/ASCAD.h5 must exist in your Colab runtime.

Help me debug why my matplotlib plot isn't showing

Build the data loading and Dataset pipeline

Now that you understand what the traces represent, you need to build the code that loads, normalizes, and serves them to your neural network during training.

  • Define the load_ascad() function that extracts profiling and attack data from the HDF5 file by adding this code in a new cell:
# --- Data Loading ---

def load_ascad(filepath):
    """Load ASCAD HDF5 file and return profiling/attack data."""
    f = h5py.File(filepath, "r")
    profiling_traces = np.array(f["Profiling_traces"]["traces"][:], dtype=np.float32)
    profiling_labels = np.array(f["Profiling_traces"]["labels"][:], dtype=np.int64)
    profiling_metadata = f["Profiling_traces"]["metadata"][:]
    attack_traces = np.array(f["Attack_traces"]["traces"][:], dtype=np.float32)
    attack_labels = np.array(f["Attack_traces"]["labels"][:], dtype=np.int64)
    attack_metadata = f["Attack_traces"]["metadata"][:]
    f.close()
    return (profiling_traces, profiling_labels, profiling_metadata,
            attack_traces, attack_labels, attack_metadata)

What does this function do?

The HDF5 file contains two top-level groups: Profiling_traces (50,000 traces where the key is known) and Attack_traces (10,000 traces where you will pretend the key is unknown).

Each group contains three datasets: traces (the raw EM measurements), labels (the S-box output class for each trace), and metadata (plaintexts and keys needed for the attack phase).

  • Define the normalize_traces() function that scales each trace to the [-1, 1] range by adding this code below:
def normalize_traces(traces):
    """Normalize traces to [-1, 1] range."""
    min_val = traces.min(axis=1, keepdims=True)
    max_val = traces.max(axis=1, keepdims=True)
    return 2.0 * (traces - min_val) / (max_val - min_val + 1e-8) - 1.0

Why normalize to [-1, 1]?

Raw EM traces have arbitrary amplitude ranges that vary between captures. Per-trace normalization ensures each trace occupies the same value range, which helps the CNN converge faster and prevents traces with larger amplitudes from dominating training.

The 1e-8 prevents division by zero for any constant-value traces.

  • Define the ASCADDataset class that wraps traces and labels into a PyTorch Dataset by adding this code below:
class ASCADDataset(Dataset):
    """PyTorch Dataset for ASCAD traces."""
    def __init__(self, traces, labels):
        self.traces = torch.from_numpy(traces).float()
        self.labels = torch.from_numpy(labels).long()

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        return self.traces[idx], self.labels[idx]

What does this Dataset class do?

PyTorch's Dataset class provides a standard interface for loading training samples. Each call to __getitem__ returns one trace (a 700-element tensor) and its label (an integer from 0 to 255 representing the S-box output class).

The labels use the Identity leakage model: the S-box output value itself is the class label. This gives 256 possible classes, one for each byte value the S-box can produce.

  • Use your new functions to load the data, normalize it, and create DataLoaders by running this code in a new cell:
# Load all data using our function
filepath = "ASCAD_data/ASCAD_databases/ASCAD.h5"
(profiling_traces, profiling_labels, profiling_metadata,
 attack_traces, attack_labels, attack_metadata) = load_ascad(filepath)

print(f"Profiling traces shape: {profiling_traces.shape}")
print(f"Attack traces shape: {attack_traces.shape}")

# Normalize traces to [-1, 1]
profiling_traces = normalize_traces(profiling_traces)
attack_traces = normalize_traces(attack_traces)

# Split profiling data: 45,000 training + 5,000 validation
train_dataset = ASCADDataset(profiling_traces[:45000], profiling_labels[:45000])
val_dataset = ASCADDataset(profiling_traces[45000:], profiling_labels[45000:])

# Create DataLoaders with batch size 256
train_loader = DataLoader(train_dataset, batch_size=256, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=256, shuffle=False)

print(f"Training samples: {len(train_dataset)}")
print(f"Validation samples: {len(val_dataset)}")
print(f"Number of classes: 256 (S-box output values)")

What does this code do?

  • load_ascad() reads all 50,000 profiling traces and 10,000 attack traces from the HDF5 file.
  • normalize_traces() scales each trace independently to [-1, 1].
  • The 45,000/5,000 split gives you a training set and a validation set to monitor overfitting.
  • DataLoaders serve batches of 256 traces to the CNN during training. Shuffling the training loader ensures the model sees traces in random order each epoch.

You should see output confirming 45,000 training samples and 5,000 validation samples.

Seeing a memory error?

Loading 50,000 traces as float32 requires about 140 MB of RAM. If your Colab session is low on memory, make sure you are not holding duplicate copies of the data from earlier cells. Restart the runtime and run cells in order.

Help me fix this memory error in Google Colab

✔️ Awesome, I've got everything!

Great. Double check you have all three function and class definitions saved in your notebook cells.

ⓧ I'd like to double check the full code

Here is the full cumulative code through this step. Your notebook cells should contain all of this:

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import numpy as np
import h5py
import matplotlib.pyplot as plt

# AES S-box lookup table
SBOX = np.array([
    0x63, 0x7C, 0x77, 0x7B, 0xF2, 0x6B, 0x6F, 0xC5, 0x30, 0x01, 0x67, 0x2B, 0xFE, 0xD7, 0xAB, 0x76,
    0xCA, 0x82, 0xC9, 0x7D, 0xFA, 0x59, 0x47, 0xF0, 0xAD, 0xD4, 0xA2, 0xAF, 0x9C, 0xA4, 0x72, 0xC0,
    0xB7, 0xFD, 0x93, 0x26, 0x36, 0x3F, 0xF7, 0xCC, 0x34, 0xA5, 0xE5, 0xF1, 0x71, 0xD8, 0x31, 0x15,
    0x04, 0xC7, 0x23, 0xC3, 0x18, 0x96, 0x05, 0x9A, 0x07, 0x12, 0x80, 0xE2, 0xEB, 0x27, 0xB2, 0x75,
    0x09, 0x83, 0x2C, 0x1A, 0x1B, 0x6E, 0x5A, 0xA0, 0x52, 0x3B, 0xD6, 0xB3, 0x29, 0xE3, 0x2F, 0x84,
    0x53, 0xD1, 0x00, 0xED, 0x20, 0xFC, 0xB1, 0x5B, 0x6A, 0xCB, 0xBE, 0x39, 0x4A, 0x4C, 0x58, 0xCF,
    0xD0, 0xEF, 0xAA, 0xFB, 0x43, 0x4D, 0x33, 0x85, 0x45, 0xF9, 0x02, 0x7F, 0x50, 0x3C, 0x9F, 0xA8,
    0x51, 0xA3, 0x40, 0x8F, 0x92, 0x9D, 0x38, 0xF5, 0xBC, 0xB6, 0xDA, 0x21, 0x10, 0xFF, 0xF3, 0xD2,
    0xCD, 0x0C, 0x13, 0xEC, 0x5F, 0x97, 0x44, 0x17, 0xC4, 0xA7, 0x7E, 0x3D, 0x64, 0x5D, 0x19, 0x73,
    0x60, 0x81, 0x4F, 0xDC, 0x22, 0x2A, 0x90, 0x88, 0x46, 0xEE, 0xB8, 0x14, 0xDE, 0x5E, 0x0B, 0xDB,
    0xE0, 0x32, 0x3A, 0x0A, 0x49, 0x06, 0x24, 0x5C, 0xC2, 0xD3, 0xAC, 0x62, 0x91, 0x95, 0xE4, 0x79,
    0xE7, 0xC8, 0x37, 0x6D, 0x8D, 0xD5, 0x4E, 0xA9, 0x6C, 0x56, 0xF4, 0xEA, 0x65, 0x7A, 0xAE, 0x08,
    0xBA, 0x78, 0x25, 0x2E, 0x1C, 0xA6, 0xB4, 0xC6, 0xE8, 0xDD, 0x74, 0x1F, 0x4B, 0xBD, 0x8B, 0x8A,
    0x70, 0x3E, 0xB5, 0x66, 0x48, 0x03, 0xF6, 0x0E, 0x61, 0x35, 0x57, 0xB9, 0x86, 0xC1, 0x1D, 0x9E,
    0xE1, 0xF8, 0x98, 0x11, 0x69, 0xD9, 0x8E, 0x94, 0x9B, 0x1E, 0x87, 0xE9, 0xCE, 0x55, 0x28, 0xDF,
    0x8C, 0xA1, 0x89, 0x0D, 0xBF, 0xE6, 0x42, 0x68, 0x41, 0x99, 0x2D, 0x0F, 0xB0, 0x54, 0xBB, 0x16,
], dtype=np.uint8)

TARGET_BYTE = 2


# --- Data Loading ---

def load_ascad(filepath):
    """Load ASCAD HDF5 file and return profiling/attack data."""
    f = h5py.File(filepath, "r")
    profiling_traces = np.array(f["Profiling_traces"]["traces"][:], dtype=np.float32)
    profiling_labels = np.array(f["Profiling_traces"]["labels"][:], dtype=np.int64)
    profiling_metadata = f["Profiling_traces"]["metadata"][:]
    attack_traces = np.array(f["Attack_traces"]["traces"][:], dtype=np.float32)
    attack_labels = np.array(f["Attack_traces"]["labels"][:], dtype=np.int64)
    attack_metadata = f["Attack_traces"]["metadata"][:]
    f.close()
    return (profiling_traces, profiling_labels, profiling_metadata,
            attack_traces, attack_labels, attack_metadata)


def normalize_traces(traces):
    """Normalize traces to [-1, 1] range."""
    min_val = traces.min(axis=1, keepdims=True)
    max_val = traces.max(axis=1, keepdims=True)
    return 2.0 * (traces - min_val) / (max_val - min_val + 1e-8) - 1.0


class ASCADDataset(Dataset):
    """PyTorch Dataset for ASCAD traces."""
    def __init__(self, traces, labels):
        self.traces = torch.from_numpy(traces).float()
        self.labels = torch.from_numpy(labels).long()

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        return self.traces[idx], self.labels[idx]
Understand the profiling attack framework

Before you build and train the CNN, you need to understand the two-phase structure that makes this entire attack possible.

What is a profiling attack?

A profiling attack has two distinct phases. In the profiling phase, the attacker has a clone device where the key is known. They record thousands of traces, compute the correct S-box output label for each trace (because they know the key), and train a classifier to map traces to labels.

In the attack phase, the attacker captures new traces from the target device where the key is unknown. They feed these traces through the trained model to get probability distributions over the 256 possible S-box outputs. Combined with the known plaintexts, these probabilities let them score every possible key byte candidate and identify the correct one.

This is exactly what your ASCAD dataset provides. The 45,000 profiling traces (with known key) train your CNN. The 10,000 attack traces (where you pretend the key is unknown) test whether your model can recover it.

Your data pipeline is ready. Next up, you will build the CNN architecture that learns to classify these 700-sample traces into 256 S-box output classes.

Build the CNN Model

Your data pipeline is ready. 45,000 normalized electromagnetic traces are loaded into DataLoaders, each labeled with one of 256 possible S-box output classes.

Now you need a model that can learn the relationship between the shape of an EM trace and the secret S-box output that produced it. In this step, you will build a Convolutional Neural Network specifically designed for 1D time-series signals like power or electromagnetic traces.

In this step, get ready to:
  • Define a 4-block CNN architecture tailored for side-channel analysis.
  • Instantiate the model on your GPU and verify its output shape.
Define the SCA_CNN architecture

A standard image CNN operates on 2D grids of pixels. Your electromagnetic traces are 1D sequences of 700 time samples. You need Conv1d layers that slide a kernel along the time axis to detect patterns in the signal.

  • In your Colab notebook, create a new code cell below your data loading code.
  • Define the SCA_CNN model class by pasting this code:
# --- Model Definition ---

class SCA_CNN(nn.Module):
    """CNN for Side-Channel Analysis on 700-sample traces."""
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            # Conv block 1
            nn.Conv1d(1, 64, kernel_size=11, padding=5),
            nn.BatchNorm1d(64),
            nn.ReLU(),
            nn.AvgPool1d(2),
            # Conv block 2
            nn.Conv1d(64, 128, kernel_size=11, padding=5),
            nn.BatchNorm1d(128),
            nn.ReLU(),
            nn.AvgPool1d(2),
            # Conv block 3
            nn.Conv1d(128, 256, kernel_size=11, padding=5),
            nn.BatchNorm1d(256),
            nn.ReLU(),
            nn.AvgPool1d(2),
            # Conv block 4
            nn.Conv1d(256, 512, kernel_size=11, padding=5),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.AvgPool1d(2),
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(512 * 43, 4096),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(4096, 256),
        )

    def forward(self, x):
        # x shape: (batch, 700) -> (batch, 1, 700)
        x = x.unsqueeze(1)
        x = self.features(x)
        x = self.classifier(x)
        return x

Why this architecture works for side-channel analysis

Each convolutional block has four layers working together:

  • Conv1d(kernel_size=11, padding=5) slides an 11-sample window across the trace. Large kernels capture the temporal patterns where leakage occurs.
  • BatchNorm1d normalizes activations to handle amplitude variation between different traces captured at different times.
  • ReLU introduces nonlinearity so the network can learn complex leakage patterns.
  • AvgPool1d(2) halves the sequence length at each block. After 4 pools, 700 samples become 43, reducing noise while preserving leakage features.

The filter progression (64, 128, 256, 512) lets deeper layers combine simpler patterns into complex ones. The classifier head maps the final 512 x 43 feature map down to 256 output classes (one per possible S-box value).

Why average pooling instead of max pooling?

In image tasks, max pooling keeps the strongest activation. In side-channel analysis, the leakage signal is subtle and spread across multiple time samples. Average pooling preserves this distributed information instead of discarding it.

  • Run the cell to register the class definition.

No output appears. The class is now defined and ready to be instantiated.

Instantiate and verify the model

With the architecture defined, you need to create an instance of the model, move it onto the GPU, and confirm it produces the correct output shape.

  • Create a new code cell and instantiate the model by running this code:
# Build model and move to GPU
model = SCA_CNN().to(device)
total_params = sum(p.numel() for p in model.parameters())
print(f"Total parameters: {total_params:,}")

# Verify output shape with a dummy forward pass
dummy_input = torch.randn(4, 700).to(device)
dummy_output = model(dummy_input)
print(f"Input shape: {dummy_input.shape}")
print(f"Output shape: {dummy_output.shape}")

What does this code do?

  • SCA_CNN().to(device) creates the model and moves all its weights onto your T4 GPU for fast computation.
  • The parameter count tells you how many learnable weights the network has. A larger model captures more complexity but also needs more data to avoid overfitting.
  • The dummy forward pass sends a fake batch of 4 traces (each 700 samples) through the model. If the output shape is (4, 256), the architecture is wired correctly.
  • Run the cell. You should see the total parameter count printed, followed by input shape torch.Size([4, 700]) and output shape torch.Size([4, 256]).

The output shape confirms that for each trace in a batch, the model produces 256 probability scores. One for each possible S-box output value.

Seeing a CUDA error?

  • Make sure your runtime is connected to a GPU. Go to **Runtime** > **Change runtime type** and confirm **T4 GPU** is selected.
  • If you restarted your runtime, you need to re-run all previous cells (including the one that sets device = torch.device("cuda" ...)) before running this cell.
  • help me debug a CUDA error when instantiating my SCA_CNN model in Google Colab

✔️ Awesome, I've got everything!

Great. Double check you have run both cells and can see the correct output shape.

ⓧ I'd like to double check the full code

Here is the full cumulative code up to this point:

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import numpy as np
import h5py
import matplotlib.pyplot as plt

# AES S-box lookup table
SBOX = np.array([
    0x63, 0x7C, 0x77, 0x7B, 0xF2, 0x6B, 0x6F, 0xC5, 0x30, 0x01, 0x67, 0x2B, 0xFE, 0xD7, 0xAB, 0x76,
    0xCA, 0x82, 0xC9, 0x7D, 0xFA, 0x59, 0x47, 0xF0, 0xAD, 0xD4, 0xA2, 0xAF, 0x9C, 0xA4, 0x72, 0xC0,
    0xB7, 0xFD, 0x93, 0x26, 0x36, 0x3F, 0xF7, 0xCC, 0x34, 0xA5, 0xE5, 0xF1, 0x71, 0xD8, 0x31, 0x15,
    0x04, 0xC7, 0x23, 0xC3, 0x18, 0x96, 0x05, 0x9A, 0x07, 0x12, 0x80, 0xE2, 0xEB, 0x27, 0xB2, 0x75,
    0x09, 0x83, 0x2C, 0x1A, 0x1B, 0x6E, 0x5A, 0xA0, 0x52, 0x3B, 0xD6, 0xB3, 0x29, 0xE3, 0x2F, 0x84,
    0x53, 0xD1, 0x00, 0xED, 0x20, 0xFC, 0xB1, 0x5B, 0x6A, 0xCB, 0xBE, 0x39, 0x4A, 0x4C, 0x58, 0xCF,
    0xD0, 0xEF, 0xAA, 0xFB, 0x43, 0x4D, 0x33, 0x85, 0x45, 0xF9, 0x02, 0x7F, 0x50, 0x3C, 0x9F, 0xA8,
    0x51, 0xA3, 0x40, 0x8F, 0x92, 0x9D, 0x38, 0xF5, 0xBC, 0xB6, 0xDA, 0x21, 0x10, 0xFF, 0xF3, 0xD2,
    0xCD, 0x0C, 0x13, 0xEC, 0x5F, 0x97, 0x44, 0x17, 0xC4, 0xA7, 0x7E, 0x3D, 0x64, 0x5D, 0x19, 0x73,
    0x60, 0x81, 0x4F, 0xDC, 0x22, 0x2A, 0x90, 0x88, 0x46, 0xEE, 0xB8, 0x14, 0xDE, 0x5E, 0x0B, 0xDB,
    0xE0, 0x32, 0x3A, 0x0A, 0x49, 0x06, 0x24, 0x5C, 0xC2, 0xD3, 0xAC, 0x62, 0x91, 0x95, 0xE4, 0x79,
    0xE7, 0xC8, 0x37, 0x6D, 0x8D, 0xD5, 0x4E, 0xA9, 0x6C, 0x56, 0xF4, 0xEA, 0x65, 0x7A, 0xAE, 0x08,
    0xBA, 0x78, 0x25, 0x2E, 0x1C, 0xA6, 0xB4, 0xC6, 0xE8, 0xDD, 0x74, 0x1F, 0x4B, 0xBD, 0x8B, 0x8A,
    0x70, 0x3E, 0xB5, 0x66, 0x48, 0x03, 0xF6, 0x0E, 0x61, 0x35, 0x57, 0xB9, 0x86, 0xC1, 0x1D, 0x9E,
    0xE1, 0xF8, 0x98, 0x11, 0x69, 0xD9, 0x8E, 0x94, 0x9B, 0x1E, 0x87, 0xE9, 0xCE, 0x55, 0x28, 0xDF,
    0x8C, 0xA1, 0x89, 0x0D, 0xBF, 0xE6, 0x42, 0x68, 0x41, 0x99, 0x2D, 0x0F, 0xB0, 0x54, 0xBB, 0x16,
], dtype=np.uint8)

TARGET_BYTE = 2


# --- Data Loading ---

def load_ascad(filepath):
    """Load ASCAD HDF5 file and return profiling/attack data."""
    f = h5py.File(filepath, "r")
    profiling_traces = np.array(f["Profiling_traces"]["traces"][:], dtype=np.float32)
    profiling_labels = np.array(f["Profiling_traces"]["labels"][:], dtype=np.int64)
    profiling_metadata = f["Profiling_traces"]["metadata"][:]
    attack_traces = np.array(f["Attack_traces"]["traces"][:], dtype=np.float32)
    attack_labels = np.array(f["Attack_traces"]["labels"][:], dtype=np.int64)
    attack_metadata = f["Attack_traces"]["metadata"][:]
    f.close()
    return (profiling_traces, profiling_labels, profiling_metadata,
            attack_traces, attack_labels, attack_metadata)


def normalize_traces(traces):
    """Normalize traces to [-1, 1] range."""
    min_val = traces.min(axis=1, keepdims=True)
    max_val = traces.max(axis=1, keepdims=True)
    return 2.0 * (traces - min_val) / (max_val - min_val + 1e-8) - 1.0


class ASCADDataset(Dataset):
    """PyTorch Dataset for ASCAD traces."""
    def __init__(self, traces, labels):
        self.traces = torch.from_numpy(traces).float()
        self.labels = torch.from_numpy(labels).long()

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        return self.traces[idx], self.labels[idx]


# --- Model Definition ---

class SCA_CNN(nn.Module):
    """CNN for Side-Channel Analysis on 700-sample traces."""
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            # Conv block 1
            nn.Conv1d(1, 64, kernel_size=11, padding=5),
            nn.BatchNorm1d(64),
            nn.ReLU(),
            nn.AvgPool1d(2),
            # Conv block 2
            nn.Conv1d(64, 128, kernel_size=11, padding=5),
            nn.BatchNorm1d(128),
            nn.ReLU(),
            nn.AvgPool1d(2),
            # Conv block 3
            nn.Conv1d(128, 256, kernel_size=11, padding=5),
            nn.BatchNorm1d(256),
            nn.ReLU(),
            nn.AvgPool1d(2),
            # Conv block 4
            nn.Conv1d(256, 512, kernel_size=11, padding=5),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.AvgPool1d(2),
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(512 * 43, 4096),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(4096, 256),
        )

    def forward(self, x):
        # x shape: (batch, 700) -> (batch, 1, 700)
        x = x.unsqueeze(1)
        x = self.features(x)
        x = self.classifier(x)
        return x

Your CNN is built and verified. Next up, you will train this model on 45,000 electromagnetic traces and watch it learn to classify the secret S-box output from raw signal data.

Train the Model

In the previous step, you built a CNN with over 90 million parameters designed to detect electromagnetic leakage patterns in 700-sample traces.

Now you need to teach it what leakage looks like. Training a neural network for side-channel analysis requires careful tuning. The learning rate schedule matters especially because SCA datasets contain subtle signal buried in noise, and aggressive learning can overshoot those patterns entirely.

In this step, get ready to:
  • Define the training function with CrossEntropyLoss, Adam optimizer, and OneCycleLR scheduler.
  • Train the CNN for 50 epochs on your profiling traces.
  • Plot training loss and validation accuracy curves to evaluate performance.
Define the training function

Side-channel traces contain extremely subtle signals. A standard constant learning rate often either overshoots the faint leakage patterns or converges too slowly across 256 classes.

The OneCycleLR scheduler solves this by warming up the learning rate to a peak, then annealing it down. This helps the model explore broadly early in training, then fine-tune once it locks onto the leakage signal.

  • Add the following train_model() function to a new cell below your model definition:
# --- Training ---

def train_model(model, train_loader, val_loader, device, epochs=50, max_lr=5e-3):
    """Train the SCA CNN model."""
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=max_lr)
    scheduler = torch.optim.lr_scheduler.OneCycleLR(
        optimizer, max_lr=max_lr, epochs=epochs, steps_per_epoch=len(train_loader)
    )

    train_losses = []
    val_accuracies = []

    for epoch in range(epochs):
        model.train()
        epoch_loss = 0.0
        for traces_batch, labels_batch in train_loader:
            traces_batch = traces_batch.to(device)
            labels_batch = labels_batch.to(device)

            optimizer.zero_grad()
            outputs = model(traces_batch)
            loss = criterion(outputs, labels_batch)
            loss.backward()
            optimizer.step()
            scheduler.step()
            epoch_loss += loss.item()

        avg_loss = epoch_loss / len(train_loader)
        train_losses.append(avg_loss)

        # Validation
        model.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for traces_batch, labels_batch in val_loader:
                traces_batch = traces_batch.to(device)
                labels_batch = labels_batch.to(device)
                outputs = model(traces_batch)
                _, predicted = torch.max(outputs, 1)
                total += labels_batch.size(0)
                correct += (predicted == labels_batch).sum().item()

        val_acc = correct / total
        val_accuracies.append(val_acc)

        if (epoch + 1) % 10 == 0:
            print(f"Epoch {epoch+1}/{epochs} - Loss: {avg_loss:.4f} - Val Acc: {val_acc:.4f}")

    return train_losses, val_accuracies

What does this code do?

This function orchestrates the entire training process for your side-channel attack model.

  • CrossEntropyLoss is the standard loss for multi-class classification. Each of the 256 possible S-box outputs is a class.
  • Adam optimizer handles adaptive learning rates per parameter, which helps when different convolutional filters converge at different speeds.
  • OneCycleLR ramps the learning rate up from max_lr/25 to 5e-3, then anneals it near zero. This avoids overfitting on noisy SCA data.
  • The validation loop runs after every epoch with torch.no_grad() to measure accuracy without affecting gradients.

Why OneCycleLR instead of constant learning rate?

SCA data is fundamentally different from image classification. The signal-to-noise ratio is very low, meaning the leakage information is a tiny fraction of what the model sees.

OneCycleLR's warm-up phase lets the model explore broadly before committing to a solution. The cool-down phase then refines the subtle patterns without overshooting. Research on CNN-based SCA shows this schedule consistently outperforms constant rates on the ASCAD benchmark.

  • Run the cell to define the function.

You won't see any output yet. The function is defined but not called.

Run the training loop

With the function defined, you can now train your model on the 45,000 profiling traces. The T4 GPU will process each of the 50 epochs in a few seconds.

  • Launch the training by running this in a new cell:
# Train the model on profiling traces
print("\nTraining...")
train_losses, val_accuracies = train_model(model, train_loader, val_loader, device)

What's happening during training?

Each epoch passes all 45,000 profiling traces through the CNN in batches of 256. The model learns to associate electromagnetic amplitude patterns at specific time samples with the 256 possible S-box output values.

Every 10 epochs, the function prints a progress update showing training loss and validation accuracy. Training takes approximately 5-8 minutes on a T4 GPU.

  • Wait for training to complete. You should see progress printed at epochs 10, 20, 30, 40, and 50.
  • By epoch 50, your validation accuracy should be in the range of 50-70%.

Training seems stuck or accuracy stays near 0%?

  • Check that your model is on the GPU. Run next(model.parameters()).device and confirm it says cuda:0.
  • Make sure your traces were normalized before creating the DataLoaders. Unnormalized traces can cause gradient issues.
  • If the Colab runtime disconnected, you need to re-run all previous cells from the top.

My model isn't learning from the electromagnetic traces. Help me debug the training loop.

Plot training curves

Validation accuracy of 50-70% might sound low for a classifier. But in side-channel analysis, the model does not need to be perfectly accurate on individual traces. What matters is that probabilities accumulate correctly across many traces during the key recovery attack.

  • Visualize how training progressed by running this in a new cell:
# Plot training loss and validation accuracy
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
ax1.plot(train_losses)
ax1.set_xlabel("Epoch")
ax1.set_ylabel("Loss")
ax1.set_title("Training Loss")
ax2.plot(val_accuracies)
ax2.set_xlabel("Epoch")
ax2.set_ylabel("Accuracy")
ax2.set_title("Validation Accuracy")
plt.tight_layout()
plt.show()

What should I see?

You should see two side-by-side plots. The left plot (Training Loss) should show a curve that decreases rapidly in the first 10-15 epochs, then flattens out. The right plot (Validation Accuracy) should climb and stabilize around 0.5-0.7.

Even 50% accuracy means the model is extracting meaningful leakage from the electromagnetic traces. In the next step, you will see how accumulating log-probabilities across hundreds of attack traces pushes the correct key to Rank 0.

  • Confirm you see a decreasing loss curve on the left and a rising accuracy curve on the right.

Plots not showing?

  • Make sure you ran the training cell first. The train_losses and val_accuracies variables must exist before you can plot them.
  • If you see a NameError, the runtime may have disconnected during training. Re-run all cells from the beginning.

My training curves look unusual or the plots are empty. Help me troubleshoot.

✔️ Awesome, I've got everything!

Great. Make sure your train_model() function definition cell is saved. The training results and plots exist in memory for the next step.

ⓧ I'd like to double check the full code

Here is the complete code you should have defined across all your cells so far:

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import numpy as np
import h5py
import matplotlib.pyplot as plt

# AES S-box lookup table
SBOX = np.array([
    0x63, 0x7C, 0x77, 0x7B, 0xF2, 0x6B, 0x6F, 0xC5, 0x30, 0x01, 0x67, 0x2B, 0xFE, 0xD7, 0xAB, 0x76,
    0xCA, 0x82, 0xC9, 0x7D, 0xFA, 0x59, 0x47, 0xF0, 0xAD, 0xD4, 0xA2, 0xAF, 0x9C, 0xA4, 0x72, 0xC0,
    0xB7, 0xFD, 0x93, 0x26, 0x36, 0x3F, 0xF7, 0xCC, 0x34, 0xA5, 0xE5, 0xF1, 0x71, 0xD8, 0x31, 0x15,
    0x04, 0xC7, 0x23, 0xC3, 0x18, 0x96, 0x05, 0x9A, 0x07, 0x12, 0x80, 0xE2, 0xEB, 0x27, 0xB2, 0x75,
    0x09, 0x83, 0x2C, 0x1A, 0x1B, 0x6E, 0x5A, 0xA0, 0x52, 0x3B, 0xD6, 0xB3, 0x29, 0xE3, 0x2F, 0x84,
    0x53, 0xD1, 0x00, 0xED, 0x20, 0xFC, 0xB1, 0x5B, 0x6A, 0xCB, 0xBE, 0x39, 0x4A, 0x4C, 0x58, 0xCF,
    0xD0, 0xEF, 0xAA, 0xFB, 0x43, 0x4D, 0x33, 0x85, 0x45, 0xF9, 0x02, 0x7F, 0x50, 0x3C, 0x9F, 0xA8,
    0x51, 0xA3, 0x40, 0x8F, 0x92, 0x9D, 0x38, 0xF5, 0xBC, 0xB6, 0xDA, 0x21, 0x10, 0xFF, 0xF3, 0xD2,
    0xCD, 0x0C, 0x13, 0xEC, 0x5F, 0x97, 0x44, 0x17, 0xC4, 0xA7, 0x7E, 0x3D, 0x64, 0x5D, 0x19, 0x73,
    0x60, 0x81, 0x4F, 0xDC, 0x22, 0x2A, 0x90, 0x88, 0x46, 0xEE, 0xB8, 0x14, 0xDE, 0x5E, 0x0B, 0xDB,
    0xE0, 0x32, 0x3A, 0x0A, 0x49, 0x06, 0x24, 0x5C, 0xC2, 0xD3, 0xAC, 0x62, 0x91, 0x95, 0xE4, 0x79,
    0xE7, 0xC8, 0x37, 0x6D, 0x8D, 0xD5, 0x4E, 0xA9, 0x6C, 0x56, 0xF4, 0xEA, 0x65, 0x7A, 0xAE, 0x08,
    0xBA, 0x78, 0x25, 0x2E, 0x1C, 0xA6, 0xB4, 0xC6, 0xE8, 0xDD, 0x74, 0x1F, 0x4B, 0xBD, 0x8B, 0x8A,
    0x70, 0x3E, 0xB5, 0x66, 0x48, 0x03, 0xF6, 0x0E, 0x61, 0x35, 0x57, 0xB9, 0x86, 0xC1, 0x1D, 0x9E,
    0xE1, 0xF8, 0x98, 0x11, 0x69, 0xD9, 0x8E, 0x94, 0x9B, 0x1E, 0x87, 0xE9, 0xCE, 0x55, 0x28, 0xDF,
    0x8C, 0xA1, 0x89, 0x0D, 0xBF, 0xE6, 0x42, 0x68, 0x41, 0x99, 0x2D, 0x0F, 0xB0, 0x54, 0xBB, 0x16,
], dtype=np.uint8)

TARGET_BYTE = 2


# --- Data Loading ---

def load_ascad(filepath):
    """Load ASCAD HDF5 file and return profiling/attack data."""
    f = h5py.File(filepath, "r")
    profiling_traces = np.array(f["Profiling_traces"]["traces"][:], dtype=np.float32)
    profiling_labels = np.array(f["Profiling_traces"]["labels"][:], dtype=np.int64)
    profiling_metadata = f["Profiling_traces"]["metadata"][:]
    attack_traces = np.array(f["Attack_traces"]["traces"][:], dtype=np.float32)
    attack_labels = np.array(f["Attack_traces"]["labels"][:], dtype=np.int64)
    attack_metadata = f["Attack_traces"]["metadata"][:]
    f.close()
    return (profiling_traces, profiling_labels, profiling_metadata,
            attack_traces, attack_labels, attack_metadata)


def normalize_traces(traces):
    """Normalize traces to [-1, 1] range."""
    min_val = traces.min(axis=1, keepdims=True)
    max_val = traces.max(axis=1, keepdims=True)
    return 2.0 * (traces - min_val) / (max_val - min_val + 1e-8) - 1.0


class ASCADDataset(Dataset):
    """PyTorch Dataset for ASCAD traces."""
    def __init__(self, traces, labels):
        self.traces = torch.from_numpy(traces).float()
        self.labels = torch.from_numpy(labels).long()

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        return self.traces[idx], self.labels[idx]


# --- Model Definition ---

class SCA_CNN(nn.Module):
    """CNN for Side-Channel Analysis on 700-sample traces."""
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            # Conv block 1
            nn.Conv1d(1, 64, kernel_size=11, padding=5),
            nn.BatchNorm1d(64),
            nn.ReLU(),
            nn.AvgPool1d(2),
            # Conv block 2
            nn.Conv1d(64, 128, kernel_size=11, padding=5),
            nn.BatchNorm1d(128),
            nn.ReLU(),
            nn.AvgPool1d(2),
            # Conv block 3
            nn.Conv1d(128, 256, kernel_size=11, padding=5),
            nn.BatchNorm1d(256),
            nn.ReLU(),
            nn.AvgPool1d(2),
            # Conv block 4
            nn.Conv1d(256, 512, kernel_size=11, padding=5),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.AvgPool1d(2),
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(512 * 43, 4096),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(4096, 256),
        )

    def forward(self, x):
        # x shape: (batch, 700) -> (batch, 1, 700)
        x = x.unsqueeze(1)
        x = self.features(x)
        x = self.classifier(x)
        return x


# --- Training ---

def train_model(model, train_loader, val_loader, device, epochs=50, max_lr=5e-3):
    """Train the SCA CNN model."""
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=max_lr)
    scheduler = torch.optim.lr_scheduler.OneCycleLR(
        optimizer, max_lr=max_lr, epochs=epochs, steps_per_epoch=len(train_loader)
    )

    train_losses = []
    val_accuracies = []

    for epoch in range(epochs):
        model.train()
        epoch_loss = 0.0
        for traces_batch, labels_batch in train_loader:
            traces_batch = traces_batch.to(device)
            labels_batch = labels_batch.to(device)

            optimizer.zero_grad()
            outputs = model(traces_batch)
            loss = criterion(outputs, labels_batch)
            loss.backward()
            optimizer.step()
            scheduler.step()
            epoch_loss += loss.item()

        avg_loss = epoch_loss / len(train_loader)
        train_losses.append(avg_loss)

        # Validation
        model.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for traces_batch, labels_batch in val_loader:
                traces_batch = traces_batch.to(device)
                labels_batch = labels_batch.to(device)
                outputs = model(traces_batch)
                _, predicted = torch.max(outputs, 1)
                total += labels_batch.size(0)
                correct += (predicted == labels_batch).sum().item()

        val_acc = correct / total
        val_accuracies.append(val_acc)

        if (epoch + 1) % 10 == 0:
            print(f"Epoch {epoch+1}/{epochs} - Loss: {avg_loss:.4f} - Val Acc: {val_acc:.4f}")

    return train_losses, val_accuracies

Your model has learned to detect electromagnetic leakage patterns in the profiling traces. Next up, you will turn those learned patterns into a weapon by running the actual key recovery attack on unseen traces.

Execute the Key Recovery Attack

Your CNN has learned to classify S-box outputs from electromagnetic traces. Now it's time to put that knowledge to work recovering the actual secret key byte.

The attack works by trying all 256 possible key byte values, scoring each guess against the model's predictions, and ranking them. The correct key should accumulate the highest probability score and rise to Rank 0.

In this step, get ready to:
  • Implement the key rank computation that scores all 256 key guesses.
  • Plot key rank convergence to see how quickly the correct key emerges.
  • Confirm the recovered key byte against the ground truth from metadata.
Implement the key rank function

During training, the attacker knew the key and taught the model to recognize leakage patterns. During the attack phase, the key is unknown. The model's predictions must be combined with known plaintexts to deduce which key byte value is correct.

How does key rank work?

For each of the 256 possible key byte guesses k, the attack computes what the S-box output would have been: Sbox[plaintext[2] XOR k]. It then looks up the model's log-probability for that class and accumulates it across all attack traces.

The correct key guess consistently points to the class the model is most confident about, so its accumulated score rises above all other guesses. When it has the highest score, its rank is 0.

  • Define the compute_key_rank function by adding this code in a new cell:
def compute_key_rank(model, attack_traces, attack_metadata, device, num_traces, target_byte=2):
    """Compute key rank using accumulated log-probabilities."""
    model.eval()
    correct_key = attack_metadata["key"][0][target_byte]

    # Get model predictions for the specified number of traces
    traces_tensor = torch.from_numpy(attack_traces[:num_traces]).float().to(device)

    with torch.no_grad():
        log_probs = torch.log_softmax(model(traces_tensor), dim=1).cpu().numpy()

    # For each key guess, accumulate log-probabilities
    scores = np.zeros(256)
    plaintexts = attack_metadata["plaintext"][:num_traces]

    for k_guess in range(256):
        for i in range(num_traces):
            # Hypothetical S-box output for this key guess
            hyp = SBOX[plaintexts[i][target_byte] ^ k_guess]
            scores[k_guess] += log_probs[i, hyp]

    # Rank: how many keys score higher than the correct key
    correct_score = scores[correct_key]
    rank = np.sum(scores > correct_score)
    return rank, correct_key

What does this code do?

  • correct_key = attack_metadata["key"][0][target_byte] extracts the real key byte from the metadata so we can check our answer later.
  • torch.log_softmax(model(traces_tensor), dim=1) gets the model's log-probability predictions for all 256 classes across all traces at once.
  • The nested loop tries every key guess (0 to 255) and for each trace computes SBOX[plaintext XOR k_guess] to find the hypothetical S-box output class, then accumulates the model's confidence for that class.
  • np.sum(scores > correct_score) counts how many key guesses scored higher than the correct key. Rank 0 means nothing scored higher.
  • Test the function with 200 attack traces by running this cell:
# Key Recovery Attack
print("\nExecuting key recovery attack...")
rank, correct_key = compute_key_rank(model, attack_traces, attack_metadata, device, num_traces=200)
print(f"Correct key byte: 0x{correct_key:02X}")
print(f"Key rank with 200 traces: {rank}")
  • Run the cell. You should see the correct key byte printed in hexadecimal and a key rank of 0 (or very close to 0).

Why does low accuracy still crack the key?

Your model only achieved 50-70% validation accuracy, yet it can recover the key. This works because probabilities accumulate across traces. Even a slightly-better-than-random prediction for the correct class adds up over hundreds of traces, drowning out the noise from incorrect predictions.

Key rank not reaching 0?

  • If the rank is high (above 10), your model may not have trained well. Go back to Step 4 and retrain for more epochs or confirm validation accuracy reached at least 40%.
  • Make sure you're using the normalized attack_traces (the ones passed through normalize_traces()), not the raw traces.
  • Try increasing num_traces to 500 or 1000. Some training runs need more traces to converge.

Why isn't my key rank reaching 0 even with 200 traces?

Plot key rank convergence

A single test with 200 traces confirms the attack works, but a convergence plot shows exactly how many traces are needed for the correct key to emerge as the top candidate.

  • Define the plot_key_rank_convergence function by adding this code in a new cell:
def plot_key_rank_convergence(model, attack_traces, attack_metadata, device, max_traces=2000, step=10):
    """Plot key rank vs number of attack traces."""
    trace_counts = list(range(step, max_traces + 1, step))
    ranks = []

    for n in trace_counts:
        rank, _ = compute_key_rank(model, attack_traces, attack_metadata, device, n)
        ranks.append(rank)

    plt.figure(figsize=(10, 5))
    plt.plot(trace_counts, ranks, "b-", linewidth=2)
    plt.xlabel("Number of Attack Traces")
    plt.ylabel("Key Rank")
    plt.title("Key Rank Convergence - AES Key Recovery")
    plt.axhline(y=0, color="r", linestyle="--", label="Rank 0 (Key Found)")
    plt.legend()
    plt.grid(True)
    plt.show()

    return trace_counts, ranks

What does this code do?

  • It iterates from 10 traces up to 2000 in steps of 10, computing the key rank at each point.
  • The plot shows how the key rank drops from a high initial value (random guess territory) down to 0 as more traces provide evidence.
  • The red dashed line at y=0 marks the target. Once the blue line hits it, the attack has succeeded.
  • Run the convergence plot and find the exact number of traces needed by adding this cell:
# Plot convergence
trace_counts, ranks = plot_key_rank_convergence(model, attack_traces, attack_metadata, device)

# Find minimum traces needed
for i, r in enumerate(ranks):
    if r == 0:
        print(f"\nKey recovered at Rank 0 with {trace_counts[i]} attack traces!")
        break
  • Run the cell. This will take 1-2 minutes as it evaluates the model 200 times with increasing trace counts.
  • You should see a convergence plot where the blue line drops sharply from a high rank and hits the red dashed line (Rank 0) within approximately 50-200 traces.

Plot takes too long or Colab disconnects?

  • The convergence loop calls the model 200 times. If your runtime is slow, reduce max_traces to 500 or increase step to 50.
  • If Colab disconnected, reconnect and re-run all cells from the top. The dataset will need to be re-downloaded if the runtime fully reset.

Help me speed up my key rank convergence plot in Colab.

Confirm the recovered key and interpret the results

Your output shows the correct key byte value (the ground truth from the dataset metadata) and confirms your model identified it as the top candidate. You just cracked AES-128 encryption using a neural network.

What does this mean for real-world security?

The ASCAD dataset was captured from a device running masked AES. Boolean masking is a countermeasure specifically designed to prevent side-channel attacks by randomizing intermediate values. Your CNN bypassed this protection by learning to extract leakage patterns that persist despite the masking.

This is the core methodology used by hardware security evaluators. Certification labs use profiling attacks like this to test whether cryptographic devices meet security standards before deployment in banking cards, passports, and IoT devices.

Why only one key byte?

AES-128 has 16 key bytes. You recovered byte index 2 because that is what the ASCAD dataset's extracted trace window targets. To recover the full 128-bit key, you would repeat this attack 16 times with traces captured around each byte's S-box operation. The methodology is identical for each byte.

✔️ Awesome, I've got everything!

Great. Double check you've run all cells and can see both the key rank convergence plot and the printed confirmation of Rank 0.

ⓧ I'd like to double check the full code

Here is the complete code across all cells in your notebook up to this point:

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import numpy as np
import h5py
import matplotlib.pyplot as plt

# AES S-box lookup table
SBOX = np.array([
    0x63, 0x7C, 0x77, 0x7B, 0xF2, 0x6B, 0x6F, 0xC5, 0x30, 0x01, 0x67, 0x2B, 0xFE, 0xD7, 0xAB, 0x76,
    0xCA, 0x82, 0xC9, 0x7D, 0xFA, 0x59, 0x47, 0xF0, 0xAD, 0xD4, 0xA2, 0xAF, 0x9C, 0xA4, 0x72, 0xC0,
    0xB7, 0xFD, 0x93, 0x26, 0x36, 0x3F, 0xF7, 0xCC, 0x34, 0xA5, 0xE5, 0xF1, 0x71, 0xD8, 0x31, 0x15,
    0x04, 0xC7, 0x23, 0xC3, 0x18, 0x96, 0x05, 0x9A, 0x07, 0x12, 0x80, 0xE2, 0xEB, 0x27, 0xB2, 0x75,
    0x09, 0x83, 0x2C, 0x1A, 0x1B, 0x6E, 0x5A, 0xA0, 0x52, 0x3B, 0xD6, 0xB3, 0x29, 0xE3, 0x2F, 0x84,
    0x53, 0xD1, 0x00, 0xED, 0x20, 0xFC, 0xB1, 0x5B, 0x6A, 0xCB, 0xBE, 0x39, 0x4A, 0x4C, 0x58, 0xCF,
    0xD0, 0xEF, 0xAA, 0xFB, 0x43, 0x4D, 0x33, 0x85, 0x45, 0xF9, 0x02, 0x7F, 0x50, 0x3C, 0x9F, 0xA8,
    0x51, 0xA3, 0x40, 0x8F, 0x92, 0x9D, 0x38, 0xF5, 0xBC, 0xB6, 0xDA, 0x21, 0x10, 0xFF, 0xF3, 0xD2,
    0xCD, 0x0C, 0x13, 0xEC, 0x5F, 0x97, 0x44, 0x17, 0xC4, 0xA7, 0x7E, 0x3D, 0x64, 0x5D, 0x19, 0x73,
    0x60, 0x81, 0x4F, 0xDC, 0x22, 0x2A, 0x90, 0x88, 0x46, 0xEE, 0xB8, 0x14, 0xDE, 0x5E, 0x0B, 0xDB,
    0xE0, 0x32, 0x3A, 0x0A, 0x49, 0x06, 0x24, 0x5C, 0xC2, 0xD3, 0xAC, 0x62, 0x91, 0x95, 0xE4, 0x79,
    0xE7, 0xC8, 0x37, 0x6D, 0x8D, 0xD5, 0x4E, 0xA9, 0x6C, 0x56, 0xF4, 0xEA, 0x65, 0x7A, 0xAE, 0x08,
    0xBA, 0x78, 0x25, 0x2E, 0x1C, 0xA6, 0xB4, 0xC6, 0xE8, 0xDD, 0x74, 0x1F, 0x4B, 0xBD, 0x8B, 0x8A,
    0x70, 0x3E, 0xB5, 0x66, 0x48, 0x03, 0xF6, 0x0E, 0x61, 0x35, 0x57, 0xB9, 0x86, 0xC1, 0x1D, 0x9E,
    0xE1, 0xF8, 0x98, 0x11, 0x69, 0xD9, 0x8E, 0x94, 0x9B, 0x1E, 0x87, 0xE9, 0xCE, 0x55, 0x28, 0xDF,
    0x8C, 0xA1, 0x89, 0x0D, 0xBF, 0xE6, 0x42, 0x68, 0x41, 0x99, 0x2D, 0x0F, 0xB0, 0x54, 0xBB, 0x16,
], dtype=np.uint8)

TARGET_BYTE = 2


# --- Data Loading ---

def load_ascad(filepath):
    """Load ASCAD HDF5 file and return profiling/attack data."""
    f = h5py.File(filepath, "r")
    profiling_traces = np.array(f["Profiling_traces"]["traces"][:], dtype=np.float32)
    profiling_labels = np.array(f["Profiling_traces"]["labels"][:], dtype=np.int64)
    profiling_metadata = f["Profiling_traces"]["metadata"][:]
    attack_traces = np.array(f["Attack_traces"]["traces"][:], dtype=np.float32)
    attack_labels = np.array(f["Attack_traces"]["labels"][:], dtype=np.int64)
    attack_metadata = f["Attack_traces"]["metadata"][:]
    f.close()
    return (profiling_traces, profiling_labels, profiling_metadata,
            attack_traces, attack_labels, attack_metadata)


def normalize_traces(traces):
    """Normalize traces to [-1, 1] range."""
    min_val = traces.min(axis=1, keepdims=True)
    max_val = traces.max(axis=1, keepdims=True)
    return 2.0 * (traces - min_val) / (max_val - min_val + 1e-8) - 1.0


class ASCADDataset(Dataset):
    """PyTorch Dataset for ASCAD traces."""
    def __init__(self, traces, labels):
        self.traces = torch.from_numpy(traces).float()
        self.labels = torch.from_numpy(labels).long()

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        return self.traces[idx], self.labels[idx]


# --- Model Definition ---

class SCA_CNN(nn.Module):
    """CNN for Side-Channel Analysis on 700-sample traces."""
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            # Conv block 1
            nn.Conv1d(1, 64, kernel_size=11, padding=5),
            nn.BatchNorm1d(64),
            nn.ReLU(),
            nn.AvgPool1d(2),
            # Conv block 2
            nn.Conv1d(64, 128, kernel_size=11, padding=5),
            nn.BatchNorm1d(128),
            nn.ReLU(),
            nn.AvgPool1d(2),
            # Conv block 3
            nn.Conv1d(128, 256, kernel_size=11, padding=5),
            nn.BatchNorm1d(256),
            nn.ReLU(),
            nn.AvgPool1d(2),
            # Conv block 4
            nn.Conv1d(256, 512, kernel_size=11, padding=5),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.AvgPool1d(2),
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(512 * 43, 4096),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(4096, 256),
        )

    def forward(self, x):
        # x shape: (batch, 700) -> (batch, 1, 700)
        x = x.unsqueeze(1)
        x = self.features(x)
        x = self.classifier(x)
        return x


# --- Training ---

def train_model(model, train_loader, val_loader, device, epochs=50, max_lr=5e-3):
    """Train the SCA CNN model."""
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=max_lr)
    scheduler = torch.optim.lr_scheduler.OneCycleLR(
        optimizer, max_lr=max_lr, epochs=epochs, steps_per_epoch=len(train_loader)
    )

    train_losses = []
    val_accuracies = []

    for epoch in range(epochs):
        model.train()
        epoch_loss = 0.0
        for traces_batch, labels_batch in train_loader:
            traces_batch = traces_batch.to(device)
            labels_batch = labels_batch.to(device)

            optimizer.zero_grad()
            outputs = model(traces_batch)
            loss = criterion(outputs, labels_batch)
            loss.backward()
            optimizer.step()
            scheduler.step()
            epoch_loss += loss.item()

        avg_loss = epoch_loss / len(train_loader)
        train_losses.append(avg_loss)

        # Validation
        model.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for traces_batch, labels_batch in val_loader:
                traces_batch = traces_batch.to(device)
                labels_batch = labels_batch.to(device)
                outputs = model(traces_batch)
                _, predicted = torch.max(outputs, 1)
                total += labels_batch.size(0)
                correct += (predicted == labels_batch).sum().item()

        val_acc = correct / total
        val_accuracies.append(val_acc)

        if (epoch + 1) % 10 == 0:
            print(f"Epoch {epoch+1}/{epochs} - Loss: {avg_loss:.4f} - Val Acc: {val_acc:.4f}")

    return train_losses, val_accuracies


# --- Key Recovery Attack ---

def compute_key_rank(model, attack_traces, attack_metadata, device, num_traces, target_byte=2):
    """Compute key rank using accumulated log-probabilities."""
    model.eval()
    correct_key = attack_metadata["key"][0][target_byte]

    # Get model predictions for the specified number of traces
    traces_tensor = torch.from_numpy(attack_traces[:num_traces]).float().to(device)

    with torch.no_grad():
        log_probs = torch.log_softmax(model(traces_tensor), dim=1).cpu().numpy()

    # For each key guess, accumulate log-probabilities
    scores = np.zeros(256)
    plaintexts = attack_metadata["plaintext"][:num_traces]

    for k_guess in range(256):
        for i in range(num_traces):
            # Hypothetical S-box output for this key guess
            hyp = SBOX[plaintexts[i][target_byte] ^ k_guess]
            scores[k_guess] += log_probs[i, hyp]

    # Rank: how many keys score higher than the correct key
    correct_score = scores[correct_key]
    rank = np.sum(scores > correct_score)
    return rank, correct_key


def plot_key_rank_convergence(model, attack_traces, attack_metadata, device, max_traces=2000, step=10):
    """Plot key rank vs number of attack traces."""
    trace_counts = list(range(step, max_traces + 1, step))
    ranks = []

    for n in trace_counts:
        rank, _ = compute_key_rank(model, attack_traces, attack_metadata, device, n)
        ranks.append(rank)

    plt.figure(figsize=(10, 5))
    plt.plot(trace_counts, ranks, "b-", linewidth=2)
    plt.xlabel("Number of Attack Traces")
    plt.ylabel("Key Rank")
    plt.title("Key Rank Convergence - AES Key Recovery")
    plt.axhline(y=0, color="r", linestyle="--", label="Rank 0 (Key Found)")
    plt.legend()
    plt.grid(True)
    plt.show()

    return trace_counts, ranks


# --- Main Execution ---

if __name__ == "__main__":
    # Setup
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    # Load data
    filepath = "ASCAD_data/ASCAD_databases/ASCAD.h5"
    (profiling_traces, profiling_labels, profiling_metadata,
     attack_traces, attack_labels, attack_metadata) = load_ascad(filepath)

    print(f"Profiling traces shape: {profiling_traces.shape}")
    print(f"Attack traces shape: {attack_traces.shape}")

    # Normalize
    profiling_traces = normalize_traces(profiling_traces)
    attack_traces = normalize_traces(attack_traces)

    # Create datasets and loaders
    train_dataset = ASCADDataset(profiling_traces[:45000], profiling_labels[:45000])
    val_dataset = ASCADDataset(profiling_traces[45000:], profiling_labels[45000:])
    train_loader = DataLoader(train_dataset, batch_size=256, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=256, shuffle=False)

    # Build model
    model = SCA_CNN().to(device)
    total_params = sum(p.numel() for p in model.parameters())
    print(f"Total parameters: {total_params:,}")

    # Train
    print("\nTraining...")
    train_losses, val_accuracies = train_model(model, train_loader, val_loader, device)

    # Plot training curves
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
    ax1.plot(train_losses)
    ax1.set_xlabel("Epoch")
    ax1.set_ylabel("Loss")
    ax1.set_title("Training Loss")
    ax2.plot(val_accuracies)
    ax2.set_xlabel("Epoch")
    ax2.set_ylabel("Accuracy")
    ax2.set_title("Validation Accuracy")
    plt.tight_layout()
    plt.show()

    # Key Recovery Attack
    print("\nExecuting key recovery attack...")
    rank, correct_key = compute_key_rank(model, attack_traces, attack_metadata, device, num_traces=200)
    print(f"Correct key byte: 0x{correct_key:02X}")
    print(f"Key rank with 200 traces: {rank}")

    # Plot convergence
    trace_counts, ranks = plot_key_rank_convergence(model, attack_traces, attack_metadata, device)

    # Find minimum traces needed
    for i, r in enumerate(ranks):
        if r == 0:
            print(f"\nKey recovered at Rank 0 with {trace_counts[i]} attack traces!")
            break

You just proved that a neural network can bypass cryptographic protections by exploiting physical electromagnetic leakage. Next up, you can push this further by attacking traces with real-world countermeasures in the Secret Mission.

Secret mission

Attack the Desynchronized Dataset

Your CNN cracked perfectly aligned traces. But real-world devices fight back with jitter countermeasures that randomly shift the leakage window up to 50 samples. Can your model still find the key when nothing lines up?

Clean Up Your Resources

Clean Up Your Resources

This project runs entirely on Google Colab with free GPU access, so there are no ongoing costs. Decide whether to keep your notebook saved, let Colab handle cleanup automatically, or remove everything yourself.

Resources you used:

  • Google Colab notebook (free, auto-deletes inactive runtimes after 90 minutes of idle time).
  • Downloaded ASCAD dataset files in Colab's ephemeral storage (ASCAD_data/ASCAD_databases/ASCAD.h5 and ASCAD_data/ASCAD_databases/ASCAD_desync50.h5).

Keep everything running

No action needed right now. Choose this if you plan to continue experimenting with different CNN architectures or hyperparameters.

  • Save your notebook to Google Drive by clicking File then Save a copy in Drive.
  • The ASCAD dataset will need to be re-downloaded if your Colab runtime disconnects (after 90 minutes of inactivity or 12 hours of total runtime).
  • Your Google account remains active. Colab is a free service tied to your existing Google account.

Pause - I'll come back to this later

Colab runtimes disconnect automatically after idle timeout. Your notebook code is preserved, but the downloaded dataset and any in-memory variables are not.

  • Save your notebook to Google Drive by clicking File then Save a copy in Drive.
  • When you return, you will need to reconnect to a GPU runtime and re-run the dataset download cell. The ASCAD dataset is public and free to re-download at any time.
  • To manually disconnect now, click Runtime then Disconnect and delete runtime.

Delete - I don't want to use this again

Remove all project resources. Colab storage is ephemeral, so most cleanup happens automatically.

  • Disconnect and delete your runtime by clicking Runtime then Disconnect and delete runtime. This removes all downloaded files (the ASCAD dataset) and frees the GPU.
  • Delete the notebook from Google Drive. Open Google Drive, find the notebook file, right-click it, and select Move to trash.
  • No other cleanup is needed. The ASCAD dataset was stored in Colab's ephemeral disk and is already gone once the runtime disconnects.

Nice Work!

Nice Work!

You just cracked AES-128 encryption using a neural network trained on electromagnetic emissions. This is the exact methodology used by hardware security evaluators to assess whether cryptographic devices leak secrets through their physical behavior.

You've learned how to:

  • Build a complete profiling side-channel attack pipeline using PyTorch and the ASCAD benchmark dataset, from data loading through key recovery.
  • Train a Convolutional Neural Network that classifies electromagnetic trace leakage into 256 S-box output classes and accumulates log-probabilities across attack traces to recover the secret key byte at Key Rank 0.
  • Understand how AES-128 leaks through physical electromagnetic emissions and how deep learning exploits the statistical dependency between secret key values and observable power consumption patterns.
  • Secret Mission: Attacked desynchronized traces with up to 50 samples of jitter, experiencing how real-world countermeasures make attacks harder and how CNN architecture choices like kernel size and network depth affect robustness.

Ready to quiz yourself?