# coding=utf8

# =============================================================================
# Python PRIDE implementation
# Version: 1.2
# Date: 29/04/2015
#
# =============================================================================
#
# Python implementation of the PRIDE cipher
# Copyright (C) 2015 Camil Staps (info@camilstaps.nl)
#
# Some general implementation ideas were taken from Cristophe Osterlynck and 
# Philippe Teuwen's PRESENT implementation:
# http://www.lightweightcrypto.org/downloads/implementations/pypresent.py
#
# This implementation is slightly (~10%) faster using Python3 as opposed to 
# using Python2. If you're concerned with performance, consider Python3.
#
# =============================================================================
#
# 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
>>> import binascii

Create a Pride object:
-----------------------
>>> key = binascii.unhexlify("00000000000000000000000000000000")
>>> cipher = Pride(key)

Encryption:
------------
>>> plain = binascii.unhexlify("0000000000000000")
>>> encrypted = cipher.encrypt(plain)
>>> binascii.hexlify(encrypted)
b'82b4109fcc70bd1f'

Decryption:
------------
>>> decrypted = cipher.decrypt(encrypted)
>>> binascii.hexlify(decrypted)
b'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.
"""

import binascii, sys

class Pride:

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

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

        if len(key) != 16:
            raise ValueError("Key must be a 128-bit raw string")
        self.whitening_key = string2number(key[:8])
        self.roundkeys = [roundKey(key[8:], i) for i in range(0,21)]

    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 = permute_dec(state)
        state = addRoundKey(state, self.whitening_key)
        # 19 rounds R
        for i in range (1,20):
            state = addRoundKey(state, self.roundkeys[i])
            state = apply_SBox(state)
            state = linear_layer(state)
        # Last round R'
        state = addRoundKey(state, self.roundkeys[20])
        state = apply_SBox(state)
        # Post-whitening & final permutation
        state = addRoundKey(state, self.whitening_key)
        state = permute(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 = permute_dec(state)
        state = addRoundKey(state, self.whitening_key)
        # Last round R'
        state = apply_SBox_dec(state)
        state = addRoundKey(state, self.roundkeys[20])
        # 19 rounds R
        for i in range(19,0,-1):
            state = linear_layer_dec(state)
            state = apply_SBox_dec(state)
            state = addRoundKey(state, self.roundkeys[i])
        # Pre-whitening & initial permutation
        state = addRoundKey(state, self.whitening_key)
        state = permute(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 range(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 range(64)]

def swap(byte):
    """Swap byte nibbles"""
    return ((byte & 0xf0) >> 4) | ((byte & 0x0f) << 4)

def rol(byte):
    """Rotate left (bit 0 := bit 7)"""
    return ((byte << 1) | (byte >> 7)) & 0xff

def ror(byte):
    """Rotate right (bit 7 := bit 0)"""
    return (byte >> 1) | ((byte & 0x01) << 7)

def apply_L0(state):
    """Apply matrix L0 using the method described in appendix H of the paper"""
    s0 = state >> 8
    s1 = state & 0x00ff
    temp = swap(s1 ^ s0)
    return ((s1 ^ temp) << 8) | (s0 ^ temp)

def apply_L1(state):
    """Apply matrix L1 using the method described in appendix H of the paper"""
    s2 = state >> 8
    s3 = swap(state & 0x00ff)
    temp = s2 ^ ror(s3)
    return ((rol(s2) ^ temp) << 8) | (s3 ^ temp)

def apply_L1_inv(state):
    """Apply matrix L1^-1 using the method described in appendix H of the paper"""
    s2 = ror(state >> 8)
    s3 = state & 0x00ff
    temp = s2 ^ ror(s3)
    return ((ror(temp) ^ s2) << 8) | swap(temp ^ s3)

def apply_L2(state):
    """Apply matrix L2 using the method described in appendix H of the paper"""
    s4 = swap(state >> 8)
    s5 = state & 0x00ff
    temp = s4 ^ ror(s5)
    return ((temp ^ rol(s4)) << 8) | (temp ^ s5)

def apply_L2_inv(state):
    """Apply matrix L2^-1 using the method described in appendix H of the paper"""
    s4 = ror(state >> 8)
    s5 = state & 0x00ff
    temp = s4 ^ ror(s5)
    return (swap(s4 ^ ror(temp)) << 8) | (s5 ^ temp)

def apply_L3(state):
    """Apply matrix L3 using the method described in appendix H of the paper"""
    s6 = state >> 8
    s7 = state & 0x00ff
    temp = swap(s6 ^ s7)
    return ((s6 ^ temp) << 8) | (s7 ^ temp)

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

    key = string_bytes(key)
    key_parts = [key[0], (key[1] + 193 * i) % 256, key[2], (key[3] + 165 * i) % 256, key[4], (key[5] + 81 * i) % 256, key[6], (key[7] + 197 * i) % 256]
    return permute_dec(sum([(1 << (8 * (7-i))) * k for i, k in enumerate(key_parts)]))

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

def apply_SBox(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 range(16)])

def apply_SBox_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 range(16)])

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

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

    return sum ([((state >> i) & 1) << PBox[i] for i in range(64)])

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

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

    return sum ([((state >> i) & 1) << PBox_inv[i] for i in range(64)])

def linear_layer(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 = permute(state)
    state = (apply_L0((state >> 48) & 0xffff) << 48) + (
        apply_L1((state >> 32) & 0xffff) << 32) + (
        apply_L2((state >> 16) & 0xffff) << 16) + (
        apply_L3(state & 0xffff))
    return permute_dec(state)

def linear_layer_dec(state):
    """Linear 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 = permute(state)
    state = (apply_L0((state >> 48) & 0xffff) << 48) + (
        apply_L1_inv((state >> 32) & 0xffff) << 32) + (
        apply_L2_inv((state >> 16) & 0xffff) << 16) + (
        apply_L3(state & 0xffff))
    return permute_dec(state)

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

    Input: string (big-endian)
    Output: long or integer
    """
    return int(binascii.hexlify(i), 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 binascii.unhexlify(s)

def string_bytes(s):
    return [ord(c) for c in s] if sys.version_info.major == 2 else s