/*
 * This file is part of the Flowee project
 * Copyright (C) 2016 The Bitcoin Core developers
 *
 * 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 3 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, see <http://www.gnu.org/licenses/>.
 */

#include "aes.h"
#include "common.h"

#include <cassert>
#include <cstring>

extern "C" {
#include "ctaes/ctaes.c"
}

AES128Encrypt::AES128Encrypt(const char key[])
{
    AES128_init(&ctx, reinterpret_cast<const uint8_t*>(key));
}

AES128Encrypt::~AES128Encrypt()
{
    memset(&ctx, 0, sizeof(ctx));
}

void AES128Encrypt::encrypt(char ciphertext[], const char plaintext[]) const
{
    AES128_encrypt(&ctx, 1, reinterpret_cast<uint8_t*>(ciphertext), reinterpret_cast<const uint8_t*>(plaintext));
}

AES128Decrypt::AES128Decrypt(const char key[])
{
    AES128_init(&ctx, reinterpret_cast<const uint8_t*>(key));
}

AES128Decrypt::~AES128Decrypt()
{
    memset(&ctx, 0, sizeof(ctx));
}

void AES128Decrypt::decrypt(char plaintext[], const char ciphertext[]) const
{
    AES128_decrypt(&ctx, 1, reinterpret_cast<uint8_t*>(plaintext), reinterpret_cast<const uint8_t*>(ciphertext));
}

AES256Encrypt::AES256Encrypt(const char key[32])
{
    AES256_init(&ctx, reinterpret_cast<const uint8_t*>(key));
}

AES256Encrypt::~AES256Encrypt()
{
    memset(&ctx, 0, sizeof(ctx));
}

void AES256Encrypt::encrypt(char ciphertext[], const char plaintext[]) const
{
    AES256_encrypt(&ctx, 1, reinterpret_cast<uint8_t*>(ciphertext), reinterpret_cast<const uint8_t*>(plaintext));
}

AES256Decrypt::AES256Decrypt(const char key[32])
{
    AES256_init(&ctx, reinterpret_cast<const uint8_t*>(key));
}

AES256Decrypt::~AES256Decrypt()
{
    memset(&ctx, 0, sizeof(ctx));
}

void AES256Decrypt::decrypt(char plaintext[], const char ciphertext[]) const
{
    AES256_decrypt(&ctx, 1, reinterpret_cast<uint8_t*>(plaintext), reinterpret_cast<const uint8_t*>(ciphertext));
}


template <typename T>
static int CBCEncrypt(const T& enc, const unsigned char iv[AES_BLOCKSIZE], const unsigned char* data, int size, bool pad, char* out)
{
    int written = 0;
    int padsize = size % AES_BLOCKSIZE;
    char mixed[AES_BLOCKSIZE];

    if (!data || !size || !out)
        return 0;

    if (!pad && padsize != 0)
        return 0;

    memcpy(mixed, iv, AES_BLOCKSIZE);

    // Write all but the last block
    while (written + AES_BLOCKSIZE <= size) {
        for (int i = 0; i != AES_BLOCKSIZE; i++)
            mixed[i] ^= *data++;
        enc.encrypt(out + written, mixed);
        memcpy(mixed, out + written, AES_BLOCKSIZE);
        written += AES_BLOCKSIZE;
    }
    if (pad) {
        // For all that remains, pad each byte with the value of the remaining
        // space. If there is none, pad by a full block.
        for (int i = 0; i != padsize; i++)
            mixed[i] ^= *data++;
        for (int i = padsize; i != AES_BLOCKSIZE; i++)
            mixed[i] ^= AES_BLOCKSIZE - padsize;
        enc.encrypt(out + written, mixed);
        written += AES_BLOCKSIZE;
    }
    return written;
}

template <typename T>
static int CBCDecrypt(const T& dec, const char iv[AES_BLOCKSIZE], const char* data, int size, bool pad, char* out)
{
    unsigned char padsize = 0;
    int written = 0;
    bool fail = false;
    const char* prev = iv;

    if (!data || !size || !out)
        return 0;

    if (size % AES_BLOCKSIZE != 0)
        return 0;

    // Decrypt all data. Padding will be checked in the output.
    while (written != size) {
        dec.decrypt(out, data + written);
        for (int i = 0; i != AES_BLOCKSIZE; i++)
            *out++ ^= prev[i];
        prev = data + written;
        written += AES_BLOCKSIZE;
    }

    // When decrypting padding, attempt to run in constant-time
    if (pad) {
        // If used, padding size is the value of the last decrypted byte. For
        // it to be valid, It must be between 1 and AES_BLOCKSIZE.
        padsize = *--out;
        fail = !padsize | (padsize > AES_BLOCKSIZE);

        // If not well-formed, treat it as though there's no padding.
        padsize *= !fail;

        // All padding must equal the last byte otherwise it's not well-formed
        for (int i = AES_BLOCKSIZE; i != 0; i--)
            fail |= ((i > AES_BLOCKSIZE - padsize) & (*out-- != padsize));

        written -= padsize;
    }
    return written * !fail;
}

AES256CBCEncrypt::AES256CBCEncrypt(const char key[], const char ivIn[], bool padIn)
    : enc(key), pad(padIn)
{
    memcpy(iv, ivIn, AES_BLOCKSIZE);
}

int AES256CBCEncrypt::encrypt(const char *data, int size, char *out) const
{
    return CBCEncrypt(enc, iv, reinterpret_cast<const uint8_t*>(data), size, pad, out);
}

int AES256CBCEncrypt::encrypt(const std::vector<unsigned char> &data, char *out) const
{
    return encrypt(reinterpret_cast<const char*>(data.data()), data.size(), out);
}

int AES256CBCEncrypt::encrypt(const std::vector<char> &data, char *out) const
{
    return encrypt(data.data(), data.size(), out);
}

AES256CBCEncrypt::~AES256CBCEncrypt()
{
    memset(iv, 0, sizeof(iv));
}

AES256CBCDecrypt::AES256CBCDecrypt(const char key[], const char ivIn[], bool padIn)
    : dec(key), pad(padIn)
{
    memcpy(iv, ivIn, AES_BLOCKSIZE);
}


int AES256CBCDecrypt::decrypt(const char *data, int size, char *out) const
{
    return CBCDecrypt(dec, iv, data, size, pad, out);
}

int AES256CBCDecrypt::decrypt(const std::vector<char> &data, char *out) const
{
    return decrypt(data.data(), data.size(), out);
}

int AES256CBCDecrypt::decrypt(const std::vector<unsigned char> &data, char *out) const
{
    return decrypt(reinterpret_cast<const char*>(data.data()), data.size(), out);
}

AES256CBCDecrypt::~AES256CBCDecrypt()
{
    memset(iv, 0, sizeof(iv));
}

AES128CBCEncrypt::AES128CBCEncrypt(const char key[], const char ivIn[], bool padIn)
    : enc(key), pad(padIn)
{
    memcpy(iv, ivIn, AES_BLOCKSIZE);
}

AES128CBCEncrypt::~AES128CBCEncrypt()
{
    memset(iv, 0, AES_BLOCKSIZE);
}

int AES128CBCEncrypt::encrypt(const char *data, int size, char *out) const
{
    return CBCEncrypt(enc, iv, reinterpret_cast<const uint8_t*>(data), size, pad, out);
}

AES128CBCDecrypt::AES128CBCDecrypt(const char key[], const char ivIn[], bool padIn)
    : dec(key), pad(padIn)
{
    memcpy(iv, ivIn, AES_BLOCKSIZE);
}

AES128CBCDecrypt::~AES128CBCDecrypt()
{
    memset(iv, 0, AES_BLOCKSIZE);
}

int AES128CBCDecrypt::decrypt(const char *data, int size, char *out) const
{
    return CBCDecrypt(dec, iv, data, size, pad, out);
}