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,counter):
    return rotate_right(keystate ^ 0x3^counter, 16, 64)

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

def apply_sbox(word, sbox):
    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):
    row_0, row_1, row_2, row_3 = get_rows(word)
    
    # apply the shiftrows transformation
    row_0 = row_0 
    row_1 = rotate_left(row_1, 4, 16)
    row_2 = rotate_left(row_2, 8, 16) 
    row_3 = rotate_left(row_3, 12, 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_1 ^ row_3) << 48
    new_word |= (row_0 ^ row_2 ^ row_3) << 32   # a |= b <==> a = a | b
    new_word |= (row_2 ^ row_3) << 16
    new_word |= (row_0 ^ row_3) << 0
    
    return new_word

def round_function(word, keystate):
    sbox = [0xc,0x7,0x2,0xe,0x1,0x4,0xb,0xd,0x5,0x8,0x9,0x3,0xa,0xf,0x0,0x6]
    
    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=8):
    keystate = key
    for i in range(rounds):
        # apply the roundfunction to word 
        word = round_function(word, keystate)
        # go to the next key state
        keystate = next_keystate(keystate,i)
    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))