working version of mixnet behavior simulation

mixnet/v2/sim/main.py

@@ -2,5 +2,5 @@ from mixnet.v2.sim.simulation import Simulation
 
 if __name__ == "__main__":
     sim = Simulation()
-    sim.run(10)
+    sim.run(30)
     print("Simulation complete!")

mixnet/v2/sim/message.py

@@ -1,36 +0,0 @@
-from __future__ import annotations
-
-
-class Message:
-    def __init__(self, pubkeys: list[bytes], attachments: list[bytes], payload: bytes):
-        assert len(pubkeys) == len(attachments)
-        eph_sk, eph_pk = bytes(32), bytes(32)  # TODO: use a random x25519 key
-        node_keys = Message.node_keys(eph_pk, pubkeys)
-        self.header = Header(eph_sk, node_keys, attachments)
-        self.payload = payload  # TODO: encrypt payload
-
-    def __bytes__(self):
-        return bytes(self.header) + self.payload
-
-    @classmethod
-    def node_keys(cls, eph_sk: bytes, pubkeys: list[bytes]) -> list[bytes]:
-        return [cls.key_exchange(eph_sk, pk) for pk in pubkeys]
-
-    @classmethod
-    def key_exchange(cls, eph_sk, pubkey) -> bytes:
-        pass
-
-    # TODO: implement unwrapping the message
-
-
-class Header:
-    DUMMY_MAC = bytes(16)
-
-    def __init__(self, eph_sk: bytes, node_keys: list[bytes], attachments: list[bytes]):
-        assert len(node_keys) == len(attachments)
-        self.eph_sk = eph_sk
-        # TODO: encapsulation
-        self.attachments = attachments
-
-    def __bytes__(self):
-        return b"".join([self.eph_sk] + [bytes(att) + self.DUMMY_MAC for att in self.attachments])

mixnet/v2/sim/node.py

@@ -1,56 +1,77 @@
 import random
 
 import simpy
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey, X25519PublicKey
 
-from mixnet.v2.sim.message import Message
+from mixnet.v2.sim.sphinx import SphinxPacket, Attachment
 from mixnet.v2.sim.p2p import P2p
 
 
 class Node:
-    N_MIXES_IN_PATH = 3
+    N_MIXES_IN_PATH = 2
 
-    def __init__(self, id: str, env: simpy.Environment, p2p: P2p):
+    def __init__(self, id: int, env: simpy.Environment, p2p: P2p):
         self.id = id
         self.env = env
         self.p2p = p2p
-        self.pubkey = bytes(32)  # TODO: replace with actual x25519 pubkey
+        self.private_key = X25519PrivateKey.generate()
+        self.public_key = self.private_key.public_key()
         self.action = self.env.process(self.send_message())
 
     def send_message(self):
         """
         Creates/encapsulate a message and send it to the network through the mixnet
         """
-        while True:
+        # while True:
+        if self.id == 0:
             msg = self.create_message()
             yield self.env.timeout(2)
             print("Sending a message at time %d" % self.env.now)
             self.env.process(self.p2p.broadcast(msg))
 
-    def create_message(self) -> bytes:
+    def create_message(self) -> SphinxPacket:
         """
         Creates a message using the Sphinx format
         @return:
         """
         mixes = self.p2p.get_nodes(self.N_MIXES_IN_PATH)
-        incentive_txs = [bytes(256) for _ in mixes]  # TODO: replace with realistic tx
-        msg = Message(mixes, incentive_txs, b"Hello, world!")
-        return bytes(msg)
+        public_keys = [mix.public_key for mix in mixes]
+        # TODO: replace with realistic tx
+        incentive_txs = [Node.create_incentive_tx(mix.public_key) for mix in mixes]
+        return SphinxPacket(public_keys, incentive_txs, b"Hello, world!")
 
-    def receive_message(self, msg: bytes):
+    def receive_message(self, msg: SphinxPacket | bytes):
         """
         Receives a message from the network, processes it,
         and forwards it to the next mix or the entire network if necessary.
         @param msg: the message to be processed
         """
-        yield self.env.timeout(random.randint(0,3))
-        print("Receiving a message at time %d" % self.env.now)
-        # TODO: this is a dummy logic
-        # if msg[0] == 0x00:  # if the msg is to be relayed
-        #     if msg[1] == 0x00:  # if I'm the exit mix,
-        #         self.env.process(self.p2p.broadcast(msg))
-        #     else: # Even if not, forward it to the next mix
-        #         yield self.env.timeout(1)  # TODO: use a random delay
-        #         # Use broadcasting here too
-        #         self.env.process(self.p2p.broadcast(msg))
-        # else:  # if the msg has gone through all mixes
-        #     pass
+        # simulating network latency
+        yield self.env.timeout(random.randint(0, 3))
+
+        if isinstance(msg, SphinxPacket):
+            msg, incentive_tx = msg.unwrap(self.private_key)
+            if self.is_my_incentive_tx(incentive_tx):
+                self.log("Receiving SphinxPacket. It's mine!")
+                if msg.is_all_unwrapped():
+                    self.env.process(self.p2p.broadcast(msg.payload))
+                else:
+                    # TODO: use Poisson delay
+                    yield self.env.timeout(random.randint(0, 5))
+                    self.env.process(self.p2p.broadcast(msg))
+            else:
+                self.log("Receiving SphinxPacket, but not mine")
+        else:
+            self.log("Received original message: %s" % msg)
+
+    @classmethod
+    def create_incentive_tx(cls, mix_public_key: X25519PublicKey) -> Attachment:
+        return Attachment(
+            mix_public_key.public_bytes(encoding=serialization.Encoding.Raw, format=serialization.PublicFormat.Raw))
+
+    def is_my_incentive_tx(self, tx: Attachment) -> bool:
+        return tx == Node.create_incentive_tx(self.public_key)
+
+    def log(self, msg):
+        print("Node:%d at %d: %s" % (self.id, self.env.now, msg))

mixnet/v2/sim/p2p.py

@@ -2,6 +2,8 @@ import random
 
 import simpy
 
+from mixnet.v2.sim.sphinx import SphinxPacket
+
 
 class P2p:
     def __init__(self, env: simpy.Environment):
@@ -11,12 +13,16 @@ class P2p:
     def add_node(self, nodes):
         self.nodes.extend(nodes)
 
-    def broadcast(self, msg):
-        print("Broadcasting a message at time %d" % self.env.now)
+    # TODO: This should accept only bytes, but SphinxPacket is also accepted until we implement the Sphinx serde
+    def broadcast(self, msg: SphinxPacket | bytes):
+        self.log("Broadcasting a msg")
         yield self.env.timeout(1)
         # TODO: gossipsub or something similar
         for node in self.nodes:
             self.env.process(node.receive_message(msg))
 
     def get_nodes(self, n: int):
-        return random.choices(self.nodes, k=n)
+        return random.sample(self.nodes, n)
+
+    def log(self, msg):
+        print("P2P at %d: %s" % (self.env.now, msg))

mixnet/v2/sim/simulation.py

@@ -8,7 +8,7 @@ class Simulation:
     def __init__(self):
         self.env = simpy.Environment()
         self.p2p = P2p(self.env)
-        self.nodes = [Node(str(i), self.env, self.p2p) for i in range(2)]
+        self.nodes = [Node(i, self.env, self.p2p) for i in range(2)]
         self.p2p.add_node(self.nodes)
 
     def run(self, until):

mixnet/v2/sim/sphinx.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+from copy import deepcopy
+
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PublicKey, X25519PrivateKey
+
+
+class SphinxPacket:
+    def __init__(self, public_keys: list[X25519PublicKey], attachments: list[Attachment], payload: bytes):
+        assert len(public_keys) == len(attachments)
+        ephemeral_private_key = X25519PrivateKey.generate()
+        ephemeral_public_key = ephemeral_private_key.public_key()
+        shared_keys = [SharedSecret(ephemeral_private_key, pk) for pk in public_keys]
+        self.header = SphinxHeader(ephemeral_public_key, shared_keys, attachments)
+        self.payload = payload  # TODO: encrypt payload
+
+    def __bytes__(self):
+        return bytes(self.header) + self.payload
+
+    def size(self) -> int:
+        return len(bytes(self))
+
+    def unwrap(self, private_key: X25519PrivateKey) -> tuple[SphinxPacket, Attachment]:
+        packet = deepcopy(self)
+        attachment = packet.header.unwrap_inplace(private_key)
+        # TODO: decrypt packet.payload
+        return packet, attachment
+
+    def is_all_unwrapped(self) -> bool:
+        return self.header.is_all_unwrapped()
+
+
+class SphinxHeader:
+    DUMMY_MAC = b'\xFF' * 16
+
+    def __init__(self, ephemeral_public_key: X25519PublicKey, shared_keys: list[SharedSecret],
+                 attachments: list[Attachment]):
+        assert len(shared_keys) == len(attachments)
+        self.ephemeral_public_key = ephemeral_public_key.public_bytes(encoding=serialization.Encoding.Raw, format=serialization.PublicFormat.Raw)
+        self.attachments = attachments  # TODO: encapsulation using node_keys
+
+    def __bytes__(self):
+        return b"".join([self.ephemeral_public_key] + [bytes(att) + self.DUMMY_MAC for att in self.attachments])
+
+    def unwrap_inplace(self, private_key: X25519PrivateKey) -> Attachment:
+        # TODO: shared_secret = SharedSecret(private_key, header.ephemeral_public_key)
+        attachment = self.attachments.pop(0)
+        self.attachments.append(Attachment(bytes(len(bytes(attachment)))))  # append a dummy attachment
+        return attachment
+
+    def is_all_unwrapped(self) -> bool:
+        # true if the first attachment is a dummy
+        return self.attachments[0] == Attachment(bytes(len(bytes(self.attachments[0]))))
+
+
+class SharedSecret:
+    def __init__(self, private_key: X25519PrivateKey, public_key: X25519PublicKey):
+        self.key = private_key.exchange(public_key)  # 32 bytes
+
+    def __bytes__(self):
+        return self.key
+
+
+class Attachment:
+    def __init__(self, data: bytes):
+        self.data = data
+
+    def __bytes__(self):
+        return self.data
+
+    def __eq__(self, other):
+        return bytes(self) == bytes(other)