import os
import sys
import time
import socket
from Cryptodome.PublicKey import RSA
from Cryptodome.Cipher import AES, PKCS1_OAEP

PORT = 50001
RUNNING_DIRECTORY = os.getcwd()


def receive_file(file_name, key, sock):
    """

    :param file_name:
    :param key:
    :param sock:

    alerts server for file download
    read's from the file and each time encrypts it and send's
    it away
    """
    send_small("download " + file_name, key, sock)
    print "receiving file"
    file_out = open(RUNNING_DIRECTORY + os.sep + "Recived_files" + os.sep + file_name, 'wb')
    l1 = recive_small(key, sock)
    while not l1 == "!@#":
        file_out.write(l1)
        l1 = recive_small(key, sock)
    file_out.close()
    print "finished"
    return


def send_small(text, key, sock):
    """

    :param text:
    :param key:
    :param sock:
    encrypte's a given string and sends it
    without alerting the server
    """
    print "Key: " + key
    cipher = AES.new(key, AES.MODE_EAX)
    nonce = cipher.nonce
    ciphertext, tag = Encryptsmall(cipher, text)
    sock.send(ciphertext)
    time.sleep(0.5)
    sock.send(tag)
    time.sleep(0.5)
    sock.send(nonce)


def recive_small(key, sock):
    """

    :param key:
    :param sock:
    receive's massage and decrypts it
    :return returns plain text if decryption successfully
    """
    text = sock.recv(4096)
    tag = sock.recv(4096)
    nonce = sock.recv(4096)
    cipher = AES.new(key, AES.MODE_EAX, nonce=nonce)
    return Decryptsmall(cipher, text, tag)


def Encryptsmall(cipher, data="Testing"):
    """

    :param cipher:
    :param data:
    :return: encrypted data
    just encrypts a given string
    """
    ciphertext, tag = cipher.encrypt_and_digest(data)
    return ciphertext, tag


def Decryptsmall(cipher, ciphertext, tag):
    """

    :param cipher:
    :param ciphertext:
    :param tag:
    :return: plain text of received message
    """
    plaintext = cipher.decrypt(ciphertext)
    try:
        cipher.verify(tag)
        print("The message is authentic:", plaintext)
    except ValueError:
        print("Key incorrect or message corrupted")

    return plaintext


def send_file(sending_file, key, sock):
    """

    :param sending_file:
    :param key:
    :param sock:
    alerts server for an upload encrypts and sends the
    encrypted data
    """
    send_small("upload " + sending_file.split(os.sep)[-1], key, sock)
    file_to_send = open(sending_file, "rb")
    l = file_to_send.read(4096)
    while l:
        send_small(l, key, sock)
        l = file_to_send.read(4096)
        print "."
    send_small("!@#", key, sock)
    print "Finished uploading"


def decrypt_data(opened_file):
    """

    :param opened_file:
    :gets encrypted data
    using rsa decrypts it used for the key transfer
    """
    file_out = open("temp2.bin", "wb")
    private_key = RSA.import_key(open('private.pem').read())

    print "."

    enc_session_key, nonce, tag, ciphertext = \
        [opened_file.read(x) for x in (private_key.size_in_bytes(), 16, 16, -1)]

    print "."

    # Decrypt the session key with the private RSA key
    cipher_rsa = PKCS1_OAEP.new(private_key)
    session_key = cipher_rsa.decrypt(enc_session_key)

    print "."

    # Decrypt the data with the AES session key
    cipher_aes = AES.new(session_key, AES.MODE_EAX, nonce)
    data = cipher_aes.decrypt_and_verify(ciphertext, tag)
    file_out.write(data)
    print "finished decrypting"


def generate_keys():
    '''
    Generate an RSA keypair with an exponent of 65537 in PEM format
    param: bits The key length in bits
    Return private key and public key
    '''

    key = RSA.generate(2048)
    private_key = key.export_key()
    file_out = open("Private.pem", "wb+")
    file_out.write(private_key)
    file_out.close()

    public_key = key.publickey().export_key()
    file_out = open("Public.pem", "wb+")
    file_out.write(public_key)
    file_out.close()


def starting_up():
    """
    makes sure all the folder's that are needed exists
    including the private and public key
    also connects the client to the server by asking for ip
    :return connected socket and the encryption key
    """

    Pass = False

    if not (os.path.exists(RUNNING_DIRECTORY + '/Public.pem')):
        print ("Generating Private and Public key.......")
        generate_keys()

    else:
        print ("Private and Public keys found")

    if not (os.path.exists(RUNNING_DIRECTORY + '/Recived_files')):
        os.mkdir(RUNNING_DIRECTORY + '/Recived_files')

    print "Enter IP Of the Server"

    while not Pass:
        ip = raw_input()
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        address = (ip, 50001)
        try:
            sock.connect(address)
            Pass = True
        except Exception, e:
            print (e)
            print "wrong ip\server isn't open try again"

    f = open("Public.pem", 'rb')
    sock.send(f.read(4096))

    f = open("temp.bin", 'wb')  # opens file for encrypted data
    f.write(sock.recv(4096))
    f.close()
    f = open("temp.bin", 'rb')  # opens file in read mode

    decrypt_data(f)
    f.close()

    os.remove("temp.bin")  # removes uneccesery files

    f = open('temp2.bin')
    key = f.read(4096)
    f.close()

    os.remove('temp2.bin')

    print key

    return key, sock


def main():

    command = ""
    password = ""
    username = ""
    server_response = ""
    key, sock = starting_up()

    repassword = "!"

    while True:

        print 'would you like to:'
        print '1) sign in'
        print '2) sign up'
        print '3) exit'

        response = raw_input()

        if response == "1":
            print ("Username: ")
            username = raw_input()
            print ("Password: ")
            password = raw_input()
            send_small("SI " + username + " " + password, key, sock)

            server_response = recive_small(key, sock)

            if server_response == "Cool":
                print "Welcome back " + username
                break
            else:
                print "wrong username/password"
                continue

        elif response == "2":
            print ("Username: ")
            username = raw_input()
            while password != repassword:
                print ("Password: ")
                password = raw_input()
                print "Repeat password"
                repassword = raw_input()
                if password == repassword:
                    send_small("S_U " + username + " " + password, key, sock)

                    server_response = recive_small(key, sock)

                    if server_response == "Cool":
                        print "welcome " + username
                        break
                else:
                    print "passwords don't match try again"
                    continue

            if server_response == "Cool":
                break
            else:
                print "username already taken"
                continue

        elif response == "3":
            send_small("bye", key, sock)
            sock.close()
            sys.exit()

        else:
            continue

    while command != "disconnect":
        print "command:"
        command = raw_input()

        if command.split(" ")[0] == "list":
            send_small("list", key, sock)
            print recive_small(key, sock)

        elif command.split(" ")[0] == "upload":
            send_file(command.split(" ")[1], key, sock)

        elif command.split(" ")[0] == 'download':
            receive_file(command.split(" ")[1], key, sock)

        elif command.split(" ")[0] == "rm":
            send_small("rm " + command.split(" ")[1], key, sock)
            print recive_small(key, sock)
        elif command.split(" ")[0] == "help":
            print "list: shows the files names in ur storage \n"
            print "upload + file path will send to the cloud the \n" \
                  "file that you gave the path to \n"
            print "download + file name will download the file that\n" \
                  "you desire REMEMBER TO ADD THE FILE TYPE\n"
            print "rm + file will remove the file you desire from the cloud \n"
            print "disconnect will literally disconnect you\n"
        else:
            print "no such command " + command + " use 'help'"
    send_small("bye", key, sock)


main()