rsa.py

import random
from itertools import count

# Extended Euclidean Algorithm
def egcd(a, b):
    x,y, u,v = 0,1, 1,0
    while a != 0:
        q, r = b//a, b%a
        m, n = x-u*q, y-v*q
        b,a, x,y, u,v = a,r, u,v, m,n
    gcd = b
    return gcd, x, y

def inverse(a, m):
    gcd, x, y = egcd(a, m)
    if gcd != 1:
        return None
    else:
        return x % m

# Todo: Optimize
def is_prime(num):
    if num == 2:
        return True
    if num < 2 or num % 2 == 0:
        return False
    for n in xrange(3, int(num**0.5)+2, 2):
        if num % n == 0:
            return False
    return True

def find_prime(l, u):
    r = random.randrange(l, u)
    while not is_prime(r):
        r += 1
    return r

def gen_key(p, q):
    if p == q:
        raise ValueError('p and q should be different')

    m = (p-1) * (q-1)
    n = p * q

    # Ensure that e & m are co-prime
    g = 0
    while g != 1:
        e = random.randrange(1, m)
        g, _, _ = egcd(e, m)

    d = inverse(e, m)

    return ((e, n), (d, n))

def encrypt(key, msg):
    k, n = key
    cipher = []

    for c in msg:
        cipher.append(pow(ord(c), k, n))

    return cipher

def decrypt(key, msg):
    k, n = key
    plain = []

    for c in msg:
        plain.append(chr(pow(c, k, n)))

    return ''.join(plain)


if __name__ == '__main__':
    l = 10 ** 12
    u = 10 ** 13

    p = find_prime(l, u)
    q = find_prime(l, u)

    public, private = gen_key(p, q)

    print "Public key: ", public
    print "Private key: ", private

    message = raw_input("Enter message: ")

    ciphertext = encrypt(private, message)

    print "Ciphertext: ", ''.join(map(lambda x: str(x), ciphertext))
    print "Plaintext (recovered): ", decrypt(public, ciphertext)