Implementation of an homomorphic encryption scheme with numpy based on the ring learning with error problem

rlwe_he_scheme.py

"""A basic homomorphic encryption scheme inspired from BFV https://eprint.iacr.org/2012/144.pdf

You can read my blog post explaining the implementation details here: https://www.ayoub-benaissa.com/blog/build-he-scheme-from-scratch-python/

Disclaimer: This implementation doesn’t neither claim to be secure nor does it follow software engineering best practices,
it is designed as simple as possible for the reader to understand the concepts behind homomorphic encryption schemes.
"""

import numpy as np
from numpy.polynomial import polynomial as poly

# Functions for random polynomial generation


def gen_binary_poly(size):
    """Generates a polynomial with coeffecients in [0, 1]

    Args:
        size: number of coeffcients, size-1 being the degree of the
            polynomial.

    Returns:
        array of coefficients with the coeff[i] being 
        the coeff of x ^ i.
    """
    return np.random.randint(0, 2, size, dtype=np.int64)


def gen_uniform_poly(size, modulus):
    """Generates a polynomial with coeffecients being integers in Z_modulus

    Args:
        size: number of coeffcients, size-1 being the degree of the
            polynomial.

    Returns:
        array of coefficients with the coeff[i] being 
        the coeff of x ^ i.
    """
    return np.random.randint(0, modulus, size, dtype=np.int64)


def gen_normal_poly(size):
    """Generates a polynomial with coeffecients in a normal distribution
    of mean 0 and a standard deviation of 2, then discretize it.

    Args:
        size: number of coeffcients, size-1 being the degree of the
            polynomial.

    Returns:
        array of coefficients with the coeff[i] being 
        the coeff of x ^ i.
    """
    return np.int64(np.random.normal(0, 2, size=size))


# Functions for polynomial evaluation in Z_q[X]/(X^N + 1)

def polymul(x, y, modulus, poly_mod):
    """Multiply two polynoms

    Args:
        x, y: two polynoms to be multiplied.
        modulus: coefficient modulus.
        poly_mod: polynomial modulus.

    Returns:
        A polynomial in Z_modulus[X]/(poly_mod).
    """
    return np.int64(
        np.round(poly.polydiv(poly.polymul(x, y) %
                              modulus, poly_mod)[1] % modulus)
    )


def polyadd(x, y, modulus, poly_mod):
    """Add two polynoms

    Args:
        x, y: two polynoms to be added.
        modulus: coefficient modulus.
        poly_mod: polynomial modulus.

    Returns:
        A polynomial in Z_modulus[X]/(poly_mod).
    """
    return np.int64(
        np.round(poly.polydiv(poly.polyadd(x, y) %
                              modulus, poly_mod)[1] % modulus)
    )


# Functions for keygen, encryption and decryption

def keygen(size, modulus, poly_mod):
    """Generate a public and secret keys

    Args:
        size: size of the polynoms for the public and secret keys.
        modulus: coefficient modulus.
        poly_mod: polynomial modulus.

    Returns:
        Public and secret key.
    """
    s = gen_binary_poly(size)
    a = gen_uniform_poly(size, modulus)
    e = gen_normal_poly(size)
    b = polyadd(polymul(-a, s, modulus, poly_mod), -e, modulus, poly_mod)

    return (b, a), s


def encrypt(pk, size, q, t, poly_mod, pt):
    """Encrypt an integer.

    Args:
        pk: public-key.
        size: size of polynomials.
        q: ciphertext modulus.
        t: plaintext modulus.
        poly_mod: polynomial modulus.
        pt: integer to be encrypted.

    Returns:
        Tuple representing a ciphertext.      
    """
    # encode the integer into a plaintext polynomial
    m = np.array([pt] + [0] * (size - 1), dtype=np.int64) % t
    delta = q // t
    scaled_m = delta * m
    e1 = gen_normal_poly(size)
    e2 = gen_normal_poly(size)
    u = gen_binary_poly(size)
    ct0 = polyadd(
        polyadd(
            polymul(pk[0], u, q, poly_mod),
            e1, q, poly_mod),
        scaled_m, q, poly_mod
    )
    ct1 = polyadd(
        polymul(pk[1], u, q, poly_mod),
        e2, q, poly_mod
    )
    return (ct0, ct1)


def decrypt(sk, size, q, t, poly_mod, ct):
    """Decrypt a ciphertext

    Args:
        sk: secret-key.
        size: size of polynomials.
        q: ciphertext modulus.
        t: plaintext modulus.
        poly_mod: polynomial modulus.
        ct: ciphertext.

    Returns:
        Integer representing the plaintext.
    """
    scaled_pt = polyadd(
        polymul(ct[1], sk, q, poly_mod),
        ct[0], q, poly_mod
    )
    delta = q // t
    decrypted_poly = np.round(scaled_pt / delta) % t
    return int(decrypted_poly[0])


# Function for adding and multiplying encrypted values

def add_plain(ct, pt, q, t, poly_mod):
    """Add a ciphertext and a plaintext.

    Args:
        ct: ciphertext.
        pt: integer to add.
        q: ciphertext modulus.
        t: plaintext modulus.
        poly_mod: polynomial modulus.

    Returns:
        Tuple representing a ciphertext.
    """
    size = len(poly_mod) - 1
    # encode the integer into a plaintext polynomial
    m = np.array([pt] + [0] * (size - 1), dtype=np.int64) % t
    delta = q // t
    scaled_m = delta * m
    new_ct0 = polyadd(ct[0], scaled_m, q, poly_mod)
    return (new_ct0, ct[1])


def add_cipher(ct1, ct2, q, poly_mod):
    """Add a ciphertext and a ciphertext.

    Args:
        ct1, ct2: ciphertexts.
        q: ciphertext modulus.
        poly_mod: polynomial modulus.

    Returns:
        Tuple representing a ciphertext.
    """
    new_ct0 = polyadd(ct1[0], ct2[0], q, poly_mod)
    new_ct1 = polyadd(ct1[1], ct2[1], q, poly_mod)
    return (new_ct0, new_ct1)


def mul_plain(ct, pt, q, t, poly_mod):
    """Multiply a ciphertext and a plaintext.

    Args:
        ct: ciphertext.
        pt: integer to multiply.
        q: ciphertext modulus.
        t: plaintext modulus.
        poly_mod: polynomial modulus.

    Returns:
        Tuple representing a ciphertext.
    """
    size = len(poly_mod) - 1
    # encode the integer into a plaintext polynomial
    m = np.array([pt] + [0] * (size - 1), dtype=np.int64) % t
    new_c0 = polymul(ct[0], m, q, poly_mod)
    new_c1 = polymul(ct[1], m, q, poly_mod)
    return (new_c0, new_c1)


if __name__ == "__main__":
    # Scheme's parameters
    # polynomial modulus degree
    n = 2**4
    # ciphertext modulus
    q = 2**15
    # plaintext modulus
    t = 2**8
    # polynomial modulus
    poly_mod = np.array([1] + [0] * (n - 1) + [1])

    # Keygen
    pk, sk = keygen(n, q, poly_mod)

    # Encryption
    pt1, pt2 = 73, 20
    cst1, cst2 = 7, 5

    ct1 = encrypt(pk, n, q, t, poly_mod, pt1)
    ct2 = encrypt(pk, n, q, t, poly_mod, pt2)

    print("[+] Ciphertext ct1({}):".format(pt1))
    print("")
    print("\t ct1_0:", ct1[0])
    print("\t ct1_1:", ct1[1])
    print("")
    print("[+] Ciphertext ct2({}):".format(pt2))
    print("")
    print("\t ct1_0:", ct2[0])
    print("\t ct1_1:", ct2[1])
    print("")

    # Evaluation
    ct3 = add_plain(ct1, cst1, q, t, poly_mod)
    ct4 = mul_plain(ct2, cst2, q, t, poly_mod)
    # ct5 = ct1 + 7 + 3 * ct2
    ct5 = add_cipher(ct3, ct4, q, poly_mod)

    # Decryption
    decrypted_ct3 = decrypt(sk, n, q, t, poly_mod, ct3)
    decrypted_ct4 = decrypt(sk, n, q, t, poly_mod, ct4)
    decrypted_ct5 = decrypt(sk, n, q, t, poly_mod, ct5)
    print("[+] Decrypted ct3(ct1 + {}): {}".format(cst1, decrypted_ct3))
    print("[+] Decrypted ct4(ct2 * {}): {}".format(cst2, decrypted_ct4))
    print("[+] Decrypted ct5(ct1 + {} + {} * ct2): {}".format(cst1, cst2, decrypted_ct5))

but it's randint(0,1) in the code, need to change to randint(0,2).

Oups, thanks for pointing out! Will update that in a second

comments of polyadd and polymul are set the other way around

Oh yeah, thanks!

Nice code! Do you have the further code about the linearization of cipher multiplication?

No, but there is a team from Bitdefneder who did an awesome job building the entire FV scheme here: https://bit-ml.github.io/blog/post/homomorphic-encryption-toy-implementation-in-python/

can we perform the addition of 2 vectors whose coefficients are negative for example, v1[-6,0,5] and v2[-2,6,1]?