Merge pull request #37 from J08nY/feat/lm

Add leakage models to core

13 files changed, 274 insertions, 12 deletions

diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
index 0d02c48..a947379 100644
--- a/.github/workflows/lint.yml
+++ b/.github/workflows/lint.yml
@@ -33,6 +33,9 @@ jobs:
       - name: Install system dependencies
         run: |
           sudo apt-get install -y $PS_PACKAGES $OTHER_PACKAGES $GMP_PACKAGES
+      - name: Install numpy
+        run: |
+          pip install "numpy<1.25"
       - name: Install picoscope bindings
         run: |
           git clone https://github.com/colinoflynn/pico-python && cd pico-python && python setup.py install && cd ..
diff --git a/.github/workflows/perf.yml b/.github/workflows/perf.yml
index 6534118..4eff89e 100644
--- a/.github/workflows/perf.yml
+++ b/.github/workflows/perf.yml
@@ -43,6 +43,9 @@ jobs:
         run: |
           sudo apt-get install -y $PS_PACKAGES $OTHER_PACKAGES
           if [ $USE_GMP == 1 ]; then sudo apt-get install -y $GMP_PACKAGES; fi
+      - name: Install numpy
+        run: |
+          pip install "numpy<1.25"
       - name: Install picoscope bindings
         run: |
           git clone https://github.com/colinoflynn/pico-python && cd pico-python && python setup.py install && cd ..
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 42eb8ba..d7aed43 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -43,6 +43,9 @@ jobs:
         run: |
           sudo apt-get install -y $PS_PACKAGES $OTHER_PACKAGES
           if [ $USE_GMP == 1 ]; then sudo apt-get install -y $GMP_PACKAGES; fi
+      - name: Install numpy
+        run: |
+          pip install "numpy<1.25"
       - name: Install picoscope bindings
         run: |
           git clone https://github.com/colinoflynn/pico-python && cd pico-python && python setup.py install && cd ..
diff --git a/Makefile b/Makefile
index 547797d..1ae4371 100644
--- a/Makefile
+++ b/Makefile
@@ -4,7 +4,7 @@ ec.test_mult ec.test_naf ec.test_op ec.test_point ec.test_signature ec.test_tran
 
 SCA_TESTS = sca.test_align sca.test_combine sca.test_edit sca.test_filter sca.test_match sca.test_process \
 sca.test_sampling sca.test_target sca.test_test sca.test_trace sca.test_traceset sca.test_plot sca.test_rpa \
-sca.test_stacked_combine
+sca.test_stacked_combine sca.test_leakage_models
 
 TESTS = ${EC_TESTS} ${SCA_TESTS}
 
diff --git a/notebook b/notebook
-Subproject 06b53bea04b38564b1205e48d802601aa0d255a
+Subproject 75d7523ba04a0a8242266a39776fae81adbc885
diff --git a/pyecsca/ec/context.py b/pyecsca/ec/context.py
index 49c886d..dd609f8 100644
--- a/pyecsca/ec/context.py
+++ b/pyecsca/ec/context.py
@@ -15,7 +15,7 @@ A :py:class:`NullContext` does not trace any actions and is the default context.
 from abc import abstractmethod, ABC
 from collections import OrderedDict
 from copy import deepcopy
-from typing import List, Optional, ContextManager, Any, Tuple, Sequence
+from typing import List, Optional, ContextManager, Any, Tuple, Sequence, Callable
 
 from public import public
 
@@ -131,6 +131,17 @@ class Tree(OrderedDict):
                 result += "\t" * depth + str(key) + ":" + str(value) + "\n"
         return result
 
+    def walk(self, callback: Callable[[Any], None]) -> None:
+        """
+        Walk the tree, depth-first, with the callback.
+
+        :param callback: The callback to call for all values in the tree.
+        """
+        for key, val in self.items():
+            callback(key)
+            if isinstance(val, Tree):
+                val.walk(callback)
+
     def __repr__(self):
         return self.repr()
 
diff --git a/pyecsca/ec/formula.py b/pyecsca/ec/formula.py
index ec8f9c0..317788b 100644
--- a/pyecsca/ec/formula.py
+++ b/pyecsca/ec/formula.py
@@ -29,6 +29,8 @@ class OpResult:
     value: Mod
 
     def __init__(self, name: str, value: Mod, op: OpType, *parents: Any):
+        if len(parents) != op.num_inputs:
+            raise ValueError(f"Wrong number of parents ({len(parents)}) to OpResult: {op} ({op.num_inputs}).")
         self.parents = tuple(parents)
         self.name = name
         self.value = value
@@ -55,6 +57,8 @@ class FormulaAction(ResultAction):
     """The input points."""
     intermediates: MutableMapping[str, List[OpResult]]
     """Intermediates computed during execution."""
+    op_results: List[OpResult]
+    """The intermediates but ordered as they were computed."""
     outputs: MutableMapping[str, OpResult]
     """The output variables."""
     output_points: List[Any]
@@ -65,19 +69,25 @@ class FormulaAction(ResultAction):
         self.formula = formula
         self.inputs = inputs
         self.intermediates = {}
+        self.op_results = []
         self.outputs = {}
         self.input_points = list(points)
         self.output_points = []
 
     def add_operation(self, op: CodeOp, value: Mod):
-        parents: List[Union[Mod, OpResult]] = []
-        for parent in {*op.variables, *op.parameters}:
-            if parent in self.intermediates:
-                parents.append(self.intermediates[parent][-1])
-            elif parent in self.inputs:
-                parents.append(self.inputs[parent])
+        parents: List[Union[int, Mod, OpResult]] = []
+        for parent in op.parents:
+            if isinstance(parent, str):
+                if parent in self.intermediates:
+                    parents.append(self.intermediates[parent][-1])
+                elif parent in self.inputs:
+                    parents.append(self.inputs[parent])
+            else:
+                parents.append(parent)
+        result = OpResult(op.result, value, op.operator, *parents)
         li = self.intermediates.setdefault(op.result, [])
-        li.append(OpResult(op.result, value, op.operator, *parents))
+        li.append(result)
+        self.op_results.append(result)
 
     def add_result(self, point: Any, **outputs: Mod):
         for k in outputs:
diff --git a/pyecsca/ec/op.py b/pyecsca/ec/op.py
index d656bc3..b83562c 100644
--- a/pyecsca/ec/op.py
+++ b/pyecsca/ec/op.py
@@ -18,7 +18,7 @@ from ast import (
 )
 from enum import Enum
 from types import CodeType
-from typing import FrozenSet, cast, Any, Optional, Union
+from typing import FrozenSet, cast, Any, Optional, Union, Tuple
 
 from public import public
 
@@ -54,6 +54,8 @@ class CodeOp:
     """The parameters used in the operation (e.g. `a`, `b`)."""
     variables: FrozenSet[str]
     """The variables used in the operation (e.g. `X1`, `Z2`)."""
+    constants: FrozenSet[int]  # TODO: Might not be only int? See issue in Formula eval.
+    """The constants used in the operation."""
     code: Module
     """The code of the operation."""
     operator: OpType
@@ -126,6 +128,15 @@ class CodeOp:
             return OpType.Pow
         return OpType.Id
 
+    @property
+    def parents(self) -> Tuple[Union[str, int]]:
+        if self.operator == OpType.Inv or self.operator == OpType.Neg:
+            return self.right,  # type: ignore
+        elif self.operator == OpType.Sqr or self.operator == OpType.Id:
+            return self.left,  # type: ignore
+        else:
+            return self.left, self.right  # type: ignore
+
     def __str__(self):
         return f"{self.result} = {self.left if self.left is not None else ''}{self.operator.op_str}{self.right if self.right is not None else ''}"
 
diff --git a/pyecsca/sca/__init__.py b/pyecsca/sca/__init__.py
index 5b359b8..89e8527 100644
--- a/pyecsca/sca/__init__.py
+++ b/pyecsca/sca/__init__.py
@@ -6,3 +6,4 @@ from .target import *
 from .trace import *
 from .trace_set import *
 from .stacked_traces import *
+from .attack import *
diff --git a/pyecsca/sca/attack/__init__.py b/pyecsca/sca/attack/__init__.py
new file mode 100644
index 0000000..23e27cb
--- /dev/null
+++ b/pyecsca/sca/attack/__init__.py
@@ -0,0 +1 @@
+from .leakage_model import *
diff --git a/pyecsca/sca/attack/leakage_model.py b/pyecsca/sca/attack/leakage_model.py
new file mode 100644
index 0000000..e589940
--- /dev/null
+++ b/pyecsca/sca/attack/leakage_model.py
@@ -0,0 +1,100 @@
+import abc
+import sys
+from typing import Literal, ClassVar
+
+from numpy.random import default_rng
+from public import public
+
+if sys.version_info[0] < 3 or sys.version_info[0] == 3 and sys.version_info[1] < 10:
+    def hw(i):
+        return bin(i).count("1")
+else:
+    def hw(i):
+        return i.bit_count()
+
+
+@public
+class NormalNoice:
+    """
+    https://www.youtube.com/watch?v=SAfq55aiqPc
+    """
+
+    def __init__(self, mean: float, sdev: float):
+        self.rng = default_rng()
+        self.mean = mean
+        self.sdev = sdev
+
+    def __call__(self, *args, **kwargs) -> float:
+        return args[0] + self.rng.normal(self.mean, self.sdev)
+
+
+@public
+class LeakageModel(abc.ABC):
+    num_args: ClassVar[int]
+
+    @abc.abstractmethod
+    def __call__(self, *args, **kwargs) -> int:
+        raise NotImplementedError
+
+
+@public
+class Identity(LeakageModel):
+    num_args = 1
+
+    def __call__(self, *args, **kwargs) -> int:
+        return int(args[0])
+
+
+@public
+class Bit(LeakageModel):
+    num_args = 1
+
+    def __init__(self, which: int):
+        if which < 0:
+            raise ValueError("which must be >= 0.")
+        self.which = which
+        self.mask = 1 << which
+
+    def __call__(self, *args, **kwargs) -> Literal[0, 1]:
+        return (int(args[0]) & self.mask) >> self.which  # type: ignore
+
+
+@public
+class Slice(LeakageModel):
+    num_args = 1
+
+    def __init__(self, begin: int, end: int):
+        if begin > end:
+            raise ValueError("begin must be <= than end.")
+        self.begin = begin
+        self.end = end
+        self.mask = 0
+        for i in range(begin, end):
+            self.mask |= 1 << i
+
+    def __call__(self, *args, **kwargs) -> int:
+        return (int(args[0]) & self.mask) >> self.begin
+
+
+@public
+class HammingWeight(LeakageModel):
+    num_args = 1
+
+    def __call__(self, *args, **kwargs) -> int:
+        return hw(int(args[0]))
+
+
+@public
+class HammingDistance(LeakageModel):
+    num_args = 2
+
+    def __call__(self, *args, **kwargs) -> int:
+        return hw(int(args[0]) ^ int(args[1]))
+
+
+@public
+class BitLength(LeakageModel):
+    num_args = 1
+
+    def __call__(self, *args, **kwargs) -> int:
+        return int(args[0]).bit_length()
diff --git a/setup.py b/setup.py
index 1a23727..d2a304f 100644
--- a/setup.py
+++ b/setup.py
@@ -30,7 +30,7 @@ setup(
     # install_package_data=True,
     python_requires='>=3.8',
     install_requires=[
-        "numpy==1.24.3",
+        "numpy==1.24.4",
         "scipy",
         "sympy>=1.7.1",
         "atpublic",
@@ -44,7 +44,7 @@ setup(
         "datashader",
         "xarray",
         "astunparse",
-        "numba==0.57.0"
+        "numba==0.57.1"
     ],
     extras_require={
         "picoscope_sdk": ["picosdk"],
diff --git a/test/sca/test_leakage_models.py b/test/sca/test_leakage_models.py
new file mode 100644
index 0000000..e9da42c
--- /dev/null
+++ b/test/sca/test_leakage_models.py
@@ -0,0 +1,119 @@
+from unittest import TestCase
+
+from pyecsca.ec.context import local, DefaultContext
+from pyecsca.ec.formula import FormulaAction, OpResult
+from pyecsca.ec.mod import Mod
+from pyecsca.ec.mult import LTRMultiplier
+from pyecsca.ec.op import OpType
+from pyecsca.ec.params import get_params
+from pyecsca.sca.attack.leakage_model import Identity, Bit, Slice, HammingWeight, HammingDistance, BitLength
+
+
+class LeakageModelTests(TestCase):
+
+    def test_identity(self):
+        val = Mod(3, 7)
+        lm = Identity()
+        self.assertEqual(lm(val), 3)
+
+    def test_bit(self):
+        val = Mod(3, 7)
+        lm = Bit(0)
+        self.assertEqual(lm(val), 1)
+        lm = Bit(4)
+        self.assertEqual(lm(val), 0)
+        with self.assertRaises(ValueError):
+            Bit(-3)
+
+    def test_slice(self):
+        val = Mod(0b11110000, 0xf00)
+        lm = Slice(0, 4)
+        self.assertEqual(lm(val), 0)
+        lm = Slice(1, 5)
+        self.assertEqual(lm(val), 0b1000)
+        lm = Slice(4, 8)
+        self.assertEqual(lm(val), 0b1111)
+        with self.assertRaises(ValueError):
+            Slice(7, 1)
+
+    def test_hamming_weight(self):
+        val = Mod(0b11110000, 0xf00)
+        lm = HammingWeight()
+        self.assertEqual(lm(val), 4)
+
+    def test_hamming_distance(self):
+        a = Mod(0b11110000, 0xf00)
+        b = Mod(0b00010000, 0xf00)
+        lm = HammingDistance()
+        self.assertEqual(lm(a, b), 3)
+
+    def test_bit_length(self):
+        a = Mod(0b11110000, 0xf00)
+        lm = BitLength()
+        self.assertEqual(lm(a), 8)
+
+
+class ModelTraceTests(TestCase):
+
+    def setUp(self):
+        self.secp128r1 = get_params("secg", "secp128r1", "projective")
+        self.base = self.secp128r1.generator
+        self.coords = self.secp128r1.curve.coordinate_model
+        self.add = self.coords.formulas["add-1998-cmo"]
+        self.dbl = self.coords.formulas["dbl-1998-cmo"]
+        self.neg = self.coords.formulas["neg"]
+        self.scale = self.coords.formulas["z"]
+
+    def test_mult_hw(self):
+        scalar = 0x123456789
+        mult = LTRMultiplier(
+            self.add,
+            self.dbl,
+            self.scale,
+            always=True,
+            complete=False,
+            short_circuit=True,
+        )
+        with local(DefaultContext()) as ctx:
+            mult.init(self.secp128r1, self.base)
+            mult.multiply(scalar)
+
+        lm = HammingWeight()
+        trace = []
+
+        def callback(action):
+            if isinstance(action, FormulaAction):
+                for intermediate in action.op_results:
+                    leak = lm(intermediate.value)
+                    trace.append(leak)
+
+        ctx.actions.walk(callback)
+        self.assertGreater(len(trace), 0)
+
+    def test_mult_hd(self):
+        scalar = 0x123456789
+        mult = LTRMultiplier(
+            self.add,
+            self.dbl,
+            self.scale,
+            always=True,
+            complete=False,
+            short_circuit=True,
+        )
+        with local(DefaultContext()) as ctx:
+            mult.init(self.secp128r1, self.base)
+            mult.multiply(scalar)
+
+        lm = HammingDistance()
+        trace = []
+
+        def callback(action):
+            if isinstance(action, FormulaAction):
+                for intermediate in action.op_results:
+                    if intermediate.op == OpType.Mult:
+                        values = list(map(lambda v: v.value if isinstance(v, OpResult) else v, intermediate.parents))
+                        leak = lm(*values)
+                        trace.append(leak)
+
+        ctx.actions.walk(callback)
+        self.assertGreater(len(trace), 0)