framework: add configurable seeds/activation to XmssKeyManager (#129)

packages/testing/src/consensus_testing/keys.py

@@ -1,6 +1,8 @@
 """XMSS key management utilities for testing."""
 
-from typing import NamedTuple, Optional
+from __future__ import annotations
+
+from typing import Any, NamedTuple, Optional
 
 from lean_spec.subspecs.containers import Attestation, Signature
 from lean_spec.subspecs.containers.slot import Slot
@@ -10,7 +12,7 @@
     TEST_SIGNATURE_SCHEME,
     GeneralizedXmssScheme,
 )
-from lean_spec.types import ValidatorIndex
+from lean_spec.types import Uint64, ValidatorIndex
 
 
 class KeyPair(NamedTuple):
@@ -23,11 +25,11 @@ class KeyPair(NamedTuple):
     """The validator's secret key (used for signing)."""
 
 
-_KEY_CACHE: dict[tuple[int, int], KeyPair] = {}
+_KEY_CACHE: dict[tuple[int, int, int], KeyPair] = {}
 """
 Cache keys across tests to avoid regenerating them for the same validator/lifetime combo.
 
-Key: (validator_index, num_active_epochs) -> KeyPair
+Key: (validator_index, activation_epoch, num_active_epochs) -> KeyPair
 """
 
 
@@ -36,11 +38,13 @@ class XmssKeyManager:
 
     DEFAULT_MAX_SLOT = Slot(100)
     """Default maximum slot horizon if not specified."""
+    DEFAULT_ACTIVATION_EPOCH = Uint64(0)
 
     def __init__(
         self,
         max_slot: Optional[Slot] = None,
         scheme: GeneralizedXmssScheme = TEST_SIGNATURE_SCHEME,
+        default_activation_epoch: Optional[Uint64] = None,
     ) -> None:
         """
         Initialize the key manager.
@@ -53,6 +57,8 @@ def __init__(
         scheme : GeneralizedXmssScheme, optional
             The XMSS scheme to use.
             Defaults to `TEST_SIGNATURE_SCHEME`.
+        default_activation_epoch : Uint64, optional
+            Activation epoch used when none is provided for key generation.
 
         Notes:
         -----
@@ -61,7 +67,67 @@ def __init__(
         """
         self.max_slot = max_slot if max_slot is not None else self.DEFAULT_MAX_SLOT
         self.scheme = scheme
+        self.default_activation_epoch = (
+            default_activation_epoch
+            if default_activation_epoch is not None
+            else self.DEFAULT_ACTIVATION_EPOCH
+        )
         self._key_pairs: dict[ValidatorIndex, KeyPair] = {}
+        self._key_metadata: dict[ValidatorIndex, dict[str, Any]] = {}
+
+    @property
+    def default_num_active_epochs(self) -> int:
+        """Default lifetime derived from the configured max_slot."""
+        return self.max_slot.as_int() + 1
+
+    def create_and_store_key_pair(
+        self,
+        validator_index: ValidatorIndex,
+        *,
+        activation_epoch: Optional[Uint64] = None,
+        num_active_epochs: Optional[Uint64] = None,
+    ) -> KeyPair:
+        """
+        Generate and store a key pair with explicit control over key generation.
+
+        Parameters
+        ----------
+        validator_index : ValidatorIndex
+            The validator for whom a key pair should be generated.
+        activation_epoch : Uint64, optional
+            First epoch for which the key is valid. Defaults to `default_activation_epoch`.
+        num_active_epochs : Uint64, optional
+            Number of consecutive epochs the key should remain active.
+            Defaults to `max_slot + 1` (to include genesis).
+        """
+        activation_epoch_val = (
+            activation_epoch if activation_epoch is not None else self.default_activation_epoch
+        )
+        num_active_epochs_val = (
+            num_active_epochs
+            if num_active_epochs is not None
+            else Uint64(self.default_num_active_epochs)
+        )
+
+        cache_key = (
+            int(validator_index),
+            int(activation_epoch_val),
+            int(num_active_epochs_val),
+        )
+
+        if cache_key in _KEY_CACHE:
+            key_pair = _KEY_CACHE[cache_key]
+        else:
+            pk, sk = self.scheme.key_gen(activation_epoch_val, num_active_epochs_val)
+            key_pair = KeyPair(public=pk, secret=sk)
+            _KEY_CACHE[cache_key] = key_pair
+
+        self._key_pairs[validator_index] = key_pair
+        self._key_metadata[validator_index] = {
+            "activation_epoch": int(activation_epoch_val),
+            "num_active_epochs": int(num_active_epochs_val),
+        }
+        return key_pair
 
     def __getitem__(self, validator_index: ValidatorIndex) -> KeyPair:
         """
@@ -83,37 +149,10 @@ def __getitem__(self, validator_index: ValidatorIndex) -> KeyPair:
         - Keys are deterministic for testing (`seed=0`).
         - Lifetime = `max_slot + 1` to include the genesis slot.
         """
-        # Return cached keys if they exist.
         if validator_index in self._key_pairs:
             return self._key_pairs[validator_index]
 
-        # Generate New Key Pair
-        #
-        # XMSS requires knowing the total number of signatures in advance.
-        # We use max_slot + 1 as the lifetime since:
-        # - Validators may sign once per slot (attestations)
-        # - We include slot 0 (genesis) in the count
-        num_active_epochs = self.max_slot.as_int() + 1
-
-        # Check global cache first (keys are reused across tests)
-        cache_key = (int(validator_index), num_active_epochs)
-        if cache_key in _KEY_CACHE:
-            key_pair = _KEY_CACHE[cache_key]
-            self._key_pairs[validator_index] = key_pair
-            return key_pair
-
-        # Generate the key pair using the default XMSS scheme.
-        #
-        # The seed is set to 0 for deterministic test keys.
-        from lean_spec.types import Uint64
-
-        pk, sk = self.scheme.key_gen(Uint64(0), Uint64(num_active_epochs))
-
-        # Store as a cohesive unit and return.
-        key_pair = KeyPair(public=pk, secret=sk)
-        _KEY_CACHE[cache_key] = key_pair  # Cache globally for reuse across tests
-        self._key_pairs[validator_index] = key_pair
-        return key_pair
+        return self.create_and_store_key_pair(validator_index)
 
     def sign_attestation(self, attestation: Attestation) -> Signature:
         """
@@ -177,16 +216,8 @@ def sign_attestation(self, attestation: Attestation) -> Signature:
         # Generate the XMSS signature using the validator's (now prepared) secret key.
         xmss_sig = self.scheme.sign(sk, epoch, message)
 
-        # Convert the signature to the wire format (byte array).
-        signature_bytes = xmss_sig.to_bytes(self.scheme.config)
-
-        # Ensure the signature meets the consensus spec length (3100 bytes).
-        #
-        # This is necessary when using TEST_CONFIG (796 bytes) vs PROD_CONFIG.
-        # Padding with zeros on the right maintains compatibility.
-        padded_bytes = signature_bytes.ljust(Signature.LENGTH, b"\x00")
-
-        return Signature(padded_bytes)
+        # Convert to the consensus Signature container (handles padding internally).
+        return Signature.from_xmss(xmss_sig, self.scheme)
 
     def get_public_key(self, validator_index: ValidatorIndex) -> PublicKey:
         """
@@ -225,3 +256,29 @@ def __contains__(self, validator_index: ValidatorIndex) -> bool:
     def __len__(self) -> int:
         """Return the number of validators with generated keys."""
         return len(self._key_pairs)
+
+    def export_test_vectors(self, include_private_keys: bool = False) -> list[dict[str, Any]]:
+        """
+        Export generated keys in a JSON-serializable structure for downstream clients.
+
+        Parameters
+        ----------
+        include_private_keys : bool
+            When True, include the full secret key dump; otherwise only public data.
+        """
+        vectors: list[dict[str, Any]] = []
+        for validator_index, key_pair in self._key_pairs.items():
+            meta = self._key_metadata.get(validator_index, {})
+            entry: dict[str, Any] = {
+                "validator_index": int(validator_index),
+                "activation_epoch": meta.get("activation_epoch"),
+                "num_active_epochs": meta.get("num_active_epochs"),
+                "public_key": key_pair.public.to_bytes(self.scheme.config).hex(),
+            }
+            if include_private_keys:
+                # Pydantic models are JSON-serializable; keep the raw dump for full fidelity.
+                entry["secret_key"] = key_pair.secret.model_dump(mode="json")
+
+            vectors.append(entry)
+
+        return vectors

src/lean_spec/subspecs/containers/signature.py

@@ -28,3 +28,22 @@ def verify(
             return scheme.verify(public_key, epoch, message, signature)
         except Exception:
             return False
+
+    @classmethod
+    def from_xmss(
+        cls, xmss_signature: XmssSignature, scheme: GeneralizedXmssScheme = TEST_SIGNATURE_SCHEME
+    ) -> "Signature":
+        """
+        Create a consensus `Signature` container from an XMSS signature object.
+
+        Handles padding to the fixed 3100-byte length required by the consensus layer,
+        delegating all encoding details to the XMSS container itself.
+        """
+        raw = xmss_signature.to_bytes(scheme.config)
+        if len(raw) > cls.LENGTH:
+            raise ValueError(
+                f"XMSS signature length {len(raw)} exceeds container size {cls.LENGTH}"
+            )
+
+        # Pad on the right to the fixed-length container expected by consensus.
+        return cls(raw.ljust(cls.LENGTH, b"\x00"))

tests/lean_spec/subspecs/containers/test_signature.py

@@ -0,0 +1,123 @@
+"""Tests for consensus Signature container."""
+
+import pytest
+
+from lean_spec.subspecs.containers import Signature
+from lean_spec.subspecs.koalabear import Fp
+from lean_spec.subspecs.xmss.constants import TEST_CONFIG
+from lean_spec.subspecs.xmss.containers import HashTreeOpening
+from lean_spec.subspecs.xmss.containers import Signature as XmssSignature
+from lean_spec.subspecs.xmss.interface import TEST_SIGNATURE_SCHEME
+
+
+class TestSignatureFromXmss:
+    """Tests for Signature.from_xmss conversion method."""
+
+    def test_from_xmss_basic_conversion(self) -> None:
+        """Test that from_xmss correctly converts an XMSS signature to consensus format."""
+        # Create a valid XMSS signature
+        path = HashTreeOpening(
+            siblings=[[Fp(value=i) for i in range(TEST_CONFIG.HASH_LEN_FE)]]
+            * TEST_CONFIG.LOG_LIFETIME
+        )
+        rho = [Fp(value=i) for i in range(TEST_CONFIG.RAND_LEN_FE)]
+        hashes = [
+            [Fp(value=i + j) for i in range(TEST_CONFIG.HASH_LEN_FE)]
+            for j in range(TEST_CONFIG.DIMENSION)
+        ]
+        xmss_sig = XmssSignature(path=path, rho=rho, hashes=hashes)
+
+        # Convert to consensus signature
+        consensus_sig = Signature.from_xmss(xmss_sig, TEST_SIGNATURE_SCHEME)
+
+        # Verify it's the correct type and length
+        assert isinstance(consensus_sig, Signature)
+        assert len(consensus_sig) == Signature.LENGTH
+        assert len(consensus_sig) == 3100
+
+    def test_from_xmss_padding(self) -> None:
+        """Test that from_xmss correctly pads to 3100 bytes."""
+        # Create a minimal XMSS signature
+        path = HashTreeOpening(
+            siblings=[[Fp(value=0)] * TEST_CONFIG.HASH_LEN_FE] * TEST_CONFIG.LOG_LIFETIME
+        )
+        rho = [Fp(value=0)] * TEST_CONFIG.RAND_LEN_FE
+        hashes = [[Fp(value=0)] * TEST_CONFIG.HASH_LEN_FE] * TEST_CONFIG.DIMENSION
+        xmss_sig = XmssSignature(path=path, rho=rho, hashes=hashes)
+
+        consensus_sig = Signature.from_xmss(xmss_sig, TEST_SIGNATURE_SCHEME)
+
+        # The signature should be padded with zeros
+        raw_xmss = xmss_sig.to_bytes(TEST_CONFIG)
+        expected_padding = 3100 - len(raw_xmss)
+
+        # Verify the last bytes are zeros (padding)
+        assert consensus_sig[-expected_padding:] == b"\x00" * expected_padding
+
+    def test_from_xmss_preserves_data(self) -> None:
+        """Test that from_xmss preserves the XMSS signature data."""
+        # Create an XMSS signature with distinct values
+        path = HashTreeOpening(
+            siblings=[
+                [Fp(value=i * j) for i in range(TEST_CONFIG.HASH_LEN_FE)]
+                for j in range(TEST_CONFIG.LOG_LIFETIME)
+            ]
+        )
+        rho = [Fp(value=i * 10) for i in range(TEST_CONFIG.RAND_LEN_FE)]
+        hashes = [
+            [Fp(value=i + j * 100) for i in range(TEST_CONFIG.HASH_LEN_FE)]
+            for j in range(TEST_CONFIG.DIMENSION)
+        ]
+        xmss_sig = XmssSignature(path=path, rho=rho, hashes=hashes)
+
+        # Convert to consensus format
+        consensus_sig = Signature.from_xmss(xmss_sig, TEST_SIGNATURE_SCHEME)
+
+        # The beginning of the consensus signature should match the XMSS bytes
+        raw_xmss = xmss_sig.to_bytes(TEST_CONFIG)
+        assert bytes(consensus_sig)[: len(raw_xmss)] == raw_xmss
+
+    def test_from_xmss_roundtrip_with_verify(self) -> None:
+        """Test that a signature created via from_xmss can be verified."""
+        from lean_spec.subspecs.xmss.interface import TEST_SIGNATURE_SCHEME
+        from lean_spec.types import Uint64
+
+        # Generate a test key pair
+        pk, sk = TEST_SIGNATURE_SCHEME.key_gen(Uint64(0), Uint64(10))
+
+        # Create a test message (must be exactly 32 bytes)
+        message = b"test message for signing123456\x00\x00"  # 32 bytes
+        assert len(message) == 32
+        epoch = Uint64(0)
+
+        # Sign the message
+        xmss_sig = TEST_SIGNATURE_SCHEME.sign(sk, epoch, message)
+
+        # Convert to consensus signature
+        consensus_sig = Signature.from_xmss(xmss_sig, TEST_SIGNATURE_SCHEME)
+
+        # Verify using the consensus signature's verify method
+        assert consensus_sig.verify(pk, epoch, message, TEST_SIGNATURE_SCHEME)
+
+    def test_from_xmss_different_signatures_produce_different_results(self) -> None:
+        """Test that different XMSS signatures produce different consensus signatures."""
+        # Create two different XMSS signatures
+        path1 = HashTreeOpening(
+            siblings=[[Fp(value=i) for i in range(TEST_CONFIG.HASH_LEN_FE)]]
+            * TEST_CONFIG.LOG_LIFETIME
+        )
+        path2 = HashTreeOpening(
+            siblings=[[Fp(value=i + 1) for i in range(TEST_CONFIG.HASH_LEN_FE)]]
+            * TEST_CONFIG.LOG_LIFETIME
+        )
+        rho = [Fp(value=0)] * TEST_CONFIG.RAND_LEN_FE
+        hashes = [[Fp(value=0)] * TEST_CONFIG.HASH_LEN_FE] * TEST_CONFIG.DIMENSION
+
+        xmss_sig1 = XmssSignature(path=path1, rho=rho, hashes=hashes)
+        xmss_sig2 = XmssSignature(path=path2, rho=rho, hashes=hashes)
+
+        consensus_sig1 = Signature.from_xmss(xmss_sig1, TEST_SIGNATURE_SCHEME)
+        consensus_sig2 = Signature.from_xmss(xmss_sig2, TEST_SIGNATURE_SCHEME)
+
+        # Different XMSS signatures should produce different consensus signatures
+        assert bytes(consensus_sig1) != bytes(consensus_sig2)