From a5b77d03066cb095b9bfcd9531b41a3e13f4cdf4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Torma=20Krist=C3=B3f?= <tormakristof@tormakristof.eu>
Date: Sun, 25 Apr 2021 00:42:36 +0200
Subject: [PATCH] finish netwrapper of server and executor improvements

---
 server/executor.py   | 22 +++++++++++++++-------
 server/netwrapper.py | 16 ++++++++++------
 2 files changed, 25 insertions(+), 13 deletions(-)

diff --git a/server/executor.py b/server/executor.py
index 4620657..9bbf60f 100644
--- a/server/executor.py
+++ b/server/executor.py
@@ -29,7 +29,8 @@ class Executor:
     def removeDirectory(self, dirName: str) -> str:
         dirName = self.sanitizeDirectory(dirName)
         actualDirName = os.path.join(self.currentDirectory, dirName)
-        os.rmdir(actualDirName)
+        if actualDirName:
+            os.rmdir(actualDirName)
         return actualDirName
 
     def getCurrentDirectory(self) -> str:
@@ -40,16 +41,19 @@ class Executor:
             if self.currentDirectory == self.baseDir:
                 return self.currentDirectory
             else:
-                directories = self.currentDirectory.split("/")
+                directories = self.currentDirectory.split(os.path.sep)
                 strdirectory = ""
                 for dir in directories:
-                    strdirectory += dir + "/"
+                    strdirectory += dir + os.path.sep
                     strdirectory = strdirectory[:-3]
-                self.currentDirectory = strdirectory
+                if os.path.exists(strdirectory):
+                    self.currentDirectory = strdirectory
                 return self.currentDirectory
         else:
             dirName = self.sanitizeDirectory(dirName)
-            self.currentDirectory = os.path.join(self.currentDirectory, dirName)
+            joinedDir = os.path.join(self.currentDirectory, dirName)
+            if os.path.exists(joinedDir):
+                self.currentDirectory = joinedDir
             return self.currentDirectory
 
     def listCurrentDirectoryContent(self) -> str:
@@ -71,10 +75,14 @@ class Executor:
     def getFileInCurrentDirectory(self, file: str):
         file = self.sanitizeFile(file)
         currentfile = os.path.join(self.currentDirectory, file)
-        return open(currentfile, "r")
+        if os.path.exists(currentfile):
+            return open(currentfile, "r")
+        else:
+            raise Exception('File not found')
 
     def removeFileInCurrentDirectory(self, file: str) -> str:
         file = self.sanitizeFile(file)
         currentfile = os.path.join(self.currentDirectory, file)
-        os.remove(currentfile)
+        if os.path.exists(currentfile):
+            os.remove(currentfile)
         return currentfile
diff --git a/server/netwrapper.py b/server/netwrapper.py
index a7919f9..496caa7 100644
--- a/server/netwrapper.py
+++ b/server/netwrapper.py
@@ -25,6 +25,8 @@ class NetWrapper:
 
     def serverIdentify(self, msg: bytes) -> None:
         incommingJson = json.loads(msg.decode('UTF-8'))
+        if incommingJson['type'] != "IDY":
+            raise Exception('Wrong message type encountered')
         self.clientAddr = incommingJson['source']
         self.currentUser = incommingJson['username']
         self.currentClientPublicKey = self.clientPublicKey[self.currentUser]
@@ -56,7 +58,7 @@ class NetWrapper:
         jsonmsg = json.dumps({'type': 'DH', 'source': self.network.own_addr, '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'):
+        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.')
@@ -68,7 +70,7 @@ class NetWrapper:
 
     def login(self) -> bool:
         b64 = {'source': '', 'type': ''}
-        while not (b64['source'] == self.clientAddr and  b64['type'] == 'AUT'):
+        while not (b64['source'] == self.clientAddr and b64['type'] == 'AUT'):
             status, msg = self.network.receive_msg(blocking=True)
             if not status:
                 raise Exception('Network error during connection.')
@@ -78,7 +80,8 @@ class NetWrapper:
             retciphertext = b64decode(b64['message'])
             retcipher = ChaCha20.new(self.cipherkey, nonce=retnonce)
             plaintext = retcipher.decrypt(retciphertext).decode('UTF-8').split(' ')
-            linsuccess = (not (len(plaintext) != 3 or plaintext[0] != "LIN" or plaintext[1] != self.currentUser)) and self.serverInstance.login(plaintext[1],plaintext[2])
+            linsuccess = (not (len(plaintext) != 3 or plaintext[0] != "LIN" or plaintext[
+                1] != self.currentUser)) and self.serverInstance.login(plaintext[1], plaintext[2])
             if linsuccess:
                 message = "OK".encode('UTF-8')
             else:
@@ -87,8 +90,9 @@ class NetWrapper:
             ciphertext = cipher.encrypt(message)
             nonce = b64encode(cipher.nonce).decode('UTF-8')
             ct = b64encode(ciphertext).decode('UTF-8')
-            sendjson = json.dumps({'type': 'AUT', 'source': self.network.own_addr, 'nonce': nonce, 'message': ct}).encode(
-              'UTF-8')
+            sendjson = json.dumps(
+                {'type': 'AUT', 'source': self.network.own_addr, 'nonce': nonce, 'message': ct}).encode(
+                'UTF-8')
             self.network.send_msg(self.clientAddr, sendjson)
             return linsuccess
         except Exception:
@@ -102,7 +106,7 @@ class NetWrapper:
             if success:
                 return b"LINOK"
             else:
-               self.logout()
+                self.logout()
         except Exception:
             self.logout()