From 781ccbddc8d00c32eb7d8596a8fa1639f9e825aa Mon Sep 17 00:00:00 2001 From: Youngjoon Lee <5462944+youngjoon-lee@users.noreply.github.com> Date: Sun, 12 May 2024 19:30:04 +0900 Subject: [PATCH] working version of mixnet behavior simulation --- mixnet/v2/sim/main.py | 2 +- mixnet/v2/sim/message.py | 36 ------------------ mixnet/v2/sim/node.py | 65 ++++++++++++++++++++++----------- mixnet/v2/sim/p2p.py | 12 ++++-- mixnet/v2/sim/simulation.py | 2 +- mixnet/v2/sim/sphinx.py | 73 +++++++++++++++++++++++++++++++++++++ 6 files changed, 127 insertions(+), 63 deletions(-) delete mode 100644 mixnet/v2/sim/message.py create mode 100644 mixnet/v2/sim/sphinx.py diff --git a/mixnet/v2/sim/main.py b/mixnet/v2/sim/main.py index f12f198..7224379 100644 --- a/mixnet/v2/sim/main.py +++ b/mixnet/v2/sim/main.py @@ -2,5 +2,5 @@ from mixnet.v2.sim.simulation import Simulation if __name__ == "__main__": sim = Simulation() - sim.run(10) + sim.run(30) print("Simulation complete!") \ No newline at end of file diff --git a/mixnet/v2/sim/message.py b/mixnet/v2/sim/message.py deleted file mode 100644 index 9922539..0000000 --- a/mixnet/v2/sim/message.py +++ /dev/null @@ -1,36 +0,0 @@ -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]) \ No newline at end of file diff --git a/mixnet/v2/sim/node.py b/mixnet/v2/sim/node.py index 03a134f..b2d35b8 100644 --- a/mixnet/v2/sim/node.py +++ b/mixnet/v2/sim/node.py @@ -1,56 +1,77 @@ import random import simpy +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey, X25519PublicKey -from mixnet.v2.sim.message import Message +from mixnet.v2.sim.sphinx import SphinxPacket, Attachment from mixnet.v2.sim.p2p import P2p class Node: - N_MIXES_IN_PATH = 3 + N_MIXES_IN_PATH = 2 - def __init__(self, id: str, env: simpy.Environment, p2p: P2p): + def __init__(self, id: int, 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.private_key = X25519PrivateKey.generate() + self.public_key = self.private_key.public_key() self.action = self.env.process(self.send_message()) def send_message(self): """ Creates/encapsulate a message and send it to the network through the mixnet """ - while True: + # while True: + if self.id == 0: msg = self.create_message() 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: + def create_message(self) -> SphinxPacket: """ Creates a message using the Sphinx format @return: """ 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) + 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] + return SphinxPacket(public_keys, incentive_txs, b"Hello, world!") - def receive_message(self, msg: bytes): + def receive_message(self, msg: SphinxPacket | bytes): """ Receives a message from the network, processes it, 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.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 + # simulating network latency + yield self.env.timeout(random.randint(0, 3)) + + if isinstance(msg, SphinxPacket): + msg, incentive_tx = msg.unwrap(self.private_key) + if self.is_my_incentive_tx(incentive_tx): + self.log("Receiving SphinxPacket. It's mine!") + if msg.is_all_unwrapped(): + self.env.process(self.p2p.broadcast(msg.payload)) + else: + # TODO: use Poisson delay + yield self.env.timeout(random.randint(0, 5)) + self.env.process(self.p2p.broadcast(msg)) + else: + self.log("Receiving SphinxPacket, but not mine") + else: + self.log("Received original message: %s" % msg) + + @classmethod + def create_incentive_tx(cls, mix_public_key: X25519PublicKey) -> Attachment: + return Attachment( + mix_public_key.public_bytes(encoding=serialization.Encoding.Raw, format=serialization.PublicFormat.Raw)) + + def is_my_incentive_tx(self, tx: Attachment) -> bool: + 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 diff --git a/mixnet/v2/sim/p2p.py b/mixnet/v2/sim/p2p.py index 5d1fd62..9397830 100644 --- a/mixnet/v2/sim/p2p.py +++ b/mixnet/v2/sim/p2p.py @@ -2,6 +2,8 @@ import random import simpy +from mixnet.v2.sim.sphinx import SphinxPacket + class P2p: def __init__(self, env: simpy.Environment): @@ -11,12 +13,16 @@ class P2p: def add_node(self, nodes): self.nodes.extend(nodes) - def broadcast(self, msg): - print("Broadcasting a message at time %d" % self.env.now) + # TODO: This should accept only bytes, but SphinxPacket is also accepted until we implement the Sphinx serde + def broadcast(self, msg: SphinxPacket | bytes): + self.log("Broadcasting a msg") yield self.env.timeout(1) # TODO: gossipsub or something similar for node in self.nodes: self.env.process(node.receive_message(msg)) def get_nodes(self, n: int): - return random.choices(self.nodes, k=n) + 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 diff --git a/mixnet/v2/sim/simulation.py b/mixnet/v2/sim/simulation.py index 87cbbb4..1d3a422 100644 --- a/mixnet/v2/sim/simulation.py +++ b/mixnet/v2/sim/simulation.py @@ -8,7 +8,7 @@ class Simulation: def __init__(self): self.env = simpy.Environment() self.p2p = P2p(self.env) - self.nodes = [Node(str(i), self.env, self.p2p) for i in range(2)] + self.nodes = [Node(i, self.env, self.p2p) for i in range(2)] self.p2p.add_node(self.nodes) def run(self, until): diff --git a/mixnet/v2/sim/sphinx.py b/mixnet/v2/sim/sphinx.py new file mode 100644 index 0000000..dafe75c --- /dev/null +++ b/mixnet/v2/sim/sphinx.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +from copy import deepcopy + +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PublicKey, X25519PrivateKey + + +class SphinxPacket: + def __init__(self, public_keys: list[X25519PublicKey], attachments: list[Attachment], payload: bytes): + assert len(public_keys) == len(attachments) + ephemeral_private_key = X25519PrivateKey.generate() + ephemeral_public_key = ephemeral_private_key.public_key() + shared_keys = [SharedSecret(ephemeral_private_key, pk) for pk in public_keys] + self.header = SphinxHeader(ephemeral_public_key, shared_keys, attachments) + self.payload = payload # TODO: encrypt payload + + def __bytes__(self): + return bytes(self.header) + self.payload + + def size(self) -> int: + return len(bytes(self)) + + def unwrap(self, private_key: X25519PrivateKey) -> tuple[SphinxPacket, Attachment]: + packet = deepcopy(self) + attachment = packet.header.unwrap_inplace(private_key) + # TODO: decrypt packet.payload + return packet, attachment + + def is_all_unwrapped(self) -> bool: + return self.header.is_all_unwrapped() + + +class SphinxHeader: + DUMMY_MAC = b'\xFF' * 16 + + def __init__(self, ephemeral_public_key: X25519PublicKey, shared_keys: list[SharedSecret], + attachments: list[Attachment]): + assert len(shared_keys) == len(attachments) + self.ephemeral_public_key = ephemeral_public_key.public_bytes(encoding=serialization.Encoding.Raw, format=serialization.PublicFormat.Raw) + self.attachments = attachments # TODO: encapsulation using node_keys + + def __bytes__(self): + return b"".join([self.ephemeral_public_key] + [bytes(att) + self.DUMMY_MAC for att in self.attachments]) + + def unwrap_inplace(self, private_key: X25519PrivateKey) -> Attachment: + # TODO: shared_secret = SharedSecret(private_key, header.ephemeral_public_key) + attachment = self.attachments.pop(0) + self.attachments.append(Attachment(bytes(len(bytes(attachment))))) # append a dummy attachment + return attachment + + def is_all_unwrapped(self) -> bool: + # true if the first attachment is a dummy + return self.attachments[0] == Attachment(bytes(len(bytes(self.attachments[0])))) + + +class SharedSecret: + def __init__(self, private_key: X25519PrivateKey, public_key: X25519PublicKey): + self.key = private_key.exchange(public_key) # 32 bytes + + def __bytes__(self): + return self.key + + +class Attachment: + def __init__(self, data: bytes): + self.data = data + + def __bytes__(self): + return self.data + + def __eq__(self, other): + return bytes(self) == bytes(other) \ No newline at end of file