nomos-simulations/deprecated/mixnet-overall/protocol/test_utils.py

47 lines
1.4 KiB
Python

import random
from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey
from protocol.config import (
GlobalConfig,
MixMembership,
NodeConfig,
NodeInfo,
)
from protocol.gossip import GossipConfig
from protocol.nomssip import TemporalMixConfig
from protocol.temporalmix import TemporalMixType
def init_mixnet_config(
num_nodes: int,
max_message_size: int = 512,
max_mix_path_length: int = 3,
) -> tuple[GlobalConfig, list[NodeConfig], dict[bytes, X25519PrivateKey]]:
gossip_config = GossipConfig(peering_degree=6)
node_configs = [
NodeConfig(
X25519PrivateKey.generate(),
max_mix_path_length,
gossip_config,
TemporalMixConfig(TemporalMixType.PURE_COIN_FLIPPING, 3, random.Random()),
)
for _ in range(num_nodes)
]
global_config = GlobalConfig(
MixMembership(
[
NodeInfo(node_config.private_key.public_key())
for node_config in node_configs
]
),
transmission_rate_per_sec=3,
max_message_size=max_message_size,
max_mix_path_length=max_mix_path_length,
)
key_map = {
node_config.private_key.public_key().public_bytes_raw(): node_config.private_key
for node_config in node_configs
}
return (global_config, node_configs, key_map)