refactor adversary

mixnet/v2/sim/adversary.py

@@ -0,0 +1,38 @@
+import math
+from collections import defaultdict
+
+import simpy
+from simpy.core import SimTime
+
+from config import Config
+from sphinx import SphinxPacket
+
+
+class Adversary:
+    def __init__(self, env: simpy.Environment, config: Config):
+        self.env = env
+        self.config = config
+        self.message_sizes = []
+        self.senders_around_interval = defaultdict(int)
+        self.mixed_msgs_per_window = []
+        self.env.process(self.update_observation_window())
+
+    def inspect_message_size(self, msg: SphinxPacket | bytes):
+        self.message_sizes.append(len(msg))
+
+    def observe_incoming_message(self, node):
+        self.mixed_msgs_per_window[-1][node] += 1
+
+    def observe_outgoing_message(self, node):
+        self.mixed_msgs_per_window[-1][node] -= 1
+        if self.is_around_message_interval(self.env.now):
+            self.senders_around_interval[node] += 1
+
+    def is_around_message_interval(self, time: SimTime):
+        now_frac, now_int = math.modf(time)
+        return now_int % self.config.message_interval == 0 and now_frac <= self.config.max_message_prep_time
+
+    def update_observation_window(self):
+        while True:
+            self.mixed_msgs_per_window.append(defaultdict(int))
+            yield self.env.timeout(self.config.io_observation_window)

mixnet/v2/sim/analysis.py

@@ -15,13 +15,13 @@ class Analysis:
         self.mixed_messages_per_node_over_time()
 
     def message_size_distribution(self):
-        df = pd.DataFrame(self.sim.p2p.message_sizes, columns=["message_size"])
+        df = pd.DataFrame(self.sim.p2p.adversary.message_sizes, columns=["message_size"])
         print(df.describe())
 
     def messages_emitted_around_interval(self):
         df = pd.DataFrame(
             [(node.id, cnt, node.id < len(self.sim.config.real_message_prob_weights))
-             for node, cnt in self.sim.p2p.senders_around_interval.items()],
+             for node, cnt in self.sim.p2p.adversary.senders_around_interval.items()],
             columns=["node_id", "msg_count", "expected"]
         )
         plt.figure(figsize=(10, 6))
@@ -34,7 +34,7 @@ class Analysis:
 
     def mixed_messages_per_node_over_time(self):
         dataframes = []
-        for mixed_msgs_per_node in self.sim.p2p.mixed_msgs_per_window:
+        for mixed_msgs_per_node in self.sim.p2p.adversary.mixed_msgs_per_window:
             df = pd.DataFrame([(node.id, cnt) for node, cnt in mixed_msgs_per_node.items()],
                               columns=["node_id", "msg_count"])
             dataframes.append(df)

mixnet/v2/sim/main.py

@@ -1,7 +1,7 @@
 import argparse
 
 from config import Config
-from mixnet.v2.sim.analysis import Analysis
+from analysis import Analysis
 from simulation import Simulation
 
 if __name__ == "__main__":

mixnet/v2/sim/p2p.py

@@ -1,10 +1,8 @@
-import math
 import random
-from collections import defaultdict
 
 import simpy
-from simpy.core import SimTime
 
+from adversary import Adversary
 from config import Config
 from sphinx import SphinxPacket
 
@@ -14,12 +12,7 @@ class P2p:
         self.env = env
         self.config = config
         self.nodes = []
-        # The followings are for an adversary.
-        # TODO: Move these to a separate class `Adversary`.
-        self.message_sizes = []
-        self.senders_around_interval = defaultdict(int)
-        self.mixed_msgs_per_window = []
-        self.env.process(self.update_observation_window())
+        self.adversary = Adversary(env, config)
 
     def add_node(self, nodes):
         self.nodes.extend(nodes)
@@ -33,10 +26,8 @@ class P2p:
         self.log("Broadcasting a msg: %d bytes" % len(msg))
 
         # Adversary
-        self.message_sizes.append(len(msg))
-        self.mixed_msgs_per_window[-1][sender] -= 1
-        if self.is_around_message_interval(self.env.now):
-            self.senders_around_interval[sender] += 1
+        self.adversary.inspect_message_size(msg)
+        self.adversary.observe_outgoing_message(sender)
 
         # Yield 0 to ensure that the broadcast is done in the same time step.
         # Without any yield, SimPy complains that the broadcast func is not a generator.
@@ -50,19 +41,8 @@ class P2p:
         # simulate network latency
         yield self.env.timeout(random.uniform(0, self.config.max_network_latency))
 
-        self.mixed_msgs_per_window[-1][node] += 1
+        self.adversary.observe_incoming_message(node)
         self.env.process(node.receive_message(msg))
 
-    # TODO: Move to a separate class `Adversary`.
-    def is_around_message_interval(self, time: SimTime):
-        now_frac, now_int = math.modf(time)
-        return now_int % self.config.message_interval == 0 and now_frac <= self.config.max_message_prep_time
-
-    # TODO: Move to a separate class `Adversary`.
-    def update_observation_window(self):
-        while True:
-            self.mixed_msgs_per_window.append(defaultdict(int))
-            yield self.env.timeout(self.config.io_observation_window)
-
     def log(self, msg):
         print("P2P at %g: %s" % (self.env.now, msg))