A simple RSA implementation in Python

rsa.py

'''
RSA asymmetric cryptographic algorithm.
A simple proof of concept written in Python.

@author jcanepa
@url https://gist.github.com/jcanepa/37d80d4691f63097e6388bcf3dd0cf51
Forked from https://gist.github.com/djego/97db0d1bc3d16a9dcb9bab0930d277ff
'''

import random

def two_distinct_large_prime_numbers():
    return 17, 29

'''
Euclid's algorithm for determining the greatest common divisor
'''
def gcd(a, b):
    while not b == 0:
        a, b = b, a % b
    return a

'''
Euclid's extended algorithm for finding the multiplicative inverse of two numbers
'''
def multiplicative_inverse(e, phi):
    for i in range(phi):
        if ((e*i)%phi) == 1:
            return i

'''
Tests to see if a number is prime.
'''
def is_prime(num):
    if num == 2:
        return True
    if num < 2 or num % 2 == 0:
        return False
    for n in range(3, int(num**0.5)+2, 2):
        if num % n == 0:
            return False
    return True

'''
Given two prime numbers, generate a new keypair.
'''
def generate_keypair(p, q):
    # validate input
    if not (is_prime(p) and is_prime(q)):
        raise ValueError('Both numbers must be prime.')
    elif p == q:
        raise ValueError('p and q cannot be equal')

    # public key modulus
    n = p * q
    # phi (φ or ϕ) is Euler's totient function of n:
    # the number of integers k in the range 1 ≤ k ≤ n
    # for which the greatest common divisor gcd(n, k) is equal to 1
    phi = (p-1) * (q-1)
    # choose an integer e such that e and φ(n) are co-prime
    # and e is between 1 and φ(n)
    e = random.randrange(1, phi)
    # use Euclid's Algorithm to verify that e and φ(n) are co-prime
    g = gcd(e, phi)
    while g != 1:
        e = random.randrange(1, phi)
        g = gcd(e, phi)
    # use Extended Euclid's Algorithm to generate the private key
    d = multiplicative_inverse(e, phi)
    # return public and private keypair
    # public key is (e, n) and private key is (d, n)
    return ((e, n), (d, n))

'''
Encrypt a given message using the public key provided.
'''
def encrypt(public_key, message):
    # unpack the key into it's components
    key, n = public_key
    # convert each letter in the message to numbers based on the character using a^b mod m
    cipher = [
        (ord(char) ** key) % n \
            for char in message]
    # return the array of bytes
    return cipher

'''
Decrypt an cipher using the private key provided.
'''
def decrypt(private_key, cipher):
    # unpack the key into its components
    key, n = private_key
    # generate the message based on the cipher and key using a^b mod m
    message = [
        chr((char ** key) % n) \
        for char in cipher]
    # return the array of bytes as a string
    return ''. join(message)

def print_seprator() -> None:
    print('*'*55)

def print_welcome() -> None:
    print_seprator()
    print('   🔏 RSA Keypair Generator & Encrypter/Decrypter 🔏  ')
    print_seprator()

'''
Run the RSA key generator & provide a user demonstration of its use.
'''
def init() -> None:
    print_welcome()

    p, q = two_distinct_large_prime_numbers()

    public, private = generate_keypair(p, q)
    print('Public  🔑', public)
    print('Private 🔑', private)

    message = input('\nEnter a message to encrypt with your private key: \n')
    encrypted_msg = encrypt(private, message)

    print_seprator()

    print('Encrypted message:',
        ''.join(map(
            lambda x: str(x),
            encrypted_msg)))

    print('Decrypted message:',
        decrypt(public, encrypted_msg))

    print_seprator()
    print()

if __name__ == '__main__':
    init()