RSA算法实现练习

rsa.py

# Rsa Demo
import math
import random

def to_binary(n):
    if n / 2 == 0:
        return str(n)
    return to_binary(n/2) + str(n%2)

assert to_binary(4) == '100'
assert to_binary(5) == '101'


def mod(a, b, n):
    """calculate a**b % n"""
    binary_b = to_binary(b)
    power = a % n
    s = 1
    for i in binary_b[::-1]:
        if i == '1':
            s = (s * power) % n
        power = (power * power) % n
    return s

assert mod(3, 1, 7) == 3
assert mod(3, 2, 7) == 2


def is_prime(n):
    if n == 1:
        return False
    
    for i in range(1, 5):
        a = random.randint(2,n-1)
        # if a % n != a**n % n:
        if a % n != mod(a,n,n):
            return False
    return True

PRIMES = [i for i in range(3, 10000) if is_prime(i)]
def generate_prime():
    while 1:
        flag = 0
        n = random.randint(2**63, 2**64-1)
        if n << 63 == 2**63:
            n = n +1
        for p in PRIMES:
            if n % p == 0:
                flag = 1
                break
        if not flag and is_prime(n):
            return n

assert not is_prime(1)
assert not is_prime(4)
assert is_prime(generate_prime())

def extend_gcd(a,b):
    if a % b == 0:
        return (0, 1)
    x2,y2 = extend_gcd(b, a % b)
    x1 = y2
    y1 = x2 - (a/b)*y2
    return x1,y1

def gcd(a, b):
    if a % b == 0:
        return b
    return gcd(b, a % b)

def get_d_from_e_m(e, m):
    x1,y1 = extend_gcd(e, m)
    return x1


def make_keypair():
    e = 65537
    d = 0
    while d <= 0:
        p1, p2 = generate_prime(), generate_prime()
        n = p1 * p2
        #m = (p1-1)*(p2-1)
        m = n - (p1+p2-1)
        d = get_d_from_e_m(e, m)
    return (e, n), (d, n)


def encrypt(msg, e, n):
    return [mod(ord(i), e, n) for i in msg]

def decrypt(numbers, d, n):
    return ''.join([chr(mod(i,d,n)) for i in numbers])

if __name__ == '__main__':
    enc_pair, dec_pair = make_keypair()
    print 'Encrypt key: ', enc_pair
    print 'Decrypt key: ', dec_pair
    msg = 'Hello'
    
    print 'Encrypt: ', msg
    code = encrypt(msg, *enc_pair)
    print code
    print
    print 'Decrypt: ', code
    print decrypt(code, *dec_pair)