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

---
 client/netwrapper.py | 50 ++++++++++++++++++--------------------------
 1 file changed, 20 insertions(+), 30 deletions(-)

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: