From 56d77f0476aaa0f11f52425a0d76d0d180c9a09a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Torma=20Krist=C3=B3f?= <tormakristof@tormakristof.eu>
Date: Mon, 26 Apr 2021 18:13:38 +0200
Subject: [PATCH] refactor chacha

---
 server/netwrapper.py | 42 +++++++++++++++++-------------------------
 1 file changed, 17 insertions(+), 25 deletions(-)

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