Add path context.

2 files changed, 46 insertions, 3 deletions

diff --git a/pyecsca/ec/context.py b/pyecsca/ec/context.py
index 6a267e6..144ca92 100644
--- a/pyecsca/ec/context.py
+++ b/pyecsca/ec/context.py
@@ -2,7 +2,7 @@ from abc import abstractmethod, ABC
 from collections import OrderedDict
 from contextvars import ContextVar, Token
 from copy import deepcopy
-from typing import List, Optional, ContextManager, Any, Tuple
+from typing import List, Optional, ContextManager, Any, Tuple, Sequence
 
 from public import public
 
@@ -173,6 +173,36 @@ class DefaultContext(Context):
         return f"{self.__class__.__name__}({self.actions!r}, current={self.current!r})"
 
 
+@public
+class PathContext(Context):
+    """A context that traces targeted actions."""
+    path: List[int]
+    current: List[int]
+    current_depth: int
+    value: Any
+
+    def enter_action(self, action: Action) -> None:
+        if self.current_depth == len(self.current):
+            self.current.append(0)
+        else:
+            self.current[self.current_depth] += 1
+        self.current_depth += 1
+        if self.path == self.current[:self.current_depth]:
+            self.value = action
+
+    def exit_action(self, action: Action) -> None:
+        if self.current_depth != len(self.current):
+            self.current.pop()
+        self.current_depth -= 1
+
+
+    def __init__(self, path: Sequence[int]):
+        self.path = list(path)
+        self.current = []
+        self.current_depth = 0
+        self.value = None
+
+
 _actual_context: ContextVar[Context] = ContextVar("operational_context", default=NullContext())
 
 
diff --git a/test/ec/test_context.py b/test/ec/test_context.py
index be84eb2..c3112d4 100644
--- a/test/ec/test_context.py
+++ b/test/ec/test_context.py
@@ -1,7 +1,8 @@
 from unittest import TestCase
 
 from pyecsca.ec.context import (local, DefaultContext, NullContext, getcontext,
-                                setcontext, resetcontext, Tree)
+                                setcontext, resetcontext, Tree, PathContext)
+from pyecsca.ec.key_generation import KeygenAction, KeyGeneration
 from pyecsca.ec.params import get_params
 from pyecsca.ec.mod import RandomModAction
 from pyecsca.ec.mult import LTRMultiplier, ScalarMultiplicationAction
@@ -51,7 +52,8 @@ class ContextTests(TestCase):
         self.base = self.secp128r1.generator
         self.coords = self.secp128r1.curve.coordinate_model
         self.mult = LTRMultiplier(self.coords.formulas["add-1998-cmo"],
-                                  self.coords.formulas["dbl-1998-cmo"], self.coords.formulas["z"])
+                                  self.coords.formulas["dbl-1998-cmo"], self.coords.formulas["z"],
+                                  always=True)
         self.mult.init(self.secp128r1, self.base)
 
     def test_null(self):
@@ -76,6 +78,17 @@ class ContextTests(TestCase):
             with self.assertRaises(ValueError):
                 default.exit_action(RandomModAction(7))
 
+    def test_path(self):
+        with local(PathContext([0, 1])) as ctx:
+            key_generator = KeyGeneration(self.mult, self.secp128r1, True)
+            key_generator.generate()
+        self.assertIsInstance(ctx.value, ScalarMultiplicationAction)
+        with local(PathContext([0, 1, 7])) as ctx:
+            key_generator = KeyGeneration(self.mult, self.secp128r1, True)
+            key_generator.generate()
+        print(ctx.value)
+
+
     def test_str(self):
         with local(DefaultContext()) as default:
             self.mult.multiply(59)