from Crypto.Util.number import getPrime, bytes_to_long, long_to_bytes
from math import gcd
from sympy.ntheory.modular import crt
import time

# key generation
p = getPrime(512)		# 512-bit prime number
q = getPrime(512)		# 512-bit prime number
n = p * q
phi = (p - 1) * (q - 1)
e = 65537       # = 2^16 +1 = 10000000000000001
assert gcd(e, phi) == 1

d = pow(e, -1, phi)

# encryption
def encrypt(message):
    m = bytes_to_long(message.encode())
    c = pow(m, e, n)
    return c

# CRT
dp    = d % (p - 1)        # d_p = d mod (p-1)
dq    = d % (q - 1)        # d_q = d mod (q-1)
q_inv = pow(q, -1, p)      # q^{-1} mod p

def decrypt(cipher_int):
    m1 = pow(cipher_int, dp, p)   # c^dp mod p
    m2 = pow(cipher_int, dq, q)   # c^dq mod q
    
    m, mod = crt([p, q], [m1, m2])  # m ≡ m1 (mod p), m ≡ m2 (mod q)
    
    # Garner's Algorithm
    # h = (m1 - m2) * q_inv % p
    # m = (m2 + h * q) % n

    message = long_to_bytes(m).decode()
    return message

m = "Hello World"
c = encrypt(m)

start = time.perf_counter()
m_recover = decrypt(c)
end = time.perf_counter()

print("Ciphertext: ", c)

print("Message: ", m_recover)
print(f"decrypt time: {(end-start)*1e3:.3f} ms")