"""
64 bit state divided in 16 4 - bit nibbles arranged in a matrix representation
this cipher is similar to the TCO2 that includes 2 changes( sbox , mix columns)
--------------------------------------

ROUND FUNCTION :
1.Add roundkey : word = word ^ ( Key & 0Xffffffff00000000 )

2.Apply sbox :[ the sbox is taken from : "https://eprint.iacr.org/2011/218.pdf "  ]
  note that i did the applying sbox function in a different way, not in for loop.
3.AES shift rows (0 , 1 , 2 , 3)

4.Mix Columns :
1 0 0 1
0 1 1 0
0 0 1 0
1 0 1 0
i chose this matrix because this matrix has an inverse matrix
--------------------------------------
KEY SCHEDULE ALGORITHM :

The Key has a 64 - bit state, the i-th keystate is computed as follows :
k0=key
Ki = (ki-1 ^ 0x3 << 16) | (ki-1 ^ 0x3 >> 48)
"""

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):
    return (((word << n) ) | ((word>>( 16 - n ) ))) & 0xFFFF


def add_roundkey ( word , keystate ):
    return word ^ ( keystate & 0xFFFFFFFF00000000 )


def apply_sbox(word, sbox):
    """ apply sbox """

    tmp0 = (sbox[word & 0xF])
    tmp1 = ((sbox[(word >> 4) & 0xF]) << 4)
    tmp2 = ((sbox[(word >> 8) & 0xF]) << 8)
    tmp3 = ((sbox[(word >> 12) & 0xF]) << 12)
    tmp4 = ((sbox[(word >> 16) & 0xF]) << 16)
    tmp5 = ((sbox[(word >> 20) & 0xF]) << 20)
    tmp6 = ((sbox[(word >> 24) & 0xF]) << 24)
    tmp7 = ((sbox[(word >> 28) & 0xF]) << 28)
    tmp8 = ((sbox[(word >> 32) & 0xF]) << 32)
    tmp9 = ((sbox[(word >> 36) & 0xF]) << 36)
    tmpA = ((sbox[(word >> 40) & 0xF]) << 40)
    tmpB = ((sbox[(word >> 44) & 0xF]) << 44)
    tmpC = ((sbox[(word >> 48) & 0xF]) << 48)
    tmpD = ((sbox[(word >> 52) & 0xF]) << 52)
    tmpE = ((sbox[(word >> 56) & 0xF]) << 56)
    tmpF = ((sbox[(word >> 60) & 0xF]) << 60)
    NewWord = (tmp0 | tmp1 | tmp2 | tmp3 | tmp4 | tmp5 | tmp6 | tmp7 | tmp8 | tmp9 |
               tmpA | tmpB | tmpC | tmpD | tmpE | tmpF)
    return NewWord



def shift_rows(word):
    row_0, row_1, row_2, row_3 = get_rows(word)
    """ SHIFT ROWs & reconstruct the word """
    r1=rotate_left(row_1,4)
    r2=rotate_left(row_2,8)
    r3=rotate_left(row_3,12)

    new_word = (row_0 << 48) | (r1 << 32) | (r2 << 16) | ( r3 << 0)

    return new_word

def reverse_shift_rows(word):
    row_0, row_1, row_2, row_3 = get_rows(word)
    """ SHIFT ROWs & reconstruct the word """
    new_word = (row_0 << 48) | (rotate_right(row_1, 4) << 32) | (rotate_right(row_2, 8) << 16) | (rotate_right(row_3, 12) << 0)

    return new_word

def rotate_right(word,n):
    return(((word >> n) ) | ((word<<( 16 - n )))) & 0xFFFF

def mix_columns ( word ):
    """ split up the word """
    row_0, row_1, row_2, row_3 = get_rows( word )
    """ Apply mix columns & reconstruct the word """
    new_word = (row_0 ^ row_3) << 48 | ((row_1 ^ row_2) << 32) | (row_2 << 16)| ((row_0 ^ row_2) << 0)

    return new_word


def reverse_mix_columns ( word ):
    """ split up the word """
    row_0, row_1, row_2, row_3 = get_rows( word )
    """ Apply mix columns & reconstruct the word ( the inverse matrix of the mix columns matrix """
    new_word = ((row_2 ^ row_3) << 48) | ((row_1 ^ row_2) << 32) | (row_2 << 16)| ((row_0 ^ row_2 ^ row_3) << 0)

    return new_word



def round_function(word, keystate ):
    sbox = [0x0, 0x3, 0x5, 0x8, 0x6, 0xC, 0xB, 0x7,
            0xA, 0x4, 0x9, 0xE, 0xF, 0x1, 0x2, 0xD]
    word = add_roundkey(word, keystate)
    word = apply_sbox(word, sbox)
    word = shift_rows(word)
    word = mix_columns(word)

    return word


def encrypt ( word, key, rounds):
    keystate = key
    i=0
    while(i<rounds):
        word = round_function ( word , keystate )
        keystate ^=0x3
        keystate = ( (keystate>>16) | (keystate<<48) ) & 0xFFFFFFFFFFFFFFFF
        i+=1
    return word


def decrypt ( word , key , rounds):
    keystate = key
    i=0
    while(i<rounds):
        word = reverse_round_function ( word , keystate )
        keystate = (keystate<<16 | keystate>>48)& 0xFFFFFFFFFFFFFFFF
        keystate ^=0x3
        i+=1
    return word


def reverse_round_function(word, keystate ):
    dec_sbox = [0x0, 0xD, 0xE, 0x1, 0x9, 0x2, 0x4, 0x7,
            0x3, 0xA, 0x8, 0x6, 0x5, 0xF, 0xB, 0xC]
    word = reverse_mix_columns(word)
    word = reverse_shift_rows(word)
    word = apply_sbox(word, dec_sbox)
    word = add_roundkey(word, 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))