From 51f0dd7bf448802c4c8651cc54b099e89cc35dec Mon Sep 17 00:00:00 2001
From: J08nY
Date: Sat, 29 Mar 2025 19:20:34 +0100
Subject: Fix mods.

---
 pyecsca/ec/mod/base.py | 16 ++++++++--------
 pyecsca/ec/mod/gmp.py  |  4 ++--
 pyecsca/ec/mod/raw.py  |  4 ++--
 3 files changed, 12 insertions(+), 12 deletions(-)

(limited to 'pyecsca/ec')

diff --git a/pyecsca/ec/mod/base.py b/pyecsca/ec/mod/base.py
index 39243c0..c57a1b9 100644
--- a/pyecsca/ec/mod/base.py
+++ b/pyecsca/ec/mod/base.py
@@ -117,11 +117,11 @@ def square_root_inner(x: M, intwrap, mod_class) -> M:
         s += 1
 
     z = intwrap(2)
-    while mod_class(z, x.n).is_residue():
+    while mod_class(z).is_residue():
         z += 1
 
     m = s
-    c = mod_class(z, x.n) ** q
+    c = mod_class(z) ** q
     t = x ** q
     r_exp = (q + 1) // 2
     r = x ** r_exp
@@ -131,8 +131,8 @@ def square_root_inner(x: M, intwrap, mod_class) -> M:
         while not (t ** (2 ** i)) == 1:
             i += 1
         two_exp = m - (i + 1)
-        b = c ** int(mod_class(2, x.n) ** two_exp)
-        m = int(mod_class(i, x.n))
+        b = c ** int(mod_class(intwrap(2)) ** two_exp)
+        m = int(mod_class(intwrap(i)))
         c = b ** 2
         t *= c
         r *= b
@@ -141,7 +141,7 @@ def square_root_inner(x: M, intwrap, mod_class) -> M:
 
 def cube_root_inner(x: M, intwrap, mod_class) -> M:
     if x.n % 3 == 2:
-        inv3 = mod_class(intwrap(3), x.n - 1).inverse()
+        inv3 = x.__class__(intwrap(3), x.n - 1).inverse()
         return x ** int(inv3)  # type: ignore
     q = x.n - 1
     s = 0
@@ -155,12 +155,12 @@ def cube_root_inner(x: M, intwrap, mod_class) -> M:
         k = (t + 1) // 3
 
     b = intwrap(2)
-    while mod_class(b, x.n).is_cubic_residue():
+    while mod_class(b).is_cubic_residue():
         b += 1
 
-    c = mod_class(b, x.n) ** t
+    c = mod_class(b) ** t
     r = x ** t
-    h = mod_class(intwrap(1), x.n)
+    h = mod_class(intwrap(1))
     cp = c ** (3 ** (s - 1))
     c = c.inverse()
     for i in range(1, s):
diff --git a/pyecsca/ec/mod/gmp.py b/pyecsca/ec/mod/gmp.py
index a8b589e..83f0aaf 100644
--- a/pyecsca/ec/mod/gmp.py
+++ b/pyecsca/ec/mod/gmp.py
@@ -88,7 +88,7 @@ if has_gmp:
                 return GMPMod(gmpy2.mpz(0), self.n, ensure=False)
             if not self.is_residue():
                 raise_non_residue()
-            return square_root_inner(self, gmpy2.mpz, partial(GMPMod, ensure=False))
+            return square_root_inner(self, gmpy2.mpz, lambda x: GMPMod(x, self.n, ensure=False))
 
         def is_cubic_residue(self) -> bool:
             if not _gmpy_is_prime(self.n):
@@ -110,7 +110,7 @@ if has_gmp:
                 return GMPMod(gmpy2.mpz(1), self.n,  ensure=False)
             if not self.is_cubic_residue():
                 raise_non_residue()
-            return cube_root_inner(self, gmpy2.mpz, partial(GMPMod, ensure=False))
+            return cube_root_inner(self, gmpy2.mpz, lambda x: GMPMod(x, self.n, ensure=False))
 
         @_check
         def __add__(self, other) -> "GMPMod":
diff --git a/pyecsca/ec/mod/raw.py b/pyecsca/ec/mod/raw.py
index 1e1aa45..a70c627 100644
--- a/pyecsca/ec/mod/raw.py
+++ b/pyecsca/ec/mod/raw.py
@@ -47,7 +47,7 @@ class RawMod(Mod):
             return RawMod(0, self.n)
         if not self.is_residue():
             raise_non_residue()
-        return square_root_inner(self, int, RawMod)
+        return square_root_inner(self, int, lambda x: RawMod(x, self.n))
 
     def is_cubic_residue(self):
         if not miller_rabin(self.n):
@@ -69,7 +69,7 @@ class RawMod(Mod):
             return RawMod(1, self.n)
         if not self.is_cubic_residue():
             raise_non_residue()
-        return cube_root_inner(self, int, RawMod)
+        return cube_root_inner(self, int, lambda x: RawMod(x, self.n))
 
     def __bytes__(self):
         return self.x.to_bytes((self.n.bit_length() + 7) // 8, byteorder="big")
-- 
cgit v1.2.3-70-g09d2