refactor and better crypto

client/client.py

@@ -99,10 +99,10 @@ while True:
             print("Invalid command format!")
             continue
 
-        if separatedCommand[0] == 'LIN' and len(separatedCommand) == 2:
+        if separatedCommand[0] == 'LIN' and len(separatedCommand) == 3:
             network = NetWrapper(CLIENT_PRIVATE_KEY, CLIENT_ADDRESS, separatedCommand[1], SERVER_PUBLIC_KEY)
             try:
-                network.connectToServer()
+                network.connectToServer(separatedCommand[2])
             except Exception as e:
                 print("Error: "+str(e))
                 continue

client/main.py

@@ -1,4 +0,0 @@
-#!/usr/bin/env python3
-
-if __name__ == '__main__':
-    print("hi")

client/netwrapper.py

@@ -7,6 +7,7 @@ import pyDH
 from Crypto.Hash import SHA512
 from Crypto.Cipher import ChaCha20_Poly1305, PKCS1_OAEP
 from Crypto.PublicKey.RSA import RsaKey
+from Crypto.Signature import pkcs1_15
 
 from netsim import network_interface
 
@@ -14,7 +15,6 @@ from netsim import network_interface
 class NetWrapper:
     def __init__(self, privateKey: RsaKey, clientAddress: str, username: str, serverPubKey: RsaKey,
                  serverAddr: str = 'A'):
-        # Create network_interface:  network_interface(path, addr) path root is shared with network / addr is own address
         self.network = network_interface('./../../netsim/network/', clientAddress)
         self.serverAddr = serverAddr
         self.username = username
@@ -22,62 +22,93 @@ class NetWrapper:
         self.serverPubKey = serverPubKey
         self.cipherkey = "".encode('UTF-8')
 
-    def randomStringGenerator(self, str_size: int = 128,
+    def randomStringGenerator(self, str_size: int = 256,
                               allowed_chars: str = string.ascii_letters + string.punctuation) -> str:
         return ''.join(random.choice(allowed_chars) for x in range(str_size))
 
-    def identifyServer(self) -> bool:
-        randommsg = self.randomStringGenerator()
+    def ecryptRSAMessage(self, message: str) -> bytes:
         cipher_rsa = PKCS1_OAEP.new(self.serverPubKey)
-        identMsg = json.dumps(
-            {'type': 'IDY', 'source': self.network.own_addr, 'username': self.username,
-             'message': b64encode(cipher_rsa.encrypt(randommsg.encode('UTF-8'))).decode('ASCII')}).encode(
-            'UTF-8')
-        self.network.send_msg(self.serverAddr, identMsg)
-        returnJson = {'source': '', 'type': ''}
-        while not (returnJson['source'] == self.serverAddr and returnJson['type'] == 'IDY'):
-            status, msg = self.network.receive_msg(blocking=True)
-            if not status:
-                raise Exception('Network error during connection.')
-            returnJson = json.loads(msg.decode('UTF-8'))
-        cipher_rsa = PKCS1_OAEP.new(self.privateKey)
-        retmsg = cipher_rsa.decrypt(b64decode(returnJson['message'])).decode('UTF-8')
-        return retmsg == randommsg
+        encrypted_msg = cipher_rsa.encrypt(message.encode('UTF-8'))
+        return encrypted_msg
 
-    def createEncryptedChannel(self):
-        dh = pyDH.DiffieHellman()
-        cipher_rsa = PKCS1_OAEP.new(self.serverPubKey)
-        mypubkey = b64encode(cipher_rsa.encrypt(str(dh.gen_public_key()).encode('UTF-8'))).decode('ASCII')
-        jsonmsg = json.dumps({'type': 'DH', 'source': self.network.own_addr, 'message': mypubkey}).encode('UTF-8')
-        self.network.send_msg(self.serverAddr, jsonmsg)
-        decodedmsg = {'source': '', 'type': ''}
-        while not (decodedmsg['source'] == self.serverAddr and decodedmsg['type'] == 'DH'):
+    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.privateKey).sign(h)
+        return header, headersignature
+
+    def verifyRSAHeaderSignature(self, header: bytes, headersignature: bytes) -> bool:
+        h = SHA512.new(header)
+        try:
+            pkcs1_15.new(self.serverPubKey).verify(h, headersignature)
+            return True
+        except Exception:
+            return False
+
+    def decryptRSAMessage(self, message: bytes) -> bytes:
+        cipher_rsa = PKCS1_OAEP.new(self.privateKey)
+        return cipher_rsa.decrypt(message)
+
+    def recieveAndUnpackRSAMessage(self) -> (str, str):
         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.privateKey)
-        serverpubkey = int(cipher_rsa.decrypt(b64decode(decodedmsg['message'])).decode('UTF-8'))
+        return decodedmsg, json.loads(b64decode(decodedmsg['header']).decode('UTF-8'))
+
+    def identifyServer(self) -> bool:
+        randommsg = self.randomStringGenerator()
+        encrypted_msg = self.ecryptRSAMessage(randommsg)
+        header, headersignature = self.signRSAHeader('IDY', {'username': self.username})
+        identMsg = json.dumps(
+            {'header': b64encode(header).decode('UTF-8'),
+             'message': b64encode(encrypted_msg).decode('UTF-8'),
+             'headersignature': b64encode(headersignature).decode('UTF-8')}).encode(
+            'UTF-8')
+        self.network.send_msg(self.serverAddr, identMsg)
+        returnJson, header = self.recieveAndUnpackRSAMessage()
+        try:
+            if not self.verifyRSAHeaderSignature(b64decode(returnJson['header']),
+                                                 b64decode(returnJson['headersignature'])) or not (
+                    header['source'] == self.serverAddr and header['type'] == 'IDY'):
+                return False
+            retmsg = self.decryptRSAMessage(b64decode(returnJson['message'])).decode('UTF-8')
+        except Exception:
+            return False
+        return retmsg == randommsg
+
+    def createEncryptedChannel(self) -> None:
+        dh = pyDH.DiffieHellman()
+        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.serverAddr, jsonmsg)
+        decodedmsg, header = self.recieveAndUnpackRSAMessage()
+        if not self.verifyRSAHeaderSignature(b64decode(decodedmsg['header']),
+                                             b64decode(decodedmsg['headersignature'])) or not (
+                header['source'] == self.serverAddr and header['type'] == 'DH'):
+            raise Exception('Header signature error')
+        serverpubkey = int(self.decryptRSAMessage(b64decode(decodedmsg['message'])).decode('UTF-8'))
         cipherkey = dh.gen_shared_key(serverpubkey).encode('UTF-8')
         hasher = SHA512.new()
         hasher.update(cipherkey)
         self.cipherkey = (hasher.hexdigest()[:32]).encode('UTF-8')
 
-    def authenticate(self, password: str):
+    def authenticate(self, password: str) -> None:
         message = f"LIN {self.username} {password}".encode('UTF-8')
         self.sendTypedMessage(message, "AUT")
         plaintext = self.recieveTypedMessage("AUT").decode('UTF-8')
         if plaintext != "OK":
             print("Authentication error")
 
-    def connectToServer(self):
+    def connectToServer(self, password: str) -> None:
         identStatus = self.identifyServer()
         if not identStatus:
             raise Exception('Server identification faliure')
         self.createEncryptedChannel()
-        print('Please enter your password:')
-        pw = input()
-        self.authenticate(pw)
+        self.authenticate(password)
 
     def sendMessage(self, message: bytes) -> None:
         try: