# targets/ml_kem_target.py
from kyber_py.ml_kem import ML_KEM_768

class MLKEMTarget:
    """
    Test harness for ML-KEM-768 using kyber_py (pure Python, FIPS 203).
    kyber_py is intentionally NOT constant-time — real leaks will be found.
    """

    def __init__(self):
        print("  [*] Generating ML-KEM-768 key pair...")
        self.ek, self.dk = ML_KEM_768.keygen()

        # One fixed ciphertext reused as the constant test vector
        _, self.fixed_ct = ML_KEM_768.encaps(self.ek)
        print("  [*] Harness ready.")

    def decapsulate_fixed(self):
        """Always decapsulates the same ciphertext — fixed test vector."""
        ML_KEM_768.decaps(self.dk, self.fixed_ct)

    def decapsulate_random(self):
        """Decapsulates a fresh random ciphertext each call."""
        _, random_ct = ML_KEM_768.encaps(self.ek)
        ML_KEM_768.decaps(self.dk, random_ct)

    def encapsulate_fixed(self):
        """Encapsulates to the same public key every time."""
        ML_KEM_768.encaps(self.ek)

    def encapsulate_random(self):
        """Encapsulates to a freshly generated key every time."""
        random_ek, _ = ML_KEM_768.keygen()
        ML_KEM_768.encaps(random_ek)

    def keygen_repeated(self):
        """Generates a fresh key pair each call."""
        ML_KEM_768.keygen()