/*
 * aes256.hpp
 *
 * Copyright (c) 2014, Danilo Treffiletti <urban82@gmail.com>
 * All rights reserved.
 *
 *     This file is part of Aes256.
 *
 *     Aes256 is free software: you can redistribute it and/or modify
 *     it under the terms of the GNU Lesser General Public License as
 *     published by the Free Software Foundation, either version 2.1
 *     of the License, or (at your option) any later version.
 *
 *     Aes256 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 Lesser General Public
 *     License along with Aes256.
 *     If not, see <http://www.gnu.org/licenses/>.
 */

#ifndef AES256_HPP
#define AES256_HPP

#include <vector>
#include "buffer.hpp"

#define BLOCK_SIZE 16

class Aes256
{

public:
    Aes256(const std::string &key);
    ~Aes256();

    static bool encrypt(const std::string &key, const Buffer &plain, Buffer &encrypted);
    static bool decrypt(const std::string &key, const Buffer &encrypted, Buffer &plain);

    std::size_t encrypt_start(const std::size_t plain_length, Buffer &encrypted);
    std::size_t encrypt_continue(const Buffer &plain, Buffer &encrypted);
    std::size_t encrypt_end(Buffer &encrypted);

    std::size_t decrypt_start(const std::size_t encrypted_length);
    std::size_t decrypt_continue(const Buffer &encrypted, Buffer &plain);
    std::size_t decrypt_end(Buffer &plain);

private:
    std::vector<unsigned char> m_key;
    std::vector<unsigned char> m_salt;
    std::vector<unsigned char> m_rkey;

    unsigned char m_buffer[3 * BLOCK_SIZE];
    unsigned char m_buffer_pos;
    std::size_t m_remainingLength;

    bool m_decryptInitialized;

    void check_and_encrypt_buffer(Buffer &encrypted);
    void check_and_decrypt_buffer(Buffer &plain);

    void encrypt(unsigned char *buffer);
    void decrypt(unsigned char *buffer);

    void expand_enc_key(unsigned char *rc);
    void expand_dec_key(unsigned char *rc);

    void sub_bytes(unsigned char *buffer);
    void sub_bytes_inv(unsigned char *buffer);

    void copy_key();

    void add_round_key(unsigned char *buffer, const unsigned char round);

    void shift_rows(unsigned char *buffer);
    void shift_rows_inv(unsigned char *buffer);

    void mix_columns(unsigned char *buffer);
    void mix_columns_inv(unsigned char *buffer);
};

#endif /* AES256_HPP */