added trusted / untrusted peers mechanism

.gitignore

@@ -1,3 +1,7 @@
+# IDE
 .idea
 cmake-build-*
+
+# Local
 assets
+storage

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)

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

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)
         )

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 ?

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)

source/behaviors/events/base/__init__.py

@@ -1 +1,2 @@
 from .BaseEvent import BaseEvent
+from .BaseTrustedEvent import BaseTrustedEvent

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

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}")

source/error/__init__.py

@@ -0,0 +1 @@
+from .UntrustedPeerException import UntrustedPeerException

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
+        # 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

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:
+            except UntrustedPeerException:
-                break
+                print("Ignored: untrusted peer.")
 
-            except:
+            except Exception:  # NOQA
                 warnings.warn(traceback.format_exc())

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)

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,

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)