dummy message encoding

mixnet/v2/sim/main.py

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

mixnet/v2/sim/message.py

@@ -0,0 +1,36 @@
+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,10 +1,20 @@
-from mixnet.v2.sim.simulation import Simulation
+import random
+
+import simpy
+
+from mixnet.v2.sim.message import Message
+from mixnet.v2.sim.p2p import P2p
 
 
 class Node:
-    def __init__(self, sim: Simulation):
-        self.sim = sim
-        self.sim.env.process(self.send_message())
+    N_MIXES_IN_PATH = 3
+
+    def __init__(self, id: str, 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.action = self.env.process(self.send_message())
 
     def send_message(self):
         """
@@ -12,15 +22,19 @@ class Node:
         """
         while True:
             msg = self.create_message()
-            yield self.sim.env.timeout(3)
-            self.sim.env.process(self.sim.p2p.broadcast(msg))
+            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:
         """
         Creates a message using the Sphinx format
         @return:
         """
-        return b""
+        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)
 
     def receive_message(self, msg: bytes):
         """
@@ -28,13 +42,15 @@ class Node:
         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.sim.env.process(self.sim.p2p.broadcast(msg))
-            else: # Even if not, forward it to the next mix
-                yield self.sim.env.timeout(1)  # TODO: use a random delay
-                # Use broadcasting here too
-                self.sim.env.process(self.sim.p2p.broadcast(msg))
-        else:  # if the msg has gone through all mixes
-            pass
+        # 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

mixnet/v2/sim/p2p.py

@@ -1,13 +1,22 @@
-from mixnet.v2.sim.node import Node
-from mixnet.v2.sim.simulation import Simulation
+import random
+
+import simpy
 
 
 class P2p:
-    def __init__(self, sim: Simulation, nodes: list[Node]):
-        self.sim = sim
-        self.nodes = nodes
+    def __init__(self, env: simpy.Environment):
+        self.env = env
+        self.nodes = []
+
+    def add_node(self, nodes):
+        self.nodes.extend(nodes)
 
     def broadcast(self, msg):
+        print("Broadcasting a message at time %d" % self.env.now)
+        yield self.env.timeout(1)
         # TODO: gossipsub or something similar
         for node in self.nodes:
-            self.sim.env.process(node.receive_message(msg))
+            self.env.process(node.receive_message(msg))
+
+    def get_nodes(self, n: int):
+        return random.choices(self.nodes, k=n)

mixnet/v2/sim/simulation.py

@@ -1,10 +1,15 @@
 import simpy
 
+from mixnet.v2.sim.node import Node
+from mixnet.v2.sim.p2p import P2p
+
 
 class Simulation:
     def __init__(self):
         self.env = simpy.Environment()
-        self.p2p = None  # TODO: implement p2p
+        self.p2p = P2p(self.env)
+        self.nodes = [Node(str(i), self.env, self.p2p) for i in range(2)]
+        self.p2p.add_node(self.nodes)
 
     def run(self, until):
         self.env.run(until=until)