Connecting client - netwrapper

client/client.py

@@ -6,13 +6,12 @@ import sys
 from Crypto.PublicKey import RSA
 from Crypto.PublicKey.RSA import RsaKey
 
+from netwrapper import NetWrapper
+
 ABSOLUTE_PATH = os.path.abspath(os.path.dirname(sys.argv[0]))
 DOWNLOAD_LOCATION = ABSOLUTE_PATH + os.path.sep + 'download' + os.path.sep
 CONFIG_LOCATION = ABSOLUTE_PATH + os.path.sep + 'config' + os.path.sep + 'config.txt'
-PASSPHRASE = ''
 LOGGED_IN = False
-SERVER_PUBLIC_KEY = None
-CLIENT_PRIVATE_KEY = None
 
 
 def loadPublicKey() -> RsaKey:
@@ -41,6 +40,10 @@ def loadPrivateKey(passphrase: str) -> RsaKey:
 
     return rsaKey
 
+def loadAddress() -> str:
+    with open(CONFIG_LOCATION) as json_file:
+        data = json.load(json_file)
+    return data['address']
 
 def printCommand():
     print('\nInvalid command! Available commands:\n' +
@@ -81,63 +84,76 @@ if not os.path.isfile(CONFIG_LOCATION) or os.stat(CONFIG_LOCATION).st_size == 0:
     print('Invalid client config file')
     sys.exit(1)
 
-#SERVER_PUBLIC_KEY = loadPublicKey()
+SERVER_PUBLIC_KEY = loadPublicKey()
 CLIENT_PRIVATE_KEY = loadPrivateKey(PASSPHRASE)
+CLIENT_ADDRESS = loadAddress()
+
 
 while True:
-    command = input("Type a command:").split(" ")
-    if len(command) > 3 or len(command) < 1:
-        print("Invalid command format!")
-        continue
+    command = input("Type a command:")
+    separatedCommand = command.split(" ")
 
-    if command[0] == 'LIN' and len(command) == 3:
-        print('TODO: Implement!')
-        LOGGED_IN = True
-        continue
+    try:
+        if len(separatedCommand) > 3 or len(separatedCommand) < 1:
+            print("Invalid command format!")
+            continue
 
-    if command[0] == 'EXIT':
-        sys.exit(1)
+        if separatedCommand[0] == 'LIN' and len(separatedCommand) == 3:
+            network = NetWrapper(CLIENT_PRIVATE_KEY, CLIENT_ADDRESS, separatedCommand[1], SERVER_PUBLIC_KEY)
+            if not network.identifyServer():
+                print('Server identification failed!')
 
-    if not LOGGED_IN:
-        printCommandsWihtoutLogin()
-        continue
+            network.createEncryptedChannel()
+            network.authenticate(password=separatedCommand[2])
+            LOGGED_IN = True
+            continue
 
-    if command[0] == 'LOUT':
-        print('TODO: Implement!')
-        continue
+        if separatedCommand[0] == 'EXIT':
+            sys.exit(1)
 
-    if command[0] == 'MKD' and len(command) == 2:
-        print('TODO: Implement!')
-        continue
+        if not LOGGED_IN:
+            printCommandsWihtoutLogin()
+            continue
 
-    if command[0] == 'RMD' and len(command) == 2:
-        print('TODO: Implement!')
-        continue
+        if separatedCommand[0] == 'LOUT' and len(separatedCommand) == 1:
+            network.sendMessage(command.encode('UTF-8'))
+            continue
 
-    if command[0] == 'GWD' and len(command) == 1:
-        print('TODO: Implement!')
-        continue
+        if separatedCommand[0] == 'MKD' and len(separatedCommand) == 2:
+            network.sendMessage(command.encode('UTF-8'))
+            continue
 
-    if command[0] == 'CWD' and len(command) == 2:
-        print('TODO: Implement!')
-        continue
+        if separatedCommand[0] == 'RMD' and len(separatedCommand) == 2:
+            network.sendMessage(command.encode('UTF-8'))
+            continue
 
-    if command[0] == 'LST' and len(command) == 1:
-        print('TODO: Implement!')
-        continue
+        if separatedCommand[0] == 'GWD' and len(separatedCommand) == 1:
+            network.sendMessage(command.encode('UTF-8'))
+            continue
 
-    if command[0] == 'UPL' and len(command) == 2:
-        if os.path.isfile(command[1]):
-            with open(command[1], "rb") as file:
+        if separatedCommand[0] == 'CWD' and len(separatedCommand) == 2:
+            network.sendMessage(command.encode('UTF-8'))
+            continue
+
+        if separatedCommand[0] == 'LST' and len(separatedCommand) == 1:
+            network.sendMessage(command.encode('UTF-8'))
+            continue
+
+        if separatedCommand[0] == 'UPL' and len(separatedCommand) == 2:
+            if os.path.isfile(separatedCommand[1]):
+                with open(separatedCommand[1], "rb") as file:
+                    print('TODO: Implement!')
+            else:
+                print('Invalid argument for file upload: ' + separatedCommand[1])
+            continue
+
+        if separatedCommand[0] == 'DNL' and len(separatedCommand) == 2:
+            dnlFilename = separatedCommand[1].rsplit(os.apth.sep, 1)[-1]
+            with open(DOWNLOAD_LOCATION + dnlFilename, "wb") as file:
                 print('TODO: Implement!')
-        else:
-            print('Invalid argument for file upload: '+command[1])
-        continue
+            continue
 
-    if command[0] == 'DNL' and len(command) == 2:
-        dnlFilename = command[1].rsplit(os.apth.sep, 1)[-1]
-        with open(DOWNLOAD_LOCATION + dnlFilename, "wb") as file:
-            print('TODO: Implement!')
-        continue
+        printCommand()
 
-    printCommand()
+    except Exception:
+        print('Error: '+Exception.args)

client/netwrapper.py

@@ -4,17 +4,15 @@ import string
 import json
 from base64 import b64encode, b64decode
 import pyDH
-from Crypto.Cipher import PKCS1_OAEP
-from Crypto.PublicKey import RSA
 from Crypto.Cipher import ChaCha20
+from Crypto.PublicKey.RSA import RsaKey
 from Crypto.Random import get_random_bytes
 
 from netsim import network_interface
 
 
 class NetWrapper:
-
-    def __init__(self, privateKey: str, clientAddress: str, username: str, serverPubKey: str = "",
+    def __init__(self, privateKey: RsaKey, clientAddress: str, username: str, serverPubKey: RsaKey,
                  serverAddr: str = 'A'):
         # Create network_interface:  network_interface(path, addr) path root is shared with network / addr is own address
         self.network = network_interface('./', clientAddress)
@@ -24,17 +22,17 @@ class NetWrapper:
         self.serverPubKey = serverPubKey
         self.cipherkey = "".encode('UTF-8')
 
+
     def randomStringGenerator(self, str_size: int = 512,
                               allowed_chars: str = string.ascii_letters + string.punctuation) -> str:
         return ''.join(random.choice(allowed_chars) for x in range(str_size))
 
+
     def identifyServer(self) -> bool:
         randommsg = self.randomStringGenerator()
-        rsakey = RSA.import_key(self.serverPubKey)
-        cipher = PKCS1_OAEP.new(rsakey)
         identMsg = json.dumps(
             {'type': 'IDY', 'source': self.network.own_addr, 'username': self.username,
-             'message': b64encode(cipher.encrypt(randommsg.encode('UTF-8')))}).encode(
+             'message': b64encode(self.privateKey.encrypt(randommsg.encode('UTF-8')))}).encode(
             'UTF-8')
         self.network.send_msg(self.serverAddr, identMsg)
         returnJson = {'source': '', 'type': ''}
@@ -43,16 +41,13 @@ class NetWrapper:
             if not status:
                 raise Exception('Network error during connection.')
             returnJson = json.loads(msg.decode('UTF-8'))
-        myrsakey = RSA.import_key(self.privateKey)
-        mycipher = PKCS1_OAEP.new(myrsakey)
-        retmsg = mycipher.decrypt(b64decode(returnJson['message'])).decode('UTF-8')
+        retmsg = self.privateKey.decrypt(b64decode(returnJson['message'])).decode('UTF-8')
         return retmsg == randommsg
 
+
     def createEncryptedChannel(self):
         dh = pyDH.DiffieHellman()
-        rsakey = RSA.import_key(self.serverPubKey)
-        cipher = PKCS1_OAEP.new(rsakey)
-        mypubkey = b64encode(cipher.encrypt(str(dh.gen_public_key()).encode('UTF-8')))
+        mypubkey = b64encode(self.serverPubKey.encrypt(str(dh.gen_public_key()).encode('UTF-8')))
         jsonmsg = json.dumps({'type': 'DH', 'source': self.network.own_addr, 'message': mypubkey}).encode('UTF-8')
         self.network.send_msg(self.serverAddr, jsonmsg)
         decodedmsg = {'source': '', 'type': ''}
@@ -61,11 +56,10 @@ class NetWrapper:
             if not status:
                 raise Exception('Network error during connection.')
             decodedmsg = json.loads(msg.decode('UTF-8'))
-        myrsakey = RSA.import_key(self.privateKey)
-        mycipher = PKCS1_OAEP.new(myrsakey)
-        serverpubkey = int(mycipher.decrypt(b64decode(decodedmsg['message'])).decode('UTF-8'))
+        serverpubkey = int(self.privateKey.decrypt(b64decode(decodedmsg['message'])).decode('UTF-8'))
         self.cipherkey = dh.gen_shared_key(serverpubkey).encode('UTF-8')
 
+
     def authenticate(self, password: str):
         message = f"LIN {self.username} {password}".encode('UTF-8')
         cipher = ChaCha20.new(self.cipherkey, get_random_bytes(12))
@@ -91,6 +85,7 @@ class NetWrapper:
         except Exception:
             print("Incorrect decryption")
 
+
     def connectToServer(self):
         identStatus = self.identifyServer()
         if not identStatus:
@@ -100,6 +95,7 @@ class NetWrapper:
         pw = input()
         self.authenticate(pw)
 
+
     def sendMessage(self, message: bytes):
         cipher = ChaCha20.new(self.cipherkey, get_random_bytes(12))
         ciphertext = cipher.encrypt(message)
@@ -109,9 +105,10 @@ class NetWrapper:
             'UTF-8')
         self.network.send_msg(self.serverAddr, sendjson)
 
+
     def recieveMessage(self) -> bytes:
         b64 = {'source': '', 'type': ''}
-        while not (b64['source'] == self.serverAddr and  b64['type'] == 'AUT'):
+        while not (b64['source'] == self.serverAddr and b64['type'] == 'AUT'):
             status, msg = self.network.receive_msg(blocking=True)
             if not status:
                 raise Exception('Network error during connection.')