#baby block cipher built on principles from
#Introduction to Modern Cryptography for design
#of substitution-permutation network.

#The key is only 16 bits long (that's the baby part)

#The s-box was computed by finding a permutation of
#{0,..,15} such that changing one bit of x changed at
#least 2 bits of pi(x)

#The mixing permutation was computed by finding a
#permutation such that the bits in each 4-bit sub-block
#are sent to 4 different sub-blocks

#convert an integer less than 2**16 to
#a 4-nibble list

import bytestuff

def num_to_list(n):
    return [(n>>j)&15 for j in [12,8,4,0]]

#and the inverse
def list_to_num(li):
    return (li[0]<<12)|(li[1]<<8)|(li[2]<<4)|li[3]


#the five  round keys are
#just cyclic shifts of the nibbles of
#the key (so round key is same in rounds 1 and 5).
#Assume k is given as a list

def key_schedule(k):
    roundkeys=[]
    roundkeys.append(k)
    for j in [1,2,3,4]:
        k=[k[3],k[0],k[1],k[2]]
        roundkeys.append(k)
##    if decrypt:
##        roundkeys.reverse()
    return roundkeys

#use this to invert a permutation of [0,..,15]
def invert_perm(per):
    inv=[0]*16
    for j in range(16):
        inv[per[j]]=j
    return inv

#state and the return value are a list of 4 nibbles:  Each
#nibble is put through the s-box
def s_box(state,decrypt=False):
    s_perm=[11,13,7,12,3,6,10,1,14,0,9,8,15,4,2,5]
    if decrypt:
        s_perm=invert_perm(s_perm)
    return [s_perm[j] for j in state]

def mix(state,decrypt=False):
    mixing_perm=[5,8,0,15,11,7,2,13,14,10,1,4,12,9,6,3]
    if decrypt:
        mixing_perm=invert_perm(mixing_perm)
    state_bits=[]
    for j in state:
        state_bits+=[(j&m)/m for m in [8,4,2,1]]
    #print state_bits
    mixed_bits=[0]*16
    for j in range(16):
        mixed_bits[j]=state_bits[mixing_perm[j]]
    #print mixed_bits
    mixed_state=[]
    for j in range(4):
        mixed_state.append(sum([(1<<(3-s))*mixed_bits[4*j+s] for s in range(4)]))
    return mixed_state
        
        
#basic encryption function
#You can adjust the number of rounds to see
#the effect.  If you set the optional last argument
#to True you can see the effect of all the phases.

def encrypt(k,m,**kw):
    
    if 'numrounds' in kw:
        numrounds=kw['numrounds']
    else:
        numrounds=5
    if 'verbose' in kw:
        verbose=kw['verbose']
    else:
        verbose=False
    expanded_key=key_schedule(num_to_list(k))
    state=num_to_list(m)
    if verbose:
        print 'plaintext'
        for k in range(4):
            print (bytestuff.bin8(state[k]))[4:],
        print
    for j in range(numrounds):
        if verbose:
            print 'Round',j+1
            print 'Round Key',
            for k in range(4):
                print (bytestuff.bin8(expanded_key[j][k]))[4:],
            print
        state = bytestuff.xor(state,expanded_key[j])
        if verbose:
            print 'Add Key'
            for k in range(4):
                print (bytestuff.bin8(state[k]))[4:],
            print
        state = s_box(state)
        if verbose:
            print 'Sub'
            for k in range(4):
                print (bytestuff.bin8(state[k]))[4:],
            print
        state = mix(state)
        if verbose:
            print 'Mix'
            for k in range(4):
                print (bytestuff.bin8(state[k]))[4:],
            print
    return list_to_num(state)


def decrypt(k,m,**kw):
    if 'numrounds' in kw:
        numrounds=kw['numrounds']
    else:
        numrounds=5
    if 'verbose' in kw:
        verbose=kw['verbose']
    else:
        verbose=False
    expanded_key=key_schedule(num_to_list(k))
    state=num_to_list(m)
    if verbose:
        print 'ciphertext'
        for k in range(4):
            print (bytestuff.bin8(state[k]))[4:],
        print
    for j in range(numrounds):
        if verbose:
            print 'Round',j+1
            print 'Round Key',
            for k in range(4):
                print (bytestuff.bin8(expanded_key[numrounds-j-1][k]))[4:],
            print
        state=mix(state,True)
        if verbose:
            print 'Mix'
            for k in range(4):
                print (bytestuff.bin8(state[k]))[4:],
            print
        state=s_box(state,True)
        if verbose:
            print 'Sub'
            for k in range(4):
                print (bytestuff.bin8(state[k]))[4:],
            print
        state = bytestuff.xor(state,expanded_key[numrounds-j-1])
        if verbose:
            print 'Add Key'
            for k in range(4):
                print (bytestuff.bin8(state[k]))[4:],
            print
    return list_to_num(state)