# Python PRIDE implementation
# Version: 1.0
# Date: 22/04/2015
#
# =============================================================================
#
# Python implementation of the PRIDE cipher
# Copyright (C) 2015 Camil Staps (info@camilstaps.nl)
# 
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
#
# =============================================================================

class Pride:

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

        key:    the key as a 128-bit raw string"""

        if len(key) * 8 == 128:
            self.key_whitening = string2number(key[:8])
            self.key_1 = key[8:]
        else:
            raise ValueError, "Key must be a 128-bit raw string"

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

        Input:  plaintext block as raw string
        Output: ciphertext block as raw string"""

        state = string2number(block)

        # Initial permutation & pre-whitening
        state = pLayer_dec(state)
        state = addRoundKey(state, self.key_whitening)
        # 19 rounds R
        for i in xrange (1,20):
            state = addRoundKey(state, roundKey(self.key_1, i))
            state = sBoxLayer(state)
            state = lLayer(state)
        # Last round R'
        state = addRoundKey(state, roundKey(self.key_1, 20))
        state = sBoxLayer(state)
        # Post-whitening & final permutation
        state = addRoundKey(state, self.key_whitening)
        state = pLayer(state)

        return number2string_N(state,8)

    def decrypt(self,block):
        """Decrypt 1 block (8 bytes)

        Input:  ciphertext block as raw string
        Output: plaintex block as raw string"""

        state = string2number(block)

        # Final permutation & post-whitening
        state = pLayer_dec(state)
        state = addRoundKey(state, self.key_whitening)
        # Last round R'
        state = sBoxLayer_dec(state)
        state = addRoundKey(state, roundKey(self.key_1, 20))
        # 19 rounds R
        for i in xrange(19,0,-1):
            state = lLayer_dec(state)
            state = sBoxLayer_dec(state)
            state = addRoundKey(state, roundKey(self.key_1, i))
        # Pre-whitening & initial permutation
        state = addRoundKey(state, self.key_whitening)
        state = pLayer(state)

        return number2string_N(state,8)

# 4 to 4-bit S-Box and its inverse
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)]
# 64-bit permutation P and its inverse
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)]

# Matrices for permutation in the L layer
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]]
L0_inv = L0
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]]
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 = [[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]]
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 = [[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]]
L3_inv = L3

def mapXor(xs):
    """Xor elements of a list together"""
    return reduce(lambda a, b: a^b, xs, 0)

def matrixMultiply(matrix, input):
    """Multiply a vector with a binary matrix

    Input:  matrix as [[Int]];
            input as Int
    Output: Int"""

    mult = [ mapXor([c * ((input >> (15 - c_i)) & 0x1) for c_i, c in reversed(list(enumerate(r)))]) for r in matrix ]
    return sum([(1 << (15-i)) * v for i,v in enumerate(mult)])

def roundKey(key, i):
    """Calculate a round key

    Input:  the base key (second half of it) as a raw string;
            the round number
    Output: the round key as raw string"""

    return pLayer_dec(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"""

    return sum([Sbox[( state >> (i*4)) & 0xF] << (i*4) for i in xrange(16)])

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

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

    return sum([Sbox_inv[( state >> (i*4)) & 0xF] << (i*4) for i in xrange(16)])

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):
    """Perform the L layer:
    * P (permutation)
    * L0 .. L3 on all four 16-bit substrings
    * P_inv (permutation inverse)

    Input:  the current state, as raw string
    Output: the new state, as an raw string"""

    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))
    return pLayer_dec(state)

def lLayer_dec(state):
    """L layer for decryption:
    * P (permutation)
    * L0_inv .. L3_inv multiplication on the four 16-bit substrings, respectively
    * P_inv (permutation inverse)

    Input:  the current state, as raw string
    Output: the new state, as raw string"""

    state = pLayer(state)
    state = (matrixMultiply(L0_inv, (state >> 48) & 0xffff) << 48) + (
        matrixMultiply(L1_inv, (state >> 32) & 0xffff) << 32) + (
        matrixMultiply(L2_inv, (state >> 16) & 0xffff) << 16) + (
        matrixMultiply(L3_inv, state & 0xffff))
    return pLayer_dec(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')

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