RSA.py

import random

def isPrime(num):
    if num == 2: return True
    if num < 2 or num % 2 == 0: return False
    for i in range(3, int(num ** 0.5) + 2, 2):
        if num % i == 0: return False
    return True
  
def gcd(e, phi): return e if phi == 0 else gcd(phi, e % phi)
   
def multiplicativeInverse(e, phi):
    for i in range(1, phi):
        if (phi * i + 1)% e == 0: return ((phi * i + 1)// e)
    return None
  
# Function to generate public and private key pairs
# Generating pulic and private keys for large prime numbers may take some time
def generateKeyPair(p, q):
    if not (isPrime(p) and (isPrime(q))): raise ValueError("Both numbers must be prime!")
    elif p == q: raise ValueError("Numbers entered should not be the same")
    
    n = p * q
    phi = (p-1) * (q-1)
    # Get a random e such that e and phi are co-prime 
    e = random.randrange(1, phi)
    value = gcd(e, phi)
    while value != 1:
        e = random.randrange(1, phi)
        value  = gcd(e, phi)
        
    d = multiplicativeInverse(e, phi)
    # Public key is (e, n) & Private key is (d, n)
    print ("Your public key is (",e,",",n,")\nYour private key is (",d,",",n,")")
    return ((e, n), (d, n))
    
def encryption(private, plainText):
    # Seperating key and n
    key, n = private
    # Converting each character in plainText to number and using C = M^d mod n
    cipherText = [pow(ord(char), key, n) for char in plainText]
    print ("The encrpyted message is: ",' '.join(map(lambda x: str(x), cipherText)))
    return cipherText
        
def decryption(public, cipherText):
    # Seperating key and n
    key, n = public
    # Generating plainText based on cipherText and key and using M = C^d mod n
    print ("The decrypted message is: ",''.join([chr(pow(char, key, n)) for char in cipherText])) 
        
if __name__ == '__main__':
    p = int(input("Enter a prime number: "))
    q = int(input("Enter a prime number other than the previous one: "))
    publicKey, privateKey = generateKeyPair(p, q)
    message = str(input("Enter the message that you want to encrypt with your private key: "))
    encrypt = encryption(privateKey, message)
    decryption(publicKey, encrypt)