Alice and Bob's RSA implementation.

rsa.py

"""
* RSA from scratch, consistent with the implementation described in EECS 70 at UC Berkeley.
* Written by Gautam Mittal
* 10/19/2019
"""

from math import floor, log

N2C = dict(enumerate('abcdefghijklmnopqrstuvwxyz', 1))
C2N = {c: n for n, c in N2C.items()}

def str_to_int(m):
    """
    Converts a message string to an integer.

    >>> str_to_int('abc')
    123
    >>> str_to_int('xyz')
    242526
    """
    i = 0
    for c in m:
        n = C2N[c]
        i = 10 ** (floor(log(n, 10)) + 1) * i + n
    return i
        
def int_to_str(i):
    """
    Finds candidate strings given integer sequence.
    """
    s = str(i) # Strings are easier to work with
    if not i:
        return []
    if len(s) == 1:
        return [N2C[int(s)]]
    candidates = []
    for i in [1, 2]:
        n = int(s[:i])
        if n > 26 or n < 1: continue
        c = N2C[n]
        rest = s[i:]
        if rest:
            for candidate in int_to_str(int(rest)):
                candidates.append(c + candidate)
        else:
            candidates.append(c)
    return candidates
    
def gcd(x, y):
    """
    Extended Euclid's algorithm.
    Given x, y, return gcd(x, y), a, b where ax + by = gcd(x, y).
    
    >>> gcd(12, 6)
    (6, 0, 1)
    >>> gcd(1, 1)
    (1, 0, 1)
    >>> gcd(13, 5)
    (1, 2, -5)
    """
    if y == 0: return (x, 1, 0)
    d, a, b = gcd(y, x % y)
    return d, b, a - (x // y) * b

def inv(x, m):
    """
    Finds modular inverse of x (mod m).

    >>> inv(5, 13)
    -5
    >>> inv(16, 27)
    -5
    """
    d, a, b = gcd(x, m)
    assert d == 1, 'Inputs are not coprime.'
    return a

def generate_key(p, q):
    """
    Generate a public/private RSA keypair.
    
    >>> generate_key(5, 7)
    ((35, 5), 5)
    >>> generate_key(2312451, 1485863)
    ((3435985380213, 3), 1145327193967)
    """
    N = p * q
    h = (p - 1) * (q - 1)
    
    # e must be coprime to (p-1)(q-1)
    e = 3
    while gcd(e, h)[0] != 1:
        e += 2

    d = inv(e, h)
    return ((N, e), d)

def encrypt(message, pub):
    """
    Encrypt message with public key.
    """
    N, e = pub
    return (message ** e) % N

def decrypt(cipher, pub, priv):
    """
    Decrypt a ciphertext given a keypair.
    """
    N, e = pub
    d = priv
    return (cipher ** d) % N

def encrypt_str(s, pub):
    """
    Encrypt a string with a public key.
    """
    return encrypt(str_to_int(s), pub)

p = 19
q = 511
pub, priv = generate_key(p, q)

x = encrypt_str('ab', pub)
y = decrypt(x, pub, priv)