Pure Python Paillier Homomorphic Cryptosystem

paillier.py

import math
import primes

def invmod(a, p):
    '''
    http://code.activestate.com/recipes/576737-inverse-modulo-p/
    The multiplicitive inverse of a in the integers modulo p.
    Return b s.t.
    a * b == 1 mod p
    '''
    r = a
    d = 1
    while True:
        d = ((p // r + 1) * d) % p
        r = (d * a) % p
        if r == 1:
            break
    else:
        raise ValueError('%d has no inverse mod %d' % (a, p))
    return d

class PrivateKey(object):
    def __init__(self, bits):
        p = primes.generate_prime(bits / 2)
        q = primes.generate_prime(bits / 2)
        n = p * q
        self.lam = (p-1) * (q-1)
        self.pub = PublicKey(bits, n)
        self.mu = invmod(self.lam, n)

class PublicKey(object):
    def __init__(self, bits, n):
        self.bits = bits
        self.n = n
        self.n_sq = n * n
        self.g = n + 1

def encrypt(plain, pub):
    while True:
        r = primes.generate_prime(long(round(math.log(pub.n, 2))))
        if r > 0 and r < pub.n:
            break
    x = pow(r, pub.n, pub.n_sq)
    cipher = (pow(pub.g, plain, pub.n_sq) * x) % pub.n_sq
    return cipher

def e_add(pub, a, b):
    return a * b % pub.n_sq

def decrypt(cipher, priv):
    pub = priv.pub
    x = pow(cipher, priv.lam, pub.n_sq) - 1
    plain = ((x // pub.n) * priv.mu) % pub.n
    return plain

primes.py

import random
import sys

def ipow(a, b, n):
    """calculates (a**b) % n via binary exponentiation, yielding itermediate
       results as Rabin-Miller requires"""
    A = a = long(a % n)
    yield A
    t = 1L
    while t <= b:
        t <<= 1

    # t = 2**k, and t > b
    t >>= 2
    
    while t:
        A = (A * A) % n
        if t & b:
            A = (A * a) % n
        yield A
        t >>= 1

def rabin_miller_witness(test, possible):
    """Using Rabin-Miller witness test, will return True if possible is
       definitely not prime (composite), False if it may be prime."""    
    return 1 not in ipow(test, possible-1, possible)

smallprimes = (3,5,7,11,13,17,19,23,29,31,37,41,43,
               47,53,59,61,67,71,73,79,83,89,97)

def default_k(bits):
    return max(64, 2 * bits)

def is_probably_prime(possible, k=None):
    if k is None:
        k = default_k(possible.bit_length())
    for i in smallprimes:
        if possible % i == 0:
            return False
    for i in xrange(k):
        test = random.randrange(2, possible) | 1
        if rabin_miller_witness(test, possible):
            return False
    return True

def generate_prime(bits, k=None):
    """Will generate an integer of b bits that is probably prime 
       (after k trials). Reasonably fast on current hardware for 
       values of up to around 512 bits."""    
    assert bits >= 8

    if k is None:
        k = default_k(bits)

    while True:
        possible = random.randrange(2 ** (bits-1) + 1, 2 ** bits) | 1
        if is_probably_prime(possible, k):
            return possible

test.py

from paillier import *

print "Generating keypair..."
priv = PrivateKey(128)
pub = priv.pub

x = 3
print "x =", x
print "Encrypting x..."
ex = encrypt(x, pub)
print "ex =", ex

y = 5
print "y =", y
print "Encrypting y..."
ey = encrypt(y, pub)
print "ey =", ey

print "Computing ex + ey..."
er = e_add(pub, ex, ey)
print "er =", er

print "Decrypting er..."
r = decrypt(er, priv)
print "r =", r