RSA Key Generation, Encryption and Decryption example

rsatest.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
#                   -------------------------------------
#    RSA Key Generation, Encryption and Decryption example
#    Copyright (C) 2014  Andrey Arapov
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
#                   -------------------------------------
#
#
# Notes:
#   - Based on the https://en.wikipedia.org/wiki/RSA_%28cryptosystem%29#Key_generation
#   - The parameters used here are artificially small
#   - I've tried to apply KISS principle here
#
#

import random
from fractions import gcd

class bcolors:
	RED = '\033[91m'
	DRED = '\033[31m'
	GREEN = '\033[92m'
	YELLOW = '\033[93m'
	BLUE = '\033[94m'
	PURPLE = '\033[95m'
	CYAN = '\033[96m'
	ENDC = '\033[0m'


# -----------------------------------------------------------------------------
# Generating the Public/Private keypair
# -----------------------------------------------------------------------------


# Primality test
# https://en.wikipedia.org/wiki/Primality_test#Python_implementation

def is_prime(num):
    if num <= 3:
        if num <= 1:
            return False
        return True
    if not num % 2 or not num % 3:
        return False
    for i in range(5, int(num ** 0.5) + 1, 6):   
        if not num % i or not num % (i + 2):
            return False
    return True


# 1. Choose two distinct prime numbers p and q.

print "1. looking for two distinct prime numbers p and q in artificially small range..."
i = 0
while i < 2:
	rand = random.randint(100, 999)
	if is_prime(rand):
		if i == 1:
			q=rand
			break
		p=rand
		i += 1

print "p =", bcolors.DRED, p, bcolors.ENDC, "\tprime?", is_prime(p)
print "q =", bcolors.DRED, q, bcolors.ENDC, "\tprime?", is_prime(q)
print


# 2. Compute n = pq.

print "2. computing n = pq ..."
n = p * q
print "n = p * q =", bcolors.DRED, p, bcolors.ENDC, "*", bcolors.DRED, q, bcolors.ENDC, "=", \
	 bcolors.BLUE, n, bcolors.ENDC
print


# 3. Compute φ(n) = φ(p)φ(q) = (p − 1)(q − 1) = n - (p + q - 1), where φ is Euler's totient function.

print "3. computing φ(n) = φ(p)φ(q) = (p − 1)(q − 1) = n - (p + q -1), where φ is Euler's totient function ..."
f_n = n - (p + q - 1)
print "φ(n) =", bcolors.DRED, f_n, bcolors.ENDC
print


# 4. Choose an integer e such that 1 < e < φ(n) and gcd(e, φ(n)) = 1; i.e., e and φ(n) are coprime.

print "4. looking for an integer e such that 1 < e < φ(n) and gcd(e, φ(n)) = 1; i.e., e and φ(n) are coprime ..."
print "\nSetting e = 2^16+1 (65537) as per recommendation in\n"+ \
  "Dan Boneh's Twenty Years of Attacks on the RSA Cryptosystem - "+ \
  "http://crypto.stanford.edu/~dabo/pubs/papers/RSA-survey.pdf\n"
#e_gcd = 2
# TOFIX: add smth like --> while (e_gcd != 1) and (e > 3):  as e should be > than 3
#while e_gcd != 1:
#	e = random.randint(1, f_n)
#	e_gcd = gcd(e, f_n)

e = 65537
print "e =", bcolors.CYAN, e, bcolors.ENDC
print
print

# 5. Determine d as d ≡ e^−1 (mod φ(n)); i.e., d is the multiplicative inverse of e (modulo φ(n)).
# http://en.wikibooks.org/wiki/Algorithm_Implementation/Mathematics/Extended_Euclidean_algorithm#Modular_inverse

def egcd(a, b):
    if a == 0:
        return (b, 0, 1)
    else:
        g, y, x = egcd(b % a, a)
        return (g, x - (b // a) * y, y)

def modinv(a, m):
    gcd, x, y = egcd(a, m)
    if gcd != 1:
        return None  # modular inverse does not exist
    else:
        return x % m

print "5. Determining d as d ≡ e^−1 (mod φ(n)); i.e., d is the multiplicative inverse of e (modulo φ(n)) ..."
d = modinv(e, f_n)
print "d =", bcolors.RED, d, bcolors.ENDC
print

print "Public key is modulus n =", bcolors.BLUE, n, bcolors.ENDC, "and the public (or encryption) exponent e =", \
	 bcolors.CYAN, e, bcolors.ENDC
print
print "Private key is modulus n =", bcolors.BLUE, n, bcolors.ENDC, "and the private (or decryption) exponent d =", \
 	bcolors.RED, d, bcolors.ENDC, "and it must be kept secret"
print bcolors.DRED+"p"+bcolors.ENDC+","+bcolors.DRED+" q"+bcolors.ENDC+", and"+bcolors.DRED+" φ(n) "+bcolors.ENDC+ \
	"must also be kept secret because they can be used to calculate "+bcolors.RED+"d"+bcolors.ENDC+"."
print



# -----------------------------------------------------------------------------
# Encrypting: c = m^e (mod n)
# -----------------------------------------------------------------------------

print "To encrypt message m: c = m^e (mod n)"

#mstr = "hi"
cstr = ""
mstr = raw_input("Enter your message m: ")
for m in [elem.encode("hex") for elem in mstr]:
	print ":: encrypting ", bcolors.YELLOW, '"'+chr(int(m, 16))+'"', \
	 int(m, 16), bcolors.ENDC, " >>>", bcolors.YELLOW, int(m, 16), \
	 bcolors.ENDC, "^", bcolors.CYAN, e, bcolors.ENDC, \
	 "( mod", bcolors.BLUE, n, bcolors.ENDC, ")", ">>> ",
	c = ( int(m, 16) ** e ) % n
	print bcolors.PURPLE, c, bcolors.ENDC
	cstr += str(c)
	cstr += ","

cstr = cstr[:-1]
print "Your encrypted message m is now a ciphertext c  =", bcolors.YELLOW, cstr, bcolors.ENDC
print


# -----------------------------------------------------------------------------
# Decrypting: m = c^d (mod n)
# -----------------------------------------------------------------------------

print "To decrypt ciphertext c: m = c^d (mod n)"
mstr = ""
for c in cstr.split(","):
	print ":: decrypting ", bcolors.PURPLE, c, bcolors.ENDC, " >>>", \
	 bcolors.PURPLE, int(c), "^", bcolors.ENDC, bcolors.RED, \
	 d, bcolors.ENDC, "( mod", bcolors.BLUE, n, bcolors.ENDC, ")", ">>>",
	m = ( int(c) ** d ) % n
	print bcolors.YELLOW, m, '"'+chr(m)+'"', bcolors.ENDC
	mstr += chr(m)

print "Decrypted ciphertext is now m =", bcolors.YELLOW, mstr, bcolors.ENDC
print