# !/usr/bin/env python3
"""
    most function are taken from previous TC0X ciphers, and changed own design
"""


def get_rows(word):
    row_0 = (word >> 48) & 0xFFFF
    row_1 = (word >> 32) & 0xFFFF
    row_2 = (word >> 16) & 0xFFFF
    row_3 = (word >>  0) & 0xFFFF
    return row_0, row_1, row_2, row_3


def rotate_left(word, n, word_size=64):
    mask = 2 ** word_size - 1
    return ((word << n) & mask) | ((word >> (word_size - n) & mask))


def rotate_right(word, n, word_size=64):
    mask = 2 ** word_size - 1
    return ((word >> n) & mask) | ((word << (word_size - n) & mask))


def next_keystate(keystate):
    return shift_rows(rotate_left(keystate ^ 0x00000000FFFFFFFF, 16, 64))


def add_roundkey(word, keystate):
    return word ^ keystate


def apply_sbox(word, sbox):
    """ apply the sbox to every nibble """
    word_new = 0
    for i in range(16):  # 16 nibbles
        nibble = (word >> (i * 4)) & 0xF  # retrieve the ith nibble
        # insert the permuted nibble in the correct position
        word_new |= sbox[nibble] << i * 4
    return word_new


def shift_rows(word, i = 0):
    row_0, row_1, row_2, row_3 = get_rows(word)

    # apply the shiftrows transformation
    row_0 = rotate_left(row_0, (0  + 4*i) % 16, 16)
    row_1 = rotate_left(row_1, (4  + 4*i) % 16, 16)
    row_2 = rotate_left(row_2, (8  + 4*i) % 16, 16)
    row_3 = rotate_left(row_3, (12 + 4*i) % 16, 16)

    # reconstruct the word
    new_word = row_0 << 48  # a |= b <==> a = a | b
    new_word |= row_1 << 32
    new_word |= row_2 << 16
    new_word |= row_3 << 0

    return new_word


def mix_columns(word):
    row_0, row_1, row_2, row_3 = get_rows(word)  # split up the word into rows
    # Apply the mix culomns transformation and reconstruct the word
    new_word = (row_0 ^ row_2) << 48
    new_word |= (row_1 ^ row_2) << 32
    new_word |= (row_0 ^ row_2 ^ row_3) << 16
    new_word |= (row_1 ^ row_3) << 0

    return new_word


def round_function(word, keystate, round):
    sbox = [0xE, 0x4, 0xD, 0x1, 0x2, 0xF, 0xB, 0x8,
            0x3, 0xA, 0x6, 0xC, 0x5, 0x9, 0x0, 0x7] # sbox taken from cited article

    word = add_roundkey(word, keystate)
    word = apply_sbox(word, sbox)
    word = shift_rows(word, round)
    word = mix_columns(word)

    return word


def encrypt(word, key, rounds):
    keystate = key
    for i in range(rounds):
        # apply the roundfunction to word
        word = round_function(word, keystate, i)
        # go to the next key state
        keystate = next_keystate(keystate)
    return word



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))