Tutorial.ipynb

{
    "cells": [
        {
            "cell_type": "markdown",
            "id": "bbf1630f-7862-413a-9f74-db43cd34491e",
            "metadata": {},
            "source": [
                "# SCALib tutorial at CARDIS 2023\n",
                "\n",
                "In this tutorial, we will perform a side-channel attack against the [ASCAD](https://github.com/ANSSI-FR/ASCAD) (v1, variable key) data set provided by the ANSSI. This is done using the side-channel analysis library [SCALib](https://github.com/simple-crypto/SCALib). This tutorial reproduces the results from [this paper](https://eprint.iacr.org/2021/817) for which the code is provided [here](https://github.com/cassiersg/ASCAD-5minutes)."
            ]
        },
        {
            "cell_type": "markdown",
            "id": "4e3862cb-beca-4a2c-942d-ac4d8a73b1b6",
            "metadata": {
                "jp-MarkdownHeadingCollapsed": true
            },
            "source": [
                "## Setup (local computer)"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "78e0b0e9-f36a-4c30-8a91-0a94ae26be27",
            "metadata": {},
            "source": [
                "### Trace file"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "4f7ace23-6496-4b50-98ca-c6d72601c57c",
            "metadata": {},
            "source": [
                "Download the traces (1.3 GB) from a local server link or, if not available, from [here](https://drive.google.com/file/d/12_vaKztMGigkKn3YmOKqLj4E1d6ItdaQ/view?usp=sharing)."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "9c63fee2-2d60-4afa-9d5d-d55b51ab75e5",
            "metadata": {},
            "outputs": [],
            "source": [
                "!pwd\n",
                "!ls"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "099e7232-0ad5-43ba-b586-5caacf64087f",
            "metadata": {},
            "source": [
                "The previous result should include a `atmega8515-raw-traces_reduced.h5` file."
            ]
        },
        {
            "cell_type": "markdown",
            "id": "906aa6be-0aee-492a-938a-f404856e5e4c",
            "metadata": {},
            "source": [
                "### Basic setup"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "4616eeb0-e2d1-4e84-bb27-bb13e490569a",
            "metadata": {},
            "source": [
                "1. Install JupyterLab: see [here](https://jupyter.org/install).\n",
                "2. Clone/download [this git repo](https://github.com/simple-crypto/scalib-tutorial) ([zip](https://github.com/simple-crypto/scalib-tutorial/archive/refs/heads/main.zip)\n",
                "3. Start JupyterLab and open this notebook."
            ]
        },
        {
            "cell_type": "markdown",
            "id": "af7540fe-d223-4626-b6ef-7709a8f0b7d1",
            "metadata": {},
            "source": [
                "### Dependencies"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "f4c25210-1950-41d6-b266-a6cd32fd4e08",
            "metadata": {},
            "source": [
                "You need a few python packages as dependencies. We recommend working in a new virtual environment.\n",
                "\n",
                "In a shell run:\n",
                "\n",
                "- Unix-like:\n",
                "\n",
                "```sh\n",
                "python -m venv tuto_scalib_ve\n",
                "tuto_scalib_ve/bin/pip install ipykernel scalib tqdm matplotlib h5py\n",
                "tuto_scalib_ve/bin/python -m ipykernel install --user --name=tuto_scalib\n",
                "```\n",
                "\n",
                "- Windows\n",
                "\n",
                "```sh\n",
                "python -m venv tuto_scalib_ve\n",
                ".\\tuto_scalib_ve\\Scripts\\pip install ipykernel scalib tqdm matplotlib h5py\n",
                ".\\tuto_scalib_ve\\Scripts\\python -m ipykernel install --user --name=tuto_scalib\n",
                "```\n",
                "\n",
                "Then start `jupyter-lab`.\n",
                "\n",
                "Go to the `Kernel > Change Kernel...` menu, and select `tuto_scalib`."
            ]
        },
        {
            "cell_type": "markdown",
            "id": "e3b5b4eb-708c-4e03-8290-f8d8b63a354c",
            "metadata": {
                "jp-MarkdownHeadingCollapsed": true
            },
            "source": [
                "## Setup (Google colab)"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "7afc4b61-6e4b-48c7-93a9-efddf118307e",
            "metadata": {},
            "source": [
                "Follow instructions [here](https://github.com/simple-crypto/scalib-tutorial). To have your work saved, use `File > Save a copy in Drive`."
            ]
        },
        {
            "cell_type": "markdown",
            "id": "495b214c-ac86-459e-a900-5d0848aff21f",
            "metadata": {},
            "source": [
                "**Run the following cells only if running the notebook on google colab.**\n",
                "\n",
                "It is expected (and not an issue) that the last command fails with a message like\n",
                "\n",
                "`ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
                "google-colab 1.0.0 requires ipykernel==5.5.6, but you have ipykernel 6.29.5 which is incompatible.`"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "af57cb1b-30b3-4031-bafb-cfb7e58b4821",
            "metadata": {},
            "outputs": [],
            "source": [
                "!gdown --fuzzy https://drive.google.com/file/d/12_vaKztMGigkKn3YmOKqLj4E1d6ItdaQ/view?usp=sharing\n",
                "!pip install h5py numpy scalib matplotlib tqdm\n",
                "!pip install --upgrade ipykernel"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "16b30f54-87d3-4717-9f02-d82fb81361d8",
            "metadata": {},
            "source": [
                "Now, restart the kernel: `Runtime > Restart session`."
            ]
        },
        {
            "cell_type": "markdown",
            "id": "5388845c-9a9f-4647-96e5-1af179fb0019",
            "metadata": {
                "jp-MarkdownHeadingCollapsed": true
            },
            "source": [
                "## Startup"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "3d75f043-8a12-4c34-a952-c154c5f08b85",
            "metadata": {},
            "source": [
                "### Imports\n",
                "The next cell imports the required libraries."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "d29cba75-56eb-47f6-9d16-1bcd0815a00c",
            "metadata": {},
            "outputs": [],
            "source": [
                "import copy\n",
                "import collections\n",
                "import functools as ft\n",
                "import math\n",
                "import statistics\n",
                "\n",
                "import h5py\n",
                "import matplotlib.pyplot as plt\n",
                "import numpy as np\n",
                "from scalib.metrics import SNR, Ttest, MTtest\n",
                "import scalib.modeling\n",
                "import scalib.attacks\n",
                "import scalib.postprocessing\n",
                "from tqdm import tqdm"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "48366241-ac5f-4f3b-b73d-44051878855f",
            "metadata": {},
            "source": [
                "### Settings for the profiling and attacks\n",
                "Choosing parameters for the following experiments. The reduced dataset contains 5100 traces. As a result, `attack + profile` should be smaller than `5100`."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "5218bc92-3f5a-4bdd-b52e-555338ef70da",
            "metadata": {},
            "outputs": [],
            "source": [
                "class Settings:\n",
                "    attacks = 100  # Number of single trace attacks\n",
                "    profile = 5000  # Number of traces used in profiling\n",
                "    poi = 512  # Number of POIs for each intermediate variables\n",
                "    dim = 8  # Number of dimensions in LDA\n",
                "    database = \"./atmega8515-raw-traces_reduced.h5\"  # Path to the database\n",
                "\n",
                "settings = Settings()\n",
                "assert(settings.attacks + settings.profile <= 5100)"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "7036e680-76df-4c5e-8a02-732ba2216d18",
            "metadata": {},
            "source": [
                "### Helper functions\n",
                "This next cell loads helper function to load the traces, instantiate AES Sbox, generates the labels of intermediate variables, etc. This is not really relevant for the tutorial."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "acee63ce-4915-4ec8-b4ee-49e3fbc10907",
            "metadata": {},
            "outputs": [],
            "source": [
                "# Matplotlib configuration\n",
                "FIGSIZE = (15, 4)\n",
                "XLIM = [0, 250_000]\n",
                "LW = 0.5\n",
                "COLORS = [\"b\", \"r\", \"c\", \"m\", \"y\", \"k\"]\n",
                "\n",
                "# number of bytes to attack\n",
                "NBYTES = 14\n",
                "\n",
                "\n",
                "def target_variables(byte):\n",
                "    \"\"\"variables that will be profiled\"\"\"\n",
                "    return [\"rin\", \"rout\"] + [\n",
                "        f\"{base}_{byte}\" for base in (\"x0\", \"x1\", \"xrin\", \"yrout\", \"y0\", \"y1\")\n",
                "    ]\n",
                "\n",
                "\n",
                "# fmt: off\n",
                "SBOX = np.array(\n",
                "    [\n",
                "        0x63, 0x7C, 0x77, 0x7B, 0xF2, 0x6B, 0x6F, 0xC5, 0x30, 0x01, 0x67, 0x2B, 0xFE, 0xD7, 0xAB,\n",
                "        0x76, 0xCA, 0x82, 0xC9, 0x7D, 0xFA, 0x59, 0x47, 0xF0, 0xAD, 0xD4, 0xA2, 0xAF, 0x9C, 0xA4,\n",
                "        0x72, 0xC0, 0xB7, 0xFD, 0x93, 0x26, 0x36, 0x3F, 0xF7, 0xCC, 0x34, 0xA5, 0xE5, 0xF1, 0x71,\n",
                "        0xD8, 0x31, 0x15, 0x04, 0xC7, 0x23, 0xC3, 0x18, 0x96, 0x05, 0x9A, 0x07, 0x12, 0x80, 0xE2,\n",
                "        0xEB, 0x27, 0xB2, 0x75, 0x09, 0x83, 0x2C, 0x1A, 0x1B, 0x6E, 0x5A, 0xA0, 0x52, 0x3B, 0xD6,\n",
                "        0xB3, 0x29, 0xE3, 0x2F, 0x84, 0x53, 0xD1, 0x00, 0xED, 0x20, 0xFC, 0xB1, 0x5B, 0x6A, 0xCB,\n",
                "        0xBE, 0x39, 0x4A, 0x4C, 0x58, 0xCF, 0xD0, 0xEF, 0xAA, 0xFB, 0x43, 0x4D, 0x33, 0x85, 0x45,\n",
                "        0xF9, 0x02, 0x7F, 0x50, 0x3C, 0x9F, 0xA8, 0x51, 0xA3, 0x40, 0x8F, 0x92, 0x9D, 0x38, 0xF5,\n",
                "        0xBC, 0xB6, 0xDA, 0x21, 0x10, 0xFF, 0xF3, 0xD2, 0xCD, 0x0C, 0x13, 0xEC, 0x5F, 0x97, 0x44,\n",
                "        0x17, 0xC4, 0xA7, 0x7E, 0x3D, 0x64, 0x5D, 0x19, 0x73, 0x60, 0x81, 0x4F, 0xDC, 0x22, 0x2A,\n",
                "        0x90, 0x88, 0x46, 0xEE, 0xB8, 0x14, 0xDE, 0x5E, 0x0B, 0xDB, 0xE0, 0x32, 0x3A, 0x0A, 0x49,\n",
                "        0x06, 0x24, 0x5C, 0xC2, 0xD3, 0xAC, 0x62, 0x91, 0x95, 0xE4, 0x79, 0xE7, 0xC8, 0x37, 0x6D,\n",
                "        0x8D, 0xD5, 0x4E, 0xA9, 0x6C, 0x56, 0xF4, 0xEA, 0x65, 0x7A, 0xAE, 0x08, 0xBA, 0x78, 0x25,\n",
                "        0x2E, 0x1C, 0xA6, 0xB4, 0xC6, 0xE8, 0xDD, 0x74, 0x1F, 0x4B, 0xBD, 0x8B, 0x8A, 0x70, 0x3E,\n",
                "        0xB5, 0x66, 0x48, 0x03, 0xF6, 0x0E, 0x61, 0x35, 0x57, 0xB9, 0x86, 0xC1, 0x1D, 0x9E, 0xE1,\n",
                "        0xF8, 0x98, 0x11, 0x69, 0xD9, 0x8E, 0x94, 0x9B, 0x1E, 0x87, 0xE9, 0xCE, 0x55, 0x28, 0xDF,\n",
                "        0x8C, 0xA1, 0x89, 0x0D, 0xBF, 0xE6, 0x42, 0x68, 0x41, 0x99, 0x2D, 0x0F, 0xB0, 0x54, 0xBB,\n",
                "        0x16,\n",
                "    ],\n",
                "    dtype=np.uint32,\n",
                ")\n",
                "\n",
                "def var_labels(key, plaintext, masks, rin, rout):\n",
                "    \"Compute value of variables of interest based on ASCAD metadata.\"\n",
                "    x0 = key ^ plaintext ^ masks\n",
                "    x1 = masks\n",
                "    xrin = ((key ^ plaintext).T ^ rin).T\n",
                "    y0 = SBOX[key ^ plaintext].astype(np.uint16) ^ masks\n",
                "    y1 = masks\n",
                "    yrout = (SBOX[(key ^ plaintext).T].astype(np.uint16) ^ rout).T\n",
                "    labels = {}\n",
                "    for i in range(14):\n",
                "        labels[f\"k_{i}\"] = key[:, i]\n",
                "        labels[f\"p_{i}\"] = plaintext[:, i]\n",
                "        labels[f\"x0_{i}\"] = x0[:, i]\n",
                "        labels[f\"x1_{i}\"] = x1[:, i]\n",
                "        labels[f\"y0_{i}\"] = y0[:, i]\n",
                "        labels[f\"y1_{i}\"] = y1[:, i]\n",
                "        labels[f\"xrin_{i}\"] = xrin[:, i]\n",
                "        labels[f\"yrout_{i}\"] = yrout[:, i]\n",
                "    labels[f\"rout\"] = rout[:]\n",
                "    labels[f\"rin\"] = rin[:]\n",
                "    return labels\n",
                "\n",
                "@ft.lru_cache(maxsize=1)\n",
                "def load_database(settings):\n",
                "    dataset = h5py.File(settings.database, \"r\")\n",
                "    traces = dataset['traces'][...].astype(np.int16)\n",
                "    metadata = {k: x[...] for k, x in dataset['metadata'].items()}\n",
                "    key = metadata[\"key\"][:, 2:].astype(np.uint16)\n",
                "    plaintext = metadata[\"plaintext\"][:, 2:].astype(np.uint16)\n",
                "    masks = metadata[\"masks\"][:, 2:16].astype(np.uint16)\n",
                "    rin = metadata[\"masks\"][:, 16].astype(np.uint16)\n",
                "    rout = metadata[\"masks\"][:, 17].astype(np.uint16)\n",
                "    labels = var_labels(key, plaintext, masks, rin, rout)\n",
                "    return traces, labels\n",
                "\n",
                "def get_traces(settings, start, l):\n",
                "    \"\"\"Load traces and labels from ASCAD database.\"\"\"    \n",
                "    I = np.arange(start, start + l)\n",
                "    traces, labels = load_database(settings)\n",
                "    return traces[I,:], {k: v[I] for k, v in labels.items()}"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "e47778ec-642c-449b-a9c8-c47f8dc2dd2a",
            "metadata": {
                "jp-MarkdownHeadingCollapsed": true
            },
            "source": [
                "## Exercise 0: Visualize the traces\n",
                "This cell plots a single trace. This works out of the box."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "d4c397e0-1d61-4ed1-acd0-287540d01690",
            "metadata": {},
            "outputs": [],
            "source": [
                "traces, labels = get_traces(settings, start=0, l=1)\n",
                "\n",
                "plt.figure(figsize=(15, 4))\n",
                "plt.plot(traces[0, :], linewidth=LW)\n",
                "plt.xlabel(\"Time\")\n",
                "plt.ylabel(\"Power Consumption\")\n",
                "plt.xlim([0, traces.shape[1]])\n",
                "plt.show()"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "0e6716fd-9bf3-4ed5-b1b9-73070f71a425",
            "metadata": {
                "jp-MarkdownHeadingCollapsed": true
            },
            "source": [
                "# Exercise 1a: Compute the SNR for all intermediate variables\n",
                "The SNR of the leakage $L$ for a variable $X$ is defined as:\n",
                "$$\n",
                "\\mathrm{SNR}_X = \\frac{\\mathrm{Var}_{x\\leftarrow X}(\\mathrm{E}[L_x])}\n",
                "                {\\mathrm{E}_{x\\leftarrow X}(\\mathrm{Var}[L_x])\n",
                "}$$\n",
                "and is the variance between classes divided by the variance within the classes.\n",
                "It can be efficiently computed with [scalib.metrics.SNR](https://scalib.readthedocs.io/en/stable/_modules/scalib/metrics/snr.html#SNR). In this exercise, you are required to fill the following code in order to assign `snrs[v]`  to the SNR value of the variable `v`. You can use the following cell in order to visualize the resulting SNR.\n"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "b283a91c-5c80-4b4b-b4a4-4ceb70957932",
            "metadata": {},
            "outputs": [],
            "source": [
                "def compute_snr(settings):\n",
                "    \"\"\"Returns the SNR of the traces samples for each target variable.\"\"\"\n",
                "\n",
                "    snrs = dict()\n",
                "    variables = [v for i in range(NBYTES) for v in target_variables(i)]\n",
                "    traces, labels = get_traces(settings, start=0, l=settings.profile)\n",
                "\n",
                "    for v in tqdm(variables, desc=\"SNR Variables\"):\n",
                "        # Compute the snr for variable v\n",
                "        # TODO\n",
                "        snrs[v] = ...\n",
                "\n",
                "        # Avoid NaN in case of scope over-range\n",
                "        np.nan_to_num(snrs[v], nan=0.0)\n",
                "    return snrs\n",
                "\n",
                "\n",
                "snrs = compute_snr(settings)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "5c2c8d4a-3faf-40ae-b2d7-4c1a40b4251f",
            "metadata": {},
            "outputs": [],
            "source": [
                "def plot_snr(snrs, variables):\n",
                "    plt.figure(figsize=FIGSIZE)\n",
                "    for color, var in zip(COLORS, variables):\n",
                "        plt.plot(snrs[var], label=var, linewidth=LW, color=color)\n",
                "    plt.xlim(XLIM)\n",
                "    plt.xlabel(\"Time\")\n",
                "    plt.ylabel(\"SNR\")\n",
                "    plt.legend()\n",
                "    plt.show()\n",
                "\n",
                "\n",
                "plot_snr(snrs, target_variables(byte=0))"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "4626e515-f0a7-4d29-9a67-5c915107d18f",
            "metadata": {},
            "source": [
                "### Improvements\n",
                "[scalib.metrics.SNR](https://scalib.readthedocs.io/en/stable/_modules/scalib/metrics/snr.html#SNR) has the ability to compute the SNR of multiple variables, which allows to save some performance overheads. In case the data-set is too large to fit in memory, the [scalib.metrics.SNR](https://scalib.readthedocs.io/en/stable/_modules/scalib/metrics/snr.html#SNR) can be updated in an incremental fashion by providing multiple chunks of data to `fit_u()`. As an extra exercise (not during the tutorial) you can:\n",
                "1. Compute the SNR on all the variables with a single call to [scalib.metrics.SNR](https://scalib.readthedocs.io/en/stable/_modules/scalib/metrics/snr.html#SNR).\n",
                "2. Provide the traces in multiple independent chunks to `fit_u()`."
            ]
        },
        {
            "cell_type": "markdown",
            "id": "fef7f159-e217-4032-b1b8-3612b671e174",
            "metadata": {
                "jp-MarkdownHeadingCollapsed": true
            },
            "source": [
                "# Exercise 1b: Select the Points-of-Interest (POIs)\n",
                "\n",
                "In this exercise, you are asked to compute the POIs for each of the variables. The POIs for the variable `v` are the `settings.poi` indexes with the largest SNR. To compute the POIs, you can leverage [np.argsort](https://numpy.org/doc/stable/reference/generated/numpy.argsort.html#numpy-argsort) function."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "60efb471-a553-48f7-a1e2-33cb198e455f",
            "metadata": {},
            "outputs": [],
            "source": [
                "def compute_pois(settings, snrs):\n",
                "    \"\"\"Compute the POIs for all variables.\"\"\"\n",
                "    pois = dict()\n",
                "\n",
                "    # Select POIs\n",
                "    for v, snr in snrs.items():\n",
                "        # Compute the POIs for the variable v.\n",
                "        # TODO\n",
                "        poi = ...\n",
                "        \n",
                "        pois[v] = poi\n",
                "    return pois\n",
                "\n",
                "\n",
                "pois = compute_pois(settings, snrs)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "c9cf70d9-4cff-43e7-ba63-3ec725499bbb",
            "metadata": {},
            "outputs": [],
            "source": [
                "def plot_pois(snrs, pois, variables):\n",
                "    plt.figure(figsize=FIGSIZE)\n",
                "    for color, var in zip(COLORS, variables):\n",
                "        poi = pois[var]\n",
                "        snr = snrs[var]\n",
                "        plt.plot(snr, label=var, linewidth=LW, color=color)\n",
                "        plt.scatter(poi, snr[poi], color=color, alpha=0.3, marker=\"x\")\n",
                "    plt.xlim(XLIM)\n",
                "    plt.xlabel(\"Time\")\n",
                "    plt.ylabel(\"SNR and POIs\")\n",
                "    plt.legend()\n",
                "    #plt.show()\n",
                "\n",
                "plot_pois(snrs, pois, target_variables(byte=0))"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "c83716f9-660a-4af5-a47e-0ff103afcf96",
            "metadata": {
                "jp-MarkdownHeadingCollapsed": true
            },
            "source": [
                "# Exercise 2: Build Gaussian templates with LDA\n",
                "In the following, you are asked to build the Gaussian templates in order to derive the distribution of the leakage of each intermediate variables. Concretely, the leakage distribution is approximated with:\n",
                "$$\n",
                " \\mathsf{\\hat{f}}(\\mathbf{l} | x)= \\frac{1}{\\sqrt{(2\\pi)^{p} \\cdot |\\mathbf{\\Sigma} |}} \\cdot \\exp^{\\frac{1}{2} (\\mathbf{W} \\cdot \\mathbf{l} - \\mathbf{\\mu}_x)                     \\mathbf{\\Sigma}\n",
                "                    ( \\mathbf{W} \\cdot \\mathbf{l}-\\mathbf{\\mu}_x)'}\n",
                "$$\n",
                "where $\\mathbf{W}$ is a projection matrix obtained with LDA. This enables to project the leakage from a space with `settings.poi` dimensions to a linear sub-space with `settings.dim` dimensions. This projection maximizes the separation between the classes. The Gaussian templates are then built in that subspace and the matrix $\\mathbf{\\Sigma}$ is the pooled covariance in that subspace.\n",
                "\n",
                "In this exercise, you are asked to assign to `ldas[v]` a Gaussian template for variable `v`. Namely, you are asked to use [scalib.modeling.LDAClassifier](https://scalib.readthedocs.io/en/stable/source/api/scalib.modeling.LDAClassifier.html#scalib.modeling.LDAClassifier). The next cell enables to visualize the different classes in the linear subspace."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "910d316b-d2bf-4539-af78-59af5efd02ad",
            "metadata": {},
            "outputs": [],
            "source": [
                "def compute_templates(settings, pois):\n",
                "    \"\"\"Compute the POIs, LDA and gaussian template for all variables.\"\"\"\n",
                "\n",
                "    ldas = dict()\n",
                "    variables = [v for i in range(NBYTES) for v in target_variables(i)]\n",
                "    traces, labels = get_traces(settings, start=0, l=settings.profile)\n",
                "\n",
                "    for v in tqdm(variables, desc=\"LDA Variables\"):\n",
                "        # Generate the Gaussian template for the variable v.\n",
                "        # TODO\n",
                "        ldas[v] = ...\n",
                "        \n",
                "    return ldas\n",
                "\n",
                "\n",
                "ldas = compute_templates(settings, pois)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "3d44ac3d-c540-455c-870d-767ef4c31edc",
            "metadata": {
                "scrolled": true
            },
            "outputs": [],
            "source": [
                "def plot_lda(settings, pois, ldas, variable, values):\n",
                "    traces, labels = get_traces(settings, start=0, l=settings.profile)\n",
                "    projection = ldas[variable].lda.get_state()[0]\n",
                "\n",
                "    plt.figure()\n",
                "    for v in values:\n",
                "        indexes = np.where(labels[variable] == v)[0]\n",
                "        traces_v = traces[np.where(labels[variable] == v)[0], :]\n",
                "        traces_v_projected = projection.T @ traces_v[:, pois[variable]].T\n",
                "        plt.scatter(\n",
                "            traces_v_projected[0, :],\n",
                "            traces_v_projected[1, :],\n",
                "            label=f\"{variable}={v}\",\n",
                "            alpha=0.3,\n",
                "        )\n",
                "\n",
                "    plt.title(\"Distributions in linear subspace\")\n",
                "    plt.xlabel(\"First dimension\")\n",
                "    plt.ylabel(\"Second dimension\")\n",
                "    plt.legend()\n",
                "    plt.show()\n",
                "\n",
                "\n",
                "for variable in target_variables(byte=0):\n",
                "    plot_lda(settings, pois, ldas, variable, np.random.randint(0, 256, 8))"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "3dc4f054-62e3-4706-8017-26838ada7fe3",
            "metadata": {
                "editable": true,
                "slideshow": {
                    "slide_type": ""
                },
                "tags": []
            },
            "source": [
                "### Improvements\n",
                "\n",
                "The previous code can be improved in two ways:\n",
                "1. When having to compute a LDA model for multiple variables, using sequentially multiple calls to `LDAClassifier` comes with overheads that can be avoided by leveraging [scalib.modeling.MultiLDA](https://scalib.readthedocs.io/en/stable/source/api/scalib.modeling.MultiLDA.html).\n",
                "2. If the dataset is too large, it might not fit in memory at once. In such a case, it can be loaded in smaller chunks and the function `fit_u()` called for each of these chunks. "
            ]
        },
        {
            "cell_type": "markdown",
            "id": "d01d85c0-af3f-4eb3-bf5b-1aeb28d6b28b",
            "metadata": {
                "jp-MarkdownHeadingCollapsed": true
            },
            "source": [
                "# Exercise 3: A First key recovery\n",
                "\n",
                "In the previous exercises, we have built templates for the shares within the implementation. In this exercise, we will run a SASCA that exploits the factor graph described below. To do so, you have to fill the TODOs in the following code.\n",
                "\n",
                "The useful function are [scalib.attacks.FactorGraph](https://scalib.readthedocs.io/en/stable/source/api/scalib.attacks.FactorGraph.html#scalib.attacks.FactorGraph) in order to generate a factor graph. Belief propagation can be ran thanks to [scalib.attacks.BPState](https://scalib.readthedocs.io/en/stable/source/api/scalib.attacks.BPState.html#scalib.attacks.BPState) and the [bp_loopy](https://scalib.readthedocs.io/en/stable/source/api/scalib.attacks.BPState.html#scalib.attacks.BPState.bp_loopy) function. The probabilities of the internal variables can be obtained from the LDA thanks to [predict_proba](https://scalib.readthedocs.io/en/stable/source/api/scalib.modeling.LDAClassifier.html#scalib.modeling.LDAClassifier.predict_proba) and then assigned to the factor graph variables thanks to [set_evidence](https://scalib.readthedocs.io/en/stable/source/api/scalib.attacks.BPState.html#scalib.attacks.BPState.set_evidence). The produced distributions after running BP can be recovered thanks to [get_distribution](https://scalib.readthedocs.io/en/stable/source/api/scalib.attacks.BPState.html#scalib.attacks.BPState.get_distribution). We strongly encourage the have a look at the examples available in the documentation."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "8a773cf1-271e-496c-af09-1dbbe1c13b7c",
            "metadata": {},
            "outputs": [],
            "source": [
                "SASCA_GRAPH = \"\"\"\n",
                "NC 256\n",
                "TABLE sbox\n",
                "\n",
                "VAR MULTI x\n",
                "VAR MULTI x0\n",
                "VAR MULTI x1\n",
                "VAR MULTI xrin\n",
                "VAR MULTI rin\n",
                "VAR MULTI y\n",
                "PUB MULTI p\n",
                "\n",
                "VAR SINGLE k\n",
                "\n",
                "PROPERTY x = p ^ k\n",
                "PROPERTY x = x0 ^ x1\n",
                "PROPERTY x = rin ^ xrin\n",
                "\n",
                "\"\"\"\n",
                "\n",
                "def attack(sasca_graph, traces, labels, ldas, pois):\n",
                "    \"\"\"Run a SASCA attack on the given traces\n",
                "    Returns the true key and the byte-wise key distribution estimated by the attack.\n",
                "    \"\"\"\n",
                "    # correct secret key\n",
                "    secret_key = [labels[f\"k_{i}\"][0] for i in range(NBYTES)]\n",
                "    \n",
                "    # distribution for each of the key bytes\n",
                "    key_distribution = []\n",
                "    \n",
                "    # Run a SASCA for each S-Box\n",
                "    for i in range(NBYTES):\n",
                "\n",
                "        # Init the sasca with the factor graph for a single trace\n",
                "        # TODO\n",
                "        graph = scalib.attacks.FactorGraph(sasca_graph, {\"sbox\": SBOX})\n",
                "\n",
                "        # Set the labels for the plaintext byte\n",
                "        # TODO\n",
                "        bp = ...\n",
                "              \n",
                "        # Set the initial distribution for target variables 'vs' if it is in the graph\n",
                "        for v in target_variables(i):\n",
                "            vs = v.split('_')[0]\n",
                "            if vs in graph.vars():\n",
                "                # Recover and assign the distribution of vs\n",
                "                # TODO\n",
                "                pass\n",
                "\n",
                "        # Run 3 iterations of belief propagation\n",
                "        # TODO       \n",
                "\n",
                "        # Get the distribution of the secret key\n",
                "        # TODO\n",
                "        byte_distribution = ...\n",
                "      \n",
                "        key_distribution.append(byte_distribution)\n",
                "        \n",
                "    key_distribution = np.array(key_distribution)\n",
                "    return secret_key, key_distribution\n",
                "\n",
                "# Get the trace and label for the attack\n",
                "traces, labels = get_traces(settings, start=settings.profile, l=1)\n",
                "# Perform the attack\n",
                "secret_key, key_distribution = attack(SASCA_GRAPH, traces,labels,ldas,pois)\n",
                "\n",
                "print(\"Correct Key:\", \" \".join([f\"0x{k:02x}\" for k in secret_key] ))\n",
                "print(\"Guessed Key:\", \" \".join([f\"0x{np.argmax(k):02x}\" for k in key_distribution] ))"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "46d250b0-f304-4d60-a554-7862c9bdf08f",
            "metadata": {
                "jp-MarkdownHeadingCollapsed": true
            },
            "source": [
                "# Exercise 4: Evaluate the resulting rank of the full correct key.\n",
                "\n",
                "After performing a side-channel attack, an evaluator is usually interested in the remaining entropy of the key. That is, he wants to know the number of brute force an adversary needs before he will discover the correct key. This is call the key rank, and you will use SCALib to evaluate it in this exercise. To do so, you will use the function [calib.postprocessing.rank_accuracy](https://scalib.readthedocs.io/en/stable/source/api/scalib.postprocessing.rankestimation.html#module-scalib.postprocessing.rankestimation)."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "6b387744-749a-4160-a1e0-fd70a890aec5",
            "metadata": {},
            "outputs": [],
            "source": [
                "def eval_rank(secret_key, key_distribution):\n",
                "    \"\"\"Returns the rank of the true key\"\"\"\n",
                "\n",
                "    # Floor the key_distribution to avoid numerical issues\n",
                "    key_distribution[np.where(key_distribution < 1e-100)] = 1e-100\n",
                "\n",
                "    # Compute the rank of the\n",
                "    log2distribution = np . log2 ( key_distribution )\n",
                "    rmin , r , rmax = scalib . postprocessing . rank_accuracy (\n",
                "        - log2distribution , secret_key , max_nb_bin = 2 ** 20\n",
                "    )\n",
                "    return r\n",
                "\n",
                "\n",
                "print(f\"Estimated rank 2**{np.log2(eval_rank(secret_key,key_distribution)):.2f}\")"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "945bb617-822d-4424-8ab0-6aac54b4b4ed",
            "metadata": {},
            "source": [
                "Putting all together, we are now able to evaluate the success rate of the adversary given an enumeration power. Concretely, the following code prints the success rate if the adversary does not perform enumeration and if the adversary is able to enumerate up to $2^{32}$ keys."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "87ee94f9-9457-461c-a965-3953427df19a",
            "metadata": {},
            "outputs": [],
            "source": [
                "def success_rate(ranks, max_rank=1):\n",
                "    return np.mean(np.array(ranks) <= max_rank)\n",
                "\n",
                "\n",
                "def attack_statistics(sasca_graph, settings, ldas, pois):\n",
                "    ranks = []\n",
                "\n",
                "    for i in tqdm(range(settings.attacks), desc=\"Evaluating the attack\"):\n",
                "        traces, labels = get_traces(settings, start=settings.profile + i, l=1)\n",
                "        secret_key, key_distribution = attack(sasca_graph, traces, labels, ldas, pois)\n",
                "        rank = eval_rank(secret_key, key_distribution)\n",
                "        ranks.append(rank)\n",
                "\n",
                "    print(f\"Success rate (rank 1): {success_rate(ranks, max_rank=1)*100:.0f}%\")\n",
                "    print(f\"Success rate (rank 2**32): {success_rate(ranks, max_rank=2**32)*100:.0f}%\")\n",
                "\n",
                "\n",
                "sasca_graph = SASCA_GRAPH\n",
                "attack_statistics(sasca_graph, settings, ldas, pois)"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "d7fd0ab6-9e3a-438d-a55a-9e46b407c7d6",
            "metadata": {
                "jp-MarkdownHeadingCollapsed": true
            },
            "source": [
                "# Exercise 5: Improving the attack\n",
                "\n",
                "You may have noticed that the factor graph above does not include all the variables within the implementation. In order to improve the attack, and adversary can encode additional information in the factor graph. The variables `y0`, `y1`, `yrout`, `rout` can be added to the factor graph. This will improve the success rate of the attack."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "709ca4f3-f413-4d4c-ab39-1aed3b954f03",
            "metadata": {
                "editable": true,
                "slideshow": {
                    "slide_type": ""
                },
                "tags": []
            },
            "outputs": [],
            "source": [
                "SASCA_GRAPH_IMPROVED = \"\"\"\n",
                "NC 256\n",
                "TABLE sbox\n",
                "\n",
                "VAR MULTI x\n",
                "VAR MULTI x0\n",
                "VAR MULTI x1\n",
                "VAR MULTI xrin\n",
                "VAR MULTI rin\n",
                "PUB MULTI p\n",
                "\n",
                "VAR SINGLE k\n",
                "\n",
                "PROPERTY x = p ^ k\n",
                "PROPERTY x = x0 ^ x1\n",
                "PROPERTY x = rin ^ xrin\n",
                "\n",
                "# TODO: add relationships for more leaking variables.\n",
                "\n",
                "\"\"\"\n",
                "sasca_graph = SASCA_GRAPH_IMPROVED\n",
                "attack_statistics(sasca_graph, settings, ldas, pois)"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "23fe087a-031a-4e66-aa34-f3cd67fa861e",
            "metadata": {},
            "source": [
                "## Exercise 6: leakage assessment with T-test\n",
                "\n",
                "Now that we have a good attack, let's take a step back and discuss how SCALib can be used to perform leakage assessment. The most common leakage assessment methodology is TVLA, which often compares a dataset with a fix plaintext (or key) to a dataset that uses random plaintexts.\n",
                "\n",
                "Our reduced ASCAD dataset cannot be used in this manner, therefore we perform a statistical test on another target: a bit on an S-Box input (this is somewhat similar to the original single-bit DPA attack)."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "b37ba71a-d6aa-4105-b459-396fcfa27e00",
            "metadata": {
                "editable": true,
                "slideshow": {
                    "slide_type": ""
                },
                "tags": []
            },
            "outputs": [],
            "source": [
                "def ttest_traces(settings, kidx=3, bit=0):\n",
                "    traces, labels = get_traces(settings, start=0, l=settings.profile)\n",
                "    labels = ((labels['k_3'] ^ labels['p_3']) >> bit) & 0x1\n",
                "    return traces, labels"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "aa087864-6129-4bb9-93ba-fe9d1cf8e4a1",
            "metadata": {},
            "source": [
                "### Univariate T-test"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "bbe214d7-b53b-41fc-90a2-3a2f6e8de087",
            "metadata": {},
            "source": [
                "Let us first run a first- and second-order univariate T-test."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "b5092344-cace-427d-9cda-2e4c20629bbc",
            "metadata": {
                "editable": true,
                "slideshow": {
                    "slide_type": ""
                },
                "tags": []
            },
            "outputs": [],
            "source": [
                "def univariate_ttest(settings):\n",
                "    \"\"\"Returns the T-test score (first axis: order, second axis: trace).\"\"\"\n",
                "    traces, x = ttest_traces(settings)\n",
                "    # TODO\n",
                "    t = ...\n",
                "\n",
                "    return t\n",
                "\n",
                "t_uni = univariate_ttest(settings)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "f55c7b05-05a3-4685-b415-cd7716d77700",
            "metadata": {
                "editable": true,
                "slideshow": {
                    "slide_type": ""
                },
                "tags": []
            },
            "outputs": [],
            "source": [
                "def plot_ttest_uni(t, p):\n",
                "    for i in range(t.shape[0]):\n",
                "        plt.figure(figsize=FIGSIZE)\n",
                "        plt.title(f\"T-test order {i+1}\")\n",
                "        plt.plot(np.abs(t[i]))\n",
                "        plt.plot([0, len(t[i])], 4.5*np.array([1, 1]), color='red')\n",
                "        p_adj = p/traces.shape[1]\n",
                "        thresh = -statistics.NormalDist().inv_cdf(p_adj)\n",
                "        plt.plot([0, len(t[i])], thresh*np.array([1, 1]), linestyle='--', color='red')\n",
                "        plt.xlim(XLIM)\n",
                "        plt.xlabel(\"Time\")\n",
                "        plt.ylabel(\"|t-score|\")\n",
                "\n",
                "plot_ttest_uni(t_uni, 0.05)"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "b77bb207-9036-44b0-8538-bda08f8c6847",
            "metadata": {},
            "source": [
                "We display 2 threshold lines:\n",
                "\n",
                "- continuous: the commonly-used 4.5 threshold\n",
                "- dashed: a threshold corresponding to a p-value of 0.05 for the *joint* test of all samples\n",
                "\n",
                "We can see that the 4.5 threshold is exceeded at the second order, while the joint threshold is not exceeed."
            ]
        },
        {
            "cell_type": "markdown",
            "id": "d2631f58-b946-499f-87da-494de8cd045f",
            "metadata": {},
            "source": [
                "## Multivariate T-test"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "843c612d-0807-4640-8c0c-39975fc550d2",
            "metadata": {},
            "source": [
                "Since we are dealing with a software masked implementation, we expect stronger leakage in a bivariate test, let us try that.\n",
                "\n",
                "We first need to generate a list of points of interests (POIs), where each POI is a pair of indices in the trace. In the example below, we take as POIs all pairs of points up to point `n` in the trace.\n",
                "\n",
                "**Caution** this step is a bit more memory-intensive that the previous ones, uses about 8 GB when `n=5000`."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "e3ec0a34-3006-4490-9961-bac42d9c6210",
            "metadata": {},
            "outputs": [],
            "source": [
                "def bivariate_ttest(settings, start=0, n=5000):\n",
                "    traces, x = ttest_traces(settings)\n",
                "\n",
                "    traces = traces[:,start:start+n]\n",
                "\n",
                "    # Template is a Boolean matrix corresponding to the positions of the POIS.\n",
                "    template = np.fromfunction(lambda i, j: i <= j, (n, n))\n",
                "    # Build POIS array from template\n",
                "    # TODO\n",
                "\n",
                "    # Compute T-test\n",
                "    # TODO\n",
                "    t = ...\n",
                "\n",
                "    # Convert back the T-test results into matrix form.\n",
                "    t_mat = np.zeros((n, n))\n",
                "    t_mat[pois_idxs] = np.abs(t)\n",
                "    return t_mat\n",
                "\n",
                "\n",
                "t = bivariate_ttest(settings)\n",
                "\n",
                "print(f\"Max T score: {np.max(t)}\")"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "6b208eae-0f6b-4f8c-ad0b-05b50f86d9f4",
            "metadata": {},
            "outputs": [],
            "source": [
                "def plot_ttest_bi(t, pixels=100):\n",
                "    # Do a max-pooling, otherwise T-test peaks are unreadable.\n",
                "    t_red = np.array([\n",
                "        [np.max(y) for y in np.array_split(x, pixels, axis=1)]\n",
                "        for x in np.array_split(np.abs(t), pixels, axis=0)\n",
                "    ])\n",
                "    plt.figure(figsize=FIGSIZE)\n",
                "    cax = plt.matshow(t_red)\n",
                "    plt.colorbar(cax)\n",
                "\n",
                "plot_ttest_bi(t)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "cb4e8105-e840-4dee-b333-de64c750084e",
            "metadata": {},
            "outputs": [],
            "source": []
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "ascad_scalib_tuto_local",
            "language": "python",
            "name": "ascad_scalib_tuto_local"
        },
        "language_info": {
            "codemirror_mode": {
                "name": "ipython",
                "version": 3
            },
            "file_extension": ".py",
            "mimetype": "text/x-python",
            "name": "python",
            "nbconvert_exporter": "python",
            "pygments_lexer": "ipython3",
            "version": "3.12.3"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 5
}