#!/usr/bin/env python3

"""
    A python reference implementation of STABS cipher
    Code is based on the reference implementation given for TC02, which is very similar because it is also
    an SPN with block and key sizes of 64 bits/
"""
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, roundconst):
	return rotate_right(keystate ^ roundconst , 16 , 64), rotate_left(roundconst, 1, 16)

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

def apply_sbox(word, sbox):
	""" apply the sbox to every byte """
	word_new = 0
	for i in range(8): # 8 Bytes
		byte = (word >> (i*8)) & 0xFF # retrieve the ith byte
		# insert the permuted nibble in the correct position
		word_new |= sbox[byte] << i*8
	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_right(row_1, 4, 16)
	row_2 = rotate_right(row_2, 8, 16)
	row_3 = rotate_right(row_3, 12, 16)

	# reconstruct the word
	new_word = row_0 << 48
	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 columns transformation and reconstruct the word
	new_word = (row_0 ^ row_2 ^ row_3) << 48
	new_word |= row_0 << 32
	new_word |= (row_1 ^ row_2) << 16
	new_word |= (row_0 ^ row_2) << 0
	return new_word

def round_function(word , keystate):
	# The AES 8-bit sbox
	sbox = [0x63, 0x7C, 0x77, 0x7B, 0xF2, 0x6B, 0x6F, 0xC5, 0x30, 0x01, 0x67, 0x2B, 0xFE, 0xD7, 0xAB, 0x76,
            0xCA, 0x82, 0xC9, 0x7D, 0xFA, 0x59, 0x47, 0xF0, 0xAD, 0xD4, 0xA2, 0xAF, 0x9C, 0xA4, 0x72, 0xC0,
            0xB7, 0xFD, 0x93, 0x26, 0x36, 0x3F, 0xF7, 0xCC, 0x34, 0xA5, 0xE5, 0xF1, 0x71, 0xD8, 0x31, 0x15,
            0x04, 0xC7, 0x23, 0xC3, 0x18, 0x96, 0x05, 0x9A, 0x07, 0x12, 0x80, 0xE2, 0xEB, 0x27, 0xB2, 0x75,
            0x09, 0x83, 0x2C, 0x1A, 0x1B, 0x6E, 0x5A, 0xA0, 0x52, 0x3B, 0xD6, 0xB3, 0x29, 0xE3, 0x2F, 0x84,
            0x53, 0xD1, 0x00, 0xED, 0x20, 0xFC, 0xB1, 0x5B, 0x6A, 0xCB, 0xBE, 0x39, 0x4A, 0x4C, 0x58, 0xCF,
            0xD0, 0xEF, 0xAA, 0xFB, 0x43, 0x4D, 0x33, 0x85, 0x45, 0xF9, 0x02, 0x7F, 0x50, 0x3C, 0x9F, 0xA8,
            0x51, 0xA3, 0x40, 0x8F, 0x92, 0x9D, 0x38, 0xF5, 0xBC, 0xB6, 0xDA, 0x21, 0x10, 0xFF, 0xF3, 0xD2,
            0xCD, 0x0C, 0x13, 0xEC, 0x5F, 0x97, 0x44, 0x17, 0xC4, 0xA7, 0x7E, 0x3D, 0x64, 0x5D, 0x19, 0x73,
            0x60, 0x81, 0x4F, 0xDC, 0x22, 0x2A, 0x90, 0x88, 0x46, 0xEE, 0xB8, 0x14, 0xDE, 0x5E, 0x0B, 0xDB,
            0xE0, 0x32, 0x3A, 0x0A, 0x49, 0x06, 0x24, 0x5C, 0xC2, 0xD3, 0xAC, 0x62, 0x91, 0x95, 0xE4, 0x79,
            0xE7, 0xC8, 0x37, 0x6D, 0x8D, 0xD5, 0x4E, 0xA9, 0x6C, 0x56, 0xF4, 0xEA, 0x65, 0x7A, 0xAE, 0x08,
            0xBA, 0x78, 0x25, 0x2E, 0x1C, 0xA6, 0xB4, 0xC6, 0xE8, 0xDD, 0x74, 0x1F, 0x4B, 0xBD, 0x8B, 0x8A,
            0x70, 0x3E, 0xB5, 0x66, 0x48, 0x03, 0xF6, 0x0E, 0x61, 0x35, 0x57, 0xB9, 0x86, 0xC1, 0x1D, 0x9E,
            0xE1, 0xF8, 0x98, 0x11, 0x69, 0xD9, 0x8E, 0x94, 0x9B, 0x1E, 0x87, 0xE9, 0xCE, 0x55, 0x28, 0xDF,
            0x8C, 0xA1, 0x89, 0x0D, 0xBF, 0xE6, 0x42, 0x68, 0x41, 0x99, 0x2D, 0x0F, 0xB0, 0x54, 0xBB, 0x16]

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

def encrypt(word, key, rounds=20):
	keystate = key
	roundconst = 0x3
	for i in range(rounds):
		# apply the roundfunction to word
		word = round_function(word, keystate)
		# go to the next key state
		keystate, roundconst = next_keystate(keystate, roundconst)
	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))