rlwe implementation in Python
import argparse
import json

import numpy as np
from rlwe import RLWE
from discrete_gauss import DiscreteGaussian
import time

def main(args):
    n = args.n
    q = args.q
    t = args.t
    sigma = 3 # standard deviation of the discrete Gaussian distribution
    
    # Initialize the DiscreteGaussian distribution
    distribution = DiscreteGaussian(sigma)
    
    # Initialize RLWE object
    rlwe = RLWE(n, q, t, distribution)
    
    # Generate secret and public key
    secret_key = rlwe.SecretKeyGen()

    pk_gen_start_time = time.time() 

    public_key = rlwe.PublicKeyGen(secret_key)

    pk_gen_end_time = time.time() 
    pk_gen_elapsed_time = pk_gen_end_time - pk_gen_start_time

    print(f"Time to generate public key: {pk_gen_elapsed_time:.6f} seconds")

    # Extract u polynomial from the RLWE object
    u = rlwe.u

    # Add zeroes at the beginning of the u polynomial to make it the same degree as the public key
    u.coefficients = np.concatenate((np.zeros(n - len(u.coefficients)), u.coefficients))

    # Generate message (plaintext)
    message = rlwe.Rt.sample_polynomial()

    # Add zeroes at the beginning of the message polynomial to make it the same degree as the public key
    message.coefficients = np.concatenate((np.zeros(n - len(message.coefficients)), message.coefficients))

    encrypt_start_time = time.time()

    # Encrypt message
    ciphertext, error = rlwe.Encrypt(public_key, message)

    encrypt_end_time = time.time()
    encrypt_elapsed_time = encrypt_end_time - encrypt_start_time

    print(f"Time to encrypt the message: {encrypt_elapsed_time:.6f} seconds")

    (c0, c1) = ciphertext
    (e0, e1) = error

    # Add zeroes at the beginning of the e0 polynomial to make it the same degree as the public key
    e0.coefficients = np.concatenate((np.zeros(n - len(e0.coefficients)), e0.coefficients))
    e1.coefficients = np.concatenate((np.zeros(n - len(e1.coefficients)), e1.coefficients))

    decrypt_start_time = time.time()
    # Decrypt ciphertext
    dec = rlwe.Decrypt(secret_key, ciphertext, error)

    decrypt_end_time = time.time()
    decrypt_elapsed_time = decrypt_end_time - decrypt_start_time

    print(f"Time to decrypt the message: {decrypt_elapsed_time:.6f} seconds")
    # Ensure that message and dec match
    for i in range(len(message.coefficients)):
        assert message.coefficients[i] == dec.coefficients[i], "Message and dec do not match"
    
    # Convert public key to JSON and save to file
    with open(args.output, "w") as f:
        json.dump({
            'pk0': adjust_negative_coefficients([int(coeff) for coeff in public_key[0].coefficients.tolist()], q),
            'pk1': adjust_negative_coefficients([int(coeff) for coeff in public_key[1].coefficients.tolist()], q),
            'm': adjust_negative_coefficients([int(coeff) for coeff in message.coefficients.tolist()], q),
            'u': adjust_negative_coefficients([int(coeff) for coeff in u.coefficients.tolist()], q),
            'e0': adjust_negative_coefficients([int(coeff) for coeff in e0.coefficients.tolist()], q),
            'e1': adjust_negative_coefficients([int(coeff) for coeff in e1.coefficients.tolist()], q),
            'c0': adjust_negative_coefficients([int(coeff) for coeff in c0.coefficients.tolist()], q),
            'c1': adjust_negative_coefficients([int(coeff) for coeff in c1.coefficients.tolist()], q),
        }, f)

def adjust_negative_coefficients(coefficients, modulus):
    return [(modulus + coeff if coeff < 0 else coeff) for coeff in coefficients]


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Generate inputs for BFV zk proof circuit in json format')
    parser.add_argument('-n', type=int, required=True, help='Degree of f(x), must be a power of 2.')
    parser.add_argument('-q', type=int, required=True, help='Modulus q of the ciphertext space')
    parser.add_argument('-t', type=int, required=True, help='Modulus t of the plaintext space.')
    parser.add_argument('--output', type=str, default='input.json', help='Output file name')
    
    args = parser.parse_args()
    main(args)


cli.py
yuriko627/rlwe-py
import numpy as np
import scipy.stats as ss

class DiscreteGaussian:
	def __init__(self, sigma: float):
		"""
		Initialize the discrete Gaussian distribution.

		Parameters:
		sigma: Standard deviation
		"""
		self.sigma = sigma
		self.z_lower = -6*sigma
		self.z_upper = 6*sigma

		# Generate the range which includes both ends
		self.x = np.arange(self.z_lower, self.z_upper + 1)

		# Define upper and lower bounds for each integer value
		self.xU, self.xL = self.x + 0.5, self.x - 0.5

		# Calculate the probability of each integer value in the range
		self.prob = ss.norm.cdf(self.xU, scale=self.sigma) - ss.norm.cdf(self.xL, scale=self.sigma)
		self.prob = self.prob / self.prob.sum() # Normalize the probabilities so their sum is 1


	def sample(self, size) -> np.ndarray:
		"""
		Sample from the Discrete Gaussian distribution

		Parameters:
		size: Number of samples to generate

		Returns: NumPy array of generated samples
		"""
		return np.random.choice(self.x, size=size, p=self.prob)

discrete_gauss.py
yuriko627/rlwe-py
import unittest
import numpy as np
from discrete_gauss import DiscreteGaussian

class TestDiscreteGaussian(unittest.TestCase):
	def setUp(self):
		self.sigma = 3
		self.gaussian = DiscreteGaussian(self.sigma)

	def test_instantiation(self):
		# assigning correct value of sigma
		self.assertEqual(self.gaussian.sigma, self.sigma)

	def test_probability_distribution(self):
		# ensure probabilities sum up to approximately 1 (due to floating-poitn precision)
		self.assertAlmostEqual(self.gaussian.prob.sum(), 1, places=5)

	def test_sampling(self):
		# sample a large number of times and check if the mean is approximately 0 (exptected for this setup)
		samples = self.gaussian.sample(100000)
		self.assertAlmostEqual(np.mean(samples), 0, places=1)
		# check that is doesn't sample values outside the range
		self.assertTrue(np.all(samples >= self.gaussian.z_lower))
		self.assertTrue(np.all(samples <= self.gaussian.z_upper))
	

if __name__ == '__main__':
    unittest.main()
gauss_test.py
yuriko627/rlwe-py
import unittest
import numpy as np
from polynomial import PolynomialRing, Polynomial, get_centered_remainder

class TestPolynomialRing(unittest.TestCase):
	def test_init_with_n(self):
		n = 4
		R = PolynomialRing(n)
		quotient = np.array([1, 0, 0, 0, 1])
		self.assertTrue(np.array_equal(R.denominator, quotient))
		self.assertEqual(R.n, n)

	def test_init_with_n_and_q(self):
		n = 4
		q = 7
		Rq = PolynomialRing(n, q)
		quotient = np.array([1, 0, 0, 0, 1])
		self.assertTrue(np.array_equal(Rq.denominator, quotient))
		self.assertEqual(Rq.Q, q)
		self.assertEqual(Rq.n, n)

	def test_init_with_n_and_invalid_q(self):
		n = 4
		q = 1
		with self.assertRaisesRegex(AssertionError, "modulus must be > 1"): PolynomialRing(n, q)

	def test_sample_poly_from_r_error(self):
		n = 4
		R = PolynomialRing(n)
		with self.assertRaisesRegex(AssertionError, "The modulus Q must be set to sample a polynomial from R_Q"): R.sample_polynomial()

	def test_sample_poly_from_rq(self):
		n = 4
		q = 8
		Rq = PolynomialRing(n, q)
		aq1 = Rq.sample_polynomial()
		aq2 = Rq.sample_polynomial()

		# Ensure that the coefficients of the polynomial are within Z_q = [-q/2, q/2)
		for coeff in aq1.coefficients:
			self.assertTrue(coeff >= -q // 2 and coeff <= q // 2)
		for coeff in aq2.coefficients:
			self.assertTrue(coeff >= -q // 2 and coeff <= q // 2)

        # Ensure that the degree of the sampled poly is equal or less than d (it might be less if the leading coefficient sampled is 0)
		count1 = 0
		for coeff in aq1.coefficients:
			count1 += 1

		count2 = 0
		for coeff in aq2.coefficients:
			count2 += 1

		self.assertTrue(count1 <= Rq.n)
		self.assertTrue(count2 <= Rq.n)

class TestPolynomialInRingR(unittest.TestCase):

    def test_init_poly_in_ring_R(self):
        n = 4
        R = PolynomialRing(n)
        coefficients = [6, 6, 6, 4, 5]
        # a is the polynomial in R reduced by the quotient polynomial
        a = Polynomial(coefficients, R)
        self.assertTrue(np.array_equal(a.coefficients, [6, 6, 4, -1]))

        # The degree of resulting poly should be less than the degree of fx
        self.assertTrue(len(a.coefficients) < len(a.ring.denominator))

        # After reduction, the polynomial should not be multiple of the quotient
        _, remainder = np.polydiv(a.coefficients, a.ring.denominator)
        self.assertTrue(np.array_equal(remainder, a.coefficients))

    def test_add_poly_in_ring_R(self):
        n = 4
        R = PolynomialRing(n)
        coefficients_1 = [3, 3, 4, 4, 4]
        coefficients_2 = [3, 3, 2, 0, 1]

        a1 = Polynomial(coefficients_1, R)
        a2 = Polynomial(coefficients_2, R)

        # a1 + a2
        result = np.polyadd(a1.coefficients, a2.coefficients)

        # The addition is happening in the ring R.
        result = Polynomial(result, R)

        # The resulting poly is 6, 6, 6, 4, 5. After reduction, the polynomial is 6, 6, 4, -1.
        self.assertTrue(np.array_equal(result.coefficients, [6, 6, 4, -1]))

        # After reduction, the polynomial should not be multiple of the quotient
        _, remainder = np.polydiv(result.coefficients, result.ring.denominator)
        self.assertTrue(np.array_equal(remainder, result.coefficients))

        # The degree of the result poly should be the max of the two degrees
        self.assertEqual(len(result.coefficients), max(len(a1.coefficients), len(a2.coefficients)))

    def test_mul_poly_in_ring_R(self):
        n = 4
        R = PolynomialRing(n)
        coefficients_1 = [3, 0, 4]
        coefficients_2 = [2, 0, 1]

        a1 = Polynomial(coefficients_1, R)
        a2 = Polynomial(coefficients_2, R)

        result = np.polymul(a1.coefficients, a2.coefficients)

        # The multiplication is happening in the ring R.
        result = Polynomial(result, R)

        self.assertTrue(np.array_equal(result.coefficients, [11, 0, -2]))

        # After reduction, the polynomial should not be multiple of the quotient
        _, remainder = np.polydiv(result.coefficients, result.ring.denominator)
        self.assertTrue(np.array_equal(remainder, result.coefficients))

        # The degree of the result poly should less than the degree of fx
        self.assertTrue(len(result.coefficients) < len(result.ring.denominator))

    def test_scalar_mul_poly_in_ring_R(self):
        n = 4
        R = PolynomialRing(n)
        coefficients = [4, 3, 0, 4]
        a = Polynomial(coefficients, R)

        # 2 * a. The resulting poly is 8, 6, 0, 8. Reduction does not change the coefficients.
        result = np.polymul(2, a.coefficients)

        # The multiplication is happening in the ring R.
        result = Polynomial(result, R)

        self.assertTrue(np.array_equal(result.coefficients, [8, 6, 0, 8]))

        # After reduction, the polynomial should not be multiple of the quotient
        _, remainder = np.polydiv(result.coefficients, result.ring.denominator)
        self.assertTrue(np.array_equal(remainder, result.coefficients))

class TestPolynomialInRingRq(unittest.TestCase):

    def test_init_poly_in_ring_Rq(self):
        n = 4
        q = 7
        Rq = PolynomialRing(n, q)
        coefficients = [3, 1, 0]
        # fetch the Z_q set which is (-q/2, q/2]
        assert Rq.Z_Q == [-3, -2, -1, 0, 1, 2, 3]
        # aq is the polynomial in Rq reduced by the quotient polynomial and by the modulus q.
        aq = Polynomial(coefficients, Rq)
        self.assertTrue(np.array_equal(aq.coefficients, [3, 1, 0]))
        # After reduction, the polynomial should not be multiple of the quotient
        _, remainder = np.polydiv(aq.coefficients, aq.ring.denominator)
        self.assertTrue(np.array_equal(remainder, aq.coefficients))
        # All coefficients should be in Z_q
        for coeff in aq.coefficients:
            self.assertTrue(coeff in Rq.Z_Q)


    def test_add_poly_in_ring_Rq(self):
        n = 4
        q = 7
        Rq = PolynomialRing(n, q)
        coefficients_1 = [3, 3, 4, 4, 4] # r(x)

        # 1st reduction
        _, remainder = np.polydiv(coefficients_1, Rq.denominator) # r(x)/f(x). where f(x)= x^n + 1

        # 2nd reduction
        for i in range(len(remainder)):
            remainder[i] = get_centered_remainder(remainder[i], Rq.Q)

        aq1 = Polynomial(coefficients_1, Rq)

        # check that coefficents of aq1 and remainder are the same
        self.assertTrue(np.array_equal(aq1.coefficients, remainder))

        coefficients_2 = [3, 3, 2, 0, 1]

        aq2 = Polynomial(coefficients_2, Rq)

        # aq1 + aq2
        result = np.polyadd(aq1.coefficients, aq2.coefficients)

        # The addition is happening in the ring Rq.
        result = Polynomial(result, Rq)

        # The resulting poly is 6, 6, 6, 4, 5. After first reduction, the polynomial is 6, 6, 4, -1.  After second reduction (modulo q), the polynomial is -1. -1, -3, -1.
        self.assertTrue(np.array_equal(result.coefficients, [-1, -1, -3, -1]))

        # After reduction, the polynomial should not be multiple of the quotient
        _, remainder = np.polydiv(result.coefficients, result.ring.denominator)
        self.assertTrue(np.array_equal(remainder, result.coefficients))

        # The degree of the result poly should be the max of the two degrees
        self.assertEqual(len(result.coefficients), max(len(aq1.coefficients), len(aq2.coefficients)))

    def test_mul_poly_in_ring_Rq(self):
        n = 4
        q = 7
        Rq = PolynomialRing(n, q)
        coefficients_1 = [3, 0, 4]
        coefficients_2 = [2, 0, 1]

        aq1 = Polynomial(coefficients_1, Rq)
        aq2 = Polynomial(coefficients_2, Rq)

        # aq1 * aq2. After reduction, the polynomial is . After second reduction (modulo q), the polynomial is -3, 0, -2.
        result = np.polymul(aq1.coefficients, aq2.coefficients)

        # The multiplication is happening in the ring Rq.
        result = Polynomial(result, Rq)

        self.assertTrue(np.array_equal(result.coefficients, [-3, 0, -2]))

        # After reduction, the polynomial should not be multiple of the quotient
        _, remainder = np.polydiv(result.coefficients, result.ring.denominator)
        self.assertTrue(np.array_equal(remainder, result.coefficients))

        # The degree of the result poly should less than the degree of fx
        self.assertTrue(len(result.coefficients) < len(result.ring.denominator))

    def test_scalar_mul_poly_in_ring_Rq(self):
        n = 4
        q = 7
        Rq = PolynomialRing(n, q)
        coefficients = [4, 3, 0, 4]
        aq = Polynomial(coefficients, Rq)

        # The resulting poly is 8, 6, 0, 8. After modulo q reduction, the polynomial is 1, -1, 0, 1.
        result = np.polymul(2, aq.coefficients)

        # The multiplication is happening in the ring Rq.
        result = Polynomial(result, Rq)

        self.assertTrue(np.array_equal(result.coefficients, [1, -1, 0, 1]))

        # After reduction, the polynomial should not be multiple of the quotient
        _, remainder = np.polydiv(result.coefficients, result.ring.denominator)
        self.assertTrue(np.array_equal(remainder, result.coefficients))


class TestCustomModulo(unittest.TestCase):

    def test_positive_values(self):
        self.assertEqual(get_centered_remainder(7, 10), -3)  # 7 % 10. Lies in the range (-5, 5]
        self.assertEqual(get_centered_remainder(15, 10), 5) # 15 % 10 = 5, which is <= 5
        self.assertEqual(get_centered_remainder(17, 10), -3) # 17 % 10 = 7, which is > 5. So, 7 - 10 = -3

    def test_negative_values(self):
        self.assertEqual(get_centered_remainder(-7, 10), 3) # Lies in the range (-5, 5]
        self.assertEqual(get_centered_remainder(-15, 10), 5) # -15 % 10 = 5 (in Python, % returns non-negative), which is <= 5
        self.assertEqual(get_centered_remainder(-17, 10), 3) # -17 % 10 = 3, which is <= 5

    def test_boundary_values(self):
        q = 7
        self.assertEqual(get_centered_remainder(-q//2 + 1, q), -q//2 + 1) # The smallest positive number in the range
        self.assertEqual(get_centered_remainder(q//2, q), q//2) # The largest number in the range

    def test_zero(self):
        self.assertEqual(get_centered_remainder(0, 10), 0) # 0 lies in the range (-5, 5]

if __name__ == "__main__":
    unittest.main()
poly_test.py
yuriko627/rlwe-py
import numpy as np

class PolynomialRing:
	# polynomial ring R = Z[x]/f(x) where f(x)=x^n+1
	# n is a power of 2
	# If the modulus Q > 1 is specified then the ring is R_Q = Z_Q[x]/f(x). Namely, the coefficients of the polynomials are in the set Z_Q = (-Q/2, Q/2]
	def __init__(self, n, modulus=None):
		# ensure that n is a power of 2
		assert n > 0 and (n & (n-1)) == 0, "n must be a power of 2"

		fx = [1] + [0] * (n-1) + [1]

		self.denominator = fx
		self.Q = modulus
		self.n = n

		# if modulus is defined, chek that it is > 1
		if modulus is not None:
			assert modulus > 1, "modulus must be > 1"
			self.Z_Q = [j for j in range (-self.Q // 2 + 1, self. Q // 2 + 1)]
		else:
			self.Z_Q = None

	def sample_polynomial(self):
		"""
		Sample polynomial a_Q from R_Q.
		"""
		# ensure that modulus is set
		if self.Q is None:
			raise AssertionError("The modulus Q must be set to sample a polynomial from R_Q")

		a_Q_coeff = np.random.choice(self.Z_Q, size=self.n)

		return Polynomial(a_Q_coeff, self)

class Polynomial:
	def __init__(self, coefficients, ring: PolynomialRing):
		self.ring = ring

		# apply redution to the ring
		remainder = reduce_coefficients(coefficients, self.ring)
		self.coefficients = remainder

def reduce_coefficients(coefficients, ring):
	# reduce (divide) coefficients by the denominator polynomial
	_, remainder = np.polydiv(coefficients, ring.denominator)

	# if the ring is R_Q, apply reduction by taking coeff mod Q
	if ring.Q is not None:
		for i in range(len(remainder)):
			remainder[i] = get_centered_remainder(remainder[i], ring.Q)

		# ensure that the coefficients are in the set Z_Q wich is defined as (-Q/2, Q/2]
		Z_Q_set = set(j for j in range(-ring.Q//2 + 1, ring.Q//2+1))
		for value in remainder:
			assert value in Z_Q_set, "Coefficients must be in Z_Q"

	return remainder

def get_centered_remainder(x, modulus):
    # The concept of the centered remainder is that after performing the modulo operation,
    # The result is in the set (-modulus/2, ..., modulus/2], rather than [0, ..., modulus-1].
    # If r is in range [0, modulus/2] then the centered remainder is r.
    # If r is in range [modulus/2 + 1, modulus-1] then the centered remainder is r - modulus.
    # If modulus is 7, then the field is {-3, -2, -1, 0, 1, 2, 3}.
    # 10 % 7 = 3. The centered remainder is 3.
    # 11 % 7 = 4. The centered remainder is 4 - 7 = -3.
    r = x % modulus
    return r if r <= modulus / 2 else r - modulus





polynomial.py
yuriko627/rlwe-py
from polynomial import PolynomialRing, Polynomial
from discrete_gauss import DiscreteGaussian
import math
import numpy as np

class RLWE:
    def __init__(self, n, q, t, distribution: DiscreteGaussian):
        """
        Initialize the RLWE instance with a given polynomial ring and error distribution.

        Parameters:
        n: degree of the f(x) which is the denominator of the polynomial ring, must be a power of 2.
        q: modulus q
        t: modulus t of the plaintext space
        distribution: Error distribution (e.g. Gaussian).
        """
        # Ensure that the modulus of the plaintext space is smaller than the modulus of the polynomial ring
        if t > q:
            raise ValueError("The modulus of the plaintext space must be smaller than the modulus of the polynomial ring.")

        # Ensure that n is a power of 2
        assert n > 0 and (n & (n-1)) == 0, "n must be a power of 2"

        # Ensure that p and q are greater than 1
        assert q > 1, "modulus q must be > 1"
        assert t > 1, "modulus t must be > 1"

        # Ensure that t is a prime number
        assert self.is_prime(t), "modulus t must be a prime number"

        self.R = PolynomialRing(n)
        self.Rq = PolynomialRing(n, q)
        self.Rt = PolynomialRing(n, t)
        self.distribution = distribution

        # Sample error polynomial from the distribution χ Error
        self.e = self.SampleFromChiErrorDistribution() ## used in public key gen

        # Sample ephemeral key polynomial u from the distribution χ Key
        self.u = self.SampleFromChiKeyDistribution()

    def SampleFromChiKeyDistribution(self):
        """
        Sample a polynomial from the χ Key distribution.
        Namely, the coefficients are sampled uniformely from the ternary set {-1, 0, 1}. (coefficients are either of them)

    	Returns: Sampled polynomial.
        """

        # Sample n coefficients from the ternary set {-1, 0, 1}
        coefficients = np.random.choice([-1, 0, 1], size=self.R.n)

        return Polynomial(coefficients, self.R)

    def SampleFromChiErrorDistribution(self):
        """
        Sample a polynomial from the χ Error distribution.

        Returns: Sampled polynomial.
        """
        # Sample a polynomial from the χ Error distribution
        coefficients = self.distribution.sample(self.R.n)
        return Polynomial(coefficients, self.R)

    def SecretKeyGen(self):
        """
        Randomly generate a secret key.

        Returns: Generated secret key polynomial.
        """

        return self.SampleFromChiKeyDistribution()

    def PublicKeyGen(self, secret_key: Polynomial):
        """
        Generate a public key from a given secret key.

        Parameters:
        secret_key: Secret key.

        Returns: Generated public key.
        """
        # Sample a polynomial a from Rq
        a = self.Rq.sample_polynomial() # TODO: what is this distribution? is it correct?

        # a*s. The result will be in Rq
        mul = np.polymul(a.coefficients, secret_key.coefficients)

        # assert that the degree of mul is at most 2 * (n - 1). Namely, the number of coefficients is at most 2 * n
        assert len(mul) <= 2 * self.R.n, f"The degree of mul is {len(mul)} which is greater than 2 * (n - 1) = {2 * self.R.n}"

        mul = Polynomial(mul, self.Rq)

        # a*s + e. The result will be in Rq
        b = np.polyadd(mul.coefficients, self.e.coefficients)

        b = Polynomial(b, self.Rq)
        pk0 = b

        # pk1 = -a. The result will be in Rq
        pk1 = np.polymul(a.coefficients, [-1])
        pk1 = Polynomial(pk1, self.Rq)

        # public_key = (b, -a)
        public_key = (pk0, pk1)

        return public_key

    def Encrypt(self, public_key: (Polynomial, Polynomial), m: Polynomial):
        """
        Encrypt a given message m with a given public_key .

        Parameters:
        public_key: Public key.
        m: message.

        Returns:
        ciphertext: Generated ciphertext.
        error: tuple of error values used in encryption.
        """
        # Ensure that the message is in Rt
        if m.ring != self.Rt:
            raise AssertionError("The message must be in Rt.")

        q = self.Rq.Q
        t = self.Rt.Q

        # Sample polynomials e0, e1 from the distribution χ Error
        e0 = self.SampleFromChiErrorDistribution()

        e1 = self.SampleFromChiErrorDistribution()

        # delta = q/t
        delta = q / t

        # Round delta to the lower integer
        delta = math.floor(delta)

        # Compute the ciphertext.
        # delta * m
        delta_m = np.polymul(delta, m.coefficients)
        # pk0 * u
        pk0_u = np.polymul(public_key[0].coefficients, self.u.coefficients)

        # delta * m + pk0 * u + e0
        ct_0 = np.polyadd(delta_m, pk0_u)
        ct_0 = np.polyadd(ct_0, e0.coefficients)

        # ct_0 will be in Rq
        ct_0 = Polynomial(ct_0, self.Rq)

        # pk1 * u
        pk1_u = np.polymul(public_key[1].coefficients, self.u.coefficients)

        # pk1 * u + e1
        ct_1 = np.polyadd(pk1_u, e1.coefficients)

        # The result will be in Rq
        ct_1 = Polynomial(ct_1, self.Rq)

        ciphertext = (ct_0, ct_1)
        error = (e0, e1)

        return ciphertext, error

    def Decrypt(self, secret_key: Polynomial, ciphertext: (Polynomial, Polynomial), error: (Polynomial, Polynomial)):
        """
        Decrypt a given ciphertext with a given secret key.

        Parameters:
        secret_key: Secret key.
        ciphertext: Ciphertext.

        Returns: Decrypted message.
        """
        # dec = round(t/q * ((ct0 + ct1*s) mod s)
        ct0 = ciphertext[0].coefficients
        ct1 = ciphertext[1].coefficients
        s = secret_key.coefficients
        t = self.Rt.Q
        q = self.Rq.Q

        ct1_s = np.polymul(ct1, s)

        # ct0 + ct1*s
        numerator_1 = np.polyadd(ct0, ct1_s)

        # Ensure that all the errors v < q/(2t) - 1/2
        # v = u * e + e0 + s * e1
        u_e = np.polymul(self.u.coefficients, self.e.coefficients)
        s_e1 = np.polymul(s, error[1].coefficients)

        v = np.polyadd(u_e, error[0].coefficients)
        v = np.polyadd(v, s_e1)

        # fresh error v is in Rq
        v = Polynomial(v, self.Rq)

        rt_Q = q % t

        threshold = q/(2*t) - rt_Q/2


        for v in v.coefficients:
            assert abs(v) < (threshold), f"Noise {abs(v)} exceeds the threshold value {threshold}, decryption won't work"

        # Numerator 1 is in Rq.
        numerator_1 = Polynomial(numerator_1, self.Rq)

        numerator = np.polymul(t, numerator_1.coefficients)

        # For each coefficient of the numerator, divide it by q and round it to the nearest integer
        quotient = [round(coeff / q) for coeff in numerator]

        # trim leading zeros
        quotient = np.trim_zeros(quotient, 'f')

        # quotient is in Rt
        quotient = Polynomial(quotient, self.Rt)

        return quotient

    def is_prime(self, n):
        if n < 2:
            return False
        for i in range(2, n):
            if n % i == 0:
                return False
        return True

rlwe.py
yuriko627/rlwe-py
import unittest
from discrete_gauss import DiscreteGaussian
from polynomial import PolynomialRing, Polynomial
from rlwe import RLWE
import time

class TestRLWE(unittest.TestCase):

    def setUp(self):
        self.n = 4
        self.q = 700
        self.sigma = 10
        self.discrete_gaussian = DiscreteGaussian(self.sigma)
        self.t = 7
        # Note that t and q do not have to be prime nor coprime.
        self.rlwe = RLWE(self.n, self.q, self.t, self.discrete_gaussian)

    def test_rlwe_initialization(self):
        self.assertEqual(self.rlwe.R.denominator, PolynomialRing(self.n).denominator)
        self.assertEqual(self.rlwe.R.n, PolynomialRing(self.n).n)
        self.assertEqual(self.rlwe.R.Q, PolynomialRing(self.n).Q)
        self.assertEqual(self.rlwe.R.Z_Q, PolynomialRing(self.n).Z_Q)

        self.assertEqual(self.rlwe.Rq.denominator, PolynomialRing(self.n, self.q).denominator)
        self.assertEqual(self.rlwe.Rq.n, PolynomialRing(self.n, self.q).n)
        self.assertEqual(self.rlwe.Rq.Q, PolynomialRing(self.n, self.q).Q)
        self.assertEqual(self.rlwe.Rq.Z_Q, PolynomialRing(self.n, self.q).Z_Q)

        self.assertEqual(self.rlwe.Rt.denominator, PolynomialRing(self.n, self.t).denominator)
        self.assertEqual(self.rlwe.Rt.n, PolynomialRing(self.n, self.t).n)
        self.assertEqual(self.rlwe.Rt.Q, PolynomialRing(self.n, self.t).Q)
        self.assertEqual(self.rlwe.Rt.Z_Q, PolynomialRing(self.n, self.t).Z_Q)

        self.assertEqual(self.rlwe.distribution, self.discrete_gaussian)

    def test_sample_from_chi_key_distribution(self):
        key = self.rlwe.SampleFromChiKeyDistribution()
        # Ensure that the degree of the sampled poly is equal or less than d (it might be less if the leading coefficient sampled is 0)
        count = 0
        for coeff in key.coefficients:
            count += 1
        self.assertTrue(count <= self.rlwe.R.n)

        # Ensure that the key is a polynomial in ring R
        self.assertEqual(key.ring, self.rlwe.R)

        # Ensure that the coefficients of the key are within {-1, 0, 1}
        for coeff in key.coefficients:
            self.assertTrue(coeff == -1 or coeff == 0 or coeff == 1)

    def test_sample_from_chi_error_distribution(self):
        error = self.rlwe.SampleFromChiErrorDistribution()
        # Ensure that the degree of the sample polynomial is at most n-1, which means the has at most n coefficients
        count = 0
        for coeff in error.coefficients:
            count += 1
        self.assertTrue(count <= self.rlwe.R.n)

        # Ensure that the error is a polynomial in ring R
        self.assertEqual(error.ring, self.rlwe.R)

        # Ensure that the secret key is sampled from the error distribution by checking that each coefficient is within the range of the error distribution
        for coeff in error.coefficients:
            self.assertTrue(coeff >= -6*self.sigma and coeff <= 6*self.sigma)


    def test_secret_key_gen(self):
        secret_key = self.rlwe.SecretKeyGen()
        # Ensure that the degree of the sample polynomial is at most n-1, which means the has at most n coefficients
        count = 0
        for coeff in secret_key.coefficients:
            count += 1
        self.assertTrue(count <= self.rlwe.R.n)

        # Ensure that the secret key is a polynomial in ring R
        self.assertEqual(secret_key.ring, self.rlwe.R)

        # Ensure that the coefficients of the key are within {-1, 0, 1}
        for coeff in secret_key.coefficients:
            self.assertTrue(coeff == -1 or coeff == 0 or coeff == 1)

    def test_public_key_gen(self):
        secret_key = self.rlwe.SecretKeyGen()
        public_key = self.rlwe.PublicKeyGen(secret_key)

        self.assertIsInstance(public_key, tuple)
        self.assertEqual(len(public_key), 2)
        self.assertIsInstance(public_key[0], Polynomial)
        self.assertIsInstance(public_key[1], Polynomial)

        # Ensure that the coefficients of the public key are within Z_q = (-q/2, q/2]
        for coeff in public_key[0].coefficients:
            self.assertTrue(coeff > -self.q // 2 and coeff <= self.q // 2)
        for coeff in public_key[1].coefficients:
            self.assertTrue(coeff > -self.q // 2 and coeff <= self.q // 2)

        # Ensure that the public key is a polynomial in Rq
        self.assertEqual(public_key[0].ring, self.rlwe.Rq)
        self.assertEqual(public_key[1].ring, self.rlwe.Rq)

        # TODO: add public key's value check

    def test_message_sample(self):
        secret_key = self.rlwe.SecretKeyGen()
        public_key = self.rlwe.PublicKeyGen(secret_key)

        message = self.rlwe.Rt.sample_polynomial()

        # Ensure that message is a polynomial in Rt
        self.assertEqual(message.ring, self.rlwe.Rt)

    def test_wrong_message_sample(self):
        secret_key = self.rlwe.SecretKeyGen()
        public_key = self.rlwe.PublicKeyGen(secret_key)

        message = self.rlwe.Rq.sample_polynomial()

        # message must be in Rt, but I sampled it from Rq. So it must throw an error.
        with self.assertRaisesRegex(AssertionError, "The message must be in Rt."):
            self.rlwe.Encrypt(public_key, message)

    def test_valid_encryption(self):
        secret_key = self.rlwe.SecretKeyGen()
        public_key = self.rlwe.PublicKeyGen(secret_key)

        message = self.rlwe.Rt.sample_polynomial()

        ciphertext, error = self.rlwe.Encrypt(public_key, message)

        # Ensure that the coefficients of the public key are within Z_q = (-q/2, q/2]
        for coeff in ciphertext[0].coefficients:
            self.assertTrue(coeff > -self.q // 2 and coeff <= self.q // 2)
        for coeff in ciphertext[1].coefficients:
            self.assertTrue(coeff > -self.q // 2 and coeff <= self.q // 2)

        # Ensure that the ciphertext is a polynomial in Rq
        self.assertEqual(ciphertext[0].ring, self.rlwe.Rq)
        self.assertEqual(ciphertext[1].ring, self.rlwe.Rq)

        ## TODO: add ciphertext value check

    def test_valid_decryption(self):
        secret_key = self.rlwe.SecretKeyGen()

        public_key = self.rlwe.PublicKeyGen(secret_key)

        message = self.rlwe.Rt.sample_polynomial()

        ciphertext, error = self.rlwe.Encrypt(public_key, message)

        dec = self.rlwe.Decrypt(secret_key, ciphertext, error)

        # ensure that message and dec are of the same degree
        self.assertEqual(len(message.coefficients), len(dec.coefficients))

        # ensure that message and dec are of the same coefficients
        for i in range(len(message.coefficients)):
            self.assertEqual(message.coefficients[i], dec.coefficients[i])

    def test_valid_decryption_real_world_params(self):
        n = 1024
        q = 2 ** 29
        sigma = 3
        discrete_gaussian = DiscreteGaussian(sigma)
        t = 7
        rlwe = RLWE(n, q, t, discrete_gaussian)

        start_time = time.time() 
        
        secret_key = rlwe.SecretKeyGen()

        end_time = time.time() 
        elapsed_time = end_time - start_time
        
        print(f"Time to generate secret key: {elapsed_time:.6f} seconds")

        print("secret_key: ", secret_key.coefficients)

        start_time = time.time() 

        public_key = rlwe.PublicKeyGen(secret_key)

        end_time = time.time() 
        elapsed_time = end_time - start_time
        
        print(f"Time to generate public key: {elapsed_time:.6f} seconds")

        print("public_key: ", public_key[0].coefficients)

        message = rlwe.Rt.sample_polynomial()

        print("message: ", message.coefficients)

        start_time = time.time() 

        ciphertext, error = rlwe.Encrypt(public_key, message)

        end_time = time.time() 
        elapsed_time = end_time - start_time
        
        print(f"Time to encrypt the message: {elapsed_time:.6f} seconds")

        print("ciphertext: ", ciphertext[0].coefficients)

        start_time = time.time() 

        dec = rlwe.Decrypt(secret_key, ciphertext, error)

        end_time = time.time() 
        elapsed_time = end_time - start_time
        
        print(f"Time to decrypt the message: {elapsed_time:.6f} seconds")

        print("dec: ", dec.coefficients)

        # ensure that message and dec are of the same degree
        self.assertEqual(len(message.coefficients), len(dec.coefficients))

        # ensure that message and dec are of the same coefficients
        for i in range(len(message.coefficients)):
            self.assertEqual(message.coefficients[i], dec.coefficients[i])

if __name__ == "__main__":
    unittest.main()

rlwe_test.py
yuriko627/rlwe-py