Rabin Rsa Hanukkah challenge in XMAS CTF'19

rabin.py

import gmpy2,math
def egcd(a, b):
    u, u1 =  1 , 0
    v, v1 =  0 , 1
    while b:
        q = a // b
        u, u1 = u1, u - q * u1
        v, v1 = v1, v - q * v1
        a, b = b, a - q * b
    return a, u, v


def phi(p, q):
    return (p - 1) * (q - 1)


def get_d(p, n, e):
    q = n / p
    phi_v = phi (p, q)
    _gcd, d, _2 = egcd (e, phi_v)
    if d < 0:
        d += phi_v
    return d


def modular_sqrt(a, p):
    if legendre_symbol(a, p) != 1:
        return 0
    elif a ==  0 :
        return 0
    elif p ==  2 :
        return p
    elif p %  4  ==  3 :
        return pow(a, (p + 1) / 4, p)
    s = p - 1
    e = 0
    while s % 2 == 0:
        s /= 2
        e += 1
    n = 2
    while legendre_symbol(n, p) != -1:
        n += 1

    x = pow(a, (s + 1) / 2, p)
    b = pow(a, s, p)
    g = pow(n, s, p)
    r = e

    while True:
        t = b
        m = 0
        for m in xrange(r):
            if t == 1:
                break
            t = pow(t, 2, p)

        if m == 0:
            return x
        gs = pow(g, 2 ** (r - m - 1), p)
        g = (gs * gs) % p
        x = (x * gs) % p
        b = (b * g) % p
        r = m


def legendre_symbol(a, p):
    ls = pow(a, (p - 1) / 2, p)
    return -1 if ls == p - 1 else ls


class Rabin(object):
    def __init__(self, p, q):
        self.p = p
        self.q = q
        self.n = p * q

    def encrypt(self, m):
        return (m * m) % self.n

    def decrypt(self, c):
        try:
            gcd, yp, yq = egcd(self.p, self.q)
            mp = modular_sqrt(c, self.p)
            mq = modular_sqrt(c, self.q)
            assert yp * self.p + yq * self.q == 1
            assert (mp * mp) % self.p == c % self.p
            assert (mq * mq) % self.q == c % self.q
            r1 = (yp *  self .p * mq + yq *  self. q * mp) %  self .n
            s1 = (yp *  self .p * mq - yq *  self .q * mp) %  self .n
            r2 = self.n - r1
            s2 = self.n - s1
            return r1, s1, r2, s2
        except AssertionError:
            return []

n = 825321266319602503456977005474981604870402407335194099572979028339224439122246767155608828548258547874076592811333439775645799852274012447643240804287007452861599291275940862131595970247906775549656137041013432613989092491697319873901497907382123859210758943466373193369020798176192106305153278525778145033 
ct = 801050608421922967220624523903721496853411844056321773877598932155971380872263121340024512973182420871402804237809506243995703890886804092449855251892886296340338442367792297266755554172082930224889412735287102163161928535579728998850091020972410977027707699268899998522781790134147981974412918582618345868
p = 416835513771282386514568836191681760971829004269725709610710626946559368006342658194525259634969297938883155195834864676790249857624739386431798157728603
q = 1979968690413591335944201971910488364616187770281197120650875477996156998030165907455238957933094804031981974077477456579642988264126768863441500694216811
e = 2L
d = get_d(p, n, e)
rabin = Rabin(p, q)
partially_decoded_ct = [ct]
for i in range ( 1 ): # this 1 is important for e=2^1
    new_partially_decoded_ct = []
    for ct_p in partially_decoded_ct:
        new_ct_p = rabin.decrypt(ct_p)
        new_partially_decoded_ct.extend(list(new_ct_p))
    partially_decoded_ct = set(new_partially_decoded_ct)
potential_plaintext = []
for potential_rsa_ct in partially_decoded_ct:
    pt = pow(potential_rsa_ct, d, n)
    potential_plaintext.append(pt)
print(potential_plaintext)
#X-MAS{H4nukk4h_Rabb1_and_Rab1n_l0ok_4nd_s0und_v3ry_much_alik3_H4nukk4h}