#!/usr/local/bin/python3

CLIENT = 'alice'
# the file size is limited to 1 mb
BUFFER_SIZE = 1048576

import os, socket, base64, json
from Crypto.Random import get_random_bytes
from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_OAEP
from Crypto.Cipher import AES


# ==============================================================================================================
#                                           PKCS encrpytion until data sesssion
# ==============================================================================================================

def send_client_hello(sock, config):
    # get bob's public key
    key = RSA.importKey(open(os.path.join(config['pubkeys'], config['to'] + '.pem')).read())
    cipher = PKCS1_OAEP.new(key)

    # create and send client_hello payload
    nonce = base64.b64encode(get_random_bytes(32)).decode('utf-8')
    payload = cipher.encrypt(json.dumps({'from': CLIENT, 'na': nonce}).encode())
    sock.sendall(payload)
    # print debug message
    print('1. send: '+ str(payload))

def receive_server_hello(sock, config):
    # get your private key
    key = RSA.importKey(open(config['privkey']).read())
    cipher = PKCS1_OAEP.new(key)

    # receive and decode server_hello payload
    payload = cipher.decrypt(sock.recv(BUFFER_SIZE)).decode()
    server_hello = json.loads(payload)

    if(server_hello['from'] != config['to']):
        print('ALERT! Invalid sender - ' + config['to'])
        sys.exit(1)

    # print debug message
    print('2. recv: '+ str(payload))
    # update config
    config['nb'] = server_hello['nb']

def send_key(sock, config):
    # get bob's public key
    key = RSA.importKey(open(os.path.join(config['pubkeys'], config['to'] + '.pem')).read())
    cipher = PKCS1_OAEP.new(key)

    # generate a 128 bits random key
    #do something with config['na] and config['nb'] to createa a session key between AB?
    session_key = get_random_bytes(16)

    # create and send client_key
    payload = cipher.encrypt(json.dumps({'nb': config['nb'], 'key': base64.b64encode(session_key).decode('utf-8')}).encode())
    sock.sendall(payload)          
    # print debug message
    print('3. send: '+ str(payload))
    # update config
    config['key'] = session_key

def receive_authenticated(sock, config):
    # get your private key
    key = RSA.importKey(open(config['privkey']).read())
    cipher = PKCS1_OAEP.new(key)

    # receive and decode authentication payload
    payload = cipher.decrypt(sock.recv(BUFFER_SIZE)).decode()
    metada = json.loads(payload)
    # print debug message
    print('4. recv: '+ str(payload))
    # check if client is authenticated
    if not metada['authenticated']:
        print('authentication failed')
        sys.exit(1)


# ==============================================================================================================
#                                           AES encrpytion from here onwards 
# ==============================================================================================================
def send_request(sock, config):
    # generate 8 rand bytes for nonce
    nonce = get_random_bytes(8)

    # create cipher
    cipher = AES.new(config['key'], AES.MODE_EAX, nonce=nonce)

    # create and send request payload
    payload = json.dumps({'request': config['request'], 'filename': config['filename']}).encode()
    message = nonce + cipher.encrypt(payload)
    sock.sendall(message)
    # print debug message
    print('5. send: '+ str(message))

def receive_ready(sock, config):

    # receive data from server
    payload = sock.recv(BUFFER_SIZE)
    
    # create cipher using first 8 bytes from payload as nonce
    cipher = AES.new(config['key'], AES.MODE_EAX, nonce=payload[:8])
    #only need to worry about 8 bytes and onwards
    plain = cipher.decrypt(payload[8:]).decode()
    metada = json.loads(plain)
    # print debug message
    print('6. recv: '+ str(plain))
    # check if server is ready
    if not metada['ready']:
        print('server ' + config['to'] + ' cannot upload file')
        sys.exit(1)

def send_upload(sock, config):
    # generate 8 rand bytes for nonce
    nonce = get_random_bytes(8)

    # create cipher
    cipher = AES.new(config['key'], AES.MODE_EAX, nonce=nonce)

    # check if file exists
    if not os.path.exists(config['filepath']):
        print('file does not exists: ' + config['filepath']) 
        sys.exit(1)
    # read the file content
    file_out = open(config['filepath'], "rb")
    file_content = file_out.read()
    file_out.close()

    message = nonce + cipher.encrypt(file_content)
    # encrypt and send file content to server
    sock.sendall(message)

    # print debug message
    print('7. send: upload ' + config['filepath'] + ' as ' + config['filename'])

def receive_download(sock, config):
    # receive data from server
    file_content = sock.recv(BUFFER_SIZE)

    # create cipher using first 8 bytes from payload as nonce
    cipher = AES.new(config['key'], AES.MODE_EAX, nonce=file_content[:8])
    plain = cipher.decrypt(file_content[8:])


    # check if filepath exists
    dirname = os.path.dirname(config['filepath'])
    if not os.path.exists(dirname):
        os.makedirs(dirname)
    # and save the file  locally
    file_out = open(config['filepath'], "wb")
    file_out.write(plain)
    file_out.close()
    # print debug message
    print('8. recv: download ' + config['filename'] + ' into ' + config['filepath'])


def generate_keys(config):
   key = RSA.generate(2048)
   private_key = key.export_key()
   dirname = os.path.dirname(config['privkey'])
   if not os.path.exists(dirname):
       os.makedirs(dirname)
   file_out = open(config['privkey'], "wb")
   file_out.write(private_key)
   file_out.close()
   if not os.path.exists(config['pubkeys']):
       os.makedirs(config['pubkeys'])
   public_key = key.publickey().export_key()
   file_out = open(os.path.join(config['pubkeys'], CLIENT + '.pem'), "wb")
   file_out.write(public_key)
   file_out.close() 

def needham_schroeder_client(config):
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
        # handshake
        sock.connect((host, port))
        send_client_hello(sock, config)
        receive_server_hello(sock, config)
        send_key(sock, config)
        receive_authenticated(sock, config)
        # data exchange
        send_request(sock, config)
        if config['request'] == 'upload':
            receive_ready(sock, config)
            send_upload(sock, config)
        elif config['request'] == 'download':
            receive_download(sock, config)
    
if __name__ == "__main__":
    import os, sys, getopt
    def usage():
        print ('Usage:    ' + os.path.basename(__file__) + ' options filepath ')
        print ('Options:')
        print ('\t -t to, --to=to')
        print ('\t -x pubkeys, --pubkeys=pubkeys')
        print ('\t -i privkey, --privkey=privkey')
        print ('\t -u, --upload')
        print ('\t -d, --download')
        print ('\t -f filename, --filename=filename')
        sys.exit(2)
    try:
      opts, args = getopt.getopt(sys.argv[1:],"hudp:s:t:x:i:f:",["help", "upload", "download", "to=", "pubkeys=", "privkey=", "filename="])
    except getopt.GetoptError as err:
      print(err)
      usage()
    # extract parameters
    request = None
    to = None
    pubkeys = None
    privkey = None
    filename = None
    filepath = args[0] if len(args) > 0 else None
    for opt, arg in opts:
        if opt in ("-h", "--help"):
           usage()
        elif opt in ("-u", "--upload"):
           request = 'upload'
        elif opt in ("-d", "--download"):
           request = 'download'
        elif opt in ("-t", "--to"):
           to = arg
        elif opt in ("-x", "--pubkeys"):
           pubkeys = arg
        elif opt in ("-i", "--privkey"):
           privkey = arg
        elif opt in ("-f", "--filename"):
           filename = arg
    # check arguments
    if (request is None):
       print('upload/download option is missing\n')
       usage()
    if (to is None):
       print('to option is missing\n')
       usage()
    if (pubkeys is None):
       print('pubkeys option is missing\n')
       usage()
    if (privkey is None):
       print('privkey option is missing\n')
       usage()       
    if (filename is None):
       print('filename option is missing\n')
       usage()
    if (filepath is None):
       print('filepath is missing\n')
       usage()
    # create config
    config = {'request': request, 'privkey': privkey, 'pubkeys': pubkeys, 'filename': filename, 'filepath': filepath}
    # extract server information
    config['to'] = to.split("@")[0]
    host = to.split("@")[1].split(":")[0]
    port = int(to.split(":")[1])
    # generate key pairs (if not exists already)
    if not os.path.exists(privkey) or not os.path.exists(os.path.join(pubkeys, CLIENT + '.pem')) :
        generate_keys(config)
    # run the client
    needham_schroeder_client(config)
        



