Initial commit

Socket.py

@@ -0,0 +1,118 @@
+import socket
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives.asymmetric import rsa
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.asymmetric import padding
+from cryptography.fernet import Fernet
+
+
+class Socket:
+    def __init__(self):
+        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        self.key = None
+        self.fernet = None
+
+    def send(self, sock: socket.socket, message: str):
+        if type(message) != bytes:
+            message = message.encode("Utf8")
+        if self.fernet:
+            message = self.fernet.encrypt(message)
+        try:
+            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:
+            response = sock.recv(1028)
+        except socket.error:
+            raise ConnectionError("Fail to receive")
+        else:
+            if self.fernet and fernet:
+                response = self.fernet.decrypt(response)
+            if utf8:
+                response = response.decode("Utf8")
+            return response
+
+    def set_secure_connexion(self, sock: socket.socket):
+        if not self.key:
+            self.key = Fernet.generate_key()
+            self.fernet = Fernet(self.key)
+
+        public_key = serialization.load_pem_public_key(
+            self.receive(sock, False, False),
+            backend=default_backend()
+        )
+
+        sock.send(
+            public_key.encrypt(
+                self.key,
+                padding.OAEP(
+                    mgf=padding.MGF1(algorithm=hashes.SHA256()),
+                    algorithm=hashes.SHA256(),
+                    label=None
+                )
+            )
+        )
+
+        if self.receive(sock) == "Key receive":
+            self.send(sock, "Key valid")
+            return True
+        else:
+            raise ConnectionError("Invalid key !")
+
+    def get_secure_connexion(self, sock: socket.socket):
+        private_key = rsa.generate_private_key(
+            public_exponent=65537,
+            key_size=2048,
+            backend=default_backend()
+        )
+
+        public_key = private_key.public_key()
+
+        pem = public_key.public_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PublicFormat.SubjectPublicKeyInfo
+        )
+
+        self.send(sock, pem)
+
+        self.key = private_key.decrypt(
+            self.receive(sock, False, False),
+            padding.OAEP(
+                mgf=padding.MGF1(algorithm=hashes.SHA256()),
+                algorithm=hashes.SHA256(),
+                label=None
+            )
+        )
+
+        self.fernet = Fernet(self.key)
+
+        self.send(sock, "Key receive")
+        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:
+            sock.connect((host, port))
+        except socket.error:
+            raise ConnectionError("Can't connect to server !")
+        else:
+            self.get_secure_connexion(self.socket)
+            self.send(sock, f"Client service ID: {service_id}")
+            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()
+        self.set_secure_connexion(connexion)
+        if self.receive(connexion) == f"Client service ID: {service_id}":
+            self.send(connexion, f"Server service ID: {service_id}")
+            return connexion
+        else:
+            raise ConnectionError("Invalid client service ID !")