#!/usr/bin/env python3

############################################################################################
############################################################################################
############################### Breaking Ciphers - Project #################################
###############################       Alaa Kryeem      #####################################
###############################      mzcdzdamri_f_     #####################################
############################################################################################
############################################################################################

"""
get_rows - Given a 64 bit word, split it into 4 rows, 4 bytes each
* Parameters:
*   Return: row_0, row_1, row_2, row_3
*   Inputs: word - Input 64-bit word to split.
"""
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

"""
rotate_left - Rotate a given message to the left.
* Parameters:
*   Return: - rotated message.
*   Inputs: word - Input 64-bit word to perform rotation on.
*			n - rotation count
*			word_size - total word size of rotated word
"""
def rotate_left(word, n, word_size=4):
    mask = 2**word_size - 1
    return ((word << n) & mask) | ((word >> (word_size - n) & mask)) 

"""
rotate_right - Rotate a given message to the right.
* Parameters:
*   Return: - rotated message.
*   Inputs: word - Input 64-bit word to perform rotation on.
*			n - rotation count
*			word_size - total word size of rotated word
"""
def rotate_right(word, n, word_size=4):
    mask = 2**word_size - 1
    return ((word >> n) & mask) | ((word << (word_size - n) & mask)) 

"""
mix_columns - The Mix Columns (M C) layer will mix the nibbles in every column according to a matrix.  The matrix is given by:
* 1 1 0 0
* 0 0 1 1
* 1 0 0 1
* 0 1 1 1
*
* Parameters:
*   Return: - new word.
*   Inputs: word - Input 64-bit word to perform MC on.
"""
def mix_columns(word):
    ## split up the word into rows
    ##
    row_0, row_1, row_2, row_3 = get_rows(word)
    
    ## apply the mix columns transformation and reconstruct the word
    ##
    new_word = (row_0 ^ row_1) << 48
    new_word |= (row_2 ^ row_3) << 32
    new_word |= (row_0 ^ row_3) << 16
    new_word |= (row_1 ^ row_2 ^ row_3) << 0
    return new_word

"""
dmix_columns - The decryption Mix Columns (M C) layer will mix the nibbles in every column according to a matrix.  The matrix is given by:
* 1	1 0 1
* 0	1 0 1
* 1	0 1 1
* 1	1 1 1
*
* Parameters:
*   Return: - new word.
*   Inputs: word - Input 64-bit word to perform MC on.
"""
def dmix_columns(word):
    ## split up the word into rows
    ##
    row_0, row_1, row_2, row_3 = get_rows(word)
    
    ## apply the mix columns transformation and reconstruct the word
    ##
    new_word = (row_0 ^ row_1 ^ row_3) << 48
    new_word |= (row_1 ^ row_3) << 32
    new_word |= (row_0 ^ row_2 ^ row_3) << 16
    new_word |= (row_0 ^ row_1 ^ row_2 ^ row_3) << 0    
    return new_word


sbox_8 = [0x33, 0x3e, 0x35, 0x3d, 0x39, 0x3a, 0x31, 0x3f, 0x37, 0x3c,
0x30, 0x38, 0x36, 0x32, 0x3b, 0x34, 0xe3, 0xee, 0xe5, 0xed,
0xe9, 0xea, 0xe1, 0xef, 0xe7, 0xec, 0xe0, 0xe8, 0xe6, 0xe2,
0xeb, 0xe4, 0x53, 0x5e, 0x55, 0x5d, 0x59, 0x5a, 0x51, 0x5f,
0x57, 0x5c, 0x50, 0x58, 0x56, 0x52, 0x5b, 0x54, 0xd3, 0xde,
0xd5, 0xdd, 0xd9, 0xda, 0xd1, 0xdf, 0xd7, 0xdc, 0xd0, 0xd8,
0xd6, 0xd2, 0xdb, 0xd4, 0x93, 0x9e, 0x95, 0x9d, 0x99, 0x9a,
0x91, 0x9f, 0x97, 0x9c, 0x90, 0x98, 0x96, 0x92, 0x9b, 0x94,
0xa3, 0xae, 0xa5, 0xad, 0xa9, 0xaa, 0xa1, 0xaf, 0xa7, 0xac,
0xa0, 0xa8, 0xa6, 0xa2, 0xab, 0xa4, 0x13, 0x1e, 0x15, 0x1d,
0x19, 0x1a, 0x11, 0x1f, 0x17, 0x1c, 0x10, 0x18, 0x16, 0x12,
0x1b, 0x14, 0xf3, 0xfe, 0xf5, 0xfd, 0xf9, 0xfa, 0xf1, 0xff,
0xf7, 0xfc, 0xf0, 0xf8, 0xf6, 0xf2, 0xfb, 0xf4, 0x73, 0x7e,
0x75, 0x7d, 0x79, 0x7a, 0x71, 0x7f, 0x77, 0x7c, 0x70, 0x78,
0x76, 0x72, 0x7b, 0x74, 0xc3, 0xce, 0xc5, 0xcd, 0xc9, 0xca,
0xc1, 0xcf, 0xc7, 0xcc, 0xc0, 0xc8, 0xc6, 0xc2, 0xcb, 0xc4,
0x3	, 0xe	, 0x5	, 0xd	, 0x9	, 0xa	, 0x1	, 0xf	,
0x7	, 0xc	, 0x0	, 0x8	, 0x6	, 0x2	, 0xb	, 0x4	,
0x83, 0x8e, 0x85, 0x8d, 0x89, 0x8a, 0x81, 0x8f, 0x87, 0x8c,
0x80, 0x88, 0x86, 0x82, 0x8b, 0x84, 0x63, 0x6e, 0x65, 0x6d,
0x69, 0x6a, 0x61, 0x6f, 0x67, 0x6c, 0x60, 0x68, 0x66, 0x62,
0x6b, 0x64, 0x23, 0x2e, 0x25, 0x2d, 0x29, 0x2a, 0x21, 0x2f,
0x27, 0x2c, 0x20, 0x28, 0x26, 0x22, 0x2b, 0x24, 0xb3, 0xbe,
0xb5, 0xbd, 0xb9, 0xba, 0xb1, 0xbf, 0xb7, 0xbc, 0xb0, 0xb8,
0xb6, 0xb2, 0xbb, 0xb4, 0x43, 0x4e, 0x45, 0x4d, 0x49, 0x4a,
0x41, 0x4f, 0x47, 0x4c, 0x40, 0x48, 0x46, 0x42, 0x4b, 0x44]

"""
generate_8_bit_sbox - given an sbox, generate a 8-bit sbox.
* Parameters:
*   Return: uint64_t  - Diffusion result.
*   Inputs: uint64_t sbox - input sbox table.
"""
def generate_8_bit_sbox(sbox):
    word = 0
    sbox_8 = [None] * 256
    for word in range(256):
        sbox_8[word] = sbox[word & 0xF] | (sbox[(word >> 4) & 0xF] << 4);
    return  sbox_8;


"""
apply_sbox8 - Apply  the  sbox to  every  nibble.
* Parameters:
*   Return: - The message after applying sbox on it.
*   Inputs: word - The message to apply sbox on.
"""
def apply_sbox8(word, nibbles=4):
    word_new = 0
    shift = 0
    for i in range(nibbles): # 4 nibbles
        word_new |= sbox_8[word & 0xFF] << shift
        word >>= 8
        shift += 8
    return  word_new

"""
apply_sbox - Apply  the  sbox to  every  nibble.
* Parameters:
*   Return: - The message after applying sbox on it.
*   Inputs: word - The message to apply sbox on.
"""
def apply_sbox(word, nibbles=8):
    word_new = 0
    sbox = (0x3,0xe,0x5,0xd,0x9,0xa,0x1,0xf,0x7,0xc,0x0,0x8,0x6,0x2,0xb,0x4)
    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

"""
sigma - Implementing the sigma permutation on the 8 bit word.
* Parameters:
*   Return: - New permutation of the given word.
*   Inputs: word - The word to apply sigma on.
"""
def sigma(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 # E
    return new_word


##
## Implementing the F function
## MAK(l = 32bits) = (σ(ll ⊕ lr)<< 16) | σ(σ(ll ⊕ lr) ⊕ lr)
## Semi-F(l, r, Ki) = S(MAK(S(l))) ⊕ r ⊕ Ki
## F(l, r, Ki) = MC(SR((Semi-F(l, r, Ki) << 32) | right))
##
"""
* Implementing the F function
* MAK(l = 32bits) = (σ(ll ⊕ lr)<< 16) | σ(σ(ll ⊕ lr) ⊕ lr)
* Semi-F(l, r, Ki) = S(MAK(S(l))) ⊕ r ⊕ Ki
* F(l, r, Ki) = MC(SR((Semi-F(l, r, Ki) << 32) | right))
*
* Parameters:
*   Return: - F(word)
*   Inputs: word - The word to apply F on.
"""
def F(word):
    sboxed_word = apply_sbox8(word)
    init_L = sboxed_word >> 16
    init_R = sboxed_word & 0xFFFF
    L = sigma(init_L ^ init_R)
    R = (sigma(init_R ^ L))
    new_word = (L << 16) | R
    return apply_sbox8(new_word)


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

"""
compute_roundkeys - Given master key K=k0|k1 the round key k is defined as follows:
ki= ki-1 ⊕ σ(ki-2) ⊕ (σ(ki-2 >> 16) << 16) ⊕ 0xC5A1B9D2

* Parameters:
*   Return: - #rounds derived keys
*   Inputs: rounds - number of rounds (derived keys)
*           key - A 64-bit master key
"""
def compute_roundkeys(key, rounds):
    key_parts = []
    for i in range(2):
        key_parts.append(key & 0xFFFFFFFF)
        key >>= 32
    # Most significant part should be on index 0
    key_parts.reverse() 

    for i in range(2, rounds):
        rk = key_parts[i-1] ^ sigma(key_parts[i-2]) ^ (sigma((key_parts[i-2] >> 16)) << 16) ^ 0xC5A1B9D2
        key_parts.append(rk)
    
    return key_parts

"""
shift_rows - rotate the input nibbles in the rows by 0, 1, 2 and 3 places to the left.
* Parameters:
*   Return: - shifted word.
*   Inputs: word - Input 64-bit word to perform shifting on.
"""
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

"""
dshift_rows - rotate the input nibbles in the rows by 0, 1, 2 and 3 places to the right.
* Parameters:
*   Return: - shifted word.
*   Inputs: word - Input 64-bit word to perform shifting on.
"""
def dshift_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) & 0xFFFF
    row_2 = rotate_right(row_2, 8, 16) & 0xFFFF
    row_3 = rotate_right(row_3, 12, 16) & 0xFFFF

    ##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 encrypt(word, key, rounds=16):
    ##compute round keys, +2 for key whitening
    ##
    round_keys = compute_roundkeys(key, rounds + 2)

    ##WKR, WKL are the key whitening keys
    ##
    WKR = round_keys[-1]
    WKL = (round_keys[-2] << 32)
    word = word ^ WKR ^ WKL

    left = (word >> 32) & 0xFFFFFFFF
    right = word & 0xFFFFFFFF


    for i in range(rounds):
        left, right = round_function(left, right, round_keys[i], mix_columns)
        combined_word = (left << 32) | right
        combined_word = shift_rows(combined_word)
        mixed_word = mix_columns(combined_word)
        left = (mixed_word >> 32) & 0xFFFFFFFF
        right = mixed_word & 0xFFFFFFFF

    before_whitening = (left << 32) | right
    after_whitening = before_whitening ^ WKR ^ WKL

    return after_whitening

def decrypt(word, key, rounds=16):
    ##compute round keys, +2 for key whitening
    ##
    round_keys = compute_roundkeys(key, rounds + 2)
    round_keys.reverse()
    ##WKR, WKL are the key whitening keys
    ##
    WKR = round_keys[0]
    WKL = (round_keys[1] << 32)
    word = word ^ WKR ^ WKL
    
    left = word & 0xFFFFFFFF
    right = (word >> 32) & 0xFFFFFFFF
    
    for i in range(2, rounds + 2):
        combined_word = (right << 32) | left
        dmixed_word = dmix_columns(combined_word)
        dmixed_word = dshift_rows(dmixed_word)

        left = dmixed_word & 0xFFFFFFFF
        right = (dmixed_word >> 32) & 0xFFFFFFFF
        left, right = round_function(left, right, round_keys[i], dmix_columns)

    before_whitening = (right << 32) | left
    after_whitening = before_whitening ^ WKR ^ WKL
    return after_whitening

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