#!/usr/bin/env python3
from time import time

def get_time_str(start_time):
    """
    :param start_time:
    :return: given start time, returns a string with time passed from start time
    """
    hours, rem = divmod(int(time() - start_time), 3600)
    minutes, seconds = divmod(rem, 60)
    return "{:0>2}:{:0>2}:{:0>2}".format(int(hours), int(minutes), int(seconds))


def d_to_h(decimal_word, hex_length=16):
    """
    :param decimal_word:
    :param hex_length:
    :return: the string rep in Hex in given length in Hex
    """
    s = '0x{:0' + str(hex_length) + 'X}'
    return s.format(decimal_word)


class Hulk(object):
    def __init__(self):
        # self.start_time = time()
        # self.test_components()
        # self.generate_testvectors()
        return

    def __del__(self):
        # print('Total time: {}'.format(get_time_str(self.start_time)))
        return

    @staticmethod
    def rotate_right(word, n, word_size_bits=32):
        """ taken from TC02 """
        mask = 2 ** word_size_bits - 1
        return ((word >> n) & mask) | (word << (word_size_bits - n) & mask)

    @staticmethod
    def apply_sbox(word, nibbles=4):
        """
        'Present' cipher SBOX
        apply the sbox to every nibble
        """
        new_word = 0
        present_sbox = (0xC, 0x5, 0x6, 0xB, 0x9, 0x0, 0xA, 0xD, 0x3, 0xE, 0xF, 0x8, 0x4, 0x7, 0x1, 0x2)

        for i in range(nibbles):
            nibble = (word >> (i * 4)) & 0xF  # retrieve the ith nibble
            new_word |= present_sbox[nibble] << i * 4
        return new_word

    @staticmethod
    def sigma(word16):
        """
        TC05 sigma
        Implementing the sigma permutation on the 8 bit word.
        """
        new_word16 = 0
        new_word16 |= (word16 & 0b1100000000001100) >> 1  # 0, 1, C, D
        new_word16 |= (word16 & 0x2000) >> 6  # 2
        new_word16 |= (word16 & 0x1000) >> 8  # 3
        new_word16 |= (word16 & 0x0C00) >> 5  # 4, 5
        new_word16 |= (word16 & 0x0200) << 6  # 6
        new_word16 |= (word16 & 0x0100) << 4  # 7
        new_word16 |= (word16 & 0x00C0) << 3  # 8, 9
        new_word16 |= (word16 & 0x0020) >> 2  # A
        new_word16 |= (word16 & 0x0010) >> 4  # B
        new_word16 |= (word16 & 0x0002) << 10  # E
        new_word16 |= (word16 & 0x0001) << 8  # E
        return new_word16

    def f_function(self, W0_32):
        """
        build with the help of Eran
        :param W0_32: W0 in 32 bits length
        W1 = SBOX(W0_32)
        L0, R0 = left(W1), right(W1)
        L1, R1 = sigma(L0), sigma(R0)
        W2 = L1 | R1
        W3 = W2 >>> 8 // rotate 8bits right
        L2, R2 = left(W3), right(W3)
        L3, R3 = sigma(L2), sigma(R2)
        W4 = L3 | R3
        :return: W4
        """
        W1 = self.apply_sbox(W0_32, 8)
        L0 = (W1 & 0xFFFF0000) >> 16
        R0 = (W1 & 0x0000FFFF)
        L1 = self.sigma(L0)
        R1 = self.sigma(R0)
        W2 = (L1 << 16) | R1
        W3 = self.rotate_right(W2, 8)
        L2 = (W3 & 0xFFFF0000) >> 16
        R2 = (W3 & 0x0000FFFF)
        L3 = self.sigma(L2)
        R3 = self.sigma(R2)
        W4 = (L3 << 16) | R3
        return W4

    def round_function(self, left32, right32, key32):
        """ TC05 round_function """
        new_left = (self.f_function(left32) ^ right32 ^ key32)
        new_right = left32
        return new_left, new_right

    def compute_roundkeys(self, key, rounds):
        """
        Based on TC05 compute_roundkeys function
        key_parts_16b = (rounds + 1) round keys 16bits length
        for i=0 to rounds:
            key_parts_32b_i = key_parts_16b[i] | key_parts_16b[i+1]
        """
        key_parts_16b = []

        for i in range(4):
            key_parts_16b.append(key & 0xFFFF)
            key >>= 16

        # Most significant part should be on index 0
        key_parts_16b.reverse()

        for i in range(4, rounds + 1):
            rk = key_parts_16b[i - 4] ^ key_parts_16b[i - 1] ^ self.sigma(key_parts_16b[i - 2]) ^ 0xC
            key_parts_16b.append(rk)

        key_parts_32b = []
        for i in range(len(key_parts_16b) - 1):
            key_parts_32b.append((key_parts_16b[i] << 16) | key_parts_16b[i + 1])

        return key_parts_32b

    def encrypt(self, word, key, rounds=16):
        """ generalized TC05 encrypt to 64 bits words"""
        left = (word >> 32) & 0xFFFFFFFF
        right = word & 0xFFFFFFFF
        round_keys = self.compute_roundkeys(key, rounds)

        for i in range(rounds):
            left, right = self.round_function(left, right, round_keys[i])

        concat = (left << 32) | right

        return concat

    def decrypt(self, word, key, rounds=16):
        """ generalized TC05 decrypt to 64 bits words"""
        left = word & 0xFFFFFFFF
        right = (word >> 32) & 0xFFFFFFFF
        round_keys = self.compute_roundkeys(key, rounds)
        round_keys.reverse()

        for i in range(rounds):
            left, right = self.round_function(left, right, round_keys[i])

        return (right << 32) | left

    def encrypt_set(self, P, key):
        """
        :param P: list of plaintext - 64bits each
        :param key: 64 bit key
        :return: list of ciphers - 64bits each
        """
        C = []
        for p in P:
            c = self.encrypt(p, key)
            C.append(c)
        return C

    def test_components(self):
        """ checking major components in Hulk cypher """
        good_msg = '\t{:<20} - Passed'
        print('Checking components:')

        # first test the sigma function
        sigma_err = 'sigma function is not working'
        assert self.sigma(0x0000) == 0x0000, sigma_err
        assert self.sigma(0x8000) == 0x4000, sigma_err
        assert self.sigma(0xF000) == 0x6090, sigma_err
        assert self.sigma(0x0F00) == 0x9060, sigma_err
        assert self.sigma(0x00F0) == 0x0609, sigma_err
        assert self.sigma(0x000F) == 0x0906, sigma_err
        assert self.sigma(0xFFFF) == 0xFFFF, sigma_err
        print(good_msg.format('Sigma'))

        # test the sbox
        sbox_err = 'sbox function is not working'
        assert self.apply_sbox(0xFF13) == 0x225B, sbox_err
        assert self.apply_sbox(0x02DE) == 0xC671, sbox_err
        print(good_msg.format('SBOX'))

        # test the invertability
        inv_err = 'd(e(w,k)) != w'

        key, w = 0x123456789ABCDEF0, 0x123456789ABCDEF0
        assert self.decrypt(self.encrypt(w, key), key) == w, inv_err

        key, w = 0x123, 0xF4F31255
        assert self.decrypt(self.encrypt(w, key), key) == w, inv_err

        key, w = 0, 0xFFFFFFFF
        assert self.decrypt(self.encrypt(w, key), key) == w, inv_err
        print(good_msg.format('d(e(w,k))==w'))

        partial_rounds_err = 'Partial encryption'
        key, w, rounds = 0x123456789ABCDEF0, 0x123456789ABCDEF0, 6
        assert self.decrypt(self.encrypt(w, key, rounds), key, rounds) == w, partial_rounds_err

        key, w, rounds = 0x123456789ABCDEF0, 0x123456789ABCDEF0, 0
        assert self.decrypt(self.encrypt(w, key, rounds), key, rounds) == w, partial_rounds_err

        print(good_msg.format('Partial encryption'))

        return

    def generate_testvectors(self):
        test_vectors = []
        print('Generating test vectors:')
        print('\tMetadata:')
        print('\t\t|P|=|C|=|K|=16hex=64bits')

        msg = "\tE(p={},k={})=c={}"
        p, k = 0, 0
        c = self.encrypt(p, k)
        test_vectors.append((p, k, c))
        print(msg.format(d_to_h(p), d_to_h(k), d_to_h(c)))

        p, k = 0x12345678, 0x1234567890ABCDEF
        c = self.encrypt(p, k)
        print(msg.format(d_to_h(p), d_to_h(k), d_to_h(c)))

        p, k = 0x123456789ABCDEF0, 0x123456789ABCDEF0
        c = self.encrypt(p, k)
        print(msg.format(d_to_h(p), d_to_h(k), d_to_h(c)))

        test_vectors.append((p, k, c))
        return test_vectors


def main():
    hulk_handler = Hulk()

    key = 0x1234567890ABCDEF
    P = [0x0, 0x12345678, 0x1234567890ABCDEF, 0x3333333333333333]
    C = hulk_handler.encrypt_set(P, key=key)

    print('Checking encrypted set P:')
    msg = "\tE(p={},k={})=c={}"
    for i in range(len(C)):
        print(msg.format(d_to_h(P[i]), d_to_h(key), d_to_h(C[i])))

def encrypt(plaintext, key, rounds):
    hulk_handler = Hulk()
    return hulk_handler.encrypt(plaintext, key=key, rounds=rounds)


if __name__ == "__main__":
    import sys
    import random

    if len(sys.argv) not in [3, 4]:
        print("Error occured %d"%(len(sys.argv)))
        exit()
    
    key = int(sys.argv[1], 16)
    rounds = int(sys.argv[2])
    if len(sys.argv) == 4:
        delta_in = int(sys.argv[3], 16)

    nrof_pairs = 2**10

    print("# %s %d rounds"%(sys.argv[0], rounds))

    for i in range(nrof_pairs):
        word = random.getrandbits(64)
        cipher = encrypt(word, key, rounds=rounds)
        if len(sys.argv) == 3:
            print("%016X %016X"%(word, cipher))
        else:
            word_2 = word ^ delta_in
            cipher_2 = encrypt(word_2, key, rounds=rounds) 
            print("%016X %016X %016X %016X"%(word, word_2, cipher, cipher_2))