AES256.cpp

/*
 * Copyright (C) 2015,2018 Southern Storm Software, Pty Ltd.
 *
 * Permission is hereby granted, free of charge, to any person obtaining a
 * copy of this software and associated documentation files (the "Software"),
 * to deal in the Software without restriction, including without limitation
 * the rights to use, copy, modify, merge, publish, distribute, sublicense,
 * and/or sell copies of the Software, and to permit persons to whom the
 * Software is furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included
 * in all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
 * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 * DEALINGS IN THE SOFTWARE.
 */

#include "AES.h"
#include "Crypto.h"
#include <string.h>

AES256::AES256()
{
    rounds = 14;
    schedule = sched;
}

AES256::~AES256()
{
    clean(sched);
}

size_t AES256::keySize() const
{
    return 32;
}

bool AES256::setKey(const uint8_t *key, size_t len)
{
    if (len != 32)
        return false;

    // Copy the key itself into the first 32 bytes of the schedule.
    uint8_t *schedule = sched;
    memcpy(schedule, key, 32);

    // Expand the key schedule until we have 240 bytes of expanded key.
    uint8_t iteration = 1;
    uint8_t n = 32;
    uint8_t w = 8;
    while (n < 240) {
        if (w == 8) {
            // Every 32 bytes (8 words) we need to apply the key schedule core.
            keyScheduleCore(schedule + 32, schedule + 28, iteration);
            schedule[32] ^= schedule[0];
            schedule[33] ^= schedule[1];
            schedule[34] ^= schedule[2];
            schedule[35] ^= schedule[3];
            ++iteration;
            w = 0;
        } else if (w == 4) {
            // At the 16 byte mark we need to apply the S-box.
            applySbox(schedule + 32, schedule + 28);
            schedule[32] ^= schedule[0];
            schedule[33] ^= schedule[1];
            schedule[34] ^= schedule[2];
            schedule[35] ^= schedule[3];
        } else {
            // Otherwise just XOR the word with the one 32 bytes previous.
            schedule[32] = schedule[28] ^ schedule[0];
            schedule[33] = schedule[29] ^ schedule[1];
            schedule[34] = schedule[30] ^ schedule[2];
            schedule[35] = schedule[31] ^ schedule[3];
        }

        // Advance to the next word in the schedule.
        schedule += 4;
        n += 4;
        ++w;
    }

    return true;
}

// Helper macros.
#define LEFT 0
#define RIGHT 16
#define ENCRYPT(phase) \
    do { \
        AESCommon::subBytesAndShiftRows(state2, state1); \
        AESCommon::mixColumn(state1,      state2); \
        AESCommon::mixColumn(state1 + 4,  state2 + 4); \
        AESCommon::mixColumn(state1 + 8,  state2 + 8); \
        AESCommon::mixColumn(state1 + 12, state2 + 12); \
        for (posn = 0; posn < 16; ++posn) \
            state1[posn] ^= schedule[posn + (phase)]; \
    } while (0)
#define DECRYPT(phase) \
    do { \
        for (posn = 0; posn < 16; ++posn) \
            state2[posn] ^= schedule[posn + (phase)]; \
        AESCommon::inverseMixColumn(state1,      state2); \
        AESCommon::inverseMixColumn(state1 + 4,  state2 + 4); \
        AESCommon::inverseMixColumn(state1 + 8,  state2 + 8); \
        AESCommon::inverseMixColumn(state1 + 12, state2 + 12); \
        AESCommon::inverseShiftRowsAndSubBytes(state2, state1); \
    } while (0)
#define KCORE(n) \
    do { \
        AESCommon::keyScheduleCore(temp, schedule + 28, (n)); \
        schedule[0] ^= temp[0]; \
        schedule[1] ^= temp[1]; \
        schedule[2] ^= temp[2]; \
        schedule[3] ^= temp[3]; \
    } while (0)
#define KXOR(a, b) \
    do { \
        schedule[(a) * 4] ^= schedule[(b) * 4]; \
        schedule[(a) * 4 + 1] ^= schedule[(b) * 4 + 1]; \
        schedule[(a) * 4 + 2] ^= schedule[(b) * 4 + 2]; \
        schedule[(a) * 4 + 3] ^= schedule[(b) * 4 + 3]; \
    } while (0)
#define KSBOX() \
    do { \
        AESCommon::applySbox(temp, schedule + 12); \
        schedule[16] ^= temp[0]; \
        schedule[17] ^= temp[1]; \
        schedule[18] ^= temp[2]; \
        schedule[19] ^= temp[3]; \
    } while (0)

AESTiny256::AESTiny256()
{
}

AESTiny256::~AESTiny256()
{
    clean(schedule);
}

size_t AESTiny256::blockSize() const
{
    return 16;
}

size_t AESTiny256::keySize() const
{
    return 32;
}

bool AESTiny256::setKey(const uint8_t *key, size_t len)
{
    if (len == 32) {
        // Make a copy of the key - it will be expanded in encryptBlock().
        memcpy(schedule, key, 32);
        return true;
    }
    return false;
}

void AESTiny256::encryptBlock(uint8_t *output, const uint8_t *input)
{
    uint8_t schedule[32];
    uint8_t posn;
    uint8_t round;
    uint8_t state1[16];
    uint8_t state2[16];
    uint8_t temp[4];

    // Start with the key in the schedule buffer.
    memcpy(schedule, this->schedule, 32);

    // Copy the input into the state and perform the first round.
    for (posn = 0; posn < 16; ++posn)
        state1[posn] = input[posn] ^ schedule[posn];
    ENCRYPT(RIGHT);

    // Perform the next 12 rounds of the cipher two at a time.
    for (round = 1; round <= 6; ++round) {
        // Expand the next 32 bytes of the key schedule.
        KCORE(round);
        KXOR(1, 0);
        KXOR(2, 1);
        KXOR(3, 2);
        KSBOX();
        KXOR(5, 4);
        KXOR(6, 5);
        KXOR(7, 6);

        // Encrypt using the left and right halves of the key schedule.
        ENCRYPT(LEFT);
        ENCRYPT(RIGHT);
    }

    // Expand the final 16 bytes of the key schedule.
    KCORE(7);
    KXOR(1, 0);
    KXOR(2, 1);
    KXOR(3, 2);

    // Perform the final round.
    AESCommon::subBytesAndShiftRows(state2, state1);
    for (posn = 0; posn < 16; ++posn)
        output[posn] = state2[posn] ^ schedule[posn];
}

void AESTiny256::decryptBlock(uint8_t *output, const uint8_t *input)
{
    // Decryption is not supported by AESTiny256.
}

void AESTiny256::clear()
{
    clean(schedule);
}

AESSmall256::AESSmall256()
{
}

AESSmall256::~AESSmall256()
{
    clean(reverse);
}

bool AESSmall256::setKey(const uint8_t *key, size_t len)
{
    uint8_t *schedule;
    uint8_t round;
    uint8_t temp[4];

    // Set the encryption key first.
    if (!AESTiny256::setKey(key, len))
        return false;

    // Expand the key schedule up to the last round which gives
    // us the round keys to use for the final two rounds.  We can
    // then work backwards from there in decryptBlock().
    schedule = reverse;
    memcpy(schedule, key, 32);
    for (round = 1; round <= 6; ++round) {
        KCORE(round);
        KXOR(1, 0);
        KXOR(2, 1);
        KXOR(3, 2);
        KSBOX();
        KXOR(5, 4);
        KXOR(6, 5);
        KXOR(7, 6);
    }
    KCORE(7);
    KXOR(1, 0);
    KXOR(2, 1);
    KXOR(3, 2);

    // Key is ready to go.
    return true;
}

void AESSmall256::decryptBlock(uint8_t *output, const uint8_t *input)
{
    uint8_t schedule[32];
    uint8_t round;
    uint8_t posn;
    uint8_t state1[16];
    uint8_t state2[16];
    uint8_t temp[4];

    // Start with the end of the decryption schedule.
    memcpy(schedule, reverse, 32);

    // Copy the input into the state and reverse the final round.
    for (posn = 0; posn < 16; ++posn)
        state1[posn] = input[posn] ^ schedule[posn];
    AESCommon::inverseShiftRowsAndSubBytes(state2, state1);
    KXOR(3, 2);
    KXOR(2, 1);
    KXOR(1, 0);
    KCORE(7);

    // Perform the next 12 rounds of the decryption process two at a time.
    for (round = 6; round >= 1; --round) {
        // Decrypt using the right and left halves of the key schedule.
        DECRYPT(RIGHT);
        DECRYPT(LEFT);

        // Expand the next 32 bytes of the key schedule in reverse.
        KXOR(7, 6);
        KXOR(6, 5);
        KXOR(5, 4);
        KSBOX();
        KXOR(3, 2);
        KXOR(2, 1);
        KXOR(1, 0);
        KCORE(round);
    }

    // Reverse the initial round and create the output words.
    DECRYPT(RIGHT);
    for (posn = 0; posn < 16; ++posn)
        output[posn] = state2[posn] ^ schedule[posn];
}

void AESSmall256::clear()
{
    clean(reverse);
    AESTiny256::clear();
}