{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "57274f32-77ab-4de6-a298-a4f923f152a0",
   "metadata": {},
   "source": [
    "# Structural reverse-engineering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bea9b7eb-3fe4-4bd1-abf7-a470131e0dcf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import operator\n",
    "import numpy as np\n",
    "import holoviews as hv\n",
    "import tabulate\n",
    "from tqdm.auto import tqdm, trange\n",
    "from IPython.display import HTML, display\n",
    "from functools import reduce\n",
    "\n",
    "from pyecsca.ec.model import ShortWeierstrassModel\n",
    "from pyecsca.ec.coordinates import AffineCoordinateModel\n",
    "from pyecsca.ec.curve import EllipticCurve\n",
    "from pyecsca.ec.params import DomainParameters\n",
    "from pyecsca.ec.formula import FormulaAction\n",
    "from pyecsca.ec.op import OpType\n",
    "from pyecsca.ec.point import Point\n",
    "from pyecsca.ec.mod import mod\n",
    "from pyecsca.ec.mult import *\n",
    "from pyecsca.ec.context import DefaultContext, local\n",
    "from pyecsca.sca.re.rpa import MultipleContext\n",
    "from pyecsca.sca.attack.leakage_model import HammingWeight\n",
    "from pyecsca.sca.trace import Trace\n",
    "from pyecsca.sca.trace.plot import plot_trace, plot_traces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "857c08bb-e3d3-472c-bc59-e7f187f0605b",
   "metadata": {},
   "outputs": [],
   "source": [
    "hv.extension(\"bokeh\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b32ac45a-cb8d-4e52-bbdf-d54a6d094b6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = ShortWeierstrassModel()\n",
    "coordsaff = AffineCoordinateModel(model)\n",
    "coords = model.coordinates[\"projective\"]\n",
    "add = coords.formulas[\"add-2007-bl\"]\n",
    "dbl = coords.formulas[\"dbl-2007-bl\"]\n",
    "neg = coords.formulas[\"neg\"]\n",
    "\n",
    "# A 64-bit prime order curve for testing things out\n",
    "p = 0xc50de883f0e7b167\n",
    "a = mod(0x4833d7aa73fa6694, p)\n",
    "b = mod(0xa6c44a61c5323f6a, p)\n",
    "gx = mod(0x5fd1f7d38d4f2333, p)\n",
    "gy = mod(0x21f43957d7e20ceb, p)\n",
    "n = 0xc50de885003b80eb\n",
    "h = 1\n",
    "\n",
    "infty = Point(coords, X=mod(0, p), Y=mod(1, p), Z=mod(0, p))\n",
    "g = Point(coords, X=gx, Y=gy, Z=mod(1, p))\n",
    "\n",
    "curve = EllipticCurve(model, coords, p, infty, dict(a=a,b=b))\n",
    "params = DomainParameters(curve, g, n, h)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c09d9152-fc42-4ea9-8398-8feee8569870",
   "metadata": {},
   "source": [
    "## Scalar multipliers\n",
    "First select a bunch of multipliers. We will be trying to distinguish among these."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a388ddac-7f82-4541-b70a-677f34eeb241",
   "metadata": {},
   "outputs": [],
   "source": [
    "multipliers = [\n",
    "    LTRMultiplier(add, dbl, None, False, AccumulationOrder.PeqPR, True, True),\n",
    "    LTRMultiplier(add, dbl, None, True, AccumulationOrder.PeqPR, True, True),\n",
    "    RTLMultiplier(add, dbl, None, False, AccumulationOrder.PeqPR, True),\n",
    "    RTLMultiplier(add, dbl, None, True, AccumulationOrder.PeqPR, False),\n",
    "    SimpleLadderMultiplier(add, dbl, None, True, True),\n",
    "    BinaryNAFMultiplier(add, dbl, neg, None, ProcessingDirection.LTR, AccumulationOrder.PeqPR, True),\n",
    "    WindowNAFMultiplier(add, dbl, neg, 3, None, AccumulationOrder.PeqPR, True, True),\n",
    "    WindowNAFMultiplier(add, dbl, neg, 4, None, AccumulationOrder.PeqPR, True, True),\n",
    "    #WindowNAFMultiplier(add, dbl, neg, 4, None, AccumulationOrder.PeqPR, False, True),\n",
    "    SlidingWindowMultiplier(add, dbl, 3, None, ProcessingDirection.LTR, AccumulationOrder.PeqPR, True),\n",
    "    SlidingWindowMultiplier(add, dbl, 5, None, ProcessingDirection.LTR, AccumulationOrder.PeqPR, True),\n",
    "    FixedWindowLTRMultiplier(add, dbl, 4, None, AccumulationOrder.PeqPR, True),\n",
    "    FixedWindowLTRMultiplier(add, dbl, 5, None, AccumulationOrder.PeqPR, True),\n",
    "    FullPrecompMultiplier(add, dbl, None, True, ProcessingDirection.LTR, AccumulationOrder.PeqPR, True, True),\n",
    "    FullPrecompMultiplier(add, dbl, None, False, ProcessingDirection.LTR, AccumulationOrder.PeqPR, True, True),\n",
    "    #FullPrecompMultiplier(add, dbl, None, False, ProcessingDirection.RTL, AccumulationOrder.PeqPR, True, True),\n",
    "    BGMWMultiplier(add, dbl, 3, None, ProcessingDirection.LTR, AccumulationOrder.PeqPR, True),\n",
    "    BGMWMultiplier(add, dbl, 5, None, ProcessingDirection.LTR, AccumulationOrder.PeqPR, True),\n",
    "    CombMultiplier(add, dbl, 3, None, AccumulationOrder.PeqPR, True),\n",
    "    CombMultiplier(add, dbl, 5, None, AccumulationOrder.PeqPR, True)\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64c77b6e-59e4-43ba-a86c-19f189c2884c",
   "metadata": {},
   "source": [
    "Now choose a scalar and compute with it using all of the multipliers. Track the amounts of add and dbl formula applications during precomputation and the main scalar multiplication."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b68a633f-fe68-43b0-ad6e-777538c70b6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "scalar = 0b1000000000000000000000000000000000000000000000000\n",
    "scalar = 0b1111111111111111111111111111111111111111111111111\n",
    "scalar = 0b1010101010101010101010101010101010101010101010101\n",
    "scalar = 0b1111111111111111111111110000000000000000000000000\n",
    "scalar = 123456789123456789"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0b69890-6652-4910-b696-56c88a1f7a1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def count_formula_actions(ctx, formula):\n",
    "    actions = []\n",
    "    \n",
    "    def callback(action):\n",
    "        if isinstance(action, FormulaAction) and action.formula == formula:\n",
    "            actions.append(action)\n",
    "\n",
    "    ctx.actions[0].walk(callback)\n",
    "    return len(actions)\n",
    "\n",
    "def simulate_trace(ctx):\n",
    "    lm = HammingWeight()\n",
    "    trace = []\n",
    "\n",
    "    def callback(action):\n",
    "        if isinstance(action, FormulaAction):\n",
    "            for intermediate in action.op_results:\n",
    "                leak = lm(intermediate.value)\n",
    "                trace.append(leak)\n",
    "            trace.extend([0] * 20)\n",
    "\n",
    "    ctx.actions[0].walk(callback)\n",
    "    return Trace(np.array(trace))\n",
    "\n",
    "traces = []\n",
    "table = [[\"Multiplier\", \"Precomp add\", \"Precomp dbl\", \"Precomp total\", \"Multiply add\", \"Multiply dbl\", \"Multiply total\", \"Total\"]]\n",
    "\n",
    "for mult in multipliers:\n",
    "    with local(DefaultContext()) as ctx:\n",
    "        mult.init(params, g)\n",
    "    precomp_add = count_formula_actions(ctx, add)\n",
    "    precomp_dbl = count_formula_actions(ctx, dbl)\n",
    "    precomp_trace = simulate_trace(ctx)\n",
    "\n",
    "    with local(DefaultContext()) as ctx:\n",
    "        mult.multiply(scalar)\n",
    "    multiply_add = count_formula_actions(ctx, add)\n",
    "    multiply_dbl = count_formula_actions(ctx, dbl)\n",
    "    formula_count = precomp_add + precomp_dbl + multiply_add + multiply_dbl\n",
    "    multiply_trace = simulate_trace(ctx)\n",
    "\n",
    "    traces.append(multiply_trace)\n",
    "    table.append([mult, precomp_add, precomp_dbl, precomp_add + precomp_dbl, multiply_add, multiply_dbl, multiply_add + multiply_dbl, formula_count])\n",
    "\n",
    "display(HTML(tabulate.tabulate(table, tablefmt=\"html\", headers=\"firstrow\")))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23bf7c6a-18e0-4362-af73-8759c44a98c5",
   "metadata": {},
   "source": [
    "Now we can look at the distributions of the number of operations for random scalars."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c400865d-66e1-436d-ba1e-5ca605e0c9d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "scalars = [int(Mod.random(params.order)) for _ in range(100)]\n",
    "counts = {}\n",
    "for i, mult in enumerate(tqdm(multipliers)):\n",
    "    counts[mult] = []\n",
    "    for scalar in tqdm(scalars, leave=False):\n",
    "        mult.init(params, g)\n",
    "        with local(DefaultContext()) as ctx:\n",
    "            mult.multiply(int(scalar))\n",
    "        multiply_add = count_formula_actions(ctx, add)\n",
    "        multiply_dbl = count_formula_actions(ctx, dbl)\n",
    "        counts[mult].append(multiply_add + multiply_dbl)\n",
    "\n",
    "count_max = max(map(max, counts.values()))\n",
    "count_min = min(map(min, counts.values()))\n",
    "\n",
    "grams = []\n",
    "for i, item in enumerate(counts.items()):\n",
    "    mult, count = item\n",
    "    freqs, edges = np.histogram(count, range=(count_min, count_max), bins=50, density=True)\n",
    "    grams.append(hv.Histogram((edges, freqs), label=mult.__class__.__name__ + str(i)))\n",
    "\n",
    "reduce(operator.mul, grams[1:], grams[0]).opts(hv.opts.Histogram(alpha=1)).opts(responsive=True, height=500, xlabel=\"Opertion count\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8609dc0b-4439-42cb-a2a4-0fc19eb4ea0c",
   "metadata": {},
   "source": [
    "## Formulas\n",
    "We can apply similar structural techniques to examine formulas."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db4f3bfe-7530-4977-b0fd-637336dcf04a",
   "metadata": {},
   "outputs": [],
   "source": [
    "adds = list(filter(lambda formula: formula.name.startswith(\"add\"), coords.formulas.values()))\n",
    "dbls = list(filter(lambda formula: formula.name.startswith(\"dbl\"), coords.formulas.values()))\n",
    "\n",
    "def op_string(formula):\n",
    "    ops = []\n",
    "    for op in formula.code:\n",
    "        if op.operator == OpType.Mult:\n",
    "            ops.append(\"m\")\n",
    "        elif op.operator == OpType.Sqr:\n",
    "            ops.append(\"s\")\n",
    "        elif op.operator == OpType.Add:\n",
    "            ops.append(\"+\")\n",
    "        elif op.operator == OpType.Sub:\n",
    "            ops.append(\"-\")\n",
    "        elif op.operator == OpType.Pow and op.right == 3:\n",
    "            ops.append(\"sm\")\n",
    "        elif op.operator == OpType.Id:\n",
    "            pass\n",
    "        else:\n",
    "            print(op.operator, op.right)\n",
    "    return \"\".join(ops)\n",
    "\n",
    "table = [[\"Formula\", \"count\", \"ops\"]]\n",
    "for add in adds:\n",
    "    ops = op_string(add)\n",
    "    table.append([add.name, len(ops), ops])\n",
    "\n",
    "for dbl in dbls:\n",
    "    ops = op_string(dbl)\n",
    "    table.append([dbl.name, len(ops), ops])\n",
    "\n",
    "display(HTML(tabulate.tabulate(table, tablefmt=\"html\", headers=\"firstrow\", colalign=(\"left\", \"center\", \"left\"))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87fed2eb-8d8b-4e7b-8265-c696e38c4663",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}