From 326aeedb2c8cea4e92e192e2ad69967fdf0eae43 Mon Sep 17 00:00:00 2001
From: Youngjoon Lee <5462944+youngjoon-lee@users.noreply.github.com>
Date: Tue, 14 May 2024 14:36:05 +0900
Subject: [PATCH] define node params

---
 mixnet/v2/sim/main.py       |  8 +++++++-
 mixnet/v2/sim/node.py       | 17 +++++++++++++----
 mixnet/v2/sim/p2p.py        |  2 +-
 mixnet/v2/sim/simulation.py |  4 ++--
 4 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/mixnet/v2/sim/main.py b/mixnet/v2/sim/main.py
index e4e9e0d..d13998b 100644
--- a/mixnet/v2/sim/main.py
+++ b/mixnet/v2/sim/main.py
@@ -4,6 +4,7 @@ import matplotlib.pyplot as plt
 import pandas as pd
 import seaborn
 
+from node import Node
 from simulation import Simulation
 
 if __name__ == "__main__":
@@ -11,9 +12,14 @@ if __name__ == "__main__":
     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")
+    parser.add_argument("--message-interval", type=int, default=1, help="Message emission interval")
+    parser.add_argument("--message-prob", type=float, default=0.2, help="Message emission probability per interval")
+    parser.add_argument("--max-message-prep-time", type=float, default=0.3, help="Max preparation time before sending a message")
     args = parser.parse_args()
 
-    sim = Simulation(args.num_nodes, args.num_mix_layers)
+    node_params = Node.Parameters(args.num_mix_layers, args.message_interval, args.message_prob, args.max_message_prep_time)
+
+    sim = Simulation(args.num_nodes, node_params)
     sim.run(until=args.running_time)
 
     df = pd.DataFrame(sim.p2p.message_sizes, columns=["message_size"])
diff --git a/mixnet/v2/sim/node.py b/mixnet/v2/sim/node.py
index 50f0621..a1066e8 100644
--- a/mixnet/v2/sim/node.py
+++ b/mixnet/v2/sim/node.py
@@ -1,4 +1,5 @@
 import random
+from dataclasses import dataclass
 
 import simpy
 from cryptography.hazmat.primitives import serialization
@@ -14,13 +15,21 @@ class Node:
     COVER_PAYLOAD = b"COVER"
     PADDING_SEPARATOR = b'\x01'
 
-    def __init__(self, id: int, env: simpy.Environment, p2p: P2p, num_mix_layers: int):
+    @dataclass
+    class Parameters:
+        num_mix_layers: int
+        message_interval: int
+        message_prob: float
+        max_message_prep_time: int
+
+
+    def __init__(self, id: int, env: simpy.Environment, p2p: P2p, params: Parameters):
         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.params = params
         self.action = self.env.process(self.send_message())
 
     def send_message(self):
@@ -39,7 +48,7 @@ class Node:
         Creates a message using the Sphinx format
         @return:
         """
-        mixes = self.p2p.get_nodes(self.num_mix_layers)
+        mixes = self.p2p.get_nodes(self.params.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]
@@ -88,4 +97,4 @@ class Node:
         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))
\ No newline at end of file
+        print("Node:%d at %g: %s" % (self.id, self.env.now, msg))
\ No newline at end of file
diff --git a/mixnet/v2/sim/p2p.py b/mixnet/v2/sim/p2p.py
index 792e943..452aff7 100644
--- a/mixnet/v2/sim/p2p.py
+++ b/mixnet/v2/sim/p2p.py
@@ -28,4 +28,4 @@ class P2p:
         return random.sample(self.nodes, n)
 
     def log(self, msg):
-        print("P2P at %d: %s" % (self.env.now, msg))
\ No newline at end of file
+        print("P2P at %g: %s" % (self.env.now, msg))
\ No newline at end of file
diff --git a/mixnet/v2/sim/simulation.py b/mixnet/v2/sim/simulation.py
index 6eac077..663d0dc 100644
--- a/mixnet/v2/sim/simulation.py
+++ b/mixnet/v2/sim/simulation.py
@@ -5,10 +5,10 @@ from p2p import P2p
 
 
 class Simulation:
-    def __init__(self, num_nodes: int, num_mix_layers: int):
+    def __init__(self, num_nodes: int, node_params: Node.Parameters):
         self.env = simpy.Environment()
         self.p2p = P2p(self.env)
-        self.nodes = [Node(i, self.env, self.p2p, num_mix_layers) for i in range(num_nodes)]
+        self.nodes = [Node(i, self.env, self.p2p, node_params) for i in range(num_nodes)]
         self.p2p.add_node(self.nodes)
 
     def run(self, until):