Solutions for Montgomery's Ladder

solution.py

from typing import Optional, Tuple
from Crypto.Util.number import inverse
import Crypto.Random.random as csprng
    
class MtgCurve:
    """
        Represents a Montgomery curve of the form
        By^2 = x^3 + Ax^2 + x
    """
    def __init__(self, B: int, A: int, p: int) -> None:
        self._B = B
        self._A = A
        self._p = p

    def __eq__(self, other: object) -> bool:
        return self._B == other._B and self._A == other._A and self._p == other._p

    def get_params(self) -> Tuple[int, int, int]:
        return (self._B, self._A, self._p)


class MtgPoint:
    """
        Represents a point on a montgomery curve 
        using affine coordinates
    """
    def __init__(self, curve: MtgCurve, x: int, y: int, infinity: bool = False) -> None:
        self._curve = curve
        self._x = x
        self._y = y
        self._infinity = infinity

    def __eq__(self, other: object) -> bool:
        return self._curve == other._curve and self._x == other._x and self._y == other._y
    
    def __add__(self, other: object):
        assert(self != other)

        if self._infinity:
            return other
        if other._infinity:
            return self

        B, A, p = self._curve.get_params()
        alpha = ( (other._y - self._y) *  inverse((other._x - self._x), p) ) % p
        x = ( (B * alpha * alpha) - A - self._x - other._x) % p 
        y = (alpha * (self._x - x) - self._y) % p
        return MtgPoint(curve=self._curve, x=x, y=y)
    
    def double(self):
        if self._infinity:
            return MtgPoint(curve=self._curve, x=0, y=0, infinity=True)
        
        B, A, p = self._curve.get_params()
        alpha = ( ( (3 * self._x * self._x) + (2 * A * self._x) + 1) * inverse((2 * B * self._y), p) ) % p
        x = ( (B * alpha * alpha) - A - 2 * self._x) % p
        y = (alpha * (self._x - x) - self._y) % p
        return MtgPoint(curve=self._curve, x=x, y=y)

    def __repr__(self) -> str:
        return f"({self._x}, {self._y})" if not self._infinity else f"0"


def tonelli_shanks(a: int, p: int) -> Tuple[int, int]:
    """
        return a tuple with 2 values
        that represent sqrt(a) mod p if they exist
    """
    def _tonelli_shanks_internal_(a: int, k: int, p: int, z: int, zinv: int) -> Tuple[int, int]:
        # we know that p is congruent to 1 mod 4
        # and a has square roots modulo p
        # Idea: find 'odd' m such that pow(a, m, p) = 1, 
        # then +/-pow(a, m + 1 // 2, p) gives square roots
        
        m = p - 1 >> k
        while m & 1 == 0 and pow(a, m, p) == 1:
            # m is even, so continue
            m >>= 1
            k += 1
        
        # at this point, m is odd
        res = None
        if pow(a, m, p) == 1:
            # we found odd m such that pow(a, m, p) = 1
            res = pow(a, m + 1 >> 1, p)
            return (res, p - res)
        
        # a^m = -1 mod p
        z_p = 1 << (k - 1)
        z_phalf = 1 << (k - 2)
        a_next = (a * pow(z, z_p, p)) % p
        r1, _ = _tonelli_shanks_internal_(a_next, k, p, z, zinv)
        a_root = (r1 * pow(zinv, z_phalf, p)) % p
        return (a_root, p - a_root)

    euler_criterion = lambda a, p: pow(a, p-1 >> 1, p)


    # check if square root even exists
    # using Euler's criterion
    if euler_criterion(a, p) != 1:
        raise Exception(f'Error: {a} does not have a square root mod {p}')

    # check if p is congruent to 1 mod 4
    if p - 3 % 4 == 0:
        val = pow(a, p + 1 // 4, p)
        return (val, p - val)

    # find a non-square mod p
    z = None
    while True:
        z = csprng.randrange(2, p - 1)
        if euler_criterion(z, p) == p - 1:
            break
    zinv = inverse(z, p)

    return _tonelli_shanks_internal_(a, 1, p, z, zinv)
    
    
def montgomery_ladder(k: int, P: MtgPoint) -> MtgPoint:
    """
        compute scalar multiplication [k]P and return resulting
        point on the elliptic curve
    """

    # Idea: perform doubling and addition whether or not the current
    #       'bit' of the secret scalar being iterated is 0 or 1
    #       to mask any timing leaks

    R0, R1 = MtgPoint(curve=P._curve, x=0, y=0, infinity=True), P
    for i in reversed(range(k.bit_length())):
        b = (k >> i) & 1
        R0, R1 = (R0 + R1, R1.double()) if b != 0 else (R0.double(), R0 + R1)
    
    return R0
    # n = k.bit_length()
    # R0 , R1 = P, P.double()
    # for i in range(n-2, 0, -1):
    #     b = k & (1 << i)
    #     R0, R1 = (R0 + R1, R1.double()) if b != 0 else (R0.double(), R0 + R1)
    
    # return R0


if __name__ == "__main__":
    #  E: Y2 = X3 + 486662X2 + X, p: 2255 - 19 
    B = 1
    A = 486662
    p = pow(2, 255) - 19
    curve = MtgCurve(B=B, A=A, p=p)
    

    # Generator point
    x = 9
    ysq = (pow(x, 3) + A * pow(x, 2) + x) % p
    y, _ = tonelli_shanks(a=ysq, p=p)
    G = MtgPoint(curve=curve, x=x, y=y)
    print (f"G = {G}")

    k = 0x1337c0decafe
    Q = montgomery_ladder(k, G)
    print (f"Q = {Q}")