#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <stdio.h>

#define AES_BLOCK_SIZE 16
#define AES_WORD_SIZE 4

#define AES_128_KEY_SIZE 16
#define AES_192_KEY_SIZE 24
#define AES_256_KEY_SIZE 32

#define AES_128_ROUND_COUNT 10
#define AES_192_ROUND_COUNT 12
#define AES_256_ROUND_COUNT 14
#define AES_MAX_ROUND_COUNT AES_256_ROUND_COUNT

// Define a structure to hold AES encryption context:
// - The number of rounds
// - The round keys
// - The IV
typedef struct
{
    int number_of_round;
    uint8_t round_key[(AES_MAX_ROUND_COUNT + 1) * AES_WORD_SIZE][AES_WORD_SIZE];
    uint8_t initialization_vector[AES_BLOCK_SIZE];

} AesContext;

// The Rijndael S-box for the SubBytes step
static const uint8_t substitution_box[256] = {
    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};

// The round constants for the key expansion
static const uint8_t round_constant[32] = { 0x7d, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1b, 0x36};

// Functions declaration
void aes_cipher_block(AesContext* ctx, uint8_t block[AES_WORD_SIZE][AES_WORD_SIZE]);
void add_round_key(uint8_t block[AES_WORD_SIZE][AES_WORD_SIZE], uint8_t key[AES_WORD_SIZE][AES_WORD_SIZE]);
void sub_bytes(uint8_t block[AES_WORD_SIZE][AES_WORD_SIZE]);
void shift_rows(uint8_t block[AES_WORD_SIZE][AES_WORD_SIZE]);
void mix_columns(uint8_t block[AES_WORD_SIZE][AES_WORD_SIZE]);
uint8_t xtime(uint8_t byte);
void aes_key_expansion(AesContext* ctx, uint8_t* key, size_t size);
void sub_word(uint8_t word[AES_WORD_SIZE]);
void rot_word(uint8_t word[AES_WORD_SIZE]);
size_t aes_pad_plaintext(uint8_t** buffer, size_t size);
void aes_init(AesContext *ctx, uint8_t *key, size_t key_size, uint8_t *initialization_vector);
void aes_ecb_encrypt(AesContext* ctx, uint8_t* buffer, size_t size);
void aes_cbc_encrypt(AesContext* ctx, uint8_t* buffer, size_t size);

// Perform AES encryption on a block of 16 bytes
void aes_cipher_block(AesContext *ctx, uint8_t block[AES_WORD_SIZE][AES_WORD_SIZE])
{
    // Initial AddRoundKey
    add_round_key(block, ctx->round_key);

    // Iterate through each round
    for (int round = 1; round < ctx->number_of_round; round++)
    {
        sub_bytes(block);
        shift_rows(block);
        mix_columns(block);
        add_round_key(block, ctx->round_key + round * AES_WORD_SIZE);
    }

    // Final round without MixColumns
    sub_bytes(block);
    shift_rows(block);
    add_round_key(block, ctx->round_key + ctx->number_of_round * AES_WORD_SIZE);
}

// Performs the AddRoundKey step
void add_round_key(uint8_t block[AES_WORD_SIZE][AES_WORD_SIZE], uint8_t key[AES_WORD_SIZE][AES_WORD_SIZE])
{
    // XOR each byte with the key
    for (int i = 0; i < AES_WORD_SIZE; i++)
        for (int j = 0; j < AES_WORD_SIZE; j++)
            block[j][i] ^= key[j][i];
}

// Performs the SubBytes step 
void sub_bytes(uint8_t block[AES_WORD_SIZE][AES_WORD_SIZE])
{
    // Substitute each byte using the S-Box
    for (int i = 0; i < AES_WORD_SIZE; i++)
        for (int j = 0; j < AES_WORD_SIZE; j++)
            block[j][i] = substitution_box[block[j][i]];
}

// Performs the ShiftRows step 
void shift_rows(uint8_t block[AES_WORD_SIZE][AES_WORD_SIZE])
{
    uint8_t temp;

    // Shift the second row one position to the left
    temp = block[0][1];
    block[0][1] = block[1][1];
    block[1][1] = block[2][1];
    block[2][1] = block[3][1];
    block[3][1] = temp;

    // Shift the third row two positions to the left (by swapping byte 1 and 3, then byte 2 and 4)
    temp = block[0][2];
    block[0][2] = block[2][2];
    block[2][2] = temp;
    temp = block[1][2];
    block[1][2] = block[3][2];
    block[3][2] = temp;

    // Shift the fourth row three positions to the left (or one position to the right)
    temp = block[3][3];
    block[3][3] = block[2][3];
    block[2][3] = block[1][3];
    block[1][3] = block[0][3];
    block[0][3] = temp;
}

// Calculates the multiplication of a byte with 2 in the Galois field.
uint8_t xtime(uint8_t byte)
{
    if (byte & 0x80)
        return (byte << 1) ^ 0x1B;
    else 
        return byte << 1;
}

// Performs the MixColumns step 
void mix_columns(uint8_t block[AES_WORD_SIZE][AES_WORD_SIZE])
{
    uint8_t xored_column;
    uint8_t temp;

    // Iterate through each column
    for (int i = 0; i < AES_WORD_SIZE; i++)
    {
        // XOR of all bytes in the column
        xored_column = block[i][0] ^ block[i][1] ^ block[i][2] ^ block[i][3];
        temp = block[i][0];

        // Perform some magic equivalent to the matrix multiplication
        block[i][0] ^= xtime(block[i][0] ^ block[i][1]) ^ xored_column;
        block[i][1] ^= xtime(block[i][1] ^ block[i][2]) ^ xored_column;
        block[i][2] ^= xtime(block[i][2] ^ block[i][3]) ^ xored_column;
        block[i][3] ^= xtime(block[i][3] ^ temp) ^ xored_column;
    }
}

// Performs the key expansion for AES encryption
void aes_key_expansion(AesContext *ctx, uint8_t *key, size_t keysize_bytes)
{
    // Calculate the number of words in the key
    const int keysize_words = keysize_bytes / AES_WORD_SIZE;

    // Calculate the final size of the expanded key
    const int expanded_keysize = AES_WORD_SIZE * (ctx->number_of_round + 1);

    // Copy the original key to the initial part of the round key
    memcpy(ctx->round_key, key, keysize_bytes);

    for (int i = keysize_words; i < expanded_keysize; i++)
    {
        // The previous word
        uint8_t word[AES_WORD_SIZE] = {ctx->round_key[i - 1][0], ctx->round_key[i - 1][1], ctx->round_key[i - 1][2], ctx->round_key[i - 1][3]};

        if (i % keysize_words == 0)
        {
            // It's time to apply some transformation 
            rot_word(word);
            sub_word(word);
            word[0] ^= round_constant[i / keysize_words];
        }
        else if (keysize_bytes == 32 && i % keysize_words == AES_WORD_SIZE)
        {
            // Additional step for 256 bit keys
            sub_word(word);
        }

        // Generate the next word in the round key by XORing 2 previous rows
        ctx->round_key[i][0] = word[0] ^ ctx->round_key[i - keysize_words][0];
        ctx->round_key[i][1] = word[1] ^ ctx->round_key[i - keysize_words][1];
        ctx->round_key[i][2] = word[2] ^ ctx->round_key[i - keysize_words][2];
        ctx->round_key[i][3] = word[3] ^ ctx->round_key[i - keysize_words][3];
    }
}

// Substitute each byte using the S-Box
void sub_word(uint8_t word[AES_WORD_SIZE])
{
    for (int i = 0; i < AES_WORD_SIZE; i++)
        word[i] = substitution_box[word[i]];
}

// Shift the bytes to the left
void rot_word(uint8_t word[AES_WORD_SIZE])
{
    const uint8_t tmp = word[0];
    word[0] = word[1];
    word[1] = word[2];
    word[2] = word[3];
    word[3] = tmp;
}

// Initializes an AES encryption context 
void aes_init(AesContext *ctx, uint8_t *key, size_t key_size, uint8_t *initialization_vector)
{
    // Configure the number of round from the key size
    if (key_size == AES_128_KEY_SIZE)
        ctx->number_of_round = AES_128_ROUND_COUNT;
   
    else if (key_size == AES_192_KEY_SIZE)
        ctx->number_of_round = AES_192_ROUND_COUNT;
    
    else if (key_size == AES_256_KEY_SIZE)
        ctx->number_of_round = AES_256_ROUND_COUNT;
   
    else
    {
        printf("Invalid key length: %ld bytes\n", key_size);
        exit(-1);
    }

    // Copy the IV in the context
    if (initialization_vector != NULL)
        memcpy(&ctx->initialization_vector, initialization_vector, AES_BLOCK_SIZE);
    else
        memset(&ctx->initialization_vector, 0x00, AES_BLOCK_SIZE);

    // Perform the Key Expansion
    aes_key_expansion(ctx, key, key_size);
}

// Pad the plaintext in preparation for AES encryption
size_t aes_pad_plaintext(uint8_t **buffer, size_t size)
{
    // Calculate the number of padding bytes required
    uint8_t padding = AES_BLOCK_SIZE - (size % AES_BLOCK_SIZE);

    // Reallocate a new block of memory with padding
    *buffer = realloc(*buffer, size + padding);

    // Fill the padding bytes with the padding value
    memset(*buffer + size, padding, padding);

    // Return the new reallocated size
    return size + padding;
}

// AES encryption using the Electronic Code Book (ECB) mode.
void aes_ecb_encrypt(AesContext *ctx, uint8_t *buffer, size_t size)
{
    // Iterate through the data in blocks of 16 bytes
    for (int i = 0; i < size; i += AES_BLOCK_SIZE)
    {
        // Perform AES encryption on the current block
        aes_cipher_block(ctx, (uint8_t(*)[AES_WORD_SIZE])(buffer + i));
    }
}

// AES encryption using the Cipher Block Chaining (CBC) mode.
void aes_cbc_encrypt(AesContext *ctx, uint8_t *buffer, size_t size)
{
    // IV is used for the first block
    uint8_t *previous = ctx->initialization_vector;

    // Iterate through the data in blocks of 16 bytes
    for (int i = 0; i < size; i += AES_BLOCK_SIZE)
    {
        // XOR the current block of plaintext with the previous ciphertext block
        for (int j = 0; j < 16; j++)
            buffer[i + j] ^= previous[j];

        // Perform AES encryption on the current block
        aes_cipher_block(ctx, (uint8_t (*)[4])(buffer + i));

        // Update the previous pointer
        previous = buffer + i;
    }
}

// AES encryption using the Counter (CTR) mode.
void aes_ctr_encrypt(AesContext *ctx, uint8_t *buffer, size_t size)
{
    // Initialize the counter from the IV
    uint8_t *counter = ctx->initialization_vector;

    // Temporary buffer for storing the encrypted counter block
    uint8_t temp[16];

    // Iterate through the data in blocks of 16 bytes
    for (int i = 0; i < size; i += AES_BLOCK_SIZE)
    {
        // Encrypt the counter
        memcpy(&temp, counter, 16);
        aes_cipher_block(ctx, (uint8_t (*)[4])temp);

        // XOR the encrypted counter with the plaintext
        for (int j = 0; j < 16 && i + j < size; j++)
            buffer[i + j] ^= temp[j];

        // Increment the counter
        for (int i= AES_BLOCK_SIZE - 1; i>=0; i--)
        {
            counter[i]++;

            // Carry overflow to the next byte if necessary
            if (counter[i] != 0) break;
        }
    }
}

// Prints the contents of a byte buffer 
void print_buffer(uint8_t *buffer, int len)
{
    for (int i = 0; i < len; i++)
        printf("%02x ", buffer[i]);
    printf("\n");
}

int main()
{
    // Initialize the plaintext, the key and the IV
    const char* plaintext = "Lorem ipsum dolor sit amet sunt labore nulla, velit pariatur aliquip pariatur. Voluptate dolore fugiat occaecat nostrud. Anim, labore non est do sunt tempor officia enim voluptate ea. Aute id adipisicing excepteur qui est officia aliqua veniam do, quis culpa sunt eu proident sint. Nostrud ad velit esse excepteur ipsum occaecat culpa ullamco pariatur consectetur exercitation, fugiat et sit. Nulla dolore cupidatat labore ut ad reprehenderit quis Lorem exercitation fugiat labore cillum consectetur.";
    const int key_size = 32;
    uint8_t key[] = {0x5e, 0x88, 0x48, 0x98, 0xda, 0x28, 0x04, 0x71, 0x51, 0xd0, 0xe5, 0x6f, 0x8d, 0xc6, 0x29, 0x27, 0x73, 0x60, 0x3d, 0x0d, 0x6a, 0xab, 0xbd, 0xd6, 0x2a, 0x11, 0xef, 0x72, 0x1d, 0x15, 0x42, 0xd8};
    uint8_t initialization_vector[] = {0xef, 0x74, 0x9f, 0xf9, 0xa0, 0x48, 0xba, 0xd0, 0xdd, 0x80, 0x80, 0x7f, 0xc4, 0x9e, 0x1c, 0xf7};

    // Convert the plaintext into a byte buffer
    size_t len = strlen(plaintext);
    uint8_t *buffer = malloc(len);
    memcpy(buffer, plaintext, len);

    // Initialize an AES context
    AesContext context;
    aes_init(&context, key, key_size, initialization_vector);
    
    // Pad the plaintext
    len = aes_pad_plaintext(&buffer, len);
    
    // Perform the AES encryption (CBC mode)
    aes_cbc_encrypt(&context, buffer, len);

    // Print the result
    print_buffer(buffer, len);

    free(buffer);
}