From 4f6fb3077940e2b3afc76a87929cd51e8d4abd0d Mon Sep 17 00:00:00 2001
From: Youngjoon Lee <5462944+youngjoon-lee@users.noreply.github.com>
Date: Tue, 14 May 2024 21:19:46 +0900
Subject: [PATCH] config.yaml

---
 mixnet/v2/sim/README.md     | 20 +++-----------------
 mixnet/v2/sim/config.py     | 11 +++++++++++
 mixnet/v2/sim/config.yaml   |  6 ++++++
 mixnet/v2/sim/main.py       | 19 +++++++++----------
 mixnet/v2/sim/node.py       | 16 ++++------------
 mixnet/v2/sim/simulation.py | 10 ++++++----
 6 files changed, 39 insertions(+), 43 deletions(-)
 create mode 100644 mixnet/v2/sim/config.py
 create mode 100644 mixnet/v2/sim/config.yaml

diff --git a/mixnet/v2/sim/README.md b/mixnet/v2/sim/README.md
index 882ae00..ff89c70 100644
--- a/mixnet/v2/sim/README.md
+++ b/mixnet/v2/sim/README.md
@@ -2,24 +2,10 @@
 
 ## How to run
 
-Make sure that all dependencies specified in the `requirements.txt` in the project root, and run the following command to run the simulation:
+First, make sure that all dependencies specified in the `requirements.txt` in the project root.
+Then, configure parameters in the [config.yaml](./config.yaml), and run the following command:
 ```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)
+python main.py --config ./config.yaml
 ```
 
 TODO: Add more details
diff --git a/mixnet/v2/sim/config.py b/mixnet/v2/sim/config.py
new file mode 100644
index 0000000..0c6328f
--- /dev/null
+++ b/mixnet/v2/sim/config.py
@@ -0,0 +1,11 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class Config:
+    running_time: int
+    num_nodes: int
+    num_mix_layers: int
+    message_interval: int
+    message_prob: float
+    max_message_prep_time: float
diff --git a/mixnet/v2/sim/config.yaml b/mixnet/v2/sim/config.yaml
new file mode 100644
index 0000000..e692487
--- /dev/null
+++ b/mixnet/v2/sim/config.yaml
@@ -0,0 +1,6 @@
+running_time: 30
+num_nodes: 5
+num_mix_layers: 3
+message_interval: 1
+message_prob: 0.2
+max_message_prep_time: 0.3
\ No newline at end of file
diff --git a/mixnet/v2/sim/main.py b/mixnet/v2/sim/main.py
index d13998b..460c58b 100644
--- a/mixnet/v2/sim/main.py
+++ b/mixnet/v2/sim/main.py
@@ -1,26 +1,25 @@
 import argparse
 
+import dacite
 import matplotlib.pyplot as plt
 import pandas as pd
 import seaborn
+import yaml
 
-from node import Node
+from config import Config
 from simulation import Simulation
 
 if __name__ == "__main__":
     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")
-    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")
+    parser.add_argument("--config", type=str, required=True, help="Configuration file path")
     args = parser.parse_args()
 
-    node_params = Node.Parameters(args.num_mix_layers, args.message_interval, args.message_prob, args.max_message_prep_time)
+    with open(args.config, "r") as f:
+        config = yaml.safe_load(f)
+    config = dacite.from_dict(data_class=Config, data=config)
 
-    sim = Simulation(args.num_nodes, node_params)
-    sim.run(until=args.running_time)
+    sim = Simulation(config)
+    sim.run()
 
     df = pd.DataFrame(sim.p2p.message_sizes, columns=["message_size"])
     print(df.describe())
diff --git a/mixnet/v2/sim/node.py b/mixnet/v2/sim/node.py
index a1066e8..1b0b46c 100644
--- a/mixnet/v2/sim/node.py
+++ b/mixnet/v2/sim/node.py
@@ -1,10 +1,10 @@
 import random
-from dataclasses import dataclass
 
 import simpy
 from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey, X25519PublicKey
 
+from config import Config
 from sphinx import SphinxPacket, Attachment
 from p2p import P2p
 
@@ -15,21 +15,13 @@ class Node:
     COVER_PAYLOAD = b"COVER"
     PADDING_SEPARATOR = b'\x01'
 
-    @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):
+    def __init__(self, id: int, env: simpy.Environment, p2p: P2p, config: Config):
         self.id = id
         self.env = env
         self.p2p = p2p
         self.private_key = X25519PrivateKey.generate()
         self.public_key = self.private_key.public_key()
-        self.params = params
+        self.config = config
         self.action = self.env.process(self.send_message())
 
     def send_message(self):
@@ -48,7 +40,7 @@ class Node:
         Creates a message using the Sphinx format
         @return:
         """
-        mixes = self.p2p.get_nodes(self.params.num_mix_layers)
+        mixes = self.p2p.get_nodes(self.config.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]
diff --git a/mixnet/v2/sim/simulation.py b/mixnet/v2/sim/simulation.py
index 663d0dc..f82621b 100644
--- a/mixnet/v2/sim/simulation.py
+++ b/mixnet/v2/sim/simulation.py
@@ -1,15 +1,17 @@
 import simpy
 
+from config import Config
 from node import Node
 from p2p import P2p
 
 
 class Simulation:
-    def __init__(self, num_nodes: int, node_params: Node.Parameters):
+    def __init__(self, config: Config):
+        self.config = config
         self.env = simpy.Environment()
         self.p2p = P2p(self.env)
-        self.nodes = [Node(i, self.env, self.p2p, node_params) for i in range(num_nodes)]
+        self.nodes = [Node(i, self.env, self.p2p, config) for i in range(config.num_nodes)]
         self.p2p.add_node(self.nodes)
 
-    def run(self, until):
-        self.env.run(until=until)
+    def run(self):
+        self.env.run(until=self.config.running_time)