feat: add signature aggregation using python bindings

packages/testing/src/consensus_testing/keys.py

@@ -39,19 +39,21 @@
 
 from lean_spec.config import LEAN_ENV
 from lean_spec.subspecs.containers import AttestationData
-from lean_spec.subspecs.containers.attestation.types import NaiveAggregatedSignature
 from lean_spec.subspecs.containers.block.types import (
     AggregatedAttestations,
     AttestationSignatures,
 )
 from lean_spec.subspecs.containers.slot import Slot
+from lean_spec.subspecs.containers.state.types import AttestationSignatureKey
+from lean_spec.subspecs.xmss.aggregation import aggregate_signatures
 from lean_spec.subspecs.xmss.containers import KeyPair, PublicKey, Signature
 from lean_spec.subspecs.xmss.interface import (
     PROD_SIGNATURE_SCHEME,
     TEST_SIGNATURE_SCHEME,
     GeneralizedXmssScheme,
 )
 from lean_spec.types import Uint64
+from lean_spec.types.byte_arrays import LeanAggregatedSignature
 
 if TYPE_CHECKING:
     from collections.abc import Mapping
@@ -273,24 +275,40 @@ def sign_attestation_data(
     def build_attestation_signatures(
         self,
         aggregated_attestations: AggregatedAttestations,
-        signature_lookup: Mapping[tuple[Uint64, bytes], Signature] | None = None,
+        signature_lookup: Mapping[AttestationSignatureKey, Signature] | None = None,
     ) -> AttestationSignatures:
-        """Build `AttestationSignatures` for already-aggregated attestations."""
+        """
+        Build `AttestationSignatures` for already-aggregated attestations.
+
+        For each aggregated attestation, collect the participating validators' public keys and
+        signatures, then produce a single leanVM aggregated signature proof blob using
+        `xmss_aggregate_signatures` (via `aggregate_signatures`).
+        """
         lookup = signature_lookup or {}
-        return AttestationSignatures(
-            data=[
-                NaiveAggregatedSignature(
-                    data=[
-                        (
-                            lookup.get((vid, agg.data.data_root_bytes()))
-                            or self.sign_attestation_data(vid, agg.data)
-                        )
-                        for vid in agg.aggregation_bits.to_validator_indices()
-                    ]
-                )
-                for agg in aggregated_attestations
+
+        proof_blobs: list[LeanAggregatedSignature] = []
+        for agg in aggregated_attestations:
+            validator_ids = agg.aggregation_bits.to_validator_indices()
+            message = agg.data.data_root_bytes()
+            epoch = agg.data.slot
+
+            public_keys: list[PublicKey] = [self.get_public_key(vid) for vid in validator_ids]
+            signatures: list[Signature] = [
+                (lookup.get((vid, message)) or self.sign_attestation_data(vid, agg.data))
+                for vid in validator_ids
             ]
-        )
+
+            # If the caller supplied raw signatures and any are invalid,
+            # aggregation should fail with exception.
+            aggregated_signature = aggregate_signatures(
+                public_keys=public_keys,
+                signatures=signatures,
+                message=message,
+                epoch=epoch,
+            )
+            proof_blobs.append(aggregated_signature)
+
+        return AttestationSignatures(data=proof_blobs)
 
 
 def _generate_single_keypair(

packages/testing/src/consensus_testing/test_fixtures/fork_choice.py

@@ -26,6 +26,7 @@
 from lean_spec.subspecs.containers.slot import Slot
 from lean_spec.subspecs.containers.state import Validators
 from lean_spec.subspecs.containers.state.state import State
+from lean_spec.subspecs.containers.state.types import AttestationSignatureKey
 from lean_spec.subspecs.forkchoice import Store
 from lean_spec.subspecs.koalabear import Fp
 from lean_spec.subspecs.ssz import hash_tree_root
@@ -231,7 +232,10 @@ def make_fixture(self) -> ForkChoiceTest:
 
                 elif isinstance(step, AttestationStep):
                     # Process attestation from gossip (immutable)
-                    store = store.on_attestation(step.attestation, is_from_block=False)
+                    store = store.on_gossip_attestation(
+                        step.attestation,
+                        scheme=LEAN_ENV_TO_SCHEMES[self.lean_env],
+                    )
 
                 else:
                     raise ValueError(f"Step {i}: unknown step type {type(step).__name__}")
@@ -301,7 +305,12 @@ def _build_block_from_spec(
         parent_root = self._resolve_parent_root(spec, store, block_registry)
 
         # Build attestations from spec
-        attestations = self._build_attestations_from_spec(spec, store, block_registry, parent_root)
+        attestations, attestation_signatures = self._build_attestations_from_spec(
+            spec, store, block_registry, parent_root, key_manager
+        )
+
+        gossip_signatures = dict(store.gossip_signatures)
+        gossip_signatures.update(attestation_signatures)
 
         # Use State.build_block for core block building (pure spec logic)
         parent_state = store.states[parent_root]
@@ -310,6 +319,8 @@ def _build_block_from_spec(
             proposer_index=proposer_index,
             parent_root=parent_root,
             attestations=attestations,
+            gossip_signatures=gossip_signatures,
+            aggregated_payloads=store.aggregated_payloads,
         )
 
         # Create proposer attestation for this block
@@ -325,8 +336,9 @@ def _build_block_from_spec(
         )
 
         # Sign all attestations and the proposer attestation
-        attestation_signatures = key_manager.build_attestation_signatures(
-            final_block.body.attestations
+        attestation_signatures_blob = key_manager.build_attestation_signatures(
+            final_block.body.attestations,
+            attestation_signatures,
         )
 
         proposer_signature = key_manager.sign_attestation_data(
@@ -340,7 +352,7 @@ def _build_block_from_spec(
                 proposer_attestation=proposer_attestation,
             ),
             signature=BlockSignatures(
-                attestation_signatures=attestation_signatures,
+                attestation_signatures=attestation_signatures_blob,
                 proposer_signature=proposer_signature,
             ),
         )
@@ -392,34 +404,38 @@ def _build_attestations_from_spec(
         store: Store,
         block_registry: dict[str, Block],
         parent_root: Bytes32,
-    ) -> list[Attestation]:
-        """Build attestations list from BlockSpec."""
+        key_manager: XmssKeyManager,
+    ) -> tuple[list[Attestation], dict[AttestationSignatureKey, Signature]]:
+        """Build attestations list from BlockSpec and their signatures."""
         if spec.attestations is None:
-            return []
+            return [], {}
 
         parent_state = store.states[parent_root]
         attestations = []
+        signature_lookup: dict[AttestationSignatureKey, Signature] = {}
 
         for att_spec in spec.attestations:
             if isinstance(att_spec, SignedAttestationSpec):
                 signed_att = self._build_signed_attestation_from_spec(
-                    att_spec, block_registry, parent_state
-                )
-                attestations.append(
-                    Attestation(validator_id=signed_att.validator_id, data=signed_att.message)
+                    att_spec, block_registry, parent_state, key_manager
                 )
             else:
-                attestations.append(
-                    Attestation(validator_id=att_spec.validator_id, data=att_spec.message)
-                )
+                signed_att = att_spec
+
+            attestation = Attestation(validator_id=signed_att.validator_id, data=signed_att.message)
+            attestations.append(attestation)
+            signature_lookup[(attestation.validator_id, attestation.data.data_root_bytes())] = (
+                signed_att.signature
+            )
 
-        return attestations
+        return attestations, signature_lookup
 
     def _build_signed_attestation_from_spec(
         self,
         spec: SignedAttestationSpec,
         block_registry: dict[str, Block],
         state: State,
+        key_manager: XmssKeyManager,
     ) -> SignedAttestation:
         """
         Build a SignedAttestation from a SignedAttestationSpec.
@@ -466,15 +482,21 @@ def _build_signed_attestation_from_spec(
         )
 
         # Create signed attestation
+        if spec.signature is not None:
+            signature = spec.signature
+        elif spec.valid_signature:
+            signature = key_manager.sign_attestation_data(
+                attestation.validator_id, attestation.data
+            )
+        else:
+            signature = Signature(
+                path=HashTreeOpening(siblings=HashDigestList(data=[])),
+                rho=Randomness(data=[Fp(0) for _ in range(PROD_CONFIG.RAND_LEN_FE)]),
+                hashes=HashDigestList(data=[]),
+            )
+
         return SignedAttestation(
             validator_id=attestation.validator_id,
             message=attestation.data,
-            signature=(
-                spec.signature
-                or Signature(
-                    path=HashTreeOpening(siblings=HashDigestList(data=[])),
-                    rho=Randomness(data=[Fp(0) for _ in range(PROD_CONFIG.RAND_LEN_FE)]),
-                    hashes=HashDigestList(data=[]),
-                )
-            ),
+            signature=signature,
         )

packages/testing/src/consensus_testing/test_fixtures/state_transition.py

@@ -11,6 +11,7 @@
 from lean_spec.subspecs.ssz.hash import hash_tree_root
 from lean_spec.types import Bytes32, Uint64
 
+from ..keys import get_shared_key_manager
 from ..test_types import BlockSpec, StateExpectation
 from .base import BaseConsensusFixture
 
@@ -258,10 +259,24 @@ def _build_block_from_spec(self, spec: BlockSpec, state: State) -> tuple[Block,
             for vid in agg.aggregation_bits.to_validator_indices()
         ]
 
+        if plain_attestations:
+            key_manager = get_shared_key_manager(max_slot=spec.slot)
+            gossip_signatures = {
+                (att.validator_id, att.data.data_root_bytes()): key_manager.sign_attestation_data(
+                    att.validator_id,
+                    att.data,
+                )
+                for att in plain_attestations
+            }
+        else:
+            gossip_signatures = {}
+
         block, post_state, _, _ = state.build_block(
             slot=spec.slot,
             proposer_index=proposer_index,
             parent_root=parent_root,
             attestations=plain_attestations,
+            gossip_signatures=gossip_signatures,
+            aggregated_payloads={},
         )
         return block, post_state

packages/testing/src/consensus_testing/test_fixtures/verify_signatures.py

@@ -16,6 +16,7 @@
     BlockWithAttestation,
     SignedBlockWithAttestation,
 )
+from lean_spec.subspecs.containers.block.types import AttestationSignatures
 from lean_spec.subspecs.containers.checkpoint import Checkpoint
 from lean_spec.subspecs.containers.state.state import State
 from lean_spec.subspecs.koalabear import Fp
@@ -179,26 +180,26 @@ def _build_block_from_spec(
             spec, state, key_manager
         )
 
+        # Provide signatures to State.build_block so it can include attestations during
+        # fixed-point collection when available_attestations/known_block_roots are used.
+        # This might contain invalid signatures as we are not validating them here.
+        gossip_signatures = {
+            (att.validator_id, att.data.data_root_bytes()): sig
+            for att, sig in zip(attestations, attestation_signature_inputs, strict=True)
+        }
+
         # Use State.build_block for core block building (pure spec logic)
-        final_block, _, _, _ = state.build_block(
+        final_block, _, _, aggregated_signatures = state.build_block(
             slot=spec.slot,
             proposer_index=proposer_index,
             parent_root=parent_root,
             attestations=attestations,
+            gossip_signatures=gossip_signatures,
+            aggregated_payloads={},
         )
 
-        # Preserve per-attestation validity from the spec.
-        #
-        # For signature tests we must ensure that the signatures in the input spec are used
-        # for any intentionally-invalid signature from the input spec remains invalid
-        # in the produced `SignedBlockWithAttestation`.
-        signature_lookup: dict[tuple[Uint64, bytes], Signature] = {
-            (att.validator_id, att.data.data_root_bytes()): sig
-            for att, sig in zip(attestations, attestation_signature_inputs, strict=True)
-        }
-        attestation_signatures = key_manager.build_attestation_signatures(
-            final_block.body.attestations,
-            signature_lookup=signature_lookup,
+        attestation_signatures = AttestationSignatures(
+            data=aggregated_signatures,
         )
 
         # Create proposer attestation for this block

packages/testing/src/consensus_testing/test_types/store_checks.py

@@ -6,7 +6,7 @@
 from lean_spec.types import Bytes32, CamelModel, Uint64
 
 if TYPE_CHECKING:
-    from lean_spec.subspecs.containers import SignedAttestation
+    from lean_spec.subspecs.containers import AttestationData
     from lean_spec.subspecs.containers.block.block import Block
     from lean_spec.subspecs.forkchoice.store import Store
 
@@ -42,7 +42,7 @@ class AttestationCheck(CamelModel):
     """
 
     def validate_attestation(
-        self, attestation: "SignedAttestation", location: str, step_index: int
+        self, attestation: "AttestationData", location: str, step_index: int
     ) -> None:
         """Validate attestation properties."""
         fields_to_check = self.model_fields_set - {"validator", "location"}
@@ -51,31 +51,31 @@ def validate_attestation(
             expected = getattr(self, field_name)
 
             if field_name == "attestation_slot":
-                actual = attestation.message.slot
+                actual = attestation.slot
                 if actual != expected:
                     raise AssertionError(
                         f"Step {step_index}: validator {self.validator} {location} "
                         f"attestation slot = {actual}, expected {expected}"
                     )
 
             elif field_name == "head_slot":
-                actual = attestation.message.head.slot
+                actual = attestation.head.slot
                 if actual != expected:
                     raise AssertionError(
                         f"Step {step_index}: validator {self.validator} {location} "
                         f"head slot = {actual}, expected {expected}"
                     )
 
             elif field_name == "source_slot":
-                actual = attestation.message.source.slot
+                actual = attestation.source.slot
                 if actual != expected:
                     raise AssertionError(
                         f"Step {step_index}: validator {self.validator} {location} "
                         f"source slot = {actual}, expected {expected}"
                     )
 
             elif field_name == "target_slot":
-                actual = attestation.message.target.slot
+                actual = attestation.target.slot
                 if actual != expected:
                     raise AssertionError(
                         f"Step {step_index}: validator {self.validator} {location} "
@@ -442,7 +442,7 @@ def validate_against_store(
                     # An attestation votes for this fork if its head is this block or a descendant
                     weight = 0
                     for attestation in store.latest_known_attestations.values():
-                        att_head_root = attestation.message.head.root
+                        att_head_root = attestation.head.root
                         # Check if attestation head is this block or a descendant
                         if att_head_root == root:
                             weight += 1

pyproject.toml

@@ -26,7 +26,11 @@ classifiers = [
     "Programming Language :: Python :: 3.14",
 ]
 requires-python = ">=3.12"
-dependencies = ["pydantic>=2.12.0,<3", "typing-extensions>=4.4"]
+dependencies = [
+    "pydantic>=2.12.0,<3",
+    "typing-extensions>=4.4",
+    "lean-multisig-py>=0.1.0",
+]
 
 [project.license]
 file = "LICENSE"
@@ -110,6 +114,7 @@ members = ["packages/*"]
 
 [tool.uv.sources]
 lean-ethereum-testing = { workspace = true }
+lean-multisig-py = { git = "https://github.com/anshalshukla/leanMultisig-py", branch = "main" }
 
 [dependency-groups]
 test = [
@@ -118,6 +123,7 @@ test = [
     "pytest-xdist>=3.6.1,<4",
     "hypothesis>=6.138.14",
     "lean-ethereum-testing",
+    "lean-multisig-py>=0.1.0",
 ]
 lint = [
     "ty>=0.0.1a34",