【FHE test】

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

def polyadd(x,y,modulus,poly_mod):
    """
    x,y:两个相加的多项式
    modulus:系数模数
    poly_mod:多项式模数
    :return:
    """
    return np.int64(np.round(poly.polydiv(poly.polyadd(x,y) % modulus, poly_mod)[1]% modulus) )

def polymul(x,y,modulus,poly_mod):
    """
    x,y:两个相乘的多项式
    modulus:系数模数
    poly_mod:多项式模数
    :return:
    """
    return np.int64(np.round(poly.polydiv(poly.polymul(x,y) % modulus, poly_mod)[1] % modulus))


def gen_binary_poly(size):
    """
    generates sk
    generates a polynomial with coffecients in [0,1]
    :param size: number of coffcients
    polynomial's degree is size - 1
    :return: array of coefficients of the polynomial
    """
    return np.random.randint(0, 2, size, dtype=np.int64)


# res = gen_binary_poly(10)
# print(res)

def gen_uniform_poly(size, modulus):
    """
    generates a polynomial with cofficients being integers in Z_modulus
    在Zq中随机取样
    :param size:number of coffcients
    polynomial's degree is size - 1
    :param modulus:
    :return:array of coefficients of the polynomial
    """
    return np.random.randint(0, modulus, size, dtype=np.int64)


# res = gen_uniform_poly(10,2**4)
# print(res)

def gen_normal_poly(size):
    """
    Generates a polynomial with coeffecients in a normal distribution
    在正态分布中随机取样
    :param size:number of coffcients
    polynomial's degree is size - 1
    :return:array of coefficients of the polynomial
    """
    return np.int64(np.random.normal(0, 2, size=size))
    # return np.random.normal(0,2,size=size)


# res = gen_normal_poly(10)
# print(res)

def kengen(size, modulus, poly_mod):
    """

    :param size:size of the polynomial for ths pk and sk
    :param modulus: cofficient modulus
    :param poly_mod: polynomial modulus
    :return:pk and sk
    """
    sk = gen_binary_poly(size)
    a = gen_uniform_poly(size, modulus)
    e = gen_normal_poly(size)
    b = polyadd(polymul(-a,sk,modulus,poly_mod),-e,modulus,poly_mod)
    # print("sk:")
    # print(sk)
    # print("a:")
    # print(a)
    # print("e:")
    # print(e)
    # print("b:")
    # print(b)
    return (b,a),sk

# poly_mod = np.array([1]+[0]*3+[1])
# res = kengen(10,2**4,poly_mod)
# print(res)

def encrypt(pk,size,q,t,poly_mod,pt):
    """
    encrypt a integer
    加密一个整数
    :param pk: pk
    :param size: size of polynomials
    :param q: ciphertext modulus
    :param t:  plaintext modulus
    :param poly_mod: polynomial modulus
    :param pt: integer to be encrypted
    :return: 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 % q
    e1,e2 = gen_normal_poly(size),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):
    """

    :param sk: (sk0,sk1)
    :param size: size of polynomials
    :param q: ciphertext modulus
    :param t: plaintext modulus
    :param poly_mod:
    :param ct: ciphertext
    :return: Integer representing the plaintext.
    """
    scaled_pt = polyadd(
        polymul(ct[1],sk,q,poly_mod),
        ct[0],q,poly_mod
    )
    decryptd_poly = np.round(scaled_pt * t / q) % t
    return np.int64(decryptd_poly[0])

def add_plain(ct,pt,q,t,poly_mod):
    """
    add a ciphertext and a plaintext
    :param ct: cithertext
    :param pt: integer to add
    :param q: cipthertext modulus
    :param t: plaintext modulus
    :param poly_mod: polynomial modulus
    :return: 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 % q
    new_ct0 = polyadd(ct[0],scaled_m,q,poly_mod)
    return (new_ct0,ct[1])

def mul_plain(ct,pt,q,t,poly_mod):
    """
    multiply a ciphertext and  plaintext
    :param ct: ciphertext
    :param pt: plaintext
    :param q:
    :param t:
    :param poly_mod:
    :return:
    """
    size = len(poly_mod) - 1
    # encode the integer into a plaintext polynomial
    m = np.array([pt]+[0]*(size-1),dtype=np.int64) % t
    new_ct0 = polymul(ct[0],m,q,poly_mod)
    new_ct1 = polymul(ct[1],m,q,poly_mod)
    return (new_ct0,new_ct1)

# polymodulus 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])
# kengen
pk,sk = kengen(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)

#decrypted
decrypted_ct3 = decrypt(sk,n,q,t,poly_mod,ct3)
decrypted_ct4 = decrypt(sk,n,q,t,poly_mod,ct4)

print("[+] Decrypted ct3(ct1 + {}): {}".format(cst1, decrypted_ct3))
print("[+] Decrypted ct4(ct2 * {}): {}".format(cst2, decrypted_ct4))