refactor and better crypto

server/netwrapper.py

@@ -6,6 +6,7 @@ from Crypto.Hash import SHA512
 from Crypto.Cipher import PKCS1_OAEP
 from Crypto.Cipher import ChaCha20_Poly1305
 from Crypto.PublicKey.RSA import RsaKey
+from Crypto.Signature import pkcs1_15
 
 from netsim import network_interface
 from authentication import Authetication
@@ -24,20 +25,49 @@ class NetWrapper:
         self.homeDirectory = ""
         self.authenticationInstance = authenticationInstance
 
+    def ecryptRSAMessage(self, message: str) -> bytes:
+        cipher_rsa = PKCS1_OAEP.new(self.clientPublicKey)
+        encrypted_msg = cipher_rsa.encrypt(message.encode('UTF-8'))
+        return encrypted_msg
+
+    def signRSAHeader(self, type: str, extradata: dict) -> (bytes, bytes):
+        header = json.dumps({'type': type, 'source': self.network.own_addr}.update(extradata)).encode('UTF-8')
+        h = SHA512.new(header)
+        headersignature = pkcs1_15.new(self.serverPrivateKey).sign(h)
+        return header, headersignature
+
+    def verifyRSAHeaderSignature(self, header: bytes, headersignature: bytes) -> bool:
+        h = SHA512.new(header)
+        try:
+            pkcs1_15.new(self.clientPublicKey).verify(h, headersignature)
+            return True
+        except Exception:
+            return False
+
+    def decryptRSAMessage(self, message: bytes) -> bytes:
+        cipher_rsa = PKCS1_OAEP.new(self.serverPrivateKey)
+        return cipher_rsa.decrypt(message)
+
     def serverIdentify(self, msg: bytes) -> None:
         incommingJson = json.loads(msg.decode('UTF-8'))
+        header = json.loads(b64decode(incommingJson['header']).decode('UTF-8'))
         if incommingJson['type'] != "IDY":
             raise Exception('Wrong message type encountered')
-        self.clientAddr = incommingJson['source']
-        self.currentUser = incommingJson['username']
+        try:
+            if not self.verifyRSAHeaderSignature(b64decode(incommingJson['header']),
+                                                 b64decode(incommingJson['headersignature'])) or header[
+                'type'] != 'IDY':
+                raise Exception('Bad initial message')
+        except Exception:
+            raise Exception('Bad initial message')
+        self.clientAddr = header['source']
+        self.currentUser = header['username']
         self.currentClientPublicKey = self.clientPublicKey[self.currentUser]
-        cipher_rsa = PKCS1_OAEP.new(self.serverPrivateKey)
-        retmsg = cipher_rsa.decrypt(b64decode(incommingJson['message'])).decode('UTF-8')
-        cipher = PKCS1_OAEP.new(self.currentClientPublicKey)
-
+        retheader, retheadersignature = self.signRSAHeader("IDY", {})
+        retmsg = self.ecryptRSAMessage(b64decode(incommingJson['message']).decode('UTF-8'))
         identMsg = json.dumps(
-            {'type': 'IDY', 'source': self.network.own_addr,
-             'message': b64encode(cipher.encrypt(retmsg.encode('UTF-8'))).decode('UTF-8')}).encode(
+            {'header': retheader, 'headersignature': retheadersignature
+             'message': b64encode(retmsg).decode('UTF-8')}).encode(
             'UTF-8')
         self.network.send_msg(self.clientAddr, identMsg)
 
@@ -45,7 +75,7 @@ class NetWrapper:
         self.sendTypedMessage(message, "CMD")
 
     def sendTypedMessage(self, message: bytes, type: str) -> None:
-        if not (type == "IDY" or type == "DH" or type == "CMD"):
+        if not (type == "AUT" or type == "CMD"):
             raise Exception('Unknown message type')
         cipher = ChaCha20_Poly1305.new(key=self.cipherkey)
         header = json.dumps({'source': self.network.own_addr, 'type': type}).encode('UTF-8')
@@ -61,19 +91,23 @@ class NetWrapper:
 
     def keyExchange(self) -> None:
         dh = pyDH.DiffieHellman()
-        cipher = PKCS1_OAEP.new(self.currentClientPublicKey)
-        mypubkey = b64encode(cipher.encrypt(str(dh.gen_public_key()).encode('UTF-8'))).decode('UTF-8')
-        jsonmsg = json.dumps({'type': 'DH', 'source': self.network.own_addr, 'message': mypubkey}).encode('UTF-8')
+        mypubkey = self.ecryptRSAMessage(str(dh.gen_public_key()))
+        header, headersignature = self.signRSAHeader("DH")
+        jsonmsg = json.dumps(
+            {'header': b64encode(header).decode('UTF-8'), 'headersignature': b64encode(headersignature).decode('UTF-8'),
+             'message': mypubkey}).encode('UTF-8')
         self.network.send_msg(self.clientAddr, jsonmsg)
-        decodedmsg = {'source': '', 'type': ''}
-        while not (decodedmsg['source'] == self.clientAddr and decodedmsg['type'] == 'DH'):
         status, msg = self.network.receive_msg(blocking=True)
         if not status:
             raise Exception('Network error during connection.')
         decodedmsg = json.loads(msg.decode('UTF-8'))
-        cipher_rsa = PKCS1_OAEP.new(self.serverPrivateKey)
-        serverpubkey = int(cipher_rsa.decrypt(b64decode(decodedmsg['message'])).decode('UTF-8'))
-        cipherkey = dh.gen_shared_key(serverpubkey).encode('UTF-8')
+        header = json.loads(b64decode(decodedmsg['header']).decode('UTF-8'))
+        if not self.verifyRSAHeaderSignature(b64decode(decodedmsg['header']),
+                                             b64decode(decodedmsg['headersignature'])) or not (
+                header['source'] == self.clientAddr and header['type'] == 'DH'):
+            raise Exception('Header signature error')
+        clientpubkey = int(self.decryptRSAMessage(b64decode(decodedmsg['message'])).decode('UTF-8'))
+        cipherkey = dh.gen_shared_key(clientpubkey).encode('UTF-8')
         hasher = SHA512.new()
         hasher.update(cipherkey)
         self.cipherkey = (hasher.hexdigest()[:32]).encode('UTF-8')
@@ -84,7 +118,7 @@ class NetWrapper:
             if not status:
                 raise Exception('Network error during connection.')
             cleartext = self.recieveEncryptedMessage(msg, "AUT").decode('UTF-8')
-            if cleartext=="ERROR":
+            if cleartext == "ERROR":
                 return False
             else:
                 plaintext = cleartext.split(' ')
@@ -137,7 +171,7 @@ class NetWrapper:
         self.homeDirectory = ""
 
     def recieveEncryptedMessage(self, msg: bytes, type: str) -> bytes:
-        if not (type == "IDY" or type == "DH" or type == "CMD"):
+        if not (type == "AUT" or type == "CMD"):
             raise Exception('Unknown message type')
         try:
             b64 = json.loads(msg.decode('UTF-8'))