1 files changed, 36 insertions, 19 deletions

diff --git a/pyecsca/ec/mult/fake.py b/pyecsca/ec/mult/fake.py
index d3f5b09..e0c8610 100644
--- a/pyecsca/ec/mult/fake.py
+++ b/pyecsca/ec/mult/fake.py
@@ -1,18 +1,30 @@
-from typing import List, Type, Callable
+from typing import Type, Callable
+from copy import deepcopy
 
 from pyecsca.ec.formula import (
-    Formula,
     AdditionFormula,
     DifferentialAdditionFormula,
     DoublingFormula,
     LadderFormula,
     NegationFormula,
+    ScalingFormula,
 )
-from pyecsca.ec.formula.fake import FakeFormula
+from pyecsca.ec.formula.fake import FakeAdditionFormula, FakeDifferentialAdditionFormula, \
+    FakeDoublingFormula, FakeLadderFormula, FakeNegationFormula, FakeScalingFormula
 from pyecsca.ec.mult import ScalarMultiplier
 from pyecsca.ec.params import DomainParameters
 
 
+fake_map = {
+    AdditionFormula: FakeAdditionFormula,
+    DifferentialAdditionFormula: FakeDifferentialAdditionFormula,
+    DoublingFormula: FakeDoublingFormula,
+    LadderFormula: FakeLadderFormula,
+    NegationFormula: FakeNegationFormula,
+    ScalingFormula: FakeScalingFormula,
+}
+
+
 def fake_mult(
     mult_class: Type[ScalarMultiplier], mult_factory: Callable, params: DomainParameters
 ) -> ScalarMultiplier:
@@ -24,22 +36,27 @@ def fake_mult(
     :param params: The domain parameters to use.
     :return: The multiplier.
     """
-    formula_classes: List[Type[Formula]] = list(
-        filter(
-            lambda klass: klass in mult_class.requires,
-            [
-                AdditionFormula,
-                DifferentialAdditionFormula,
-                DoublingFormula,
-                LadderFormula,
-                NegationFormula,
-            ],
-        )
-    )
     formulas = []
-    for formula in formula_classes:
-        for subclass in formula.__subclasses__():
-            if issubclass(subclass, FakeFormula):
-                formulas.append(subclass(params.curve.coordinate_model))
+    for formula, fake_formula in fake_map.items():
+        if formula in mult_class.requires:
+            formulas.append(fake_formula(params.curve.coordinate_model))
     mult = mult_factory(*formulas, short_circuit=False)
     return mult
+
+
+def turn_fake(mult: ScalarMultiplier) -> ScalarMultiplier:
+    """
+    Turn a multiplier into a fake multiplier.
+
+    :param mult: The multiplier to turn into a fake multiplier.
+    :return: The multiplier with fake formulas.
+    """
+    copy = deepcopy(mult)
+    copy.short_circuit = False
+    formulas = {}
+    for key, formula in copy.formulas.items():
+        for real, fake in fake_map.items():
+            if isinstance(formula, real):
+                formulas[key] = fake(formula.coordinate_model)
+    copy.formulas = formulas
+    return copy