Fix to_affine mapping, fix affine formulas with infinity point.

6 files changed, 42 insertions, 10 deletions

diff --git a/pyecsca/ec/coordinates.py b/pyecsca/ec/coordinates.py
index f8642cf..48a8f76 100644
--- a/pyecsca/ec/coordinates.py
+++ b/pyecsca/ec/coordinates.py
@@ -1,4 +1,4 @@
-from ast import parse, Expression, Module
+from ast import parse, Module
 from os.path import join
 from typing import List, Any, MutableMapping
 
@@ -6,9 +6,8 @@ from pkg_resources import resource_listdir, resource_isdir, resource_stream
 from public import public
 
 from .formula import (Formula, EFDFormula, AdditionEFDFormula, DoublingEFDFormula,
-                      TriplingEFDFormula,
-                      DifferentialAdditionEFDFormula, LadderEFDFormula, ScalingEFDFormula,
-                      NegationEFDFormula)
+                      TriplingEFDFormula, DifferentialAdditionEFDFormula, LadderEFDFormula,
+                      ScalingEFDFormula, NegationEFDFormula)
 
 
 @public
diff --git a/pyecsca/ec/curve.py b/pyecsca/ec/curve.py
index 1514c92..722c980 100644
--- a/pyecsca/ec/curve.py
+++ b/pyecsca/ec/curve.py
@@ -6,7 +6,7 @@ from public import public
 
 from .coordinates import CoordinateModel, AffineCoordinateModel
 from .mod import Mod
-from .model import CurveModel, ShortWeierstrassModel
+from .model import CurveModel
 from .point import Point, InfinityPoint
 
 
@@ -56,12 +56,22 @@ class EllipticCurve(object):
         return Point(AffineCoordinateModel(self.model), x=locals["x"], y=locals["y"])
 
     def affine_add(self, one: Point, other: Point) -> Point:
+        if isinstance(one, InfinityPoint):
+            return other
+        if isinstance(other, InfinityPoint):
+            return one
+        if one == other:
+            return self.affine_double(one)
         return self._execute_base_formulas(self.model.base_addition, one, other)
 
     def affine_double(self, one: Point) -> Point:
+        if isinstance(one, InfinityPoint):
+            return one
         return self._execute_base_formulas(self.model.base_doubling, one)
 
     def affine_negate(self, one: Point) -> Point:
+        if isinstance(one, InfinityPoint):
+            return one
         return self._execute_base_formulas(self.model.base_negation, one)
 
     def affine_multiply(self, point: Point, scalar: int) -> Point:
@@ -69,6 +79,8 @@ class EllipticCurve(object):
             raise ValueError
         if not isinstance(point.coordinate_model, AffineCoordinateModel):
             raise ValueError
+        if isinstance(point, InfinityPoint):
+            return point
         q = copy(point)
         r = copy(point)
 
@@ -93,7 +105,7 @@ class EllipticCurve(object):
 
     @property
     def neutral_is_affine(self):
-        """Whether the neurtal point is an affine point."""
+        """Whether the neutral point is an affine point."""
         return bool(self.model.base_neutral)
 
     def is_neutral(self, point: Point) -> bool:
diff --git a/pyecsca/ec/mod.py b/pyecsca/ec/mod.py
index 4e69790..6f9b91c 100644
--- a/pyecsca/ec/mod.py
+++ b/pyecsca/ec/mod.py
@@ -2,6 +2,7 @@ import random
 import secrets
 from functools import wraps, lru_cache
 from abc import ABC, abstractmethod
+from typing import Type
 
 has_gmp = False
 try:
@@ -360,6 +361,7 @@ class Undefined(BaseMod):
     def __pow__(self, n):
         raise NotImplementedError
 
+
 if has_gmp:
 
     class GMPMod(BaseMod):
diff --git a/pyecsca/ec/point.py b/pyecsca/ec/point.py
index 2f9c29a..a0a1879 100644
--- a/pyecsca/ec/point.py
+++ b/pyecsca/ec/point.py
@@ -89,20 +89,24 @@ class Point(object):
             for op in ops:
                 try:
                     locls[op.result] = op(**locls)
-                except:
+                except Exception:
                     continue
             result = {}
-            for var in coordinate_model.variables:  
+            for var in coordinate_model.variables:
                 if var in locls:  # Try this first.
                     result[var] = locls[var]
-                elif var == "X":  #  XXX: This just works for the stuff currently in EFD.
+                elif var == "X":
                     result[var] = self.coords["x"]
+                    if coordinate_model.name == "inverted":
+                        result[var] = result[var].inverse()
                 elif var == "Y":
                     result[var] = self.coords["y"]
+                    if coordinate_model.name == "inverted":
+                        result[var] = result[var].inverse()
                 elif var.startswith("Z"):
                     result[var] = Mod(1, curve.prime)
                 elif var == "T":
-                    result[var] = Mod(int(affine_point.coords["x"] * affine_point.coords["y"]), curve.prime)
+                    result[var] = Mod(int(self.coords["x"] * self.coords["y"]), curve.prime)
                 else:
                     raise NotImplementedError
             return action.exit(Point(coordinate_model, **result))
@@ -113,6 +117,11 @@ class Point(object):
             return False
         if self.coordinate_model.curve_model != other.coordinate_model.curve_model:
             return False
+        if "z" in self.coordinate_model.formulas:
+            formula = self.coordinate_model.formulas["z"]
+            self_mapped = formula(self)
+            other_mapped = formula(other)
+            return self_mapped == other_mapped
         return self.to_affine() == other.to_affine()
 
     def __bytes__(self):
diff --git a/test/ec/test_curve.py b/test/ec/test_curve.py
index 1421398..d15228a 100644
--- a/test/ec/test_curve.py
+++ b/test/ec/test_curve.py
@@ -55,6 +55,10 @@ class CurveTests(TestCase):
                    y=Mod(0xbcdaf32a2c08fd4271228fef35070848, self.secp128r1.curve.prime))
         self.assertIsNotNone(self.secp128r1.curve.affine_add(self.affine_base, pt))
 
+        added = self.secp128r1.curve.affine_add(self.affine_base, self.affine_base)
+        doubled = self.secp128r1.curve.affine_double(self.affine_base)
+        self.assertEqual(added, doubled)
+
     def test_affine_double(self):
         self.assertIsNotNone(self.secp128r1.curve.affine_double(self.affine_base))
 
diff --git a/test/ec/test_point.py b/test/ec/test_point.py
index 2fc4ca4..b8d355a 100644
--- a/test/ec/test_point.py
+++ b/test/ec/test_point.py
@@ -70,10 +70,16 @@ class PointTests(TestCase):
                       X=Mod(0x2, self.secp128r1.curve.prime),
                       Y=Mod(0x3, self.secp128r1.curve.prime),
                       Z=Mod(1, self.secp128r1.curve.prime))
+        third = Point(self.coords,
+                      X=Mod(0x5, self.secp128r1.curve.prime),
+                      Y=Mod(0x3, self.secp128r1.curve.prime),
+                      Z=Mod(1, self.secp128r1.curve.prime))
         self.assertTrue(pt.equals(other))
         self.assertNotEqual(pt, other)
         self.assertFalse(pt.equals(2))
         self.assertNotEqual(pt, 2)
+        self.assertFalse(pt.equals(third))
+        self.assertNotEqual(pt, third)
 
         infty_one = InfinityPoint(self.coords)
         infty_other = InfinityPoint(self.coords)