A lot of comments and add encryption in option

Socket.py

@@ -8,45 +8,57 @@ from cryptography.fernet import Fernet
 
 
 class Socket:
+    """Self
+    Socket service with security system"""
     def __init__(self):
-        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        self.key = None
-        self.fernet = None
+        """Self
+        Socket service with security system"""
+        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)  # Set the socket object
+        self.key = None  # Set the key encryption
+        self.fernet = None  # Set the encryption object
 
-    def send(self, sock: socket.socket, message: str):
-        if type(message) != bytes:
+    def send(self, sock: socket.socket, message: str, encryption=True):
+        """self, sock (socket), message (str), encryption (bool)
+        Send a message with a specific socket, encrypted if available"""
+        if type(message) != bytes:  # If the message are not a bytes type, encode in Utf8
             message = message.encode("Utf8")
-        if self.fernet:
+        if self.fernet and encryption:  # Encrypt if available
             message = self.fernet.encrypt(message)
-        try:
+        try:  # Try to send, if fail raise a custom error
             sock.send(message)
         except socket.error:
             raise ConnectionError("Fail to send")
         else:
             return True
 
-    def receive(self, sock: socket.socket, utf8=True, fernet=True):
-        try:
+    def receive(self, sock: socket.socket, utf8=True, encryption=True):
+        """self, sock (socket), utf8 (bool), encryption (bool)
+        Receive a message from a socket"""
+        try:  # Try to receive, else raise a custom error
             response = sock.recv(1028)
         except socket.error:
             raise ConnectionError("Fail to receive")
         else:
-            if self.fernet and fernet:
+            if self.fernet and encryption:  # Encrypt if available
                 response = self.fernet.decrypt(response)
-            if utf8:
+            if utf8:  # decode the utf8 by default
                 response = response.decode("Utf8")
             return response
 
     def set_secure_connexion(self, sock: socket.socket):
-        if not self.key:
+        """self, sock (socket)
+        Set encryption for a connected socket"""
+        if not self.key:  # If any key generated, generate one
             self.key = Fernet.generate_key()
             self.fernet = Fernet(self.key)
 
+        # Receive the public key of the targed socket
         public_key = serialization.load_pem_public_key(
             self.receive(sock, False, False),
             backend=default_backend()
         )
 
+        # Encrypt the key and send back with private key encryption to the target socket
         sock.send(
             public_key.encrypt(
                 self.key,
@@ -58,28 +70,33 @@ class Socket:
             )
         )
 
-        if self.receive(sock) == "Key receive":
+        if self.receive(sock) == "Key receive":  # Check if the kay is valid
             self.send(sock, "Key valid")
             return True
         else:
             raise ConnectionError("Invalid key !")
 
     def get_secure_connexion(self, sock: socket.socket):
+        """self, sock (socket)
+        Get encryption for a connected socket"""
+        # Generate private key
         private_key = rsa.generate_private_key(
             public_exponent=65537,
             key_size=2048,
             backend=default_backend()
         )
 
-        public_key = private_key.public_key()
+        public_key = private_key.public_key()  # Create pubic kay
 
+        # Make a sendable public key
         pem = public_key.public_bytes(
             encoding=serialization.Encoding.PEM,
             format=serialization.PublicFormat.SubjectPublicKeyInfo
         )
 
-        self.send(sock, pem)
+        self.send(sock, pem)  # Send the public key
 
+        # Decrypt the receive key with the private key and set it
         self.key = private_key.decrypt(
             self.receive(sock, False, False),
             padding.OAEP(
@@ -89,29 +106,37 @@ class Socket:
             )
         )
 
-        self.fernet = Fernet(self.key)
+        self.fernet = Fernet(self.key)  # Set the encryption object
 
-        self.send(sock, "Key receive")
+        self.send(sock, "Key receive")  # Check if the key is valid
         if self.receive(sock) != "Key valid":
             raise ConnectionError("Invalid key !")
         return True
 
-    def connect_server(self, sock: socket.socket, host: str, port: int, service_id: int):
-        try:
+    def connect_server(self, sock: socket.socket, host: str, port: int, service_id: int, encryption=True):
+        """self, sock (socket), host (str), port (int), service_id (int), encryption (bool)
+        Connect a socket to a socket server"""
+        try:  # Try to connect, else raise a custom error
             sock.connect((host, port))
         except socket.error:
             raise ConnectionError("Can't connect to server !")
         else:
+            if encryption:  # Set secure connexion is asked
                 self.get_secure_connexion(self.socket)
-            self.send(sock, f"Client service ID: {service_id}")
+
+            self.send(sock, f"Client service ID: {service_id}")  # Check if service id match
             if self.receive(sock) != f"Server service ID: {service_id}":
                 raise ConnectionError("Server service ID")
             return True
 
-    def connect_client(self, sock: socket, service_id: int):
-        connexion, address = sock.accept()
+    def connect_client(self, sock: socket, service_id: int, encryption=True):
+        """self, sock (socket), service_id (int), encryption (bool)
+        Connect a socket server to a socket"""
+        connexion, address = sock.accept()  # Await for connexion
+        if encryption:
             self.set_secure_connexion(connexion)
-        if self.receive(connexion) == f"Client service ID: {service_id}":
+
+        if self.receive(connexion) == f"Client service ID: {service_id}":  # Check if service id match
             self.send(connexion, f"Server service ID: {service_id}")
             return connexion
         else: