diff --git a/server/netwrapper.py b/server/netwrapper.py index 72e9d5c..ba8d556 100644 --- a/server/netwrapper.py +++ b/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') + status, msg = self.network.receive_msg(blocking=True) + if not status: + raise Exception('Network error during connection.') + decodedmsg = json.loads(msg.decode('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'))