add args

mixnet/v2/sim/README.md

@@ -6,5 +6,20 @@ Make sure that all dependencies specified in the `requirements.txt` in the proje
 ```bash
 python main.py
 ```
+The following parameters can be configured:
+```
+usage: main.py [-h] [--running-time RUNNING_TIME] [--num-nodes NUM_NODES] [--num-mix-layers NUM_MIX_LAYERS]
+
+Run simulation
+
+options:
+  -h, --help            show this help message and exit
+  --running-time RUNNING_TIME
+                        Running time of the simulation (default: 30)
+  --num-nodes NUM_NODES
+                        Number of nodes in the network (default: 2)
+  --num-mix-layers NUM_MIX_LAYERS
+                        Number of mix layers in the network (default: 2)
+```
 
 TODO: Add more details

mixnet/v2/sim/main.py

@@ -1,6 +1,14 @@
+import argparse
+
 from simulation import Simulation
 
 if __name__ == "__main__":
-    sim = Simulation()
-    sim.run(30)
+    parser = argparse.ArgumentParser(description='Run simulation', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument("--running-time", type=int, default=30, help="Running time of the simulation")
+    parser.add_argument("--num-nodes", type=int, default=2, help="Number of nodes in the network")
+    parser.add_argument("--num-mix-layers", type=int, default=2, help="Number of mix layers in the network")
+    args = parser.parse_args()
+
+    sim = Simulation(args.num_nodes, args.num_mix_layers)
+    sim.run(until=args.running_time)
     print("Simulation complete!")

mixnet/v2/sim/node.py

@@ -9,18 +9,18 @@ from p2p import P2p
 
 
 class Node:
-    N_MIXES_IN_PATH = 2
     INCENTIVE_TX_SIZE = 512
     REAL_PAYLOAD = b"BLOCK"
     COVER_PAYLOAD = b"COVER"
     PADDING_SEPARATOR = b'\x01'
 
-    def __init__(self, id: int, env: simpy.Environment, p2p: P2p):
+    def __init__(self, id: int, env: simpy.Environment, p2p: P2p, num_mix_layers: int):
         self.id = id
         self.env = env
         self.p2p = p2p
         self.private_key = X25519PrivateKey.generate()
         self.public_key = self.private_key.public_key()
+        self.num_mix_layers = num_mix_layers
         self.action = self.env.process(self.send_message())
 
     def send_message(self):
@@ -31,7 +31,7 @@ class Node:
             msg = self.create_message()
             # TODO: Use the realistic cover traffic emission rate
             yield self.env.timeout(2)
-            print("Sending a message at time %d" % self.env.now)
+            self.log("Sending a message to the mixnet")
             self.env.process(self.p2p.broadcast(msg))
 
     def create_message(self) -> SphinxPacket:
@@ -39,7 +39,7 @@ class Node:
         Creates a message using the Sphinx format
         @return:
         """
-        mixes = self.p2p.get_nodes(self.N_MIXES_IN_PATH)
+        mixes = self.p2p.get_nodes(self.num_mix_layers)
         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]

mixnet/v2/sim/simulation.py

@@ -5,10 +5,10 @@ from p2p import P2p
 
 
 class Simulation:
-    def __init__(self):
+    def __init__(self, num_nodes: int, num_mix_layers: int):
         self.env = simpy.Environment()
         self.p2p = P2p(self.env)
-        self.nodes = [Node(i, self.env, self.p2p) for i in range(2)]
+        self.nodes = [Node(i, self.env, self.p2p, num_mix_layers) for i in range(num_nodes)]
         self.p2p.add_node(self.nodes)
 
     def run(self, until):