Apply leakage model to execution tree.

3 files changed, 60 insertions, 2 deletions

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..226734d 100644
--- a/pyecsca/ec/formula.py
+++ b/pyecsca/ec/formula.py
@@ -55,6 +55,8 @@ class FormulaAction(ResultAction):
     """The input points."""
     intermediates: MutableMapping[str, List[OpResult]]
     """Intermediates computed during execution."""
+    op_results: List[OpResult]
+    """"""
     outputs: MutableMapping[str, OpResult]
     """The output variables."""
     output_points: List[Any]
@@ -65,6 +67,7 @@ class FormulaAction(ResultAction):
         self.formula = formula
         self.inputs = inputs
         self.intermediates = {}
+        self.op_results = []
         self.outputs = {}
         self.input_points = list(points)
         self.output_points = []
@@ -76,8 +79,10 @@ class FormulaAction(ResultAction):
                 parents.append(self.intermediates[parent][-1])
             elif parent in self.inputs:
                 parents.append(self.inputs[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/test/sca/test_leakage_models.py b/test/sca/test_leakage_models.py
index f951083..d1e6fd1 100644
--- a/test/sca/test_leakage_models.py
+++ b/test/sca/test_leakage_models.py
@@ -1,6 +1,10 @@
 from unittest import TestCase
 
+from pyecsca.ec.context import local, DefaultContext
+from pyecsca.ec.formula import FormulaAction
 from pyecsca.ec.mod import Mod
+from pyecsca.ec.mult import LTRMultiplier
+from pyecsca.ec.params import get_params
 from pyecsca.sca.attack.leakage_model import Identity, Bit, Slice, HammingWeight, HammingDistance, BitLength
 
 
@@ -46,3 +50,41 @@ class LeakageModelTests(TestCase):
         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(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)