"""Secure client implementation

This is a skeleton file for you to build your secure file store client.

Fill in the methods for the class Client per the project specification.

You may add additional functions and classes as desired, as long as your
Client class conforms to the specification. Be sure to test against the
included functionality tests.

Hits from part1

## CreateFakeKey                                  ## FIXED
It is important to bind the name of a file to its contents. Otherwise, if
a user creates files A and B, a malicious server could swap the names and have downloads
for A return the content of B and vice versa.

## CheckEncryptionMACDifferentKey                 ## FIXED
MAC keys and encryption keys should be different,
and this test case checks just that.

## CheckPaddingOracle                             ## pretty sure it's FIXED
A padding oracle attack is a chosen-ciphertext attack: if the attacker
can convince the client to decrypt (tampered) ciphertexts chosen by the attacker,
then a padding oracle attack allows the attacker to decrypt an encrypted message. In this
attack, the attacker modifies the ciphertext in some cleverly chosen fashion, asks the client
to decrypt it, and then observes the decryption process caused an invalid-padding error. If
the attacker can observe whether such an error occurred, then this leaks partial information;
after repeating this many times, an attacker can piece together all of these clues to deduce
what the original message must have been. These attacks are very powerful and are subtle
if you don’t know about them.
Failure of this test case typically indicates that you used encryption without authentication,
or that you decrypted some ciphertext without checking the MAC or before verifying the
validity of the MAC. Defend against padding oracle attacks by always verifying the MAC
before decrypting message contents.


## NameIsNotFromHash                                ## FIXED (Using HMAC, hope it's ok)
Verifies that IDs are not generated by hashing the filename.
That has the same problems outlined [below].
IDs should not be generated by encrypting the filename with
ECB mode or CBC/CTR mode with constant IV, as this leaks partial information. In
particular, if two filenames start with the same 16 bytes, then the adversary can detect this
(by noticing that the ciphertexts start with the same 16 bytes). Also, these modes allow
dictionary attacks: an adversary who can convince you to upload files under a name of the
attacker’s choice will be able to exploit the properties of ECB/CBC/CTR to figure out the
filename of your secret file, if the attacker can identify a list of candidate names for the secret
file. Finally, CTR mode with a constant IV is especially bad: it is vulnerable to the attacks
on pad reuse with the one-time pad.

"""
from base_client import BaseClient, IntegrityError
import json
from crypto import CryptoError
from util import *


def path_join(*strings):
    """Joins a list of strings putting a "/" between each.

    :param strings: a list of strings to join
    :returns: a string
    """
    return '/'.join(strings)


class Client(BaseClient):
    def __init__(self, storage_server, public_key_server, crypto_object,
                 username):
        super().__init__(storage_server, public_key_server, crypto_object,
                         username)

        self.client_storage = {}  # Make a place for the client to keep state for efficiency.

        # deterministic location for filename signature key
        location = self.username + "/key_file"
        location = self.crypto.message_authentication_code(location, str(self.rsa_priv_key.key.d)[:64], "SHA256")

        # set key to what was previously stored
        encrypted_key = self.storage_server.get(location)

        if not encrypted_key:  # check if key exists
            # no key exists
            self.sk_n = self.crypto.get_random_bytes(32)  # key for filename
            self.sk_n2 = self.crypto.get_random_bytes(32)  # key for filename
            keys = [self.sk_n, self.sk_n2]
            keys = json.dumps(keys)
            packed_key = self.asym_pack(keys, "key_file")
            self.storage_server.put(location, packed_key)  # add new secret key to the server, with signature attached.
        else:  # a previous key exists
            # verify the key is valid
            keys = self.asym_unpack(encrypted_key, "key_file")
            keys = json.loads(keys)
            self.sk_n, self.sk_n2 = keys

    ######################################################################################
    # Helper Functions                                                                  ##
    ######################################################################################

    def hash(self, msg):
        return self.crypto.cryptographic_hash(msg, "SHA256")

    def hash2(self, msg, key):
        """

        :param str msg: what you gonna hash
        :param str key: key to hash with
        :return: HMAC of msg
        :rtype: str
        """
        return self.crypto.message_authentication_code(msg, key, "SHA256")

    def asym_pack(self, pt, filename, to=None):
        """Securely pack a plaintext for storage on the insecure server

            :param str pt: The key to pack.
            :param str to: user who's public key encrypts the msg
            :param str filename: the correct filename to avoid swap attacks

            :returns: the packed key
            :rtype: str
        """
        if not to:
            to = self.username
        if len(pt) > 256:
            raise CryptoError("plaintext too long")
        try:
            encrypted_key = self.crypto.asymmetric_encrypt(pt, self.pks.get_encryption_key(to))  # encrypt key
        except:
            raise CryptoError("bad 'to' user")
        signed = encrypted_key + self.hash(self.username + filename) + "ct"
        signed = self.crypto.asymmetric_sign(signed, self.rsa_priv_key)  # sign key
        return signed + encrypted_key

    def asym_unpack(self, ct, filename, sender=None):
        """unpack a symmetric key that has been stored on the insecure server

            :param str ct: The key to unpack.
            :param str sender: user who's public key is used to verify
            :param str filename: the correct filename to avoid swap attacks

            :returns: the symmetric key
            :rtype: str

            :raises IntegrityError: if the signature is invalid
        """
        if not ct:
            return None
        if not sender:
            sender = self.username
        # verify the key is valid
        signed, check = ct[:512], ct[512:] + self.hash(sender + filename) + "ct"
        if self.crypto.asymmetric_verify(check, signed, self.pks.get_signature_key(sender)):
            ct = ct[512:]  # if valid, extract the cipher text
            try:
                pt = self.crypto.asymmetric_decrypt(ct, self.elg_priv_key)  # decrypt
            except:
                raise IntegrityError("decryption failure")
        else:
            raise IntegrityError("Bad Signature")
        return pt

    def sym_pack(self, pt, sk_m, sk_s, filename):
        """
        symmetrically encrypts and signs a plaintext string

        :param str pt: the plaintext to be packed
        :param str sk_m: symmetric key for encrypting
        :param str sk_s: symmetric key for signing
        :param str filename: the correct filename to avoid swap attacks

        :return: the ciphertext as a string with IV and signature concatenated
        :rtype: str
        """
        iv = self.crypto.get_random_bytes(16)  # get 16-byte IV (size of digest_block)
        value = self.crypto.symmetric_encrypt(pt, sk_m, 'AES', 'CBC', iv)  # encrypt w/AES-CBC
        value += iv  # append IV to msg CT
        value += self.crypto.message_authentication_code(value + self.hash(filename), sk_s, "SHA256")
        return value

    def sym_unpack(self, ct, sk_m, sk_s, filename):
        """
        symmetrically decrypts and verifies a ciphertext string
        :param ct: ciphertext to be unpacked
        :param sk_m: symmetric key for encrypting
        :param sk_s: symmetric key for signing
        :param str filename: the correct filename to avoid swap attacks

        :return: the plaintext
        :rtype: str
        """
        if not ct:
            return None
        sign = ct[-64:]  # get signature from end of msg CT
        resp = ct[:-64]  # remove signature from msg CT
        if self.crypto.message_authentication_code(resp + self.hash(filename), sk_s, "SHA256") != sign:  # verify
            raise IntegrityError("Bad Signature")
        try:
            iv = resp[-32:]  # save the IV at the end of the message CT
            resp = resp[:(-32)]  # remove IV from message CT
            resp = self.crypto.symmetric_decrypt(resp, sk_m, "AES", "CBC", iv)  # decrypt msg w/AES-CBC
        except:  # any exception should be due to bit flipping by malicious server.
            raise IntegrityError("Data Corrupt")
        return resp

    def get_chunks(self, value, size=256):
        """ Returns a list made up of the value broken into chunks of length "size"

        :param str value: the value to break down
        :param int size: how large a given chunk is
        :return: chunks which together make up the value
        :rtype: list[str]
        """
        chunks = []
        while value:
            if len(value) > size:
                chunks.append(value[:size])
                value = value[size:]
            else:
                chunks.append(value)
                value = ""
        return chunks

    def build_hash_tree(self, chunks, sk_t):
        """
        uses a list of hashes to build a Merkle tree on the server based at Location
        :param chunks: the list of chunks which make up the file
        :return: client side copy of the tree
        """
        # make hash_list
        hashes = []
        for i in range(len(chunks)):
            hashes.append([self.hash2(chunks[i], sk_t), None, None, i])  # integrity hash

        # organize into merkle tree
        while len(hashes) > 1:
            temp = []
            for i in range(0, len(hashes), 2):  # pair up the hashes
                try:
                    parent = self.hash(hashes[i][0]+hashes[i+1][0])  # add the hash of two hashes to temp
                    temp.append([parent, hashes[i], hashes[i + 1]])
                except IndexError:
                    parent = self.hash(hashes[i][0]+hashes[i][0])  # if the num hashes is odd, hash last with itself
                    temp.append([parent, hashes[i], None])
            hashes = temp

        return hashes[0]

    def put_tree(self, node_loc, tree):
        """
        recursively places tree nodes in the server. assumes secure location
        :param node_loc: location to start placing nodes at
        :param tree: the tree to place onto the server
        :return:
        """
        if tree:
            self.storage_server.put(node_loc, tree[0])  # if hash is modified, it don't matter
            self.put_tree(node_loc + '1', tree[1])
            self.put_tree(node_loc + '2', tree[2])

    def compare_server_tree(self, location, tree):
        """
        takes the server location of a hash tree and a client side, and returns the locations in the
        chunk list where there are discrepancies.
        :param location: where to start the tree
        :param tree: where to
        :return: indexes of chunks with discrepancies
        """
        bad = []
        server_root = self.storage_server.get(location)
        if tree and tree[0] != server_root:
            if len(tree) == 4:
                return [tree[3]]
            elif not server_root:
                return self.get_leaf_indices(tree)
            else:
                bad += self.compare_server_tree(location + '1', tree[1])
                bad += self.compare_server_tree(location + '2', tree[2])
        return bad

    def update_server_tree(self, location, tree, old_tree=None):
        """
        takes the server location of a hash tree and a client side, and returns the locations in the
        chunk list where there are discrepancies, changes the server tree to match the client.
        :param location: where to start the tree
        :param tree: where to
        :return: indexes of chunks with discrepancies
        """
        bad = []
        server_root = self.storage_server.get(location)
        if old_tree:
            if tree and old_tree:
                if tree[0] != old_tree[0]:
                    self.storage_server.put(location, tree[0])
                    try:
                        return [tree[3]]
                    except IndexError:
                        bad += self.update_server_tree(location + '1', tree[1], old_tree[1])
                        bad += self.update_server_tree(location + '2', tree[2], old_tree[2])
        elif tree and tree[0] != server_root:
            self.storage_server.put(location, tree[0])
            if len(tree) == 4:
                return [tree[3]]
            else:
                bad += self.update_server_tree(location + '1', tree[1])
                bad += self.update_server_tree(location + '2', tree[2])
        return bad

    def compare_hash_tree(self, old_tree, new_tree):
        """
        takes two hash trees and returns the locations in the chunk list where there are discrepancies.
        :param old_tree:
        :param new_tree:
        :return: indexes of chunks with discrepancies
        """
        bad = []
        if new_tree and old_tree:
            if new_tree[0] != old_tree[0]:
                try:
                    return [new_tree[3]]
                except IndexError:
                    bad += self.compare_hash_tree(old_tree[1], new_tree[1])
                    bad += self.compare_hash_tree(old_tree[2], new_tree[2])
        elif new_tree:
            bad = self.get_leaf_indices(new_tree)
        return bad

    def get_leaf_indices(self, tree):
        indices = []
        if tree:
            try:
                return [tree[3]]
            except IndexError:
                indices += self.get_leaf_indices(tree[1])
                indices += self.get_leaf_indices(tree[2])
        return indices

    def resolve(self, uid):
        """
        navigates through and unpacks a series of linked pointers until it reaches a [DATA] file.

        :param uid: starting location

        :return uid: The location of the [DATA] file
        :return sk_m: the decryption , symmetric key
        :return sk_s: the signature symmetric key
        :rtype str, str, str
        """
        sk_n2 = self.sk_n2
        sk_m = None
        sk_s = None
        sk_t = None
        owner = self.username
        requester = self.username
        while True:
            res = self.storage_server.get(uid)
            if res is None or res.startswith("[DATA]"):
                return uid, sk_m, sk_s, sk_t
            elif res.startswith("[POINTER]"):
                pointer = res[10:]  # remove [POINTER] tag
                signature = pointer[-512:]
                key_list = pointer[:-512]
                if not self.crypto.asymmetric_verify(key_list + owner, signature, self.pks.get_signature_key(owner)):
                    raise IntegrityError("bad key_list verification")
                key_list = json.loads(key_list)  # Load the dictionary with names
                try:
                    pointer = key_list[self.hash2(requester, sk_n2)]  # go to your place in dictionary
                    if sk_m and sk_s:  # if you have symmetric keys, you are following the chain'
                        for i in range(6):
                            pointer[i] = self.sym_unpack(pointer[i], sk_m, sk_s, uid)
                    else:  # if you do not have symmetric keys, you should own this pointer.
                        pointer = [self.asym_unpack(p, uid) for p in pointer]
                except KeyError:
                    self.storage_server.put("invalid", "restricted")
                    return "invalid", None, None, None
                requester = owner  # you are now requesting as the owner of the last pointer
                uid, sk_m, sk_s, owner, sk_n2, sk_t = pointer[:6]  # update the location, keys and who owns the pointer
            else:
                raise IntegrityError()

    ######################################################################################
    # Member Functions                                                                  ##
    ######################################################################################

    def upload(self, name, value, revoke=False, sk_m=None, sk_s=None, sk_t=None):
        # Set UID as a hash of the username and value, use salt/signed hash
        if not revoke:
            uid = self.hash(path_join(self.username, name))  # join username and file name
            uid = self.crypto.message_authentication_code(uid, self.sk_n, "SHA256")  # use HMAC to hash/salt
            uid, sk_m, sk_s, sk_t = self.resolve(uid)  # use hashed uid to resolve
        else:
            uid = name

        # If there are no keys, then this is a new file
        if not sk_m or not sk_s:
            sk_m = self.crypto.get_random_bytes(32)  # generate a symmetric key for the message
            sk_s = self.crypto.get_random_bytes(32)  # key for signing
            sk_t = self.crypto.get_random_bytes(32)  # tree key
            location = self.crypto.get_random_bytes(64)
            new = True
        else:
            location = uid
            new = False

        # Encrypt message
        pt_chunks = self.get_chunks(value)
        ct_chunks = []
        for i in range(len(pt_chunks)):
            ct_chunks.append(self.sym_pack(pt_chunks[i], sk_m, sk_s, path_join(location,str(i))))  # you always need the chunks encrypted
        hash_tree = self.build_hash_tree(pt_chunks, sk_t)  # location/tree12212...12121

        # no state upload
        if new:
            for i in range(len(ct_chunks)):
                self.storage_server.put(path_join(location, str(i)), ct_chunks[i])  # location/###
            self.put_tree(path_join(location, "tree"), hash_tree)
        else:  # stateful upload
            if name in self.client_storage:
                state_htree = self.client_storage[name][1]
                if state_htree[0] == self.storage_server.get(path_join(location, "tree")):  # the root hashes are equal
                    # we need to only modify the chunks that are different in the update and our local
                    # push updates to the server
                    change = self.update_server_tree(path_join(location, "tree"), hash_tree, state_htree)
                else:
                    change = self.update_server_tree(path_join(location, "tree"), hash_tree)
            else:
                change = self.update_server_tree(path_join(location, "tree"), hash_tree)
            for i in change:
                self.storage_server.put(path_join(location, str(i)), ct_chunks[i])

        # pack/store data
        self.client_storage[name] = [value, hash_tree, pt_chunks]  # update client state
        value = "[DATA]"
        self.storage_server.put(location, value)

        # setup pointer to file location, if it does not yet exist
        if new:
            key_list = {}
            sk_n2 = self.crypto.get_random_bytes(32)  # key for HMAC of names in keylsit
            pointer = [location, sk_m, sk_s, self.username, sk_n2, sk_t]  # make list of keys and location
            pointer = [self.asym_pack(x, uid) for x in pointer]  # encrypt each item
            key_list[self.hash2(self.username, self.sk_n2)] = pointer
            key_list = json.dumps(key_list)
            signature = self.crypto.asymmetric_sign(key_list + self.username, self.rsa_priv_key)   # sign key_list as owner
            key_list += signature
            key_list = "[POINTER] " + key_list
            self.storage_server.put(uid, key_list)

    def download(self, name, revoke=False, uid=None, sk_m=None, sk_s=None, sk_t=None):
        # Set UID as a hash of the username and value, use salt/signed hash
        if not revoke:
            uid = self.hash(path_join(self.username, name))  # join username and file name
            uid = self.crypto.message_authentication_code(uid, self.sk_n, "SHA256")  # use HMAC to hash/salt
            uid, sk_m, sk_s, sk_t = self.resolve(uid)  # use hashed uid to resolve

        # Decrypt message
        resp = self.storage_server.get(uid)  # get CT from uid
        if resp is None or resp == "restricted":  # check if empty or revoked
            return None
        if name in self.client_storage:
            state_htree = self.client_storage[name][1]
            if state_htree[0] == self.storage_server.get(path_join(uid, "tree")):
                # the top hashes match, return the stored value
                pass
            else:
                # find the changes in the file and update the client version
                changes = self.compare_server_tree(path_join(uid, "tree"), state_htree)
                for i in changes:
                    ct_chunk = self.storage_server.get(path_join(uid, str(i)))
                    self.client_storage[name][2][i] = self.sym_unpack(ct_chunk, sk_m, sk_s, path_join(uid, str(i)))
                self.client_storage[name][0] = "".join(self.client_storage[name][2])  # update client state value
        else:  # no state exists, download all chunks
            self.client_storage[name] = ["", [], []]
            i = 0
            ct_chunk = self.storage_server.get(path_join(uid, str(i)))
            while ct_chunk:
                self.client_storage[name][2].append(self.sym_unpack(ct_chunk, sk_m, sk_s, path_join(uid, str(i))))
                i += 1
                ct_chunk = self.storage_server.get(path_join(uid, str(i)))
            self.client_storage[name][0] = "".join(self.client_storage[name][2])  # update client state value
            self.client_storage[name][1] = self.build_hash_tree(self.client_storage[name][2], sk_t)
        # print(self.username, "got", self.client_storage[name][0])
        return self.client_storage[name][0]

    def share(self, user, name):
        """
        create a new entry on your key_list file and add the user you intend to share with, then
        create a signed message which will let them find your pointer, and access their entry on
        the key_list file with a symmetric key.


        :param user: Who you are sharing with
        :param name: the file you plan to share

        :return: message, which will be sent to 'user'
        :rtype: str
        """
        uid = self.hash(path_join(self.username, name))  # join username and file name
        uid = self.crypto.message_authentication_code(uid, self.sk_n, "SHA256")  # use HMAC to hash/salt

        # open the pointer
        pointer = self.storage_server.get(uid)
        pointer = pointer[10:]  # remove [POINTER] tag
        signature = pointer[-512:]  # remove signature
        key_list = pointer[:-512]  # remove key_list

        if not self.crypto.asymmetric_verify(key_list + self.username, signature, self.pks.get_signature_key(self.username)):
            raise IntegrityError("bad key_list signature")
        key_list = json.loads(key_list)  # Load the dictionary with names

        # Get the info for the Data/Pointer that this pointer is aimed at
        try:
            pointer = key_list[self.hash2(self.username, self.sk_n2)]
            pointer = [self.asym_unpack(p, uid) for p in pointer]
        except KeyError:
            raise IntegrityError("own key entry missing")

        # make the new pointer
        sk_m = self.crypto.get_random_bytes(32)  # generate a symmetric key for the message
        sk_s = self.crypto.get_random_bytes(32)  # key for signing
        sk_n2 = pointer[4]  # key for HMAC of name (shared for the whole pointer)
        pointer = [self.sym_pack(p, sk_m, sk_s, uid) for p in pointer]
        owner_access = [self.asym_pack(sk_m, uid), self.asym_pack(sk_s, uid)]  # let owner remember keys for revocation
        pointer += owner_access  # append the owner's keys to the back of the pointer
        key_list[self.hash2(user, sk_n2)] = pointer
        key_list = json.dumps(key_list)
        signature = self.crypto.asymmetric_sign(key_list + self.username, self.rsa_priv_key)  # sign key_list as owner
        key_list += signature
        key_list = "[POINTER] " + key_list
        self.storage_server.put(uid, key_list)
        msg = [uid, sk_m, sk_s, sk_n2]
        msg = [self.asym_pack(ms, "msg", user) for ms in msg]
        msg = json.dumps(msg)
        return msg

    def receive_share(self, from_username, new_name, message):
        """
        This should just take the info and turn it into a new pointer which is located at [hash(newname)]

        :param from_username: the sharer
        :param new_name: the new filename for this user's storage
        :param message: encrypted message from sharer

        :return: None
        """
        uid = self.hash(path_join(self.username, new_name))  # join username and file name
        uid = self.crypto.message_authentication_code(uid, self.sk_n, "SHA256")  # use HMAC to hash/salt

        # verify message
        message = json.loads(message)
        message = [self.asym_unpack(msg, "msg", from_username) for msg in message]
        location, sk_m, sk_s, sk_n2 = message

        key_list = {}
        pointer = [location, sk_m, sk_s, from_username, sk_n2, ""]  # make list of keys and location
        pointer = [self.asym_pack(p, uid) for p in pointer]  # encrypt each part of the pointer
        key_list[self.hash2(self.username, self.sk_n2)] = pointer
        key_list = json.dumps(key_list)
        key_list += self.crypto.asymmetric_sign(key_list + self.username, self.rsa_priv_key)  # sign key_list as owner
        key_list = "[POINTER] " + key_list
        self.storage_server.put(uid, key_list)
        if self.storage_server.get(uid) != key_list:
            raise IntegrityError("Bad Upload")

    def revoke(self, user, name):
        """
        remove user from list of people with access rights, update other users with the new key.

        :param user: name of the user to be revoked
        :param name: file to revoke the user from

        :return: True if success, False otherwise
        :rtype: bool
        """
        uid = self.hash(path_join(self.username, name))  # join username and file name
        uid = self.crypto.message_authentication_code(uid, self.sk_n, "SHA256")  # use HMAC to hash/salt

        # load the pointer list
        pointer = self.storage_server.get(uid)
        pointer = pointer[10:]  # remove [POINTER] tag
        signature = pointer[-512:]
        key_list = pointer[:-512]
        if not self.crypto.asymmetric_verify(key_list + self.username, signature, self.pks.get_signature_key(self.username)):
            raise IntegrityError("bad key_list signature")
        key_list = json.loads(key_list)  # Load the dictionary with names

        # generate a new location and set of symmetric keys for the file
        sk_m = self.crypto.get_random_bytes(32)  # generate a symmetric key for the message
        sk_s = self.crypto.get_random_bytes(32)  # key for signing
        sk_n2 = self.crypto.get_random_bytes(32)  # key for HMAC
        sk_t = self.crypto.get_random_bytes(32)  # Key for trees
        location = self.crypto.get_random_bytes(64)
        new_pointer = [location, sk_m, sk_s, self.username, sk_n2, sk_t]

        # get the old keys
        old_pointer = key_list[self.hash2(self.username, self.sk_n2)]
        old_pointer = [self.asym_unpack(p, uid) for p in old_pointer]

        # update the file
        if self.storage_server.get(old_pointer[0]).startswith("[POINTER]"):
            print("you can't revoke if you don't own")
            return False
        file = self.download(name, True, old_pointer[0], old_pointer[1], old_pointer[2], old_pointer[5])
        self.upload(new_pointer[0], file, True, new_pointer[1], new_pointer[2], new_pointer[5])

        # clean up
        self.storage_server.delete(old_pointer[0])

        # update your keys
        pointer = [self.asym_pack(np, uid) for np in new_pointer]
        key_list[self.hash2(self.username, self.sk_n2)] = pointer

        del key_list[self.hash2(user, old_pointer[4])]  # remove the offending user

        # update everyone else's keys
        keys = list(key_list.keys())
        keys.remove(self.hash2(self.username, self.sk_n2))  # don't update the owner pointer
        for key in keys:
            pointer = key_list[key]
            share_sk_m = self.asym_unpack(pointer[-2], uid)
            share_sk_s = self.asym_unpack(pointer[-1], uid)
            foo = [self.sym_pack(np, share_sk_m, share_sk_s, uid) for np in new_pointer]
            key_list[key] = foo + pointer[-2:]

        key_list = json.dumps(key_list)
        key_list += self.crypto.asymmetric_sign(key_list + self.username, self.rsa_priv_key)  # sign key_list as owner
        key_list = "[POINTER] " + key_list
        self.storage_server.put(uid, key_list)
        if self.storage_server.get(uid) != key_list:
            raise IntegrityError("Bad Upload")

        return True


if __name__ == "__main__":
    # A basic unit test suite for the insecure Client to demonstrate
    # its functions.
    from servers import PublicKeyServer, StorageServer
    from crypto import Crypto

    print("Initializing servers and clients...")
    pks = PublicKeyServer()
    server = StorageServer()
    crypto = Crypto()
    alice = Client(server, pks, crypto, "alice")
    bob = Client(server, pks, crypto, "bob")
    carol = Client(server, pks, crypto, "carol")
    dave = Client(server, pks, crypto, "dave")

    print("testing private functions")
    a = "f"*256
    asmCT = alice.asym_pack(a, "test")
    if a == alice.asym_unpack(asmCT, "test"):
        print("simple asym enc/dec work")
    asmCT = alice.asym_pack(a, "test", 'bob')
    if a == bob.asym_unpack(asmCT, "test", 'alice'):
        print("shared asym enc/dec work")

    a = "The quick brown fox jumped over the lazy dog!1"*1000
    b = "The quick brown fox jumped over the lazy dog!1"*999 + "The quick brown fox jumped over the lazy dog!2"
    c = "The quick brown fox jumped over the lazy dog!1"*499 + "The quick brown fox jumped over the lazy dog!2" + "The quick brown fox jumped over the lazy dog!1" * 500
    print("test that the chunk/tree system works")
    chunks_a = alice.get_chunks(a)
    print("chunks work: ", "".join(chunks_a) == a)
    tree_a = alice.build_hash_tree(chunks_a)
    alice.put_tree("test_tree a", tree_a)

    chunks_b = alice.get_chunks(b)
    tree_b = alice.build_hash_tree(chunks_b)
    alice.put_tree("test_tree b", tree_b)

    chunks_c = alice.get_chunks(c)
    tree_c = alice.build_hash_tree(chunks_c)
    alice.put_tree("test_tree c", tree_c)

    for a, b in zip(chunks_a,chunks_b):
        e = compute_edits(a, b)
        if e:
            print(e)

    for b, c in zip(chunks_b,chunks_c):
        e = compute_edits(b, c)
        if e:
            print(e)

    test_key = alice.crypto.get_random_bytes(32)

    tlen = json.dumps(tree_a)
    print(len(a))
    print(len(tlen))
    print(len(alice.sym_pack(a, test_key, test_key, test_key)))
    print(alice.compare_hash_tree(tree_a, tree_b))
    print(alice.compare_server_tree("test_tree a", tree_a))
    print(alice.compare_server_tree("test_tree b", tree_b))
    print(alice.compare_server_tree("test_tree b", tree_a))

    print(len(chunks_a[0]))
    print(len(alice.sym_pack(chunks_a[0], test_key, test_key, test_key)))

    print("Testing client put and share...")
    a = "The quick brown fox jumped over the lazy dog!1"
    alice.upload("a", a)
    b = alice.download("a")
    if b == a:
        print("success!")
    else:
        print("a: ", a)
        print("b: ", b)
        print("Failed :(")

    alice.upload("a", "b")

    m = alice.share("bob", "a")
    bob.receive_share("alice", "q", m)

    m = bob.share("carol", "q")
    carol.receive_share("bob", "w", m)

    m = alice.share("dave", "a")
    dave.receive_share("alice", "e", m)

    print("Testing Bob, Carol, and Dave getting their new shares...")
    print(bob.download('q'))
    assert bob.download("q") == "b"
    assert carol.download("w") == "b"
    assert dave.download("e") == "b"

    print("Revoking Bob...")
    alice.revoke("bob", "a")
    dave.upload("e", "c")

    print("Testing Bob, Carol, and Dave getting their shares...")
    assert alice.download("a") == "c"
    assert bob.download("q") != "c"
    assert carol.download("w") != "c"
    assert dave.download("e") == "c"

    print("Testing restarting PKS and clients...")
    pks2 = PublicKeyServer()
    alice2 = Client(server, pks2, crypto, "alice")
    bob2 = Client(server, pks2, crypto, "bob")
    assert alice2.rsa_priv_key.publickey() == bob2.pks.get_signature_key("alice")
    assert alice2.elg_priv_key.publickey() == bob2.pks.get_encryption_key("alice")

    crypto._remove_keyfile("alice")
    crypto._remove_keyfile("bob")
    crypto._remove_keyfile("carol")
    crypto._remove_keyfile("dave")
    print("Basic tests passed.")



