#!/usr/bin/python
# -*- coding: utf-8 -*-
#
# Copyright (c) 2014, NewAE Technology Inc
# All rights reserved.
#
# Author: Colin O'Flynn
#
# Find this and more at newae.com - this file is part of the chipwhisperer
# project, http://www.assembla.com/spaces/chipwhisperer
#
#    This file is part of chipwhisperer.
#
#    chipwhisperer 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 3 of the License, or
#    (at your option) any later version.
#
#    chipwhisperer 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 Lesser General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with chipwhisperer.  If not, see <http://www.gnu.org/licenses/>.
#=================================================

#
# This file performs forward AND backwards key scheduling. Can work from arbitrary
# key locations (i.e. first to last, last to first, etc.)
#
# Currently only support AES-128 and AES-256
#
from chipwhisperer.common.utils.util import camel_case_deprecated


def sbox(inp):
    s =  [0x63, 0x7c, 0x77, 0x7b, 0xf2, 0x6b, 0x6f, 0xc5, 0x30, 0x01, 0x67,
            0x2b, 0xfe, 0xd7, 0xab, 0x76, 0xca, 0x82, 0xc9, 0x7d, 0xfa, 0x59,
            0x47, 0xf0, 0xad, 0xd4, 0xa2, 0xaf, 0x9c, 0xa4, 0x72, 0xc0, 0xb7,
            0xfd, 0x93, 0x26, 0x36, 0x3f, 0xf7, 0xcc, 0x34, 0xa5, 0xe5, 0xf1,
            0x71, 0xd8, 0x31, 0x15, 0x04, 0xc7, 0x23, 0xc3, 0x18, 0x96, 0x05,
            0x9a, 0x07, 0x12, 0x80, 0xe2, 0xeb, 0x27, 0xb2, 0x75, 0x09, 0x83,
            0x2c, 0x1a, 0x1b, 0x6e, 0x5a, 0xa0, 0x52, 0x3b, 0xd6, 0xb3, 0x29,
            0xe3, 0x2f, 0x84, 0x53, 0xd1, 0x00, 0xed, 0x20, 0xfc, 0xb1, 0x5b,
            0x6a, 0xcb, 0xbe, 0x39, 0x4a, 0x4c, 0x58, 0xcf, 0xd0, 0xef, 0xaa,
            0xfb, 0x43, 0x4d, 0x33, 0x85, 0x45, 0xf9, 0x02, 0x7f, 0x50, 0x3c,
            0x9f, 0xa8, 0x51, 0xa3, 0x40, 0x8f, 0x92, 0x9d, 0x38, 0xf5, 0xbc,
            0xb6, 0xda, 0x21, 0x10, 0xff, 0xf3, 0xd2, 0xcd, 0x0c, 0x13, 0xec,
            0x5f, 0x97, 0x44, 0x17, 0xc4, 0xa7, 0x7e, 0x3d, 0x64, 0x5d, 0x19,
            0x73, 0x60, 0x81, 0x4f, 0xdc, 0x22, 0x2a, 0x90, 0x88, 0x46, 0xee,
            0xb8, 0x14, 0xde, 0x5e, 0x0b, 0xdb, 0xe0, 0x32, 0x3a, 0x0a, 0x49,
            0x06, 0x24, 0x5c, 0xc2, 0xd3, 0xac, 0x62, 0x91, 0x95, 0xe4, 0x79,
            0xe7, 0xc8, 0x37, 0x6d, 0x8d, 0xd5, 0x4e, 0xa9, 0x6c, 0x56, 0xf4,
            0xea, 0x65, 0x7a, 0xae, 0x08, 0xba, 0x78, 0x25, 0x2e, 0x1c, 0xa6,
            0xb4, 0xc6, 0xe8, 0xdd, 0x74, 0x1f, 0x4b, 0xbd, 0x8b, 0x8a, 0x70,
            0x3e, 0xb5, 0x66, 0x48, 0x03, 0xf6, 0x0e, 0x61, 0x35, 0x57, 0xb9,
            0x86, 0xc1, 0x1d, 0x9e, 0xe1, 0xf8, 0x98, 0x11, 0x69, 0xd9, 0x8e,
            0x94, 0x9b, 0x1e, 0x87, 0xe9, 0xce, 0x55, 0x28, 0xdf, 0x8c, 0xa1,
            0x89, 0x0d, 0xbf, 0xe6, 0x42, 0x68, 0x41, 0x99, 0x2d, 0x0f, 0xb0,
            0x54, 0xbb, 0x16]
    return s[inp]


def invsbox(inp):
    s = [0x52, 0x09, 0x6a, 0xd5, 0x30, 0x36, 0xa5, 0x38, 0xbf, 0x40, 0xa3,
            0x9e, 0x81, 0xf3, 0xd7, 0xfb , 0x7c, 0xe3, 0x39, 0x82, 0x9b, 0x2f,
            0xff, 0x87, 0x34, 0x8e, 0x43, 0x44, 0xc4, 0xde, 0xe9, 0xcb , 0x54,
            0x7b, 0x94, 0x32, 0xa6, 0xc2, 0x23, 0x3d, 0xee, 0x4c, 0x95, 0x0b,
            0x42, 0xfa, 0xc3, 0x4e , 0x08, 0x2e, 0xa1, 0x66, 0x28, 0xd9, 0x24,
            0xb2, 0x76, 0x5b, 0xa2, 0x49, 0x6d, 0x8b, 0xd1, 0x25 , 0x72, 0xf8,
            0xf6, 0x64, 0x86, 0x68, 0x98, 0x16, 0xd4, 0xa4, 0x5c, 0xcc, 0x5d,
            0x65, 0xb6, 0x92 , 0x6c, 0x70, 0x48, 0x50, 0xfd, 0xed, 0xb9, 0xda,
            0x5e, 0x15, 0x46, 0x57, 0xa7, 0x8d, 0x9d, 0x84 , 0x90, 0xd8, 0xab,
            0x00, 0x8c, 0xbc, 0xd3, 0x0a, 0xf7, 0xe4, 0x58, 0x05, 0xb8, 0xb3,
            0x45, 0x06 , 0xd0, 0x2c, 0x1e, 0x8f, 0xca, 0x3f, 0x0f, 0x02, 0xc1,
            0xaf, 0xbd, 0x03, 0x01, 0x13, 0x8a, 0x6b , 0x3a, 0x91, 0x11, 0x41,
            0x4f, 0x67, 0xdc, 0xea, 0x97, 0xf2, 0xcf, 0xce, 0xf0, 0xb4, 0xe6,
            0x73 , 0x96, 0xac, 0x74, 0x22, 0xe7, 0xad, 0x35, 0x85, 0xe2, 0xf9,
            0x37, 0xe8, 0x1c, 0x75, 0xdf, 0x6e , 0x47, 0xf1, 0x1a, 0x71, 0x1d,
            0x29, 0xc5, 0x89, 0x6f, 0xb7, 0x62, 0x0e, 0xaa, 0x18, 0xbe, 0x1b ,
            0xfc, 0x56, 0x3e, 0x4b, 0xc6, 0xd2, 0x79, 0x20, 0x9a, 0xdb, 0xc0,
            0xfe, 0x78, 0xcd, 0x5a, 0xf4 , 0x1f, 0xdd, 0xa8, 0x33, 0x88, 0x07,
            0xc7, 0x31, 0xb1, 0x12, 0x10, 0x59, 0x27, 0x80, 0xec, 0x5f , 0x60,
            0x51, 0x7f, 0xa9, 0x19, 0xb5, 0x4a, 0x0d, 0x2d, 0xe5, 0x7a, 0x9f,
            0x93, 0xc9, 0x9c, 0xef , 0xa0, 0xe0, 0x3b, 0x4d, 0xae, 0x2a, 0xf5,
            0xb0, 0xc8, 0xeb, 0xbb, 0x3c, 0x83, 0x53, 0x99, 0x61 , 0x17, 0x2b,
            0x04, 0x7e, 0xba, 0x77, 0xd6, 0x26, 0xe1, 0x69, 0x14, 0x63, 0x55,
            0x21, 0x0c, 0x7d]
    return s[inp]


rcon = [0x8d, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1b, 0x36,
            0x6c, 0xd8, 0xab, 0x4d, 0x9a, 0x2f, 0x5e, 0xbc, 0x63, 0xc6, 0x97,
            0x35, 0x6a, 0xd4, 0xb3, 0x7d, 0xfa, 0xef, 0xc5, 0x91, 0x39, 0x72,
            0xe4, 0xd3, 0xbd, 0x61, 0xc2, 0x9f, 0x25, 0x4a, 0x94, 0x33, 0x66,
            0xcc, 0x83, 0x1d, 0x3a, 0x74, 0xe8, 0xcb, 0x8d, 0x01, 0x02, 0x04,
            0x08, 0x10, 0x20, 0x40, 0x80, 0x1b, 0x36, 0x6c, 0xd8, 0xab, 0x4d,
            0x9a, 0x2f, 0x5e, 0xbc, 0x63, 0xc6, 0x97, 0x35, 0x6a, 0xd4, 0xb3,
            0x7d, 0xfa, 0xef, 0xc5, 0x91, 0x39, 0x72, 0xe4, 0xd3, 0xbd, 0x61,
            0xc2, 0x9f, 0x25, 0x4a, 0x94, 0x33, 0x66, 0xcc, 0x83, 0x1d, 0x3a,
            0x74, 0xe8, 0xcb, 0x8d, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40,
            0x80, 0x1b, 0x36, 0x6c, 0xd8, 0xab, 0x4d, 0x9a, 0x2f, 0x5e, 0xbc,
            0x63, 0xc6, 0x97, 0x35, 0x6a, 0xd4, 0xb3, 0x7d, 0xfa, 0xef, 0xc5,
            0x91, 0x39, 0x72, 0xe4, 0xd3, 0xbd, 0x61, 0xc2, 0x9f, 0x25, 0x4a,
            0x94, 0x33, 0x66, 0xcc, 0x83, 0x1d, 0x3a, 0x74, 0xe8, 0xcb, 0x8d,
            0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1b, 0x36, 0x6c,
            0xd8, 0xab, 0x4d, 0x9a, 0x2f, 0x5e, 0xbc, 0x63, 0xc6, 0x97, 0x35,
            0x6a, 0xd4, 0xb3, 0x7d, 0xfa, 0xef, 0xc5, 0x91, 0x39, 0x72, 0xe4,
            0xd3, 0xbd, 0x61, 0xc2, 0x9f, 0x25, 0x4a, 0x94, 0x33, 0x66, 0xcc,
            0x83, 0x1d, 0x3a, 0x74, 0xe8, 0xcb, 0x8d, 0x01, 0x02, 0x04, 0x08,
            0x10, 0x20, 0x40, 0x80, 0x1b, 0x36, 0x6c, 0xd8, 0xab, 0x4d, 0x9a,
            0x2f, 0x5e, 0xbc, 0x63, 0xc6, 0x97, 0x35, 0x6a, 0xd4, 0xb3, 0x7d,
            0xfa, 0xef, 0xc5, 0x91, 0x39, 0x72, 0xe4, 0xd3, 0xbd, 0x61, 0xc2,
            0x9f, 0x25, 0x4a, 0x94, 0x33, 0x66, 0xcc, 0x83, 0x1d, 0x3a, 0x74,
            0xe8, 0xcb ]


def g_func(inp, rcon):
    #Step 1: change order
    newlist = [inp[1], inp[2], inp[3], inp[0]]

    #Step 2: s-box
    newlist = [sbox(t) for t in newlist]

    #Step 3: apply rcon
    newlist[0] ^= rcon

    return newlist


def h_func(inp):
    #Step 1: s-box
    newlist = [sbox(t) for t in inp]

    return newlist


def xor(l1, l2):
    return [l1[i] ^ l2[i] for i in range(0, len(l1))]


def key_schedule_rounds(input_key, input_round, desired_round):
    """Use key schedule to determine key for different round.

    When dealing with AES-256, inputkey is 16 bytes and inputround
    indicates round that bytes 0...15 are from.

    Args:
        input_key (list): List of bytes of starting key, 16/32 bytes
        input_round (int): Starting round number (i.e. 0 = first)
        desired_round (int): Desired round number (i.e. 10 = last for 16-byte)

    Returns:
         list: Key at desired round number. Can go forward or backwards.
    """

    #Some constants
    n = len(input_key)
    if n == 16:
        pass
    elif n == 32:
        desiredfull = desired_round
        desired_round = int(desired_round / 2)

        #Special case for inputround of 13, needed for 'final' round...
        if input_round != 13:
            if input_round % 2 == 1:
                raise ValueError("Input round must be divisible by 2")
            input_round = int(input_round / 2)
        else:
            if input_round <= desiredfull:
                if desiredfull < 13:
                    raise ValueError("Round = 13 only permissible for reverse")

                if desiredfull == 13:
                    return input_key[0:16]
                else:
                    return input_key[16:32]

    else:
        raise ValueError("Invalid keylength: %d"%n)

    rnd = input_round
    state = list(input_key)

    #Check if we are going forward or backwards
    while rnd < desired_round:
        rnd += 1

        #Forward algorithm, thus need at least one round
        state[0:4] = xor(state[0:4], g_func(state[(n-4):n], rcon[rnd]))
        for i in range(4, n, 4):
            if n == 32 and i == 16:
                inp = h_func(state[(i-4):i])
            else:
                inp = state[(i - 4):i]
            state[i:(i+4)] = xor(state[i:(i+4)], inp)

    while rnd > desired_round:
        #For AES-256 final-round is 13 as that includes 32 bytes
        #of key. Convert to round 12 then continue as usual...
        if n == 32 and rnd == 13:
            inputrnd = int(12/2)
            rnd = inputrnd
            oldstate = list(state[16:32])
            state[16:32] = state[0:16]

            for i in range(12, 0, -4):
                state[i:(i+4)] = xor(oldstate[i:(i+4)], oldstate[(i-4):i])
            state[0:4] = xor(oldstate[0:4], g_func(state[(n - 4):n], rcon[7]))

        if rnd == desired_round:
            break

        # Reverse algorithm, thus need at least one round
        for i in range(n-4, 0, -4):
            if n == 32 and i == 16:
                inp = h_func(state[(i-4):i])
            else:
                inp = state[(i - 4):i]
            state[i:(i+4)] = xor(state[i:(i+4)], inp)
        state[0:4] = xor(state[0:4], g_func(state[(n - 4):n], rcon[rnd]))
        rnd -= 1

    #For AES-256, we use half the generated key at once...
    if n == 32:
        if desiredfull % 2:
            state = state[16:32]
        else:
            state = state[0:16]

    #Return answer
    return state

keyScheduleRounds = camel_case_deprecated(key_schedule_rounds)

def test():
    #Manual tests right now - need to automate this.

    ##### AES-128 Tests
    print("**********AES-128 Tests***************")

    ik = [0]*16
    for i in range(0, 11):
        result = keyScheduleRounds(ik, 0, i)
        print((" ".join(["%2x"%d for d in result])))
        ok = result

    # 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00
    # 62 63 63 63 62 63 63 63 62 63 63 63 62 63 63 63
    # 9b 98 98 c9 f9 fb fb aa 9b 98 98 c9 f9 fb fb aa
    # 90 97 34 50 69 6c cf fa f2 f4 57 33 0b 0f ac 99
    # ee 06 da 7b 87 6a 15 81 75 9e 42 b2 7e 91 ee 2b
    # 7f 2e 2b 88 f8 44 3e 09 8d da 7c bb f3 4b 92 90
    # ec 61 4b 85 14 25 75 8c 99 ff 09 37 6a b4 9b a7
    # 21 75 17 87 35 50 62 0b ac af 6b 3c c6 1b f0 9b
    # 0e f9 03 33 3b a9 61 38 97 06 0a 04 51 1d fa 9f
    # b1 d4 d8 e2 8a 7d b9 da 1d 7b b3 de 4c 66 49 41
    #b4 ef 5b cb 3e 92 e2 11 23 e9 51 cf 6f 8f 18 8e

    print("")

    for i in range(0, 11):  # 10 Rounds
        result = keyScheduleRounds(ok, 10, i)
        print((" ".join(["%2x" % d for d in result])))

    ##### AES-256 Tests
    print("**********AES-256 Tests***************")

    ik = [0]*32
    for i in range(0, 15):  # 14 Rounds
        result = keyScheduleRounds(ik, 0, i)
        print((" ".join(["%02x"%d for d in result])))

    # 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00
    # 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00
    # 62 63 63 63 62 63 63 63 62 63 63 63 62 63 63 63
    # aa fb fb fb aa fb fb fb aa fb fb fb aa fb fb fb
    # 6f 6c 6c cf 0d 0f 0f ac 6f 6c 6c cf 0d 0f 0f ac
    # 7d 8d 8d 6a d7 76 76 91 7d 8d 8d 6a d7 76 76 91
    # 53 54 ed c1 5e 5b e2 6d 31 37 8e a2 3c 38 81 0e
    # 96 8a 81 c1 41 fc f7 50 3c 71 7a 3a eb 07 0c ab
    # 9e aa 8f 28 c0 f1 6d 45 f1 c6 e3 e7 cd fe 62 e9
    # 2b 31 2b df 6a cd dc 8f 56 bc a6 b5 bd bb aa 1e
    # 64 06 fd 52 a4 f7 90 17 55 31 73 f0 98 cf 11 19
    # 6d bb a9 0b 07 76 75 84 51 ca d3 31 ec 71 79 2f
    # e7 b0 e8 9c 43 47 78 8b 16 76 0b 7b 8e b9 1a 62
    # 74 ed 0b a1 73 9b 7e 25 22 51 ad 14 ce 20 d4 3b
    #10 f8 0a 17 53 bf 72 9c 45 c9 79 e7 cb 70 63 85

    print("")

    ik = [0x74 ,0xed ,0x0b ,0xa1 ,0x73 ,0x9b ,0x7e ,0x25 ,0x22 ,
          0x51 ,0xad ,0x14 ,0xce ,0x20 ,0xd4 ,0x3b ,0x10 ,0xf8 ,
          0x0a ,0x17 ,0x53 ,0xbf ,0x72 ,0x9c ,0x45 ,0xc9 ,0x79 ,
          0xe7 ,0xcb ,0x70 ,0x63, 0x85]

    for i in range(0, 14):
        result = keyScheduleRounds(ik, 13, i)
        print((" ".join(["%2x"%d for d in result])))

if __name__ == "__main__":
    test()
