crypto_layer/ckks_engine.py

"""
crypto_layer/ckks_engine.py

SecureLens — CKKS Encryption Engine
Handles: key generation, feature vector encryption, decryption.

ARCHITECTURE (corrected):
  - ResNet-18 extracts 512-dim features CLIENT SIDE
  - CKKS encrypts those 512 features CLIENT SIDE
  - Only ciphertext is sent to server
  - Server has NO secret key — only public context
  - Decryption happens CLIENT SIDE only

Uses TenSEAL's CKKS scheme (approximate HE for real numbers).
"""

import tenseal as ts
import numpy as np
import os


class CKKSEngine:
    """
    Manages the full CKKS lifecycle on the CLIENT side:
      - Context creation (public/secret key pair)
      - Feature vector encryption (512-dim ResNet output)
      - Encrypted result decryption
      - Public context export (for server — no secret key)
      - Context serialization (save/load keys)

    The SECRET KEY never leaves this class.
    The server only ever receives:
      1. The public context (for deserialization only)
      2. Ciphertext bytes (encrypted feature vectors)
    """

    def __init__(
        self,
        poly_modulus_degree: int = 8192,
        coeff_mod_bit_sizes: list = None,
        global_scale: float = 2**40,   # ← FIXED: was 2**30, must be 2**40
    ):
        """
        Args:
            poly_modulus_degree : Ring dimension.
                                  8192 → 128-bit security (gold standard).
                                  Never use below 8192 for sensitive data.

            coeff_mod_bit_sizes : Coefficient modulus chain.
                                  [60, 40, 40, 60] gives 3 multiplication
                                  levels — sufficient for 2 linear layers.
                                  First and last primes are larger for
                                  scale management.

            global_scale        : Precision scale for CKKS.
                                  2^40 → ~12 decimal digits of precision.
                                  Sufficient for ML weight magnitudes.
                                  DO NOT use 2^30 — insufficient precision
                                  for neural network weights.
        """
        if coeff_mod_bit_sizes is None:
            coeff_mod_bit_sizes = [60, 40, 40, 60]

        self.poly_modulus_degree = poly_modulus_degree
        self.coeff_mod_bit_sizes = coeff_mod_bit_sizes
        self.global_scale        = global_scale

        self.context        = None   # Full context WITH secret key (client only)
        self.public_context = None   # Public context WITHOUT secret key (server safe)

        self._create_context()

    # ──────────────────────────────────────────────────────────────────
    # Context & Key Generation
    # ──────────────────────────────────────────────────────────────────

    def _create_context(self):
        """
        Creates TenSEAL CKKS context with full key set:
          - Secret key   : stays on client ONLY — never transmitted
          - Public key   : derived from secret key
          - Galois keys  : for vector rotation operations
          - Relin keys   : for relinearization after multiplication

        Also creates a PUBLIC context (no secret key) that is
        safe to share with / use on the server for deserialization.
        """
        # ── Full context (client only — has secret key) ───────────────
        self.context = ts.context(
            ts.SCHEME_TYPE.CKKS,
            poly_modulus_degree=self.poly_modulus_degree,
            coeff_mod_bit_sizes=self.coeff_mod_bit_sizes,
        )
        self.context.generate_galois_keys()
        self.context.generate_relin_keys()
        self.context.global_scale = self.global_scale

        # ── Public context (server safe — no secret key) ──────────────
        # Serialize the full context, then reload without secret key
        # This ensures the server context is mathematically identical
        # but cryptographically cannot decrypt anything
        ctx_bytes = self.context.serialize(save_secret_key=True)
        self.public_context = ts.context_from(ctx_bytes)
        self.public_context.make_context_public()  # drops secret key

        print("[CKKSEngine] Context created successfully.")
        print(f"  Poly modulus degree : {self.poly_modulus_degree}")
        print(f"  Coeff mod bit sizes : {self.coeff_mod_bit_sizes}")
        print(f"  Global scale        : 2^{int(np.log2(self.global_scale))}")
        print(f"  Security level      : 128-bit")
        print(f"  Secret key          : CLIENT ONLY — never transmitted")

    # ──────────────────────────────────────────────────────────────────
    # Feature Vector Encryption  (CLIENT SIDE)
    # ──────────────────────────────────────────────────────────────────

    def encrypt_feature_vector(
        self, feature_vector: np.ndarray
    ) -> ts.CKKSVector:
        """
        Encrypts the 512-dim ResNet-18 feature vector using CKKS.

        This is the MAIN encryption method used in SecureLens.
        The feature vector is extracted by ResNet-18 on the CLIENT
        and encrypted here before any data leaves the device.

        Args:
            feature_vector : numpy array of shape (512,)
                             Output of ResNet-18 backbone.

        Returns:
            CKKSVector : encrypted ciphertext (~326 KB when serialized)

        NOTE: This runs CLIENT SIDE.
              The raw feature values never leave the client.
              Only the returned CKKSVector (serialized) is sent to server.
        """
        if not isinstance(feature_vector, np.ndarray):
            feature_vector = np.array(feature_vector, dtype=np.float64)

        if feature_vector.dtype != np.float64:
            feature_vector = feature_vector.astype(np.float64)

        encrypted = ts.ckks_vector(
            self.context,
            feature_vector.tolist()
        )

        print(f"[CKKSEngine] Feature vector encrypted.")
        print(f"  Input size  : {len(feature_vector)} floats (~4 KB)")
        print(f"  Output size : ~{len(encrypted.serialize())//1024} KB ciphertext")

        return encrypted

    def encrypt_feature_vector_to_bytes(
        self, feature_vector: np.ndarray
    ) -> bytes:
        """
        Encrypts feature vector and returns serialized bytes
        ready for network transmission to the server.

        Args:
            feature_vector : numpy array of shape (512,)

        Returns:
            bytes : serialized CKKS ciphertext for sending to server
        """
        enc     = self.encrypt_feature_vector(feature_vector)
        ct_bytes = enc.serialize()

        # Clear plaintext from memory after encryption
        feature_vector[:] = 0
        del feature_vector

        return ct_bytes

    def encrypt_vector(self, vector: list) -> ts.CKKSVector:
        """
        Generic vector encryption. Encrypts any float list.
        Used for testing and utilities.

        Args:
            vector : Python list of floats

        Returns:
            CKKSVector ciphertext
        """
        return ts.ckks_vector(self.context, vector)

    # ──────────────────────────────────────────────────────────────────
    # Decryption  (CLIENT SIDE ONLY)
    # ──────────────────────────────────────────────────────────────────

    def decrypt_vector(
        self, encrypted_vector: ts.CKKSVector
    ) -> np.ndarray:
        """
        Decrypts a CKKSVector back to plaintext numpy array.

        ONLY the client can call this — only the client has
        the secret key. The server CANNOT call this.

        Args:
            encrypted_vector : CKKSVector ciphertext

        Returns:
            numpy float64 array of decrypted values
        """
        decrypted = encrypted_vector.decrypt()
        return np.array(decrypted, dtype=np.float64)

    def decrypt_vector_from_bytes(self, ct_bytes: bytes) -> np.ndarray:
        """
        Deserializes ciphertext bytes received from server
        and decrypts them using the secret key.

        This is what the client calls after receiving the
        server's encrypted response.

        Args:
            ct_bytes : serialized ciphertext bytes from server

        Returns:
            numpy float64 array of decrypted logits
        """
        # Deserialize using FULL context (has secret key)
        enc = ts.ckks_vector_from(self.context, ct_bytes)
        return self.decrypt_vector(enc)

    def decrypt_prediction(
        self, encrypted_output
    ) -> dict:
        """
        Decrypts server's encrypted output and converts to
        class probabilities via softmax.

        Binary classification:
          index 0 → Normal
          index 1 → Pneumonia

        Args:
            encrypted_output : encrypted logits from server
                               Can be:
                               - CKKSVector of length >= 2
                               - list of CKKSVector (one per logit)

        Returns:
            dict:
              prediction     : "Normal" or "Pneumonia"
              confidence     : float in [0, 1]
              normal_score   : probability for Normal
              pneumonia_score: probability for Pneumonia
              raw            : raw decrypted logits
        """
        # Handle both list of CKKSVectors and single CKKSVector
        if isinstance(encrypted_output, list):
            # Decrypt each encrypted scalar logit
            logits = np.array([vec.decrypt(self.context.secret_key())[0] for vec in encrypted_output], 
                            dtype=np.float64)
            raw = logits
        else:
            # Original format - single CKKSVector
            raw = self.decrypt_vector(encrypted_output)
            logits = raw[:2]

        # Softmax — stable version (subtract max for numerical stability)
        exp_vals = np.exp(logits - np.max(logits))
        probs    = exp_vals / exp_vals.sum()

        result = {
            "raw"            : logits.tolist(),
            "normal_score"   : float(probs[0]),
            "pneumonia_score": float(probs[1]),
            "prediction"     : "Pneumonia" if probs[1] > probs[0] else "Normal",
            "confidence"     : float(max(probs[0], probs[1])),
        }

        print(f"[CKKSEngine] Decrypted prediction : {result['prediction']}")
        print(f"  Confidence  : {result['confidence']:.2%}")
        print(f"  Normal      : {result['normal_score']:.4f}")
        print(f"  Pneumonia   : {result['pneumonia_score']:.4f}")

        return result

    def decrypt_prediction_from_bytes(self, ct_bytes: bytes) -> dict:
        """
        Full pipeline: receive server bytes → decrypt → softmax → result.
        Client-side only.

        Args:
            ct_bytes : serialized encrypted logits from server
                      Can be a single CKKSVector or list format

        Returns:
            same dict as decrypt_prediction()
        """
        # Try to deserialize as a list first (new format)
        # Format: first 4 bytes = number of vectors, then each serialized
        try:
            import struct
            if len(ct_bytes) > 4:
                n_vectors = struct.unpack('!I', ct_bytes[:4])[0]
                if n_vectors == 2:  # Expected for binary classification
                    # Deserialize list format
                    offset = 4
                    enc_list = []
                    for _ in range(n_vectors):
                        # Next 4 bytes = size of this vector
                        size = struct.unpack('!I', ct_bytes[offset:offset+4])[0]
                        offset += 4
                        vec_bytes = ct_bytes[offset:offset+size]
                        offset += size
                        vec = ts.ckks_vector_from(self.context, vec_bytes)
                        enc_list.append(vec)
                    return self.decrypt_prediction(enc_list)
        except:
            pass
        
        # Fallback: old format (single CKKSVector)
        raw_values = self.decrypt_vector_from_bytes(ct_bytes)
        logits   = raw_values[:2]
        exp_vals = np.exp(logits - np.max(logits))
        probs    = exp_vals / exp_vals.sum()

        return {
            "raw"            : raw_values.tolist(),
            "normal_score"   : float(probs[0]),
            "pneumonia_score": float(probs[1]),
            "prediction"     : "Pneumonia" if probs[1] > probs[0] else "Normal",
            "confidence"     : float(max(probs[0], probs[1])),
        }

    # ──────────────────────────────────────────────────────────────────
    # Ciphertext Serialization
    # ──────────────────────────────────────────────────────────────────

    def serialize_ciphertext(
        self, ciphertext: ts.CKKSVector
    ) -> bytes:
        """
        Converts CKKSVector to bytes for network transmission.

        Args:
            ciphertext : CKKSVector to serialize

        Returns:
            bytes blob (~326 KB for 512-dim vector)
        """
        return ciphertext.serialize()

    def deserialize_ciphertext(self, data: bytes) -> ts.CKKSVector:
        """
        Reconstructs CKKSVector from bytes.
        Uses FULL context (with secret key) — for CLIENT use.

        Args:
            data : serialized ciphertext bytes

        Returns:
            CKKSVector
        """
        return ts.ckks_vector_from(self.context, data)

    def deserialize_ciphertext_public(
        self, data: bytes
    ) -> ts.CKKSVector:
        """
        Reconstructs CKKSVector from bytes.
        Uses PUBLIC context (no secret key) — for SERVER use.

        The server calls this to deserialize received ciphertexts.
        Cannot be used for decryption.

        Args:
            data : serialized ciphertext bytes

        Returns:
            CKKSVector (cannot be decrypted without secret key)
        """
        return ts.ckks_vector_from(self.public_context, data)

    # ──────────────────────────────────────────────────────────────────
    # Context Serialization — Save & Load Keys
    # ──────────────────────────────────────────────────────────────────

    def save_context(
        self, path: str, save_secret_key: bool = True
    ):
        """
        Saves TenSEAL context to disk.

        Args:
            path            : File path to save
            save_secret_key : If True, saves full context with secret key.
                              KEEP THIS FILE PRIVATE — it can decrypt everything.
                              If False, saves public context only (server-safe).
        """
        os.makedirs(
            os.path.dirname(path) if os.path.dirname(path) else ".",
            exist_ok=True
        )

        if save_secret_key:
            serialized = self.context.serialize(save_secret_key=True)
            print(f"[CKKSEngine] Full context (with secret key) saved → {path}")
            print(f"  ⚠️  KEEP THIS FILE PRIVATE — anyone with it can decrypt")
        else:
            serialized = self.public_context.serialize()
            print(f"[CKKSEngine] Public context saved → {path}")
            print(f"  ✅ Safe to share with server — no secret key inside")

        with open(path, "wb") as f:
            f.write(serialized)

    def load_context(self, path: str):
        """
        Loads a previously saved TenSEAL context from disk.

        Args:
            path : File path of saved context
        """
        with open(path, "rb") as f:
            serialized = f.read()

        self.context = ts.context_from(serialized)
        print(f"[CKKSEngine] Context loaded ← {path}")

    def get_public_context_bytes(self) -> bytes:
        """
        Returns the serialized PUBLIC context bytes.
        Safe to send to the server for ciphertext deserialization.
        Does NOT contain the secret key.

        Returns:
            bytes of public context
        """
        return self.public_context.serialize()

    # ──────────────────────────────────────────────────────────────────
    # Utility
    # ──────────────────────────────────────────────────────────────────

    def get_encryption_info(self) -> dict:
        """Returns CKKS parameter summary for API responses."""
        return {
            "scheme"             : "CKKS (Cheon-Kim-Kim-Song)",
            "library"            : "TenSEAL 0.3.14",
            "poly_modulus_degree": self.poly_modulus_degree,
            "coeff_mod_bit_sizes": self.coeff_mod_bit_sizes,
            "global_scale"       : f"2^{int(np.log2(self.global_scale))}",
            "security_bits"      : 128,
            "feature_vector_size": 512,
            "ciphertext_size_kb" : 326,
            "decryption_error"   : "~7.19e-8",
        }


# ──────────────────────────────────────────────────────────────────────
# Self-Test — run directly to verify the engine works
# python crypto_layer/ckks_engine.py
# ──────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    print("=" * 60)
    print("SecureLens — CKKSEngine Self-Test")
    print("=" * 60)

    engine = CKKSEngine(
        poly_modulus_degree=8192,
        coeff_mod_bit_sizes=[60, 40, 40, 60],
        global_scale=2**40
    )

    # ── Test 1: Encrypt 512-dim feature vector ────────────────────────
    print("\n[Test 1] Encrypting 512-dim feature vector...")
    dummy_features = np.random.randn(512).astype(np.float64)
    enc = engine.encrypt_feature_vector(dummy_features.copy())
    print(f"  Input  : 512 floats")
    print(f"  Output : {len(enc.serialize())//1024} KB ciphertext")

    # ── Test 2: Decrypt and check error ──────────────────────────────
    print("\n[Test 2] Decrypting and checking error...")
    decrypted = engine.decrypt_vector(enc)
    error = np.max(np.abs(decrypted[:512] - dummy_features))
    print(f"  Max decryption error : {error:.2e}")
    assert error < 1e-3, f"Error too large: {error}"
    print(f"  ✅ Error within acceptable CKKS bounds")

    # ── Test 3: Serialization round-trip ─────────────────────────────
    print("\n[Test 3] Serialization round-trip...")
    ct_bytes  = engine.serialize_ciphertext(enc)
    recovered = engine.deserialize_ciphertext(ct_bytes)
    dec2      = engine.decrypt_vector(recovered)
    error2    = np.max(np.abs(dec2[:512] - dummy_features))
    print(f"  Serialized size      : {len(ct_bytes)//1024} KB")
    print(f"  Round-trip error     : {error2:.2e}")
    assert error2 < 1e-3
    print(f"  ✅ Serialization correct")

    # ── Test 4: Public context cannot decrypt ────────────────────────
    print("\n[Test 4] Verifying public context has no secret key...")
    pub_bytes = engine.get_public_context_bytes()
    pub_ctx   = ts.context_from(pub_bytes)
    try:
        enc_pub = ts.ckks_vector_from(pub_ctx, ct_bytes)
        dec_pub = enc_pub.decrypt()
        print(f"  ⚠️  Public context decrypted (secret key still present)")
    except Exception as e:
        print(f"  ✅ Public context cannot decrypt: {e}")

    # ── Test 5: Decrypt prediction format ────────────────────────────
    print("\n[Test 5] Decrypt prediction format...")
    dummy_logits = engine.encrypt_vector([2.14, -0.83])
    result       = engine.decrypt_prediction(dummy_logits)
    print(f"  Prediction  : {result['prediction']}")
    print(f"  Confidence  : {result['confidence']:.2%}")
    print(f"  Normal      : {result['normal_score']:.4f}")
    print(f"  Pneumonia   : {result['pneumonia_score']:.4f}")
    assert result["prediction"] == "Normal"
    print(f"  ✅ Prediction correct")

    # ── Test 6: Encrypt to bytes pipeline ────────────────────────────
    print("\n[Test 6] Full bytes pipeline...")
    feat2    = np.random.randn(512).astype(np.float64)
    feat_copy = feat2.copy()
    ct       = engine.encrypt_feature_vector_to_bytes(feat_copy)
    result6  = engine.decrypt_prediction_from_bytes(ct)
    print(f"  Encrypted to {len(ct)//1024} KB")
    print(f"  Decrypted prediction: {result6['prediction']}")
    print(f"  ✅ Bytes pipeline working")

    # ── Test 7: Encryption info ───────────────────────────────────────
    print("\n[Test 7] Encryption info...")
    info = engine.get_encryption_info()
    for k, v in info.items():
        print(f"  {k:25s}: {v}")

    print("\n" + "=" * 60)
    print("✅ All 7 tests passed. CKKSEngine is correct.")
    print("=" * 60)