Add equality check via homogenity.

6 files changed, 258 insertions, 9 deletions

diff --git a/pyecsca/ec/mod/base.py b/pyecsca/ec/mod/base.py
index b7aa48a..7a275c0 100644
--- a/pyecsca/ec/mod/base.py
+++ b/pyecsca/ec/mod/base.py
@@ -3,7 +3,7 @@ import secrets
 from functools import lru_cache, wraps
 
 from public import public
-from typing import Tuple, Any, Dict, Type
+from typing import Tuple, Any, Dict, Type, Set
 
 from pyecsca.ec.context import ResultAction
 from pyecsca.misc.cfg import getconfig
@@ -65,6 +65,44 @@ def jacobi(x: int, n: int) -> int:
 
 
 @public
+def square_roots(x: "Mod") -> Set["Mod"]:
+    """
+
+    :param x:
+    :return:
+    """
+    if not x.is_residue():
+        return set()
+    sqrt = x.sqrt()
+    return {sqrt, -sqrt}
+
+
+@public
+def cube_roots(x: "Mod") -> Set["Mod"]:
+    """
+
+    :param x:
+    :return:
+    """
+    if not x.is_cubic_residue():
+        return set()
+    cube_root = x.cube_root()
+    if (x.n - 1) % 3 != 0:
+        # gcd(3, p - 1) = 1
+        return {cube_root}
+    else:
+        # gcd(3, p - 1) = 3
+        m = (x.n - 1) // 3
+        # Find 3rd root of unity
+        while True:
+            z = x.__class__(random.randrange(2, x.n - 1), x.n)
+            r = z ** m
+            if r != 1:
+                break
+        return {cube_root, cube_root * r, cube_root * (r ** 2)}
+
+
+@public
 @lru_cache
 def miller_rabin(n: int, rounds: int = 50) -> bool:
     """Miller-Rabin probabilistic primality test."""
@@ -215,6 +253,20 @@ class Mod:
         """
         raise NotImplementedError
 
+    def is_cubic_residue(self) -> bool:
+        """
+        Whether this element is a cubic residue (only implemented for prime modulus).
+        """
+        raise NotImplementedError
+
+    def cube_root(self) -> "Mod":
+        """
+        Compute the cube root of this element (only implemented for prime modulus).
+
+        Uses the Adleman-Manders-Miller algorithm (which is adapted Tonelli-Shanks).
+        """
+        raise NotImplementedError
+
     @_check
     def __mul__(self, other) -> "Mod":
         return self.__class__((self.x * other.x) % self.n, self.n)
diff --git a/pyecsca/ec/mod/raw.py b/pyecsca/ec/mod/raw.py
index 90e0f9d..0d76594 100644
--- a/pyecsca/ec/mod/raw.py
+++ b/pyecsca/ec/mod/raw.py
@@ -77,6 +77,64 @@ class RawMod(Mod):
             r *= b
         return r
 
+    def is_cubic_residue(self):
+        if not miller_rabin(self.n):
+            raise NotImplementedError
+        if self.x in (0, 1):
+            return True
+        if self.n % 3 == 2:
+            return True
+        pm1 = self.n - 1
+        r = self ** (pm1 // 3)
+        return r == 1
+
+    def cube_root(self) -> "RawMod":
+        if not miller_rabin(self.n):
+            raise NotImplementedError
+        if self.x == 0:
+            return RawMod(0, self.n)
+        if self.x == 1:
+            return RawMod(1, self.n)
+        if not self.is_cubic_residue():
+            raise_non_residue()
+        if self.n % 3 == 2:
+            inv3 = RawMod(3, self.n - 1).inverse()
+            return self ** int(inv3)
+        q = self.n - 1
+        s = 0
+        while q % 3 == 0:
+            q //= 3
+            s += 1
+        t = q
+        if t % 3 == 1:
+            k = (t - 1) // 3
+        else:
+            k = (t + 1) // 3
+
+        b = 2
+        while RawMod(b, self.n).is_cubic_residue():
+            b += 1
+
+        c = RawMod(b, self.n) ** t
+        r = self ** t
+        h = RawMod(1, self.n)
+        cp = c ** (3 ** (s - 1))
+        c = c.inverse()
+        for i in range(1, s):
+            d = r ** (3 ** (s - i - 1))
+            if d == cp:
+                h *= c
+                r *= c ** 3
+            elif d != 1:
+                h *= c ** 2
+                r *= c ** 6
+            c = c ** 3
+        x: RawMod = (self ** k) * h
+        if t % 3 == 1:
+            return x.inverse()
+        else:
+            return x
+
     def __bytes__(self):
         return self.x.to_bytes((self.n.bit_length() + 7) // 8, byteorder="big")
 
diff --git a/pyecsca/ec/mod/symbolic.py b/pyecsca/ec/mod/symbolic.py
index 57e7233..c3d006b 100644
--- a/pyecsca/ec/mod/symbolic.py
+++ b/pyecsca/ec/mod/symbolic.py
@@ -49,16 +49,22 @@ class SymbolicMod(Mod):
     def __neg__(self) -> "SymbolicMod":
         return self.__class__(-self.x, self.n)
 
-    def bit_length(self):
+    def bit_length(self) -> int:
         raise NotImplementedError
 
     def inverse(self) -> "SymbolicMod":
         return self.__class__(self.x ** (-1), self.n)
 
+    def is_residue(self) -> bool:
+        raise NotImplementedError
+
     def sqrt(self) -> "SymbolicMod":
         raise NotImplementedError
 
-    def is_residue(self):
+    def is_cubic_residue(self) -> bool:
+        raise NotImplementedError
+
+    def cube_root(self) -> "SymbolicMod":
         raise NotImplementedError
 
     def __invert__(self) -> "SymbolicMod":
diff --git a/pyecsca/ec/point.py b/pyecsca/ec/point.py
index 162b9aa..1833f46 100644
--- a/pyecsca/ec/point.py
+++ b/pyecsca/ec/point.py
@@ -1,13 +1,15 @@
 """Provides a :py:class:`.Point` class and a special :py:class:`.InfinityPoint` class for the point at infinity."""
 
 from copy import copy
-from typing import Mapping, TYPE_CHECKING
+from operator import itemgetter
+from typing import Mapping, Set, TYPE_CHECKING
 
 from public import public
 
 from pyecsca.ec.context import ResultAction
 from pyecsca.ec.coordinates import AffineCoordinateModel, CoordinateModel
-from pyecsca.ec.mod import Mod, Undefined, mod
+from pyecsca.ec.mod import Mod, Undefined, mod, square_roots, cube_roots
+from pyecsca.ec.error import NonResidueError
 from pyecsca.ec.op import CodeOp
 
 
@@ -141,11 +143,13 @@ class Point:
             if randomized:
                 lmbd = Mod.random(curve.prime)
                 for var, value in result.items():
-                    result[var] = value * (lmbd ** coordinate_model.homogweights[var])
+                    weight = coordinate_model.homogweights[var]
+                    lpow = lmbd ** weight
+                    result[var] = value * lpow
             return action.exit(Point(coordinate_model, **result))
 
     def equals_affine(self, other: "Point") -> bool:
-        """Test whether this point is equal to :paramref:`~.equals_affine.other` irrespective of the coordinate model (in the affine sense)."""
+        """Test whether this point is equal to :paramref:`~.equals_affine.other` in the affine sense."""
         if not isinstance(other, Point) or isinstance(other, InfinityPoint):
             return False
         if self.coordinate_model.curve_model != other.coordinate_model.curve_model:
@@ -158,7 +162,7 @@ class Point:
 
         The "z" scaling formula maps the projective class to a single representative.
 
-        :param other: The point to compare
+        :param other: The point to compare.
         :raises ValueError: If the "z" formula is not available for the coordinate system.
         :return: Whether the points are equal.
         """
@@ -174,6 +178,60 @@ class Point:
         else:
             raise ValueError("No scaling formula available.")
 
+    def equals_homog(self, other: "Point") -> bool:
+        """
+        Test whether this point is equal to :paramref:`~.equals_homog.other` in the coordinate system.
+
+        :param other: The point to compare.
+        :return: Whether the points are equal.
+        """
+        if not isinstance(other, Point) or isinstance(other, InfinityPoint):
+            return False
+        if self.coordinate_model.curve_model != other.coordinate_model.curve_model:
+            return False
+        weights = sorted(self.coordinate_model.homogweights.items(), key=itemgetter(1))
+        lambdas: Set[Mod] = set()
+        for var, weight in weights:
+            var1 = self.coords[var]
+            var2 = other.coords[var]
+            if var1 == 0 and var2 == 0:
+                continue
+            if var1 == 0 or var2 == 0:
+                return False
+            val = var2 / var1
+            if not lambdas:
+                if weight == 1:
+                    lambdas.add(val)
+                elif weight == 2:
+                    if not val.is_residue():
+                        return False
+                    lambdas.update(square_roots(val))
+                elif weight == 3:
+                    if not val.is_cubic_residue():
+                        return False
+                    lambdas.update(cube_roots(val))
+                elif weight == 4:
+                    if not val.is_residue():
+                        return False
+                    first = val.sqrt()
+                    try:
+                        lambdas.update(square_roots(first))
+                    except NonResidueError:
+                        pass
+                    try:
+                        lambdas.update(square_roots(-first))
+                    except NonResidueError:
+                        pass
+                else:
+                    raise NotImplementedError(
+                        f"Equality checking does not support {weight} weight."
+                    )
+            else:
+                lambdas = set(filter(lambda candidate: candidate ** weight == val, lambdas))
+                if not lambdas:
+                    return False
+        return True
+
     def equals(self, other: "Point") -> bool:
         """Test whether this point is equal to `other` irrespective of the coordinate model (in the affine sense)."""
         return self.equals_affine(other)
@@ -240,6 +298,9 @@ class InfinityPoint(Point):
     def equals_scaled(self, other: "Point") -> bool:
         return self == other
 
+    def equals_homog(self, other: "Point") -> bool:
+        return self == other
+
     def equals(self, other: "Point") -> bool:
         return self == other
 
diff --git a/test/ec/test_mod.py b/test/ec/test_mod.py
index 18eeeac..724ea42 100644
--- a/test/ec/test_mod.py
+++ b/test/ec/test_mod.py
@@ -17,7 +17,7 @@ from pyecsca.ec.mod import (
     miller_rabin,
     RawMod,
     SymbolicMod,
-    jacobi,
+    jacobi, cube_roots,
 )
 from pyecsca.ec.mod.gmp import has_gmp
 from pyecsca.ec.mod.flint import has_flint
@@ -116,6 +116,42 @@ def test_sqrt():
     getconfig().ec.non_residue_action = "error"
 
 
+def test_is_cubic_residue():
+    # p is 2 mod 3 so all are residues
+    p = 11
+    for i in range(11):
+        assert mod(i, p).is_cubic_residue()
+    p = 13
+    assert not mod(4, p).is_cubic_residue()
+    assert mod(5, p).is_cubic_residue()
+    assert not mod(6, p).is_cubic_residue()
+    assert not mod(7, p).is_cubic_residue()
+
+
+def test_cube_root():
+    p = 11
+    for i in range(11):
+        assert len(cube_roots(mod(i, p))) == 1
+    assert mod(2, p).cube_root() == 7
+    assert mod(3, p).cube_root() == 9
+    assert mod(4, p).cube_root() == 5
+    assert mod(5, p).cube_root() == 3
+    assert mod(6, p).cube_root() == 8
+    assert mod(7, p).cube_root() == 6
+
+    p = 13
+    assert mod(1, p).cube_root() == 1
+    assert mod(0, p).cube_root() == 0
+    assert mod(5, p).cube_root() == 8
+    assert mod(8, p).cube_root() == 5
+    assert mod(12, p).cube_root() == 12
+
+    assert cube_roots(mod(0, p)) == {mod(0, p)}
+    assert cube_roots(mod(1, p)) == {mod(1, p), mod(3, p), mod(9, p)}
+    assert cube_roots(mod(5, p)) == {mod(8, p), mod(11, p), mod(7, p)}
+    assert cube_roots(mod(8, p)) == {mod(5, p), mod(2, p), mod(6, p)}
+
+
 def test_eq():
     assert mod(1, 7) == 1
     assert mod(1, 7) != "1"
diff --git a/test/ec/test_point.py b/test/ec/test_point.py
index 6dfced5..789b4f1 100644
--- a/test/ec/test_point.py
+++ b/test/ec/test_point.py
@@ -8,6 +8,7 @@ from pyecsca.ec.mod import mod
 from pyecsca.ec.model import ShortWeierstrassModel, MontgomeryModel
 from pyecsca.ec.params import get_params
 from pyecsca.ec.point import Point, InfinityPoint
+from pyecsca.ec.error import UnsatisfiedAssumptionError
 
 
 @pytest.fixture()
@@ -129,15 +130,23 @@ def test_equals(secp128r1, secp128r1_coords):
     assert pt.equals_affine(other)
     assert not pt.equals_scaled(third)
 
+    assert pt.equals_homog(pt)
+    assert pt.equals_homog(other)
+    assert other.equals_homog(pt)
+    assert not pt.equals_homog(third)
+    assert not third.equals_homog(pt)
+
     infty_one = InfinityPoint(secp128r1_coords)
     infty_other = InfinityPoint(secp128r1_coords)
     assert infty_one.equals(infty_other)
     assert infty_one.equals_affine(infty_other)
     assert infty_one.equals_scaled(infty_other)
+    assert infty_one.equals_homog(infty_other)
     assert infty_one == infty_other
     assert not pt.equals(infty_one)
     assert not pt.equals_affine(infty_one)
     assert not pt.equals_scaled(infty_one)
+    assert not pt.equals_homog(infty_one)
 
     mont = MontgomeryModel()
     different = Point(
@@ -152,6 +161,33 @@ def test_equals(secp128r1, secp128r1_coords):
     assert pt != different
 
 
+def test_homog():
+    model = ShortWeierstrassModel()
+    for coords_name, coords in model.coordinates.items():
+        try:
+            params = get_params("secg", "secp128r1", coords_name, infty=True)
+        except UnsatisfiedAssumptionError:
+            continue
+        infty = params.curve.neutral
+        rand_aff = params.curve.affine_random()
+        one1 = rand_aff.to_model(coords, params.curve)
+        one2 = rand_aff.to_model(coords, params.curve, randomized=True)
+        one3 = rand_aff.to_model(coords, params.curve, randomized=True)
+        assert one1.equals_homog(one2)
+        assert one1.equals_homog(one3)
+        assert one2.equals_homog(one3)
+        assert not one1.equals_homog(infty)
+        assert infty.equals_homog(infty)
+        while True:
+            other_aff = params.curve.affine_random()
+            if other_aff != rand_aff:
+                break
+        other = other_aff.to_model(coords, params.curve)
+        assert not one1.equals_homog(other)
+        assert not one2.equals_homog(other)
+        assert not one3.equals_homog(other)
+
+
 def test_bytes(secp128r1, secp128r1_coords):
     pt = Point(
         secp128r1_coords,