import os
from dilithium_py.ml_dsa import ML_DSA_65

class MLDSATarget:

    FIXED_MSG = b"fixed test message for side channel analysis" * 4

    def __init__(self):
        print("  [*] Generating ML-DSA-65 key pair...")
        self.pk, self.sk = ML_DSA_65.keygen()

        # Pre-generate a FIXED random seed to pass into sign
        # so the rejection sampling uses a deterministic path
        self.fixed_sig = ML_DSA_65.sign(self.sk, self.FIXED_MSG)

        # Pre-generate a set of random signatures for verify tests
        # so verify_random doesn't have to call sign at all
        print("  [*] Pre-generating verify test vectors...")
        self.random_msgs = [os.urandom(64) for _ in range(3000)]
        self.random_sigs = [ML_DSA_65.sign(self.sk, m) for m in self.random_msgs]
        self._verify_idx = 0
        print("  [*] ML-DSA harness ready.")

    def sign_fixed(self):
        """Sign fixed message — tests if timing varies with message content."""
        ML_DSA_65.sign(self.sk, self.FIXED_MSG)

    def sign_random(self):
        """Sign random message — same key, different message each time."""
        ML_DSA_65.sign(self.sk, os.urandom(64))

    def verify_fixed(self):
        """Verify the fixed pre-computed signature."""
        ML_DSA_65.verify(self.pk, self.FIXED_MSG, self.fixed_sig)

    def verify_random(self):
        """
        Verify a pre-computed random signature.
        Pre-generated at init so sign() overhead doesn't contaminate timing.
        """
        i = self._verify_idx % len(self.random_msgs)
        ML_DSA_65.verify(self.pk, self.random_msgs[i], self.random_sigs[i])
        self._verify_idx += 1

    def keygen_repeated(self):
        ML_DSA_65.keygen()