4 files changed, 185 insertions, 8 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