import hh
import pickle
import re
import random
from phe import paillier





"""""""""""""""""""""""""""""""""""""""""""""""""""
""""""""""""""""""""""""""""""""""""""""""""""""""
#   · • ●· • ●· • ●· • ●· • ●· • ●· • ● • ·   #
#                                             #
#                                             #
#         The following functions             #
#                save/load                    #
#      Python Objects to/from a file          #
#                                             #
#                                             #
#                                             #
#   · • ●· • ●· • ●· • ●· • ●· • ●· • ● • ·   #
"""""""""""""""""""""""""""""""""""""""""""""""""""

def toFile(filename, obj):
    """Returns a python Object as a file."""

    with open(filename, 'wb') as output:
        pickle.dump(obj, output, pickle.HIGHEST_PROTOCOL)


def fromFile(filename):
    """Loads and returns a python Object from a file."""

    with open(filename, 'rb') as input:
        return pickle.load(input)






"""""""""""""""""""""""""""""""""""""""""""""""""""
""""""""""""""""""""""""""""""""""""""""""""""""""
#   · • ●· • ●· • ●· • ●· • ●· • ●· • ● • ·   #
#                                             #
#                                             #
#         The following functions             #
#                   are                       #
#     analyzer/verifier methods/classes       #
#                                             #
#                                             #
#                                             #
#   · • ●· • ●· • ●· • ●· • ●· • ●· • ● • ·   #
"""""""""""""""""""""""""""""""""""""""""""""""""""


class CipherSum(object):
    """"Cipher Object, which calculates a Cipher Sum (self.cipher) based on the CipherList
	which can be obtained by :meth: get_ciphers. In the end one will get a structure with
	Cipher Sum and a list, which consists of the ID's of all encrypted summands .
	Pub must be a PK-Object generated by phe/paillier.py."""

    def __init__(self, CipherList, index, pub):
        self.EN = paillier.EncryptedNumber(pub, 0)
        self.cipher = addEncryptedNumbersList(self.EN, [i[0] for i in CipherList])
        self.summands = [i[1] for i in CipherList]
        self.index = index

    def __repr__(self):
        return '<Cipher: %s || Summands: %s>' % (self.cipher, self.summands)

class HashSum(object):
    """Equivalent to Cipher Object. Calculates the Hash Sum based on a HashList. HashList is
       also obtained by :meth: get_ciphers. Endresult will be the Hash Sum and ID's of all
       hashed summands. Key must be the RSA PK-Object generated by e.g. pythonic rsa library.
    """

    def __init__(self, HashList, index, key):
        self.hash = hh.addHomomorphList(key, [i[0] for i in HashList])
        self.summands = [i[1] for i in HashList]
        self.index = HashList.index

    def __repr__(self):
        return '<Hash: %s || Summands: %s>' % (self.hash, self.summands)


def addEncryptedNumbersList(encryptedNumber, cipherList):

    product = 1
    for cipher in cipherList:
        product *= cipher
        product %= encryptedNumber.public_key.nsquare
    return product

def checkEquality(CipherSum, HashSum, negatives, pri, key):
    """Is H(Dec(Enc(CipherSum))) == H(h1+h2+h3+...+hn)?
       If true, Enc(CipherSum) will be replaced with Dec(Enc(Ciphersum)) in :attr: cipher.
       Else :attr: cipher = None.
    """
    factor = getSubtractionFactor(CipherSum, negatives)
    x = pri.raw_decrypt(CipherSum.cipher)
    if x >= 0 and hh.hash(key, x + factor * key.n) == HashSum.hash:
        CipherSum.cipher = x
    elif x < 0 and hh.hash(key, x + (factor - 1) * key.n) == HashSum.hash:
        CipherSum.cipher = x
    else:
        CipherSum.cipher = None
    return CipherSum

def derandomizeSum(sums, randoms):
    """
    If plain data was randomized, results in the end can be derandomized by substracting all
    values, which were added in the beginning based on the IDs the encrypted sums was made of.
    """
    randoms = fromFile(randoms)

    for sum in sums:
        relevantItems = []
        for id in sum.summands:
            item = [item for item in randoms[sum.index - 1] if item[0] == id - 1]
            if item:
                relevantItems.append(item)
        for item in relevantItems:
            try:
                sum.cipher -= [x[1] for x in item][0]
            except TypeError:
                pass

def getCiphers(filename, index, range=(-1)):
    """Returns a list of ciphers of encrypted data. Index is need for column specification
       Range is by default (-1), which implies that all ciphers from the column will be listed.
       If specified by a list, only ID's within the list will be considered."""

    if index > 0:
        CipherList = []
        with open(filename, 'rb') as f:
            next(f)
            if range == (-1):
                for line in f:
                    CipherList.append((int(line.split()[index]),
                                       int(line.split()[0])))
            else:
                for line in f:
                    if int(line.split(None, 1)[0]) in range:
                        CipherList.append((int(line.split()[index]),
                                           int(line.split()[0])))
        return CipherList
    else:
        print('illegal index')

def getHashSum(CipherSum, filename, key):
    """ Creates a HashSum Object through a given CipherSum Object. Loads the hash summands from
    a given file, which contains the hashed data, based on the ID's within the Cipher Object. """

    hashes = getCiphers(filename, CipherSum.index, CipherSum.summands)
    return HashSum(hashes, CipherSum.index, key)


def getIndex(filename, title):
    """Returns the index of a given column title (e.g. 'Age') within a file"""

    with open(filename, 'r') as f:
        words = f.readline().split()
        try:
            index = words.index(title)
            return index
        except ValueError:
            print("could not find title.")
            return -1

def getSubtractionFactor(sum, negatives):
    """Given a cipher sum and IDs, this method returns how many summands are negative."""

    counter = 0
    negatives = fromFile(negatives)
    for id in sum.summands:
        if id in negatives[sum.index - 1]:
            counter += 1
    return counter

def verify(keyfile, sumfile, hdata, negatives, threshold=0):
    """ Gets at least one CipherSum as an input. All Sums will then be checked with
        :meth: check_equality. threshold specifies how many summands a sum must me made of
        at least. """

    pri, key = fromFile(keyfile)
    sums = fromFile(sumfile)
    s = []

    try:
        for sum in sums:
            if threshold <= len(sum.summands):
                hashsum = getHashSum(sum, hdata, key)
                s.append(checkEquality(sum, hashsum, negatives, pri, key))
            else:
                sum.cipher = None
                s.append(sum)
    except TypeError:
        sum = sums
        if threshold <= len(sum.summands):
            hashsum = getHashSum(sum, hdata, key)
            s.append(checkEquality(sum, hashsum, negatives, pri, key))
        else:
            print(threshold)
            sum.cipher = None
            s.append(sum)
    toFile('response.pkl', s)



"""""""""""""""""""""""""""""""""""""""""""""""""""
""""""""""""""""""""""""""""""""""""""""""""""""""
#   · • ●· • ●· • ●· • ●· • ●· • ●· • ● • ·   #
#                                             #
#                                             #
#         The following functions are         #
#               setup methods                 #
#            used by the Data Holder          #
#                                             #
#                                             #
#   · • ●· • ●· • ●· • ●· • ●· • ●· • ● • ·   #

"""""""""""""""""""""""""""""""""""""""""""""""""""

def createNegativeList(input):
    """Since we also operate with negative numbers, we need to know how many negative numbers
       are used. This is because our negative hashes are represented e.g. like this (-4) -> (n-4).
       Therefore it is needed to know, how many times key.n must be added to the decrypted sum x.
       This method creates a list, stating which number is negative in the respective index.
    """

    with open(input, 'r') as infile:
        firstline = infile.readline()
        randoms = [[] for i in range(len(firstline.split()) - 1)]
        for id, line in enumerate(infile):
            data = line.split()
            del data[0]
            for idx, plain in enumerate(data):
                if int(plain) < 0:
                    randoms[idx].append((id + 1))
    toFile('negatives.pkl', randoms)

def encryptData(filename1, filename2, pub):
    """Encrypts a plain data textfile with pub, where first row contains index titles
       and first column contains ID's. Replaces all plain data except index titles and ID's
       with the encrypted version. """

    with open(filename1, 'r') as infile, open(filename2, 'w') as outfile:
        outfile.write(infile.readline())
        for line in infile:

            data = line.split()
            tmp = [data[0]]

            del data[0]

            line = " ".join(data)

            for plain in data:
                if int(plain) > 0:
                    line = re.sub(r'\b%s\b' % plain, str(pub.encrypt(int(plain)).ciphertext()), line, 1)
                else:
                    line = re.sub(r'\B%s\b' % plain, str(pub.encrypt(int(plain)).ciphertext()), line, 1)
            tmp.append(line + "\n")
            outfile.write(" ".join(tmp))


def hashData(filename1, filename2, key):
    """Same as encrypt data. Replaces all plain data except index titles and ID's
       with the hashed version. """

    with open(filename1, 'r') as infile, open(filename2, 'w') as outfile:
        outfile.write(infile.readline())
        for line in infile:
            data = line.split()
            tmp = [data[0]]
            del data[0]
            line = " ".join(data)
            for plain in data:
                if int(plain) > 0:
                    line = re.sub(r'\b%s\b' % plain, str(hh.hash(key, int(plain))), line, 1)
                else:
                    line = re.sub(r'\B%s\b' % plain, str(hh.hash(key, int(plain))), line, 1)
            tmp.append(line + "\n")
            outfile.write(" ".join(tmp))

def randomizeData(input, output, p=1, max_number=250000):
    """
    Randomizes plain data by adding a random number from 1 to max_number to the plain value.
    Creates a Python Object as a file, which contains which random number was added to which ID.
    """

    with open(input, 'r') as infile, open(output, 'w') as outfile:
        firstline = infile.readline()
        outfile.write(firstline)
        randoms = [[] for i in range(len(firstline.split()) - 1)]
        for id, line in enumerate(infile):
            data = line.split()
            del data[0]
            for idx, plain in enumerate(data):
                if random.random() <= p:
                    r = random.randint(1, max_number)
                    if int(plain) > 0:
                        line = re.sub(" " + r'\b%s\b' % plain, " " + str(int(plain) + r), line, 1)
                    else:
                        line = re.sub(r'\B%s\b' % plain, str(int(plain) + r), line, 1)
                    randoms[idx].append((id, r))
            outfile.write(line)
    toFile('randoms.pkl', randoms)


def setup(logfile,threshold=0,random=True):
    pub,pri = paillier.generate_paillier_keypair()
    key     = hh.generateKey()
    if random == True:
        randomizeData(logfile,'rlog.txt')
    try:
        encryptData('rlog.txt','enclog.txt',pub)
        hashData('rlog.txt','hashlog.txt',key)
    except (OSError,IOError):
        encryptData(logfile,'enclog.txt',pub)
        hashData(logfile,'hashlog.txt',key)

    createNegativeList('enclog.txt')
    keys = pri,key
    toFile("publickey.pkl",pub)
    toFile("keys.pkl",keys)
    toFile("threshold.pkl",threshold)
