polyhe 1.0.0__py3-none-any.whl

polyhe/__init__.py

@@ -0,0 +1,31 @@
+"""PolyHE package"""
+
+import enum
+import importlib
+
+from polyhe.core import Options
+from polyhe.core.context import Context
+from polyhe.core.ciphertext import Ciphertext
+
+
+__all__ = (
+    "new",
+    "Backend",
+    "Options",
+    "Context",
+    "Ciphertext"
+)
+
+
+class Backend(enum.StrEnum):
+    """Encryption backend"""
+    OPENFHE = enum.auto()
+    TENSEAL = enum.auto()
+    UARCHFHE = enum.auto()
+
+
+def new(backend: Backend = Backend.OPENFHE, options: Options = Options()) -> Context:
+    """Generate encryption context"""
+    module = importlib.import_module(f"{__name__}.{backend}")
+    cls: type[Context] = getattr(module, "Context")
+    return cls(options)

polyhe/core/__init__.py

@@ -0,0 +1,80 @@
+"""Structures package"""
+
+from dataclasses import dataclass
+
+import numpy as np
+
+
+__all__ = (
+    "Options",
+)
+
+
+# COEFF_MODULUS[security_level][poly_degree]
+# source: sealapi.CoeffModulus.BFVDefault
+COEFF_MODULUS = {
+    128: {
+        10: [27],
+        11: [54],
+        12: [36, 36, 37],
+        13: [43, 43, 44, 44, 44],
+        14: [48, 48, 48, 49, 49, 49, 49, 49, 49],
+        15: [55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 56]
+    },
+    192: {
+        10: [19],
+        11: [37],
+        12: [25, 25, 25],
+        13: [38, 38, 38, 38],
+        14: [50, 50, 50, 50, 50, 50],
+        15: [54, 54, 54, 54, 54, 55, 55, 55, 55, 55, 55]
+    },
+    256: {
+        10: [14],
+        11: [29],
+        12: [58],
+        13: [39, 39, 40],
+        14: [47, 47, 47, 48, 48],
+        15: [52, 53, 53, 53, 53, 53, 53, 53, 53]
+    }
+}
+
+
+@dataclass(repr=False, slots=True, frozen=True)
+class Descriptor:
+    """Data descriptor"""
+    type: type
+    shape: tuple[int, ...]
+
+    @property
+    def numpy(self) -> bool:
+        """Is numpy-like"""
+        return issubclass(self.type, np.number)
+
+    def __repr__(self) -> str:
+        """Descriptor representation"""
+        return f"<{self.__class__.__name__} type={self.type.__name__} shape={self.shape}>"
+
+
+@dataclass(slots=True, frozen=True)
+class Options:
+    """Encryption options"""
+    slots: int = 13
+    scale: int = 40
+    security: int = 128
+
+
+def numpy2python(scalar: np.number) -> type:
+    """Transform numpy scalar to python scalar"""
+    return type(scalar().tolist())  # type: ignore
+
+
+def materialize(obj):
+    """Generate description and array"""
+    is_np = isinstance(obj, np.ndarray)
+    obj = np.asarray(obj)
+    type = obj.dtype.type
+    if not is_np:
+        type = numpy2python(type)  # type: ignore
+    shape = obj.shape
+    return Descriptor(type, shape), obj

polyhe/core/ciphertext.py

@@ -0,0 +1,115 @@
+"""Ciphertext package"""
+from __future__ import annotations
+
+import typing
+import operator
+import itertools
+from dataclasses import dataclass
+
+import numpy as np
+
+from polyhe.core import Descriptor
+
+if typing.TYPE_CHECKING:
+    from polyhe.core.context import Context
+
+
+__all__ = (
+    "Ciphertext",
+)
+
+
+@dataclass(eq=False, order=False, slots=True, frozen=True)
+class Ciphertext[T]:
+    """Abstract ciphertext"""
+    _context: Context
+    _descriptor: Descriptor
+    _chunks: tuple
+
+    def _new(self, /, *args, **kwds):
+        """Create new operable ciphertext"""
+        return self.__class__(
+            _context=self._context,
+            _descriptor=self._descriptor,
+            *args, **kwds
+        )
+
+    def _operable(self, other):
+        """Ensure ciphertext is operable"""
+        if not isinstance(other, Ciphertext):
+
+            # Try scalar to array
+            try:
+                is_np = isinstance(other, np.ndarray)
+                other = np.full(self._descriptor.shape, other, self._descriptor.type)
+                if not is_np:
+                    other = other.tolist()
+            except Exception:
+                pass
+
+            # Try data to plaintext
+            try:
+                other = self._context._plaintext(other)
+            except Exception:
+                raise NotImplementedError(f"Incompatible type ({other.__class__})")  # from None
+
+        if other._context.options != self._context.options:
+            raise ValueError(f"Different underlying contexts ({other._context} != {self._context})")
+
+        if other._descriptor != self._descriptor:
+            raise ValueError(f"Different underlying ciphertext ({other} != {self})")
+
+        return other
+
+    def _link(self, context):
+        """Link ciphertext to context"""
+
+    def _neg(self, a):
+        """Negate two ciphertext"""
+        return operator.neg(a)
+
+    def _add(self, a, b):
+        """Add two ciphertexts"""
+        return operator.add(a, b)
+
+    def _sub(self, a, b):
+        """Substract two ciphertexts"""
+        return operator.sub(a, b)
+
+    def _mul(self, a, b):
+        """Multiply two ciphertexts"""
+        return operator.mul(a, b)
+
+    def _op_unary(self, op):
+        """Execute a unary operation"""
+        self._link(self._context)
+        chunks = tuple(map(op, self._chunks))
+        return self._new(_chunks=chunks)
+
+    def _op_binary(self, op, other):
+        """Execute a binary operation"""
+        other = self._operable(other)
+        self._link(self._context)
+        other._link(self._context)
+        chunks = tuple(itertools.starmap(op, zip(self._chunks, other._chunks)))
+        return self._new(_chunks=chunks)
+
+    def __neg__(self):
+        """Negate a ciphertext"""
+        return self._op_unary(self._neg)
+
+    def __add__(self, other):
+        """Add two ciphertexts"""
+        return self._op_binary(self._add, other)
+
+    def __sub__(self, other):
+        """Substract two ciphertexts"""
+        return self._op_binary(self._sub, other)
+
+    def __mul__(self, other):
+        """Multiply two ciphertexts"""
+        return self._op_binary(self._mul, other)
+
+    def __repr__(self) -> str:
+        """Ciphertext representation"""
+        return f"<{self.__class__.__name__} format={self._descriptor.type.__name__} shape={self._descriptor.shape}>"

polyhe/core/context.py

@@ -0,0 +1,125 @@
+"""Context package"""
+
+import copy
+import math
+import typing
+import itertools
+from abc import ABC, abstractmethod
+from collections.abc import Iterable
+from functools import cached_property
+
+import numpy as np
+
+from polyhe.core.ciphertext import Ciphertext
+from polyhe.core import COEFF_MODULUS, Options, materialize
+
+
+__all__ = (
+    "Context",
+)
+
+
+class Context(ABC):
+    """Abstract encryption context"""
+    _cls: type[Ciphertext]
+
+    def __init__(self, options: Options = Options()) -> None:
+        """Initialize context"""
+        self.options = options
+
+    @property
+    def _slots_exp(self) -> int:
+        """Slot exponent"""
+        return self.options.slots
+
+    @property
+    def _slots(self) -> int:
+        """Slot count"""
+        return 2 ** self._slots_exp
+
+    @property
+    def _scale_exp(self) -> int:
+        """Scale exponent"""
+        return self.options.scale
+
+    @property
+    def _security_level(self) -> int:
+        """Security level"""
+        return self.options.security
+
+    @property
+    def _poly_exp(self) -> int:
+        """Polynomial exponent"""
+        return self._slots_exp + 1
+
+    @property
+    def _poly(self) -> int:
+        """Polynomial size"""
+        return 2 ** self._poly_exp
+
+    @property
+    def _coeff_modulus(self) -> list[int]:
+        """Coefficient modulus"""
+        return COEFF_MODULUS[self._security_level][self._poly_exp]
+
+    def __getstate__(self) -> dict:
+        """Get serializable state"""
+        state: dict = dict(super().__getstate__())  # type: ignore
+        state.pop("_public", None)
+        return state
+
+    @cached_property
+    def public(self) -> typing.Self:
+        """Get public context"""
+        return copy.deepcopy(self)
+
+    def _new(self, /, *args, **kwds) -> Ciphertext:
+        """Create new operable ciphertext"""
+        return self._cls(_context=self.public, *args, **kwds)
+
+    def _partition(self, array: np.ndarray) -> Iterable[np.ndarray]:
+        """Transform array into flat batches"""
+        parts = (array.size + self._slots - 1) // self._slots
+        space = np.zeros(parts * self._slots, dtype=np.float64)
+        space[:array.size] = array.ravel()
+        return np.split(space, parts)
+
+    def _encode(self, chunk):
+        """Encode data to plaintext"""
+        return chunk
+
+    @abstractmethod
+    def _encrypt(self, chunk):
+        """Encrypt plaintext to ciphertext"""
+
+    @abstractmethod
+    def _decrypt(self, chunk):
+        """Decrypt cypertext to plaintext"""
+
+    def _decode(self, chunk):
+        """Decode plaintext to data"""
+        return chunk
+
+    def _plaintext[T](self, data: T) -> Ciphertext[T]:
+        """Encode data to plaintext"""
+        descriptor, obj = materialize(data)
+        chunks = tuple(map(self._encode, self._partition(obj)))
+        return Ciphertext(_context=self.public, _descriptor=descriptor, _chunks=chunks)
+
+    def encrypt[T](self, data: T) -> Ciphertext[T]:
+        """Encrypt data to ciphertext"""
+        descriptor, obj = materialize(data)
+        chunks = tuple(map(self._encrypt, map(self._encode, self._partition(obj))))
+        return self._new(_descriptor=descriptor, _chunks=chunks)
+
+    def decrypt[T](self, data: Ciphertext[T]) -> T:
+        """Decrypt cypertext to data"""
+        data._link(self)
+        descriptor = data._descriptor
+        result = itertools.chain.from_iterable(map(self._decode, map(self._decrypt, data._chunks)))
+        result = np.fromiter(iter=result, dtype=descriptor.type, count=math.prod(descriptor.shape)).reshape(descriptor.shape)
+        return result if descriptor.numpy else result.tolist()  # type: ignore
+
+    def __repr__(self) -> str:
+        """Context representation"""
+        return f"{self.__class__.__name__}(options={self.options!r})"

polyhe/openfhe.py

@@ -0,0 +1,160 @@
+"""OpenFHE encryption"""
+
+import sys
+import typing
+import copyreg
+from dataclasses import dataclass
+from functools import cached_property
+
+from polyhe import core
+from polyhe.core import context
+from polyhe.core import ciphertext
+
+# Make sure global package is not confused with current package
+_pkg = sys.path.pop(0)
+try:
+    import openfhe
+finally:
+    sys.path.insert(0, _pkg)
+
+
+__all__ = (
+    "Context",
+)
+
+
+SECURITY_LEVEL = {
+    # 0: openfhe.SecurityLevel.HEStd_NotSet,
+    128: openfhe.SecurityLevel.HEStd_128_classic,
+    192: openfhe.SecurityLevel.HEStd_192_classic,
+    256: openfhe.SecurityLevel.HEStd_256_classic
+}
+
+
+@dataclass(repr=False, eq=False, order=False, slots=True, frozen=True)
+class Ciphertext(ciphertext.Ciphertext):
+    """OpenFHE ciphertext"""
+    _context: "Context"
+
+    def _neg(self, a: openfhe.Ciphertext) -> openfhe.Ciphertext:
+        """Negate a ciphertext"""
+        return self._context._context.EvalNegate(a)
+
+    def _add(self, a: openfhe.Ciphertext, b: openfhe.Ciphertext) -> openfhe.Ciphertext:
+        """Add two ciphertexts"""
+        return self._context._context.EvalAdd(a, b)
+
+    def _sub(self, a: openfhe.Ciphertext, b: openfhe.Ciphertext) -> openfhe.Ciphertext:
+        """Substract two ciphertexts"""
+        return self._context._context.EvalSub(a, b)
+
+    def _mul(self, a: openfhe.Ciphertext, b: openfhe.Ciphertext) -> openfhe.Ciphertext:
+        """Multiply two ciphertexts"""
+        return self._context._context.EvalMult(a, b)
+
+
+class Context(context.Context):
+    """OpenFHE context"""
+    _cls = Ciphertext
+
+    def __init__(self, options: core.Options = core.Options()) -> None:
+        """Initialize context"""
+        super().__init__(options)
+
+        # Context
+        level = SECURITY_LEVEL[self._security_level]
+        parameters = openfhe.CCParamsCKKSRNS()
+        parameters.SetRingDim(self._poly)
+        parameters.SetSecurityLevel(level)
+        parameters.SetScalingModSize(self._scale_exp)
+        self._context = openfhe.GenCryptoContext(parameters)
+        self._context.Enable(openfhe.PKESchemeFeature.PKE)
+        self._context.Enable(openfhe.PKESchemeFeature.LEVELEDSHE)
+
+        # Keys
+        keys = self._context.KeyGen()
+        self._public_key = keys.publicKey
+        self._private_key = keys.secretKey
+        self._context.EvalMultKeyGen(self._private_key)
+
+    @cached_property
+    def public(self) -> typing.Self:
+        """Get public context"""
+        context = super().public
+        try:
+            del context._private_key
+        except AttributeError:
+            pass
+        return context
+
+    def _encode(self, chunk):
+        """Ecnode data to plaintext"""
+        return self._context.MakeCKKSPackedPlaintext(chunk)
+
+    def _encrypt(self, chunk: openfhe.Plaintext) -> openfhe.Ciphertext:
+        """Encrypt plaintext to ciphertext"""
+        return self._context.Encrypt(self._public_key, chunk)
+
+    def _decrypt(self, chunk: openfhe.Ciphertext) -> openfhe.Plaintext:
+        """Decrypt cypertext to plaintext"""
+        return self._context.Decrypt(chunk, self._private_key)
+
+    def _decode(self, chunk: openfhe.Plaintext):
+        """Decode plaintext to data"""
+        return chunk.GetRealPackedValue()
+
+
+# Serialization
+def deserialize_context(str: bytes) -> openfhe.CryptoContext:
+    """OpenFHE context deserializer"""
+    return openfhe.DeserializeCryptoContextString(str, openfhe.BINARY)
+
+
+def deserialize_ciphertext(str: bytes) -> openfhe.Ciphertext:
+    """OpenFHE cipher text deserializer"""
+    return openfhe.DeserializeCiphertextString(str, openfhe.BINARY)
+
+
+def deserialize_publickey(str: bytes) -> openfhe.PublicKey:
+    """OpenFHE public key deserializer"""
+    return openfhe.DeserializePublicKeyString(str, openfhe.BINARY)
+
+
+def deserialize_privatekey(str: bytes) -> openfhe.PrivateKey:
+    """OpenFHE private key deserializer"""
+    return openfhe.DeserializePrivateKeyString(str, openfhe.BINARY)
+
+
+# Pickle support
+def context_reducer(context: openfhe.CryptoContext):
+    """OpenFHE context pickle reducer"""
+    cls = deserialize_context
+    args = (openfhe.Serialize(context, openfhe.BINARY),)
+    return (cls, args)
+
+
+def ciphertext_reducer(ciphertext: openfhe.Ciphertext):
+    """OpenFHE cipher text pickle reducer"""
+    cls = deserialize_ciphertext
+    args = (openfhe.Serialize(ciphertext, openfhe.BINARY),)
+    return (cls, args)
+
+
+def public_key_reducer(public_key: openfhe.PublicKey):
+    """OpenFHE public key pickle reducer"""
+    cls = deserialize_publickey
+    args = (openfhe.Serialize(public_key, openfhe.BINARY),)
+    return (cls, args)
+
+
+def private_key_reducer(private_key: openfhe.PrivateKey):
+    """OpenFHE private key pickle reducer"""
+    cls = deserialize_privatekey
+    args = (openfhe.Serialize(private_key, openfhe.BINARY),)
+    return (cls, args)
+
+
+copyreg.pickle(openfhe.CryptoContext, context_reducer)
+copyreg.pickle(openfhe.Ciphertext, ciphertext_reducer)
+copyreg.pickle(openfhe.PublicKey, public_key_reducer)
+copyreg.pickle(openfhe.PrivateKey, private_key_reducer)

polyhe/tenseal.py

@@ -0,0 +1,105 @@
+"""TenSEAL encryption"""
+
+import sys
+import typing
+import copyreg
+from dataclasses import dataclass
+from functools import cached_property
+
+from polyhe import core
+from polyhe.core import context
+from polyhe.core import ciphertext
+
+# Make sure global package is not confused with current package
+_pkg = sys.path.pop(0)
+try:
+    import tenseal
+    from tenseal import sealapi
+    from tenseal.tensors import CKKSVector
+    from tenseal.enc_context import Context as SealContext
+finally:
+    sys.path.insert(0, _pkg)
+
+
+__all__ = (
+    "Context",
+)
+
+
+SECURITY_LEVEL = {
+    # 0: sealapi.SEC_LEVEL_TYPE.NONE,
+    128: sealapi.SEC_LEVEL_TYPE.TC128,
+    192: sealapi.SEC_LEVEL_TYPE.TC192,
+    256: sealapi.SEC_LEVEL_TYPE.TC256
+}
+
+
+@dataclass(repr=False, eq=False, order=False, slots=True, frozen=True)
+class Ciphertext(ciphertext.Ciphertext):
+    """TenSEAL ciphertext"""
+    _context: "Context"
+
+    def _link(self, context: "Context") -> None:
+        """Link ciphertext to context"""
+        for chunk in self._chunks:
+            chunk.link_context(context._context)
+
+
+class Context(context.Context):
+    """TenSEAL context"""
+    _cls = Ciphertext
+
+    def __init__(self, options: core.Options = core.Options()) -> None:
+        """Initialize context"""
+        super().__init__(options)
+
+        # Context
+        self._context = tenseal.context(
+            scheme=tenseal.SCHEME_TYPE.CKKS,
+            poly_modulus_degree=self._poly,
+            coeff_mod_bit_sizes=self._coeff_modulus
+        )
+
+        # Keys
+        self._context.global_scale = 2 ** self._scale_exp
+        self._context.generate_galois_keys()
+        self._context.generate_relin_keys()
+
+    @cached_property
+    def public(self) -> typing.Self:
+        """Get public context"""
+        context = super().public
+        context._context.make_context_public()
+        return context
+
+    def _encode(self, chunk) -> tenseal.PlainTensor:
+        """Encrypt plaintext to ciphertext"""
+        return tenseal.plain_tensor(chunk)
+
+    def _encrypt(self, chunk: tenseal.PlainTensor) -> CKKSVector:
+        """Encrypt plaintext to ciphertext"""
+        return tenseal.ckks_vector(self.public._context, chunk)
+
+    def _decrypt(self, chunk: CKKSVector):
+        """Decrypt cypertext to list"""
+        return chunk.decrypt()
+
+
+# Pickle support
+def context_reducer(context: SealContext):
+    """TenSEAL context pickle reducer"""
+    # FIXME: Context deserialization is very slow
+    cls = context.load
+    args = (context.serialize(save_secret_key=True),)
+    return (cls, args)
+
+
+def ckks_vector_reducer(vector: CKKSVector):
+    """TenSEAL CKKS vector pickle reducer"""
+    cls = vector.lazy_load
+    args = (vector.serialize(),)
+    return (cls, args)
+
+
+copyreg.pickle(SealContext, context_reducer)
+copyreg.pickle(CKKSVector, ckks_vector_reducer)  # type: ignore (wrong inferred typing)

polyhe/test.py

@@ -0,0 +1,50 @@
+"""PolyHE test"""
+
+from argparse import ArgumentParser, Namespace
+
+import numpy as np
+
+import polyhe
+
+
+__all__ = ()
+
+
+def main(config: Namespace) -> None:
+    ctx = polyhe.new(config.backend)
+
+    # Numbers
+    p1 = 1
+    p2 = -1
+    c1 = ctx.encrypt(p1)
+    c2 = ctx.encrypt(p2)
+    c3 = c1 * c2 + p1
+    p3 = ctx.decrypt(c3)
+    ref = 0
+    assert np.isclose(p3, ref).all(), f"numbers error; got {p3}, expect {ref}"
+
+    # Lists
+    p1 = [0.3, 0.4]
+    p2 = [0.4, 0.3]
+    c1 = ctx.encrypt(p1)
+    c2 = ctx.encrypt(p2)
+    c3 = -c1 + c2
+    p3 = ctx.decrypt(c3)
+    ref = [0.1, -0.1]
+    assert np.isclose(p3, ref).all(), f"lists error; got {p3}, expect {ref}"
+
+    # NumPy
+    p1 = np.full((2, 3, 4), 0.4, np.float64)
+    p2 = np.full((2, 3, 4), 0.3, np.float64)
+    c1 = ctx.encrypt(p1)
+    c2 = ctx.encrypt(p2)
+    c3 = c1 - c2
+    p3 = ctx.decrypt(c3)
+    ref = np.full((2, 3, 4), 0.1, np.float64)
+    assert np.isclose(p3, ref).all(), f"numpy error; got {p3}, expect {ref}"
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser(prog="polyhe-test", description="PolyHE test")
+    parser.add_argument("backend", choices=list(polyhe.Backend), help="Operational backend")
+    main(parser.parse_args())