# Python PRIDE implementation
# Version: 1.0
# Date: 22/04/2015
#
# =============================================================================
# Copyright (c) 2015 Camil Staps (info@camilstaps.nl)
# =============================================================================

class Pride:

    def __init__(self,key):
        """Create a PRIDE cipher object

        key:    the key as a 128-bit rawstring
        """
        if len(key) * 8 == 128:
            self.key_whitening = string2number(key[:8])
            self.key_1 = key[8:]
            #self.roundkeys = generateRoundkeys(key[8:])
        else:
            raise ValueError, "Key must be a 128-bit rawstring"

        # print self.key_whitening
        #print self.roundkeys

    def encrypt(self,block):
        """Encrypt 1 block (8 bytes)

        Input:  plaintext block as raw string
        Output: ciphertext block as raw string
        """
        state = string2number(block)
        state = pLayer_dec(state)
        state = addRoundKey(state, self.key_whitening)
        for i in xrange (1,20):
            state = addRoundKey(state, pLayer_dec(roundKey(self.key_1, i)))
            state = sBoxLayer(state)
            state = lLayer(state)
        state = addRoundKey(state, pLayer_dec(roundKey(self.key_1, 20)))
        state = sBoxLayer(state)
        state = addRoundKey(state, self.key_whitening)
        state = pLayer(state)
        return number2string_N(state,8)

#        0   1   2   3   4   5   6   7   8   9   a   b   c   d   e   f
Sbox= [0x0,0x4,0x8,0xf,0x1,0x5,0xe,0x9,0x2,0x7,0xa,0xc,0xb,0xd,0x6,0x3]
Sbox_inv = [Sbox.index(x) for x in xrange(16)]
PBox = [0,16,32,48,1,17,33,49,2,18,34,50,3,19,35,51,4,20,36,52,5,21,37,53,6,22,38,54,7,23,39,55,8,24,40,56,9,25,41,57,10,26,42,58,11,27,43,59,12,28,44,60,13,29,45,61,14,30,46,62,15,31,47,63]
PBox_inv = [PBox.index(x) for x in xrange(64)]

L0 = [[0,0,0,0,1,0,0,0,1,0,0,0,1,0,0,0],
    [0,0,0,0,0,1,0,0,0,1,0,0,0,1,0,0],
    [0,0,0,0,0,0,1,0,0,0,1,0,0,0,1,0],
    [0,0,0,0,0,0,0,1,0,0,0,1,0,0,0,1],
    [1,0,0,0,0,0,0,0,1,0,0,0,1,0,0,0],
    [0,1,0,0,0,0,0,0,0,1,0,0,0,1,0,0],
    [0,0,1,0,0,0,0,0,0,0,1,0,0,0,1,0],
    [0,0,0,1,0,0,0,0,0,0,0,1,0,0,0,1],
    [1,0,0,0,1,0,0,0,0,0,0,0,1,0,0,0],
    [0,1,0,0,0,1,0,0,0,0,0,0,0,1,0,0],
    [0,0,1,0,0,0,1,0,0,0,0,0,0,0,1,0],
    [0,0,0,1,0,0,0,1,0,0,0,0,0,0,0,1],
    [1,0,0,0,1,0,0,0,1,0,0,0,0,0,0,0],
    [0,1,0,0,0,1,0,0,0,1,0,0,0,0,0,0],
    [0,0,1,0,0,0,1,0,0,0,1,0,0,0,0,0],
    [0,0,0,1,0,0,0,1,0,0,0,1,0,0,0,0]]

L1 = [[1,1,0,0,0,0,0,0,0,0,0,1,0,0,0,0],
    [0,1,1,0,0,0,0,0,0,0,0,0,1,0,0,0],
    [0,0,1,1,0,0,0,0,0,0,0,0,0,1,0,0],
    [0,0,0,1,1,0,0,0,0,0,0,0,0,0,1,0],
    [0,0,0,0,1,1,0,0,0,0,0,0,0,0,0,1],
    [0,0,0,0,0,1,1,0,1,0,0,0,0,0,0,0],
    [0,0,0,0,0,0,1,1,0,1,0,0,0,0,0,0],
    [1,0,0,0,0,0,0,1,0,0,1,0,0,0,0,0],
    [1,0,0,0,0,0,0,0,0,0,0,1,1,0,0,0],
    [0,1,0,0,0,0,0,0,0,0,0,0,1,1,0,0],
    [0,0,1,0,0,0,0,0,0,0,0,0,0,1,1,0],
    [0,0,0,1,0,0,0,0,0,0,0,0,0,0,1,1],
    [0,0,0,0,1,0,0,0,1,0,0,0,0,0,0,1],
    [0,0,0,0,0,1,0,0,1,1,0,0,0,0,0,0],
    [0,0,0,0,0,0,1,0,0,1,1,0,0,0,0,0],
    [0,0,0,0,0,0,0,1,0,0,1,1,0,0,0,0]]

L2 = [[0,0,0,0,1,1,0,0,0,0,0,0,0,0,0,1],
    [0,0,0,0,0,1,1,0,1,0,0,0,0,0,0,0],
    [0,0,0,0,0,0,1,1,0,1,0,0,0,0,0,0],
    [1,0,0,0,0,0,0,1,0,0,1,0,0,0,0,0],
    [1,1,0,0,0,0,0,0,0,0,0,1,0,0,0,0],
    [0,1,1,0,0,0,0,0,0,0,0,0,1,0,0,0],
    [0,0,1,1,0,0,0,0,0,0,0,0,0,1,0,0],
    [0,0,0,1,1,0,0,0,0,0,0,0,0,0,1,0],
    [0,0,0,0,1,0,0,0,1,0,0,0,0,0,0,1],
    [0,0,0,0,0,1,0,0,1,1,0,0,0,0,0,0],
    [0,0,0,0,0,0,1,0,0,1,1,0,0,0,0,0],
    [0,0,0,0,0,0,0,1,0,0,1,1,0,0,0,0],
    [1,0,0,0,0,0,0,0,0,0,0,1,1,0,0,0],
    [0,1,0,0,0,0,0,0,0,0,0,0,1,1,0,0],
    [0,0,1,0,0,0,0,0,0,0,0,0,0,1,1,0],
    [0,0,0,1,0,0,0,0,0,0,0,0,0,0,1,1]]

L3 = [[1,0,0,0,1,0,0,0,0,0,0,0,1,0,0,0],
    [0,1,0,0,0,1,0,0,0,0,0,0,0,1,0,0],
    [0,0,1,0,0,0,1,0,0,0,0,0,0,0,1,0],
    [0,0,0,1,0,0,0,1,0,0,0,0,0,0,0,1],
    [1,0,0,0,1,0,0,0,1,0,0,0,0,0,0,0],
    [0,1,0,0,0,1,0,0,0,1,0,0,0,0,0,0],
    [0,0,1,0,0,0,1,0,0,0,1,0,0,0,0,0],
    [0,0,0,1,0,0,0,1,0,0,0,1,0,0,0,0],
    [0,0,0,0,1,0,0,0,1,0,0,0,1,0,0,0],
    [0,0,0,0,0,1,0,0,0,1,0,0,0,1,0,0],
    [0,0,0,0,0,0,1,0,0,0,1,0,0,0,1,0],
    [0,0,0,0,0,0,0,1,0,0,0,1,0,0,0,1],
    [1,0,0,0,0,0,0,0,1,0,0,0,1,0,0,0],
    [0,1,0,0,0,0,0,0,0,1,0,0,0,1,0,0],
    [0,0,1,0,0,0,0,0,0,0,1,0,0,0,1,0],
    [0,0,0,1,0,0,0,0,0,0,0,1,0,0,0,1]]

L0_inv = L0

L1_inv = [[0,0,0,0,0,0,1,1,0,0,0,0,0,0,1,0],
    [1,0,0,0,0,0,0,1,0,0,0,0,0,0,0,1],
    [1,1,0,0,0,0,0,0,1,0,0,0,0,0,0,0],
    [0,1,1,0,0,0,0,0,0,1,0,0,0,0,0,0],
    [0,0,1,1,0,0,0,0,0,0,1,0,0,0,0,0],
    [0,0,0,1,1,0,0,0,0,0,0,1,0,0,0,0],
    [0,0,0,0,1,1,0,0,0,0,0,0,1,0,0,0],
    [0,0,0,0,0,1,1,0,0,0,0,0,0,1,0,0],
    [0,0,0,1,0,0,0,0,0,0,0,1,1,0,0,0],
    [0,0,0,0,1,0,0,0,0,0,0,0,1,1,0,0],
    [0,0,0,0,0,1,0,0,0,0,0,0,0,1,1,0],
    [0,0,0,0,0,0,1,0,0,0,0,0,0,0,1,1],
    [0,0,0,0,0,0,0,1,1,0,0,0,0,0,0,1],
    [1,0,0,0,0,0,0,0,1,1,0,0,0,0,0,0],
    [0,1,0,0,0,0,0,0,0,1,1,0,0,0,0,0],
    [0,0,1,0,0,0,0,0,0,0,1,1,0,0,0,0]]

L2_inv = [[0,0,1,1,0,0,0,0,0,0,1,0,0,0,0,0],
    [0,0,0,1,1,0,0,0,0,0,0,1,0,0,0,0],
    [0,0,0,0,1,1,0,0,0,0,0,0,1,0,0,0],
    [0,0,0,0,0,1,1,0,0,0,0,0,0,1,0,0],
    [0,0,0,0,0,0,1,1,0,0,0,0,0,0,1,0],
    [1,0,0,0,0,0,0,1,0,0,0,0,0,0,0,1],
    [1,1,0,0,0,0,0,0,1,0,0,0,0,0,0,0],
    [0,1,1,0,0,0,0,0,0,1,0,0,0,0,0,0],
    [0,0,0,0,0,0,0,1,1,0,0,0,0,0,0,1],
    [1,0,0,0,0,0,0,0,1,1,0,0,0,0,0,0],
    [0,1,0,0,0,0,0,0,0,1,1,0,0,0,0,0],
    [0,0,1,0,0,0,0,0,0,0,1,1,0,0,0,0],
    [0,0,0,1,0,0,0,0,0,0,0,1,1,0,0,0],
    [0,0,0,0,1,0,0,0,0,0,0,0,1,1,0,0],
    [0,0,0,0,0,1,0,0,0,0,0,0,0,1,1,0],
    [0,0,0,0,0,0,1,0,0,0,0,0,0,0,1,1]]

L3_inv = L3

def matrixMultiply(matrix, input):
    r_outs = []
    for r in matrix:
        r_out = 0;
        for c_i, c in reversed(list(enumerate(r))):
            r_out ^= c * ((input >> (15 - c_i)) & 0x1)
        r_outs.append(r_out)
    return sum([(1 << (15-i)) * v for i,v in enumerate(r_outs)])

def roundKey(key, i):
    return string2number(key[0] + chr((ord(key[1]) + 193 * i) % 256) + key[2] + chr((ord(key[3]) + 165 * i) % 256) + key[4] + chr((ord(key[5]) + 81 * i) % 256) + key[6] + chr((ord(key[7]) + 197 * i) % 256))

def addRoundKey(state,roundkey):
    return state ^ roundkey

def sBoxLayer(state):
    """SBox function for encryption

    Input:  64-bit integer
    Output: 64-bit integer"""

    output = 0
    for i in xrange(16):
        output += Sbox[( state >> (i*4)) & 0xF] << (i*4)
    return output

def sBoxLayer_dec(state):
    """Inverse SBox function for decryption

    Input:  64-bit integer
    Output: 64-bit integer"""

    output = 0
    for i in xrange(16):
        output += Sbox_inv[( state >> (i*4)) & 0xF] << (i*4)
    return output

def pLayer(state):
    """Permutation layer for encryption

    Input:  64-bit integer
    Output: 64-bit integer"""

    return sum ([((state >> i) & 0x01) << PBox[i] for i in xrange(64)])

def pLayer_dec(state):
    """Permutation layer for decryption

    Input:  64-bit integer
    Output: 64-bit integer"""

    return sum ([((state >> i) & 0x01) << PBox_inv[i] for i in xrange(64)])

def lLayer(state):
    state = pLayer(state)
    state = (matrixMultiply(L0, (state >> 48) & 0xffff) << 48) + (matrixMultiply(L1, (state >> 32) & 0xffff) << 32) + (matrixMultiply(L2, (state >> 16) & 0xffff) << 16) + matrixMultiply(L3, state & 0xffff)
    state = pLayer_dec(state)
    return state

def string2number(i):
    """ Convert a string to a number

    Input: string (big-endian)
    Output: long or integer
    """
    return int(i.encode('hex'),16)

def number2string_N(i, N):
    """Convert a number to a string of fixed size

    i: long or integer
    N: length of string
    Output: string (big-endian)
    """
    s = '%0*x' % (N*2, i)
    return s.decode('hex')

def _test():
    import doctest
    doctest.testmod()

if __name__ == "__main__":
    print Pride("0000000000000000fedcba9876543210".decode('hex')).encrypt("0123456789abcdef".decode('hex')).encode('hex')