add sphinx packet construction/deconstruction

mixnet/mixnet.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import random
 from dataclasses import dataclass
 from typing import List, TypeAlias
 
@@ -43,24 +44,26 @@ def build_topology(
             layers.append(layer)
         return MixnetTopology(layers)
 
+    def choose_mixnode(self) -> MixNode:
+        return random.choice(self.mix_nodes)
+
 
 @dataclass
 class MixNode:
-    identity_public_key: BlsPublicKey
-    encryption_public_key: X25519PublicKey
+    identity_private_key: BlsPrivateKey
+    encryption_private_key: X25519PrivateKey
     addr: NodeAddress
 
-    def __init__(
-        self,
-        identity_private_key: BlsPrivateKey,
-        encryption_private_key: X25519PrivateKey,
-        addr: NodeAddress,
-    ):
-        self.identity_public_key = identity_private_key.get_g1()
-        self.encryption_public_key = encryption_private_key.public_key()
-        self.addr = addr
+    def identity_public_key(self) -> BlsPublicKey:
+        return self.identity_private_key.get_g1()
+
+    def encryption_public_key(self) -> X25519PublicKey:
+        return self.encryption_private_key.public_key()
 
 
 @dataclass
 class MixnetTopology:
     layers: List[List[MixNode]]
+
+    def generate_route(self) -> list[MixNode]:
+        return [random.choice(layer) for layer in self.layers]

mixnet/sphinx/const.py

@@ -0,0 +1,22 @@
+# k in the Sphinx paper
+SECURITY_PARAMETER = 16
+# r in the Sphinx paper
+# In this specification, the max number of mix nodes in a route is limited to this value.
+MAX_PATH_LENGTH = 5
+# The length of node address which contains an IP address and a port.
+NODE_ADDRESS_LENGTH = 2 * SECURITY_PARAMETER
+# The length of flag that represents the type of routing information (forward-hop or final-hop)
+FLAG_LENGTH = 1
+
+VERSION_LENGTH = 3
+VERSION = b"\x00\x00\x00"
+
+# In our architecture, SURB is not used.
+# But, for the consistency with Nym's Sphinx implementation, keep this field in the Sphinx header.
+SURB_IDENTIFIER_LENGTH = SECURITY_PARAMETER
+SURB_IDENTIFIER = b"\x01\x01\x01\x01\x01\x01\x01\x01\x01\x01\x01\x01\x01\x01\x01\x01"
+
+# In our architecture, delays are determined by each mix node (not by a packet sender).
+# But, for the consistency with Nym's Sphinx implementation, keep the delay field in the Sphinx header.
+DELAY_LENGTH = 8
+DELAY = b"\x00\x00\x00\x00\x00\x00\x00\x00"

mixnet/sphinx/crypto.py

@@ -0,0 +1,23 @@
+from cryptography.hazmat.primitives import hashes, hmac
+from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+
+
+def aes128ctr(data: bytes, key: bytes, nonce: bytes) -> bytes:
+    encryptor = Cipher(algorithms.AES128(key), modes.CTR(nonce)).encryptor()
+    return encryptor.update(data) + encryptor.finalize()
+
+
+def compute_hmac_sha256(data: bytes, key: bytes) -> bytes:
+    h = hmac.HMAC(key, hashes.SHA256())
+    h.update(data)
+    return h.finalize()
+
+
+def lioness_encrypt(data: bytes, key: bytes) -> bytes:
+    # TODO: Couldn't find a lioness package that works with the latest Python. Implement it.
+    return data
+
+
+def lioness_decrypt(data: bytes, key: bytes) -> bytes:
+    # TODO: Couldn't find a lioness package that works with the latest Python. Implement it.
+    return data

mixnet/sphinx/header/header.py

@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import List, Self, Tuple
+
+from cryptography.hazmat.primitives.asymmetric.x25519 import (
+    X25519PrivateKey,
+    X25519PublicKey,
+)
+
+from mixnet.mixnet import MixNode, NodeAddress
+from mixnet.sphinx.header.keys import KeyMaterial, RoutingKeys
+from mixnet.sphinx.header.routing import EncapsulatedRoutingInformation, Filler
+
+
+@dataclass
+class SphinxHeader:
+    """
+    A Sphinx header contains an encapsulated routing information
+    and a shared secret that can be used to unwrap one layer of the encapsulated routing information.
+    """
+
+    shared_pubkey: X25519PublicKey
+    routing_info: EncapsulatedRoutingInformation
+
+    @classmethod
+    def build(
+        cls,
+        initial_ephemeral_privkey: X25519PrivateKey,
+        route: List[MixNode],
+        destination: MixNode,
+    ) -> Tuple[Self, List[bytes]]:
+        """
+        Construct a SphinxHeader by encapsulating all routing information
+        and keys that can be used to encrypt a payload.
+        """
+        key_material = KeyMaterial.derive(initial_ephemeral_privkey, route)
+        filler = Filler.build(key_material.routing_keys)
+        routing_info = EncapsulatedRoutingInformation.build(
+            route, destination, key_material.routing_keys, filler
+        )
+        payload_keys = [
+            routing_key.payload_key for routing_key in key_material.routing_keys
+        ]
+        return (cls(key_material.initial_ephemeral_pubkey, routing_info), payload_keys)
+
+    def process(
+        self, private_key: X25519PrivateKey
+    ) -> ProcessedForwardHopHeader | ProcessedFinalHopHeader:
+        """
+        Unwrap one layer of encapsulated routing information using private_key.
+
+        If there are other encapsulated layers left after being unwrapped, this method returns ProcessedForwardHopHeader.
+        If not, this returns ProcessedFinalHopHeader.
+        """
+        routing_keys = self.compute_routing_keys(self.shared_pubkey, private_key)
+
+        assert self.routing_info.integrity_mac.verify(
+            self.routing_info.encrypted_routing_info.value,
+            routing_keys.header_integrity_hmac_key,
+        )
+
+        routing_info_and_addr = self.routing_info.encrypted_routing_info.unwrap(
+            routing_keys.stream_cipher_key
+        )
+        encapsulated_routing_info = routing_info_and_addr[0]
+        next_node_address = routing_info_and_addr[1]
+
+        if encapsulated_routing_info is not None:
+            new_shared_pubkey = KeyMaterial.blind_shared_pubkey(
+                self.shared_pubkey, routing_keys.blinding_factor
+            )
+            return ProcessedForwardHopHeader(
+                SphinxHeader(new_shared_pubkey, encapsulated_routing_info),
+                next_node_address,
+                routing_keys.payload_key,
+            )
+        else:
+            return ProcessedFinalHopHeader(next_node_address, routing_keys.payload_key)
+
+    @staticmethod
+    def compute_routing_keys(
+        shared_pubkey: X25519PublicKey, private_key: X25519PrivateKey
+    ) -> RoutingKeys:
+        """
+        Derive RoutingKeys from a shared key created by Diffie-Hellman key exchange between shared_pubkey and private_key.
+        """
+        dh_shared_key = private_key.exchange(shared_pubkey)
+        return RoutingKeys.derive(dh_shared_key)
+
+
+@dataclass
+class ProcessedForwardHopHeader:
+    """
+    A forward-hop header unwrapped from SphinxHeader
+
+    This class contains another SphinxHeader to be forwarded to the next mix node,
+    and a payload key for the current mix node to decrypt one layer of payload encryption.
+    """
+
+    next_header: SphinxHeader
+    next_node_address: NodeAddress
+    payload_key: bytes
+
+
+@dataclass
+class ProcessedFinalHopHeader:
+    """
+    A final-hop header unwrapped from SphinxHeader
+
+    This class contains a payload key for the current mix node to decrypt the last layer of payload encryption,
+    and a destination address to which the decrypted payload will be delivered.
+    """
+
+    destination_address: NodeAddress
+    payload_key: bytes

mixnet/sphinx/header/keys.py

@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import List, Self
+
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.asymmetric.x25519 import (
+    X25519PrivateKey,
+    X25519PublicKey,
+)
+from cryptography.hazmat.primitives.kdf.hkdf import HKDF
+
+from mixnet.mixnet import MixNode
+
+
+@dataclass
+class KeyMaterial:
+    """
+    Contain a list of RoutingKeys for all mix nodes in the route,
+    and a shared secret that will be contained in a SphinxHeader for the first mix node in the route.
+    """
+
+    initial_ephemeral_pubkey: X25519PublicKey
+    routing_keys: List[RoutingKeys]
+
+    @classmethod
+    def derive(
+        cls, initial_ephemeral_privkey: X25519PrivateKey, route: List[MixNode]
+    ) -> Self:
+        """
+        Derive KeyMaterial for route using initial_ephemeral_privkey provided.
+        """
+        initial_ephemeral_pubkey = initial_ephemeral_privkey.public_key()
+
+        routing_keys = []
+        accumulated_privkey = initial_ephemeral_privkey
+        for node in route:
+            dh_shared_key = accumulated_privkey.exchange(node.encryption_public_key())
+            node_routing_keys = RoutingKeys.derive(dh_shared_key)
+
+            # TODO: find a proper library for Ristretto operations
+            # https://github.com/nymtech/sphinx/blob/ca107d94360cdf8bbfbdb12fe5320ed74f80e40c/src/header/keys.rs#L128-L128
+            # blinding_factor_scalar = Scalar.from_bytes_mod_order(node_routing_keys.blinding_factor)
+            # accumulated_privkey = product(accumulated_privkey, blinding_factor_scalar)
+
+            routing_keys.append(node_routing_keys)
+
+        return cls(initial_ephemeral_pubkey, routing_keys)
+
+    @staticmethod
+    def blind_shared_pubkey(
+        shared_pubkey: X25519PublicKey, blinding_factor: bytes
+    ) -> X25519PublicKey:
+        """
+        Blind shared_pubkey to derive a next public key.
+        """
+        # TODO: find a proper library for Ristretto operations
+        # https://github.com/nymtech/sphinx/blob/ca107d94360cdf8bbfbdb12fe5320ed74f80e40c/src/header/mod.rs#L236-L236
+        # For now, we're skipping blinding because we don't accumulate a private key using blinding factor
+        # when deriving RoutingKeys.
+        return shared_pubkey
+
+
+# Adopted from https://github.com/nymtech/sphinx/blob/ca107d94360cdf8bbfbdb12fe5320ed74f80e40c/src/constants.rs#L26-L26
+HKDF_INPUT_SEED = b"Dwste mou enan moxlo arketa makru kai ena upomoxlio gia na ton topothetisw kai tha kinisw thn gh."
+
+
+@dataclass
+class RoutingKeys:
+    """
+    Contain all keys for a mix node in the route.
+    """
+
+    # For Sphinx header encryption (AES-128)
+    stream_cipher_key: bytes
+    # For HMAC integrity authentication
+    header_integrity_hmac_key: bytes
+    # For payload encryption (ChaCha20)
+    payload_key: bytes
+    # For deriving a shared key for a next mix node, combining with the previous ephemeral private key
+    blinding_factor: bytes
+
+    @classmethod
+    def derive(cls, dh_shared_key: bytes) -> Self:
+        """
+        Derive all keys from dh_shared_key using HKDF-SHA256.
+        """
+        derived_key = HKDF(
+            algorithm=hashes.SHA256(), length=256, salt=None, info=HKDF_INPUT_SEED
+        ).derive(dh_shared_key)
+        assert len(derived_key) == 256
+
+        stream_cipher_key = derived_key[0:16]  # 16bytes == 128bits
+        header_integrity_hmac_key = derived_key[16:32]  # 16bytes
+        payload_key = derived_key[32:224]  # 192bytes
+        blinding_factor = derived_key[224:]  # 32bytes
+        return cls(
+            stream_cipher_key, header_integrity_hmac_key, payload_key, blinding_factor
+        )

mixnet/sphinx/header/mac.py

@@ -0,0 +1,36 @@
+from dataclasses import dataclass
+from typing import Self
+
+from mixnet.sphinx.const import SECURITY_PARAMETER
+from mixnet.sphinx.crypto import compute_hmac_sha256
+
+
+@dataclass
+class IntegrityHmac:
+    """
+    This class represents a HMAC-SHA256 that can be used for integrity authentication.
+    """
+
+    value: bytes
+
+    def __init__(self, value: bytes):
+        """Override the default constructor to assert the size of value"""
+        assert len(value) == IntegrityHmac.size()
+        self.value = value
+
+    @staticmethod
+    def size() -> int:
+        return SECURITY_PARAMETER
+
+    @classmethod
+    def compute(cls, data: bytes, key: bytes) -> Self:
+        """
+        Build IntegrityHmac using data and key.
+        """
+        return cls(compute_hmac_sha256(data, key)[: cls.size()])
+
+    def verify(self, data: bytes, key: bytes) -> bool:
+        """
+        Verify a HMAC computed from data and key matches with the expected HMAC.
+        """
+        return self.value == compute_hmac_sha256(data, key)[: self.size()]