#!/usr/bin/env python3

def apply_sbox(word, nibbles=4):
    """ apply the sbox to every nibble """
    word_new = 0
    sbox = (0xC, 5, 6, 0xB, 9, 0, 0xA, 0xD, 3, 0xE, 0xF, 8, 4, 7, 1, 2)
    for i in range(nibbles):  # 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 sigma(word):
    """
    Implementing the sigma permutation on the 8 bit word.
    """
    new_word = 0
    # first move the two most significant bits of nibble 0 and 3
    new_word |= (word & 0b1100000000001100) >> 1  # 0, 1, C, D

    # now move the rest of the bits
    new_word |= (word & 0x2000) >> 6  # 2
    new_word |= (word & 0x1000) >> 8  # 3
    new_word |= (word & 0x0C00) >> 5  # 4, 5
    new_word |= (word & 0x0200) << 6  # 6
    new_word |= (word & 0x0100) << 4  # 7
    new_word |= (word & 0x00C0) << 3  # 8, 9
    new_word |= (word & 0x0020) >> 2  # A
    new_word |= (word & 0x0010) >> 4  # B
    new_word |= (word & 0x0002) << 10  # E
    new_word |= (word & 0x0001) << 8  # F

    return new_word


def F1(word):
    return sigma(apply_sbox(word))

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

def F2(word):
     return rotate_left(word, 3, 16) ^ rotate_left(word, 8, 16) ^ rotate_left(word, 14, 16)

def F(word):
    upper = (word >> 16) & 0xFFFF
    upper = F1(upper)
    lower = word & 0xFFFF
    lower = F2(lower)
    return (lower << 16) | upper


def round_function(left, right, key):
    return ((F(left) ^ right ^ key), left)


def compute_roundkeys(key, rounds):
    key_parts = []
    key_parts.append(key & 0xFFFFFFFF)
    key >>= 32
    key_parts.append(key & 0xFFFFFFFF)

    for i in range(2, rounds):
        rk = key_parts[i - 1] ^ sigma(key_parts[i - 2]) ^ 0x3
        key_parts.append(rk)


    return key_parts


def encrypt(word, key, rounds=12):
    left = (word >> 32) & 0xFFFFFFFF
    right = word & 0xFFFFFFFF

    round_keys = compute_roundkeys(key, rounds)

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

    return (left << 32) | right


def decrypt(word, key, rounds=12):
    left = word & 0xFFFFFFFF
    right = (word >> 32) & 0xFFFFFFFF

    round_keys = compute_roundkeys(key, rounds)
    round_keys.reverse()

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

    return (right << 32) | left


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