Init script extension and key loading for netwrapper

server/authentication.py

@@ -7,13 +7,15 @@ from base64 import b64encode
 
 from Crypto.Hash import SHA256
 from Crypto.Protocol.KDF import bcrypt, bcrypt_check
+from Crypto.PublicKey import RSA
+from Crypto.PublicKey.RSA import RsaKey
 
 auth_logger = logging.getLogger('AUTH APPLICATION ')
 auth_logger.root.setLevel(logging.INFO)
 
 class Authetication:
     ABSOLUTE_PATH = os.path.abspath(os.path.dirname(sys.argv[0]))
-    HOME_DIRECTORY_LOCATION = ABSOLUTE_PATH + os.path.sep +"home"
+    HOME_DIRECTORY_LOCATION = ABSOLUTE_PATH + os.path.sep + "home"
     CONFIG_DIRECTORY_LOCATION = ABSOLUTE_PATH + os.path.sep + "config"
     CONFIG_FILE_LOCATION = ABSOLUTE_PATH + os.path.sep + "config" + os.path.sep + "config.txt"
     PRIVATE_KEY_DIRECTORY_LOCATION = CONFIG_DIRECTORY_LOCATION + os.path.sep + "private_keys"
@@ -111,3 +113,34 @@ class Authetication:
                 auth_logger.debug("User NOT saved! Home directory already exists")
 
             return True
+
+
+    def loadUserPublicKeys(self) -> dict:
+        with open(self.CONFIG_FILE_LOCATION) as json_file:
+            data = json.load(json_file)
+
+        dictionary: dict
+
+        for user in data['user']:
+            key = data['publicKey']
+            key = bytes.fromhex(key)
+            try:
+                rsaKey = RSA.import_key(key)
+                dictionary[user['username']] = rsaKey
+            except ValueError:
+                print('Invalid server public key!')
+
+        return dictionary
+
+
+    def loadServerPrivateKey(self) -> RsaKey:
+        with open(self.CONFIG_FILE_LOCATION) as json_file:
+            data = json.load(json_file)
+
+        key = data['serverPrivateKey']
+        key = bytes.fromhex(key)
+        try:
+            rsaKey = RSA.import_key(key)
+        except ValueError:
+            print('Invalid server private key!')
+        return rsaKey

server/config_init.py

@@ -6,7 +6,7 @@ from Crypto.PublicKey import RSA
 from authentication import Authetication
 
 
-def generatePrivateKeyForUser(auth: Authetication, username: str, user_passphrase: str, public_server_key='') -> bool:
+def generatePrivateKeyForUser(auth: Authetication,username: str, user_passphrase: str, public_server_key: str) -> bool:
     if auth.checkUserExists(username):
         with open(auth.CONFIG_FILE_LOCATION) as json_file:
             data = json.load(json_file)
@@ -17,9 +17,10 @@ def generatePrivateKeyForUser(auth: Authetication, username: str, user_passphras
         public_key_value = bytes.hex(public_key.exportKey('DER', pkcs=8))
 
         ##Save private key in separate file
-        user_privatekey = {'passphrase': user_passphrase, 'privateClientKey': private_key_value,
+        user_privatekey = {'privateClientKey': private_key_value,
                            'publicServerKey': public_server_key}
-        with open(auth.PRIVATE_KEY_DIRECTORY_LOCATION + os.path.sep + str(data['index']) + '.txt', 'w+') as outfile:
+        with open(auth.PRIVATE_KEY_DIRECTORY_LOCATION + os.path.sep + str(data['index']) + '.txt',
+                  'w+') as outfile:
             json.dump(user_privatekey, outfile)
         outfile.close()
 
@@ -27,7 +28,7 @@ def generatePrivateKeyForUser(auth: Authetication, username: str, user_passphras
         for user in data['user']:
             if username == user['username']:
                 user['publicKey'] = public_key_value
-                with open(auth.CONFIG_FILE_LOCATION, 'w') as outfile:
+                with open(auth.CONFIG_FILE_LOCATION, 'w+') as outfile:
                     json.dump(data, outfile)
                 break
         outfile.close()
@@ -36,10 +37,28 @@ def generatePrivateKeyForUser(auth: Authetication, username: str, user_passphras
         return False
 
 
+def generatePrivateKeyForServer(auth: Authetication,passphrase: str) -> str:
+    with open(auth.CONFIG_FILE_LOCATION) as json_file:
+        data = json.load(json_file)
+    json_file.close()
+
+    private_key = RSA.generate(2048)
+    public_key = private_key.publickey()
+    private_key_value = bytes.hex(private_key.exportKey('DER', passphrase=passphrase, pkcs=8))
+    public_key_value = bytes.hex(public_key.exportKey('DER', pkcs=8))
+
+    data['serverPrivateKey'] = private_key_value
+    with open(auth.CONFIG_FILE_LOCATION, 'w+') as outfile:
+        json.dump(data, outfile)
+
+    return public_key_value
+
+
 if __name__ == '__main__':
     auth = Authetication()
     auth.initConfig()
-    auth.saveUser('alma' ,'alma')
-    generatePrivateKeyForUser('alma', 'amla')
-    auth.saveUser('citrom' ,'citrom')
-    generatePrivateKeyForUser('citrom', 'mortic')
+    serverPublicKey = generatePrivateKeyForServer(auth, 'admin')
+    auth.saveUser('alma', 'alma')
+    generatePrivateKeyForUser(auth, 'alma', 'amla', serverPublicKey)
+    auth.saveUser('citrom', 'citrom')
+    generatePrivateKeyForUser(auth, 'citrom', 'mortic', serverPublicKey)

server/netwrapper.py

@@ -5,6 +5,7 @@ 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
@@ -13,7 +14,7 @@ from server.server import Server
 
 class NetWrapper:
 
-    def __init__(self, clientPublicKey: dict, serverPrivateKey: str, serverInstance: Server):
+    def __init__(self, clientPublicKey: dict, serverPrivateKey: RsaKey, serverInstance: Server):
         self.clientPublicKey = clientPublicKey
         self.currentClientPublicKey = "".encode('UTF-8')
         self.serverPrivateKey = serverPrivateKey
@@ -30,9 +31,7 @@ class NetWrapper:
         self.clientAddr = incommingJson['source']
         self.currentUser = incommingJson['username']
         self.currentClientPublicKey = self.clientPublicKey[self.currentUser]
-        myrsakey = RSA.import_key(self.serverPrivateKey)
-        mycipher = PKCS1_OAEP.new(myrsakey)
-        retmsg = mycipher.decrypt(b64decode(incommingJson['message'])).decode('UTF-8')
+        retmsg = self.serverPrivateKey.decrypt(b64decode(incommingJson['message'])).decode('UTF-8')
         rsakey = RSA.import_key(self.currentClientPublicKey)
         cipher = PKCS1_OAEP.new(rsakey)
         identMsg = json.dumps(
@@ -63,9 +62,7 @@ class NetWrapper:
             if not status:
                 raise Exception('Network error during connection.')
             decodedmsg = json.loads(msg.decode('UTF-8'))
-        myrsakey = RSA.import_key(self.serverPrivateKey)
-        mycipher = PKCS1_OAEP.new(myrsakey)
-        serverpubkey = int(mycipher.decrypt(b64decode(decodedmsg['message'])).decode('UTF-8'))
+        serverpubkey = int(self.serverPrivateKey.decrypt(b64decode(decodedmsg['message'])).decode('UTF-8'))
         self.cipherkey = dh.gen_shared_key(serverpubkey).encode('UTF-8')
 
     def login(self) -> bool: