diff --git a/client/netwrapper.py b/client/netwrapper.py index 0f7ce56..03e3fb3 100644 --- a/client/netwrapper.py +++ b/client/netwrapper.py @@ -7,7 +7,6 @@ import pyDH from Crypto.Hash import SHA512 from Crypto.Cipher import ChaCha20_Poly1305, PKCS1_OAEP from Crypto.PublicKey.RSA import RsaKey -from Crypto.Random import get_random_bytes from netsim import network_interface @@ -66,32 +65,10 @@ class NetWrapper: def authenticate(self, password: str): message = f"LIN {self.username} {password}".encode('UTF-8') - cipher = ChaCha20_Poly1305.new(key=self.cipherkey) - header = json.dumps({'type': 'AUT', 'source': self.network.own_addr}).encode('UTF-8') - cipher.update(header) - ciphertext, tag = cipher.encrypt_and_digest(message) - nonce = b64encode(cipher.nonce).decode('UTF-8') - ct = b64encode(ciphertext).decode('UTF-8') - b64tag = b64encode(tag).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.serverAddr, sendjson) - try: - status, msg = self.network.receive_msg(blocking=True) - if not status: - raise Exception('Network error during connection.') - b64 = json.loads(msg.decode('UTF-8')) - retnonce = b64decode(b64['nonce']) - retciphertext = b64decode(b64['message']) - retcipher = ChaCha20_Poly1305.new(key=self.cipherkey, nonce=retnonce) - retcipher.update(b64decode(b64['header'])) - retheader = json.loads(b64decode(b64['header']).decode('UTF-8')) - plaintext = retcipher.decrypt_and_verify(retciphertext, b64decode(b64['tag'])).decode('UTF-8') - if plaintext != "OK" or not (retheader['source'] == self.serverAddr and retheader['type'] == 'AUT'): - raise Exception('Authentication error') - except Exception as e: - print(e) + self.sendTypedMessage(message, "AUT") + plaintext = self.recieveTypedMessage("AUT").decode('UTF-8') + if plaintext != "OK": + print("Authentication error") def connectToServer(self): identStatus = self.identifyServer() @@ -102,9 +79,17 @@ class NetWrapper: pw = input() self.authenticate(pw) - def sendMessage(self, message: bytes): + def sendMessage(self, message: bytes) -> None: + try: + self.sendTypedMessage(message, "CMD") + except Exception: + print("Could not send message") + + def sendTypedMessage(self, message: bytes, type: str) -> None: + if not (type == "AUT" or type == "CMD"): + raise Exception('Unknown message type') cipher = ChaCha20_Poly1305.new(key=self.cipherkey) - header = json.dumps({'type': 'CMD', 'source': self.network.own_addr}).encode('UTF-8') + header = json.dumps({'type': type, 'source': self.network.own_addr}).encode('UTF-8') cipher.update(header) ciphertext, tag = cipher.encrypt_and_digest(message) nonce = b64encode(cipher.nonce).decode('UTF-8') @@ -115,6 +100,11 @@ class NetWrapper: self.network.send_msg(self.serverAddr, sendjson) def recieveMessage(self) -> bytes: + return self.recieveTypedMessage("CMD") + + def recieveTypedMessage(self, type: str) -> bytes: + if not (type == "AUT" or type == "CMD"): + raise Exception('Unknown message type') try: status, msg = self.network.receive_msg(blocking=True) if not status: @@ -126,7 +116,7 @@ class NetWrapper: retcipher.update(b64decode(b64['header'])) plaintext = retcipher.decrypt_and_verify(retciphertext, b64decode(b64['tag'])) retheader = json.loads(b64decode(b64['header']).decode('UTF-8')) - if not (retheader['source'] == self.serverAddr and retheader['type'] == 'CMD'): + if not (retheader['source'] == self.serverAddr and retheader['type'] == type): return "ERROR".encode('UTF-8') return plaintext except Exception: