diff --git a/server/netwrapper.py b/server/netwrapper.py index e54eb50..72e9d5c 100644 --- a/server/netwrapper.py +++ b/server/netwrapper.py @@ -42,8 +42,13 @@ class NetWrapper: self.network.send_msg(self.clientAddr, identMsg) def sendMessage(self, message: bytes) -> None: + self.sendTypedMessage(message, "CMD") + + def sendTypedMessage(self, message: bytes, type: str) -> None: + if not (type == "IDY" or type == "DH" or type == "CMD"): + raise Exception('Unknown message type') cipher = ChaCha20_Poly1305.new(key=self.cipherkey) - header = json.dumps({'source': self.network.own_addr, 'type': 'CMD'}).encode('UTF-8') + header = json.dumps({'source': self.network.own_addr, 'type': type}).encode('UTF-8') cipher.update(header) ciphertext, tag = cipher.encrypt_and_digest(message) nonce = b64encode(cipher.nonce).decode('UTF-8') @@ -78,16 +83,11 @@ class NetWrapper: status, msg = self.network.receive_msg(blocking=True) if not status: raise Exception('Network error during connection.') - b64 = json.loads(msg.decode('UTF-8')) - retheader = json.loads(b64decode(b64['header']).decode('UTF-8')) - retnonce = b64decode(b64['nonce']) - retciphertext = b64decode(b64['message']) - rettag = b64decode(b64['tag']) - retcipher = ChaCha20_Poly1305.new(key=self.cipherkey, nonce=retnonce) - retcipher.update(b64decode(b64['header'])) - plaintext = retcipher.decrypt_and_verify(retciphertext, rettag).decode('UTF-8').split(' ') - if not (retheader['source'] == self.clientAddr and retheader['type'] == 'AUT'): + cleartext = self.recieveEncryptedMessage(msg, "AUT").decode('UTF-8') + if cleartext=="ERROR": return False + else: + plaintext = cleartext.split(' ') self.homeDirectory = self.authenticationInstance.login(plaintext[1], plaintext[2]) linsuccess = (not (len(plaintext) != 3 or plaintext[0] != "LIN" or plaintext[ 1] != self.currentUser)) and self.homeDirectory @@ -95,20 +95,10 @@ class NetWrapper: message = "OK".encode('UTF-8') else: message = "ERROR".encode('UTF-8') - header = json.dumps({'source': self.network.own_addr, 'type': 'AUT'}).encode('UTF-8') - cipher = ChaCha20_Poly1305.new(key=self.cipherkey) - cipher.update(header) - ciphertext, tag = cipher.encrypt_and_digest(message) - b64tag = b64encode(tag).decode('UTF-8') - nonce = b64encode(cipher.nonce).decode('UTF-8') - ct = b64encode(ciphertext).decode('UTF-8') - sendjson = json.dumps( - {'header': b64encode(header).decode('UTF-8'), 'nonce': nonce, 'message': ct, 'tag': b64tag}).encode( - 'UTF-8') - self.network.send_msg(self.clientAddr, sendjson) + self.sendTypedMessage(message, "AUT") return linsuccess except Exception: - print("Incorrect decryption") + print("Login failed") return False def initClientConnection(self, msg: bytes) -> bytes: @@ -137,7 +127,7 @@ class NetWrapper: if not self.clientAddr: return self.initClientConnection(msg) else: - return self.recieveEncryptedMessage(msg) + return self.recieveEncryptedMessage(msg, "CMD") def logout(self) -> None: self.clientAddr = "" @@ -146,7 +136,9 @@ class NetWrapper: self.currentUser = "" self.homeDirectory = "" - def recieveEncryptedMessage(self, msg: bytes) -> bytes: + def recieveEncryptedMessage(self, msg: bytes, type: str) -> bytes: + if not (type == "IDY" or type == "DH" or type == "CMD"): + raise Exception('Unknown message type') try: b64 = json.loads(msg.decode('UTF-8')) retheader = json.loads(b64decode(b64['header']).decode('UTF-8')) @@ -156,7 +148,7 @@ class NetWrapper: retcipher = ChaCha20_Poly1305.new(key=self.cipherkey, nonce=retnonce) retcipher.update(b64decode(b64['header'])) plaintext = retcipher.decrypt_and_verify(retciphertext, rettag) - if not (retheader['source'] == self.clientAddr and retheader['type'] == 'CMD'): + if not (retheader['source'] == self.clientAddr and retheader['type'] == type): return "ERROR".encode('UTF-8') else: return plaintext