Fix sign treatment in switch sign

2 files changed, 29 insertions, 15 deletions

diff --git a/pyecsca/ec/formula/graph.py b/pyecsca/ec/formula/graph.py
index e2f10ea..1473858 100644
--- a/pyecsca/ec/formula/graph.py
+++ b/pyecsca/ec/formula/graph.py
@@ -9,7 +9,7 @@ from ..op import CodeOp, OpType
 import matplotlib.pyplot as plt
 import networkx as nx
 from ast import parse
-from typing import Dict, List, Tuple, Set, Optional, MutableMapping
+from typing import Dict, List, Tuple, Set, Optional, MutableMapping, Any
 from copy import deepcopy
 from abc import ABC, abstractmethod
 
@@ -235,9 +235,11 @@ class EFDFormulaGraph:
     input_nodes: MutableMapping[str, InputNode]
     output_names: Set[str]
     roots: List[Node]
+    coordinate_model: Any
 
     def __init__(self, formula: EFDFormula, rename=True):
         self._formula = formula  # TODO remove, its here only for to_EFDFormula
+        self.coordinate_model = formula.coordinate_model
         self.output_names = formula.outputs
         self.input_nodes = {v: InputNode(v) for v in formula_input_variables(formula)}
         self.roots = list(self.input_nodes.values())
diff --git a/pyecsca/ec/formula/switch_sign.py b/pyecsca/ec/formula/switch_sign.py
index 66ac4b4..2e917ab 100644
--- a/pyecsca/ec/formula/switch_sign.py
+++ b/pyecsca/ec/formula/switch_sign.py
@@ -4,6 +4,8 @@ from ..op import OpType, CodeOp
 from .graph import EFDFormulaGraph, ConstantNode, Node, CodeOpNode
 from itertools import chain, combinations
 from .efd import EFDFormula
+from ..point import Point
+from ..mod import Mod
 
 
 def generate_switched_formulas(
@@ -38,23 +40,33 @@ def switch_sign(graph: EFDFormulaGraph, node_combination) -> EFDFormulaGraph:
         node, variable = queue.pop()
         queue = switch_sign_propagate(node, variable, output_signs) + queue
 
-    # TODO rewrite this hacky solution:
-    if graph._formula.coordinate_model.name.startswith("jacobian"):
-        output_signs = {out[0]: sign for out, sign in output_signs.items()}
-        X, Y, Z = (output_signs[var] for var in ["X", "Y", "Z"])
-        correct_output = X / (Z**2) == 1 and Y / (Z**3) == 1
-    elif graph._formula.coordinate_model.name.startswith("modified"):
-        output_signs = {out[0]: sign for out, sign in output_signs.items()}
-        X, Y, Z, T = (output_signs[var] for var in ["X", "Y", "Z", "T"])
-        correct_output = X / (Z**2) == 1 and Y / (Z**3) == 1
-        correct_output &= T == 1
-    else:
-        correct_output = len(set(output_signs.values())) == 1
-    if not correct_output:
-        raise BadSignSwitch
+    sign_test(output_signs, graph.coordinate_model)
     return graph
 
 
+def sign_test(output_signs, coordinate_model):
+    scale = coordinate_model.formulas.get("z", None)
+    if scale is None:
+        scale = coordinate_model.formulas.get("scale", None)
+    result_signs = {}
+    p = 7
+    out_inds = set(map(lambda x: "".join([o for o in x if o.isdigit()]), output_signs))
+    for ind in out_inds:
+        point = {}
+        for out, sign in output_signs.items():
+            out_var = out.removesuffix(ind)
+            if not out_var.isalpha():
+                continue
+            point[out_var] = Mod(sign, p)
+        point = Point(coordinate_model, **point)
+        try:
+            apoint = point.to_affine()
+        except NotImplementedError:
+            apoint = scale(p, point)[0]
+        if set(apoint.coords.values()) != set([Mod(1, p)]):
+            raise BadSignSwitch
+
+
 class BadSignSwitch(Exception):
     pass