From f28454c2b2da7aee95c10695863ca2b14cfa9efe Mon Sep 17 00:00:00 2001 From: study-faraphel Date: Sun, 5 Jan 2025 12:46:21 +0100 Subject: [PATCH] added trusted / untrusted peers mechanism --- .gitignore | 4 + source/behaviors/events/AudioEvent.py | 4 +- source/behaviors/events/KeyEvent.py | 4 +- source/behaviors/events/PeerEvent.py | 8 ++ source/behaviors/events/RequestKeyEvent.py | 6 +- .../behaviors/events/base/BaseTrustedEvent.py | 18 +++++ source/behaviors/events/base/__init__.py | 1 + source/behaviors/roles/MasterRole.py | 2 +- source/error/UntrustedPeerException.py | 8 ++ source/error/__init__.py | 1 + source/managers/CommunicationManager.py | 77 ++++++++++++++++++- source/managers/EventManager.py | 7 +- source/managers/Manager.py | 7 +- source/packets/PeerPacket.py | 3 + source/structures/Peer.py | 3 + 15 files changed, 140 insertions(+), 13 deletions(-) create mode 100644 source/behaviors/events/base/BaseTrustedEvent.py create mode 100644 source/error/UntrustedPeerException.py create mode 100644 source/error/__init__.py diff --git a/.gitignore b/.gitignore index 5d95bde..477f53b 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,7 @@ +# IDE .idea cmake-build-* + +# Local assets +storage diff --git a/source/behaviors/events/AudioEvent.py b/source/behaviors/events/AudioEvent.py index 3507dbc..26e5551 100644 --- a/source/behaviors/events/AudioEvent.py +++ b/source/behaviors/events/AudioEvent.py @@ -2,11 +2,13 @@ from source import packets from source.behaviors.events import base -class AudioEvent(base.BaseEvent): +class AudioEvent(base.BaseTrustedEvent): """ Event reacting to receiving audio data. """ def handle(self, packet: packets.AudioPacket, address: tuple): + super().handle(packet, address) + # add the audio chunk to the buffer of audio packet to play self.manager.audio.add_audio(packet) diff --git a/source/behaviors/events/KeyEvent.py b/source/behaviors/events/KeyEvent.py index d38c04d..22ee651 100644 --- a/source/behaviors/events/KeyEvent.py +++ b/source/behaviors/events/KeyEvent.py @@ -3,12 +3,14 @@ from source.behaviors import roles from source.behaviors.events import base -class KeyEvent(base.BaseEvent): +class KeyEvent(base.BaseTrustedEvent): """ Event reacting to a machine sending us their secret key """ def handle(self, packet: packets.KeyPacket, address: tuple): + super().handle(packet, address) + # check if we are a slave if not isinstance(self.manager.role.current, roles.SlaveRole): return diff --git a/source/behaviors/events/PeerEvent.py b/source/behaviors/events/PeerEvent.py index 881e51c..9540a11 100644 --- a/source/behaviors/events/PeerEvent.py +++ b/source/behaviors/events/PeerEvent.py @@ -8,8 +8,16 @@ class PeerEvent(base.BaseEvent): """ def handle(self, packet: packets.PeerPacket, address: tuple): + # ignore peers with a banned key + if self.manager.communication.is_peer_banned(packet.public_key): + return + + # TODO(Faraphel): SHOULD NOT BE TRUSTED AUTOMATICALLY ! + self.manager.communication.trust_peer(packet.public_key) + # update our peers database to add new peer information self.manager.peer.peers[address] = structures.Peer( public_key=packet.public_key, master=packet.master, + trusted=self.manager.communication.is_peer_trusted(packet.public_key) ) diff --git a/source/behaviors/events/RequestKeyEvent.py b/source/behaviors/events/RequestKeyEvent.py index ee5c284..7ce64c3 100644 --- a/source/behaviors/events/RequestKeyEvent.py +++ b/source/behaviors/events/RequestKeyEvent.py @@ -4,12 +4,14 @@ from source.behaviors.events import base from source.utils.crypto.type import CipherType -class RequestKeyEvent(base.BaseEvent): +class RequestKeyEvent(base.BaseTrustedEvent): """ Event reacting to a machine trying to get our secret symmetric key for secure communication """ def handle(self, packet: packets.RequestKeyPacket, address: tuple): + super().handle(packet, address) + # check if we are a master if not isinstance(self.manager.role.current, roles.MasterRole): return @@ -18,5 +20,3 @@ class RequestKeyEvent(base.BaseEvent): packet = packets.KeyPacket(self.manager.role.current.secret_key) # send it back to the slave self.manager.communication.send(packet, CipherType.RSA, address) - - # TODO(Faraphel): check if we trust the slave ? in a list of trusted public / private key ? diff --git a/source/behaviors/events/base/BaseTrustedEvent.py b/source/behaviors/events/base/BaseTrustedEvent.py new file mode 100644 index 0000000..05c26fe --- /dev/null +++ b/source/behaviors/events/base/BaseTrustedEvent.py @@ -0,0 +1,18 @@ +import abc + +from source import packets +from source.behaviors.events.base import BaseEvent +from source.error import UntrustedPeerException + + +class BaseTrustedEvent(BaseEvent, abc.ABC): + """ + Event that can only be triggered if the distant peer is trusted + """ + + def handle(self, packet: packets.base.BasePacket, address: tuple) -> None: + # get the peer that sent the message + peer = self.manager.peer.peers.get(address) + # check if it is trusted + if peer is None or not peer.trusted: + raise UntrustedPeerException(peer) diff --git a/source/behaviors/events/base/__init__.py b/source/behaviors/events/base/__init__.py index 1525289..cfe96e8 100644 --- a/source/behaviors/events/base/__init__.py +++ b/source/behaviors/events/base/__init__.py @@ -1 +1,2 @@ from .BaseEvent import BaseEvent +from .BaseTrustedEvent import BaseTrustedEvent diff --git a/source/behaviors/roles/MasterRole.py b/source/behaviors/roles/MasterRole.py index 54d24b6..f8d70ae 100644 --- a/source/behaviors/roles/MasterRole.py +++ b/source/behaviors/roles/MasterRole.py @@ -27,7 +27,7 @@ class MasterRole(base.BaseActiveRole): # prepare the audio file that will be streamed # TODO(Faraphel): use another audio source - self.audio = pydub.AudioSegment.from_file("../assets/Queen - Another One Bites the Dust.mp3") + self.audio = pydub.AudioSegment.from_file("./assets/Queen - Another One Bites the Dust.mp3") self.play_time = datetime.now() # calculate the number of bytes per milliseconds in the audio diff --git a/source/error/UntrustedPeerException.py b/source/error/UntrustedPeerException.py new file mode 100644 index 0000000..2eb9c00 --- /dev/null +++ b/source/error/UntrustedPeerException.py @@ -0,0 +1,8 @@ +import typing + +from source import structures + + +class UntrustedPeerException(Exception): + def __init__(self, peer: typing.Optional[structures.Peer]): + super().__init__(f"Peer not trusted: {peer}") diff --git a/source/error/__init__.py b/source/error/__init__.py new file mode 100644 index 0000000..d569e8b --- /dev/null +++ b/source/error/__init__.py @@ -0,0 +1 @@ +from .UntrustedPeerException import UntrustedPeerException diff --git a/source/managers/CommunicationManager.py b/source/managers/CommunicationManager.py index a511da8..b358ae8 100644 --- a/source/managers/CommunicationManager.py +++ b/source/managers/CommunicationManager.py @@ -1,3 +1,5 @@ +import hashlib +import json import socket import typing import zlib @@ -35,8 +37,27 @@ class CommunicationManager: # create a dictionary to hold the types of packets and their headers. self.packet_types: bidict.bidict[bytes, typing.Type[packets.base.BasePacket]] = bidict.bidict() - # create a private and public key for RSA communication - self.private_key, self.public_key = rsa_create_key_pair() + # load or create a private and public key for asymmetric communication + private_key_path = self.manager.storage / "private_key.der" + public_key_path = self.manager.storage / "public_key.der" + + if public_key_path.exists() and private_key_path.exists(): + self.private_key = private_key_path.read_bytes() + self.public_key = public_key_path.read_bytes() + else: + self.private_key, self.public_key = rsa_create_key_pair() + private_key_path.write_bytes(self.private_key) + public_key_path.write_bytes(self.public_key) + + self._trusted_peers_path = self.manager.storage / "trusted-peers.json" + self._trusted_peers: set[str] = set() + if self._trusted_peers_path.exists(): + self._trusted_peers = set(json.loads(self._trusted_peers_path.read_text())) + + self._banned_peers_path = self.manager.storage / "banned-peers.json" + self._banned_peers: set[str] = set() + if self._banned_peers_path.exists(): + self._banned_peers = set(json.loads(self._banned_peers_path.read_text())) def __del__(self): # close the socket @@ -229,3 +250,55 @@ class CommunicationManager: # no matching interfaces have been found return False + + def save_trusted_peers(self) -> None: + """ + Save the list of trusted peers + """ + + self._trusted_peers_path.write_text(json.dumps(list(self._trusted_peers))) + + def save_banned_peers(self) -> None: + """ + Save the list of banned peers + """ + + self._banned_peers_path.write_text(json.dumps(list(self._banned_peers))) + + def trust_peer(self, public_key: bytes) -> None: + """ + Mark a peer as trusted for future connexion + Automatically save it to a file + :param public_key: the public key of the peer + """ + + self._trusted_peers.add(hashlib.sha256(public_key).hexdigest()) + self.save_trusted_peers() + + def ban_peer(self, public_key: bytes) -> None: + """ + Ban a peer from being used for any future connexion + Automatically save it to a file + :param public_key: the public key of the peer + """ + + self._banned_peers.add(hashlib.sha256(public_key).hexdigest()) + self.save_banned_peers() + + def is_peer_trusted(self, public_key: bytes) -> bool: + """ + Determinate is a peer is trusted or not + :param public_key: the public key of the peer + :return: True if the peer is trusted, False otherwise + """ + + return hashlib.sha256(public_key).hexdigest() in self._trusted_peers + + def is_peer_banned(self, public_key: bytes) -> bool: + """ + Determinate if a peer is banned or not + :param public_key: the public key of the peer + :return: True if the peer is banned, False otherwise + """ + + return hashlib.sha256(public_key).hexdigest() in self._banned_peers diff --git a/source/managers/EventManager.py b/source/managers/EventManager.py index a3ddf35..fed3677 100644 --- a/source/managers/EventManager.py +++ b/source/managers/EventManager.py @@ -4,6 +4,7 @@ import warnings from source import packets from source.behaviors import events +from source.error import UntrustedPeerException from source.managers import Manager @@ -55,8 +56,8 @@ class EventManager: # give it to the event handler self.manager.event.handle(packet, address) - except KeyboardInterrupt: - break + except UntrustedPeerException: + print("Ignored: untrusted peer.") - except: + except Exception: # NOQA warnings.warn(traceback.format_exc()) diff --git a/source/managers/Manager.py b/source/managers/Manager.py index 1cc494c..8f52462 100644 --- a/source/managers/Manager.py +++ b/source/managers/Manager.py @@ -1,8 +1,8 @@ import threading +from pathlib import Path -from source import packets, structures +from source import packets from source.behaviors import events -from source.utils.dict import TimestampedDict class Manager: @@ -13,6 +13,9 @@ class Manager: def __init__(self, interface: str): from . import CommunicationManager, EventManager, RoleManager, AudioManager, PeerManager + self.storage = Path("./storage/") + self.storage.mkdir(exist_ok=True) + # communication manager self.communication = CommunicationManager(self, interface) self.communication.register_packet_type(b"DISC", packets.DiscoveryPacket) diff --git a/source/packets/PeerPacket.py b/source/packets/PeerPacket.py index 80903c1..5acc5e5 100644 --- a/source/packets/PeerPacket.py +++ b/source/packets/PeerPacket.py @@ -17,6 +17,9 @@ class PeerPacket(base.BasePacket): # is the machine a master master: bool = dataclasses.field() + # TODO(Faraphel): share our trusted / banned peers with the other peer so that only one machine need to trust / ban it + # to propagate it to the whole network ? + def pack(self) -> bytes: return msgpack.packb(( self.public_key, diff --git a/source/structures/Peer.py b/source/structures/Peer.py index da4dce0..ade4989 100644 --- a/source/structures/Peer.py +++ b/source/structures/Peer.py @@ -13,5 +13,8 @@ class Peer: # secret symmetric key secret_key: Optional[bytes] = dataclasses.field(default=None) + # is the machine trusted + trusted: bool = dataclasses.field(default=False) + # when did the peer last communication with us occurred last_interaction: datetime = dataclasses.field(default_factory=datetime.now)