client/client/client.py

import getopt
import json
import os
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'
LOGGED_IN = False


def loadPublicKey() -> RsaKey:
    with open(CONFIG_LOCATION) as json_file:
        data = json.load(json_file)
    key = data['publicServerKey']
    key = bytes.fromhex(key)
    try:
        rsaKey = RSA.import_key(key)
    except ValueError:
        print('Invalid server public key!')
        sys.exit(1)
    return rsaKey


def loadPrivateKey(passphrase: str) -> RsaKey:
    with open(CONFIG_LOCATION) as json_file:
        data = json.load(json_file)
    key = data['privateClientKey']
    key = bytes.fromhex(key)
    try:
        rsaKey = RSA.import_key(key, passphrase)
    except ValueError:
        print('Invalid client key!')
        sys.exit(1)

    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' +
          '     Create directory                        -> MKD <directory name> \n' +
          '     Remove directory                        -> RMV <directory name> \n' +
          '     Get current directory                   -> GWD \n' +
          '     Change current directory                -> CWD <path> \n' +
          '     List content of current directory       -> LST \n' +
          '     Upload file to current directory        -> UPL <filename> \n' +
          '     Download file from current directory    -> DNL <filename> \n' +
          '     Login                                   -> LIN <username> <password> \n' +
          '     Logout                                  -> LOUT \n')


def printCommandsWihtoutLogin():
    print('\nYou must log in before issuing other commads!\n',
          '     Login    -> LIN <username> <password> \n',
          '     Exit     -> EXIT\n')


try:
    opts, args = getopt.getopt(sys.argv[1:], 'hp:')
except getopt.GetoptError:
    print('Error: Unknown option detected.')
    sys.exit(1)

for opt, arg in opts:
    if opt in '-p':
        PASSPHRASE = arg

if PASSPHRASE == '':
    print('Key required to start client!')
    print('Usage:')
    print('  client.py -p <passphrase>')
    sys.exit(1)

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()
CLIENT_PRIVATE_KEY = loadPrivateKey(PASSPHRASE)
CLIENT_ADDRESS = loadAddress()

while True:
    command = input("Type a command:")
    separatedCommand = command.split(" ")

    try:
        if len(separatedCommand) > 3 or len(separatedCommand) < 1:
            print("Invalid command format!")
            continue

        if separatedCommand[0] == 'LIN' and len(separatedCommand) == 3:
            network = NetWrapper(CLIENT_PRIVATE_KEY, CLIENT_ADDRESS, separatedCommand[1], SERVER_PUBLIC_KEY)
            print('Start')
            if not network.identifyServer():
                print('Server identification failed!')
                continue
            print('Identified')
            network.createEncryptedChannel()
            print('Channel created')
            network.authenticate(password=separatedCommand[2])
            print('Authenticated')
            LOGGED_IN = True
            continue

        if separatedCommand[0] == 'EXIT':
            sys.exit(1)

        if not LOGGED_IN:
            printCommandsWihtoutLogin()
            continue

        if separatedCommand[0] == 'LOUT' and len(separatedCommand) == 1:
            network.sendMessage(command.encode('UTF-8'))
            print(network.recieveMessage().decode('UTF-8'))
            continue

        if separatedCommand[0] == 'MKD' and len(separatedCommand) == 2:
            network.sendMessage(command.encode('UTF-8'))
            print(network.recieveMessage().decode('UTF-8'))
            continue

        if separatedCommand[0] == 'RMD' and len(separatedCommand) == 2:
            network.sendMessage(command.encode('UTF-8'))
            print(network.recieveMessage().decode('UTF-8'))
            continue

        if separatedCommand[0] == 'GWD' and len(separatedCommand) == 1:
            network.sendMessage(command.encode('UTF-8'))
            print(network.recieveMessage().decode('UTF-8'))
            continue

        if separatedCommand[0] == 'CWD' and len(separatedCommand) == 2:
            network.sendMessage(command.encode('UTF-8'))
            print(network.recieveMessage().decode('UTF-8'))
            continue

        if separatedCommand[0] == 'LST' and len(separatedCommand) == 1:
            network.sendMessage(command.encode('UTF-8'))
            print(network.recieveMessage().decode('UTF-8'))
            continue

        if separatedCommand[0] == 'UPL' and len(separatedCommand) == 2:
            if os.path.isfile(separatedCommand[1]):
                cmd = 'UPL ' + separatedCommand[1].split(os.path.sep)[-1]
                network.sendMessage(cmd.encode('UTF-8'))

                with open(separatedCommand[1], "rb") as file:
                    network.sendMessage(file.readlines())

                network.sendMessage('EOF'.encode('UTF-8'))
                response = network.recieveMessage().decode('UTF-8')
                print(response)

            else:
                print('Invalid argument for file upload: ' + separatedCommand[1])
            continue

        if separatedCommand[0] == 'DNL' and len(separatedCommand) == 2:
            dnlFilename = separatedCommand[1].split(os.path.sep)[-1]
            cmd = 'DNL ' + dnlFilename
            network.sendMessage(cmd.encode('UTF-8'))
            file = network.recieveMessage().decode('UTF-8')
            response = network.recieveMessage().decode('UTF-8')

            if response == 'OK':
                with open(DOWNLOAD_LOCATION + dnlFilename, "wb+") as file:
                    file.writelines(file)
            else:
                print(response)

            continue

        printCommand()

    except Exception as e:
        print('Error: ' + str(e))