Make Mod a dynamic class.

Fixes #6.

2 files changed, 80 insertions, 19 deletions

diff --git a/pyecsca/ec/mod.py b/pyecsca/ec/mod.py
index 3f97637..65018f7 100644
--- a/pyecsca/ec/mod.py
+++ b/pyecsca/ec/mod.py
@@ -1,7 +1,7 @@
 import random
 import secrets
 from functools import wraps, lru_cache
-from abc import ABC, abstractmethod
+from abc import abstractmethod
 from public import public
 
 from .error import NonInvertibleError, NonResidueError
@@ -14,7 +14,7 @@ try:
 
     has_gmp = True
 except ImportError:
-    pass
+    gmpy2 = None
 
 
 @public
@@ -104,7 +104,19 @@ class RandomModAction(ResultAction):
         return f"{self.__class__.__name__}({self.order:x})"
 
 
-class BaseMod(ABC):
+_mod_classes = []
+
+
+@public
+class Mod(object):
+
+    def __new__(cls, *args, **kwargs):
+        if cls != Mod:
+            return cls.__new__(cls, *args, **kwargs)
+        if not _mod_classes:
+            raise ValueError("Cannot find a working Mod class.")
+        return _mod_classes[-1].__new__(_mod_classes[-1], *args, **kwargs)
+
     def __init__(self, x, n):
         self.x = x
         self.n = n
@@ -129,12 +141,32 @@ class BaseMod(ABC):
         return self.__class__(self.n - self.x, self.n)
 
     @abstractmethod
-    def inverse(self):
+    def inverse(self) -> "Mod":
+        """
+        Invert the element.
+
+        :return: The inverse.
+        :raises: :py:class:`NonInvertibleError` if the element is not invertible.
+        """
         ...
 
     def __invert__(self):
         return self.inverse()
 
+    @abstractmethod
+    def is_residue(self) -> bool:
+        """Whether this element is a quadratic residue (only implemented for prime modulus)."""
+        ...
+
+    @abstractmethod
+    def sqrt(self) -> "Mod":
+        """
+        The modular square root of this element (only implemented for prime modulus).
+
+        Uses the `Tonelli-Shanks <https://en.wikipedia.org/wiki/Tonelli–Shanks_algorithm>`_ algorithm.
+        """
+        ...
+
     @check
     def __mul__(self, other):
         return self.__class__((self.x * other.x) % self.n, self.n)
@@ -172,18 +204,42 @@ class BaseMod(ABC):
         q, r = divmod(self.x, divisor.x)
         return self.__class__(q, self.n), self.__class__(r, self.n)
 
+    @abstractmethod
+    def __bytes__(self):
+        ...
+
+    @abstractmethod
+    def __int__(self):
+        ...
+
     @classmethod
     def random(cls, n: int):
+        """
+        Generate a random :py:class:`Mod` in ℤₙ.
+
+        :param n: The order.
+        :return: The random :py:class:`Mod`.
+        """
         with RandomModAction(n) as action:
             return action.exit(cls(secrets.randbelow(n), n))
 
+    @abstractmethod
+    def __pow__(self, n):
+        ...
+
+    def __str__(self):
+        return str(self.x)
+
 
 @public
-class RawMod(BaseMod):
+class RawMod(Mod):
     """An element x of ℤₙ."""
     x: int
     n: int
 
+    def __new__(cls, *args, **kwargs):
+        return object.__new__(cls)
+
     def __init__(self, x: int, n: int):
         super().__init__(x % n, n)
 
@@ -196,7 +252,6 @@ class RawMod(BaseMod):
         return RawMod(x, self.n)
 
     def is_residue(self):
-        """Whether this element is a quadratic residue (only implemented for prime modulus)."""
         if not miller_rabin(self.n):
             raise NotImplementedError
         if self.x == 0:
@@ -207,11 +262,6 @@ class RawMod(BaseMod):
         return legendre_symbol == 1
 
     def sqrt(self):
-        """
-        The modular square root of this element (only implemented for prime modulus).
-
-        Uses the `Tonelli-Shanks <https://en.wikipedia.org/wiki/Tonelli–Shanks_algorithm>`_ algorithm.
-        """
         if not miller_rabin(self.n):
             raise NotImplementedError
         if self.x == 0:
@@ -280,8 +330,14 @@ class RawMod(BaseMod):
         return RawMod(pow(self.x, n, self.n), self.n)
 
 
+_mod_classes.append(RawMod)
+
+
 @public
-class Undefined(BaseMod):
+class Undefined(Mod):
+    def __new__(cls, *args, **kwargs):
+        return object.__new__(cls)
+
     def __init__(self):
         super().__init__(None, None)
 
@@ -303,6 +359,12 @@ class Undefined(BaseMod):
     def inverse(self):
         raise NotImplementedError
 
+    def sqrt(self):
+        raise NotImplementedError
+
+    def is_residue(self) -> bool:
+        raise NotImplementedError
+
     def __invert__(self):
         raise NotImplementedError
 
@@ -355,11 +417,14 @@ class Undefined(BaseMod):
 if has_gmp:
 
     @public
-    class GMPMod(BaseMod):
+    class GMPMod(Mod):
         """An element x of ℤₙ. Implemented by GMP."""
         x: gmpy2.mpz
         n: gmpy2.mpz
 
+        def __new__(cls, *args, **kwargs):
+            return object.__new__(cls)
+
         def __init__(self, x: int, n: int):
             super().__init__(gmpy2.mpz(x % n), gmpy2.mpz(n))
 
@@ -462,8 +527,4 @@ if has_gmp:
             return GMPMod(gmpy2.powmod(self.x, gmpy2.mpz(n), self.n), self.n)
 
 
-    Mod = GMPMod
-else:
-    Mod = RawMod
-
-public(Mod=Mod)
+    _mod_classes.append(GMPMod)
diff --git a/test/ec/test_mod.py b/test/ec/test_mod.py
index 71a5a74..bf08cd9 100644
--- a/test/ec/test_mod.py
+++ b/test/ec/test_mod.py
@@ -89,7 +89,7 @@ class ModTests(TestCase):
     def test_undefined(self):
         u = Undefined()
         for k, meth in u.__class__.__dict__.items():
-            if k in ("__module__", "__init__", "__doc__", "__hash__", "__abstractmethods__", "_abc_impl"):
+            if k in ("__module__", "__new__", "__init__", "__doc__", "__hash__", "__abstractmethods__", "_abc_impl"):
                 continue
             args = [5 for _ in range(meth.__code__.co_argcount - 1)]
             if k == "__repr__":