# coding=utf8

# =============================================================================
# Python PRIDE implementation
# Version: 1.0
# Date: 23/04/2015
#
# =============================================================================
#
# Python implementation of the PRIDE cipher
# Copyright (C) 2015 Camil Staps (info@camilstaps.nl)
#
# =============================================================================
#
# PRIDE is a modern (2014) lightweight block cipher optimized for 8-bit 
# microcontrollers, that "significantly outperforms all existing block ciphers
# of similar key-sizes, with the exception of SIMON and SPECK". Its design 
# rationale is described in Block Ciphers -- Focus On The Linear Layer (feat. 
# PRIDE); Martin R. Albrecht, Benedikt Driessen, Elif Bilge Kavun, Gregor 
# Leander, Christof Paar, Tolga Yalçın: https://eprint.iacr.org/2014/453
# 
# =============================================================================

"""
PRIDE block cipher implementation

USAGE EXAMPLE:
---------------
Importing:
-----------
>>> from pypride import Pride

Create a Pride object:
-----------------------
>>> key = "00000000000000000000000000000000".decode('hex')
>>> cipher = Pride(key)

Encryption:
------------
>>> plain = "0000000000000000".decode('hex')
>>> encrypted = cipher.encrypt(plain)
>>> encrypted.encode('hex')
'82b4109fcc70bd1f'

Decryption:
------------
>>> decrypted = cipher.decrypt(encrypted)
>>> decrypted.encode('hex')
'0000000000000000'

This implementation is fully based on the report PRIDE was presented in (https://eprint.iacr.org/2014/453; specifically section 5.4).
Test vectors can be found in test-vectors.py and were taken from appendix J of that paper.
"""

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