// Example program
#include <iostream>
#include <string>
#include <stdio.h>
#include <math.h>
#include <bitset>

static const int sbox[16][16] =
 {{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}};

static const unsigned char rcon[10] = { 1, 2, 4, 8, 16, 32, 64, 128, 27, 54 };

void DrawBits(unsigned char blyat)
{
    std::bitset<8> x(blyat);
    std::cout << x << '\n';
}

void DrawMatrix(unsigned char block[4][4])
{
    for(int i = 0; i < 4; i++)
    {
        for(int j = 0; j < 4; j++)
        {
            //std::cout << block[i][j] << " ";
            std::cout << std::hex << (int)block[i][j] << " ";
            //DrawBits(block[i][j]);
            //std::cout << ((block[i][j] >> (4*0)) & 0xff) << " " << ((block[i][j] >> (4*1)) & 0xff);
        }
        std::cout << std::endl;
    }
    std::cout << std::endl;
}

void DrawRoundMatrix(unsigned char** block)
{
    for(int i = 0; i < 4; i++)
    {
        for(int j = 0; j < 4; j++)
        {
            std::cout << std::hex << (int)block[i][j] << " ";
        }
        std::cout << std::endl;
    }
    std::cout << std::endl;
}

unsigned char bitxor(unsigned char hex1, unsigned char hex2)
{
    //DrawBits(hex1);
    //DrawBits(hex2);
    //std::cout << "----------------" << std::endl;
    //DrawBits(hex1 ^ hex2);
    //std::cout << std::endl;
    return hex1 ^ hex2;
}

unsigned char** KeySchedule(unsigned char key[4][4], int roundcount)
{
    unsigned char** newkey = 0;
    newkey = new unsigned char*[4];
    unsigned char rotWord[4];
    rotWord[0] = key[1][3];
    rotWord[1] = key[2][3];
    rotWord[2] = key[3][3];
    rotWord[3] = key[0][3];

    for(int j = 0; j < 4; j++)
    {
        std::bitset<8> x(rotWord[j]);
        std::bitset<4> first(0);
        std::bitset<4> second(0);
        for(int i = 0; i < 8; i++)
        {
            if(i < 4)
                second[i] = x[i];
            else
                first[i-4] = x[i];
        }
        rotWord[j] = sbox[(int)(first.to_ulong())][(int)(second.to_ulong())];
    }

    unsigned char word[4] = { key[0][0], key[1][0], key[2][0], key[3][0]};
    unsigned char rc[4] = { rcon[roundcount], 0, 0, 0};

    for(int i = 0; i < 4; i++)
    {
        newkey[i] = new unsigned char[4];
        newkey[i][0] = (word[i] ^ rotWord[i] ^ rc[i]);
    }

    word[0] = key[0][1];
    word[1] = key[1][1];
    word[2] = key[2][1];
    word[3] = key[3][1];
    for(int i = 0; i < 4; i++)
    {
        newkey[i][1] = (word[i] ^ newkey[i][0]);
    }

    word[0] = key[0][2];
    word[1] = key[1][2];
    word[2] = key[2][2];
    word[3] = key[3][2];
    for(int i = 0; i < 4; i++)
    {
        newkey[i][2] = (word[i] ^ newkey[i][1]);
    }

    word[0] = key[0][3];
    word[1] = key[1][3];
    word[2] = key[2][3];
    word[3] = key[3][3];
    for(int i = 0; i < 4; i++)
    {
        newkey[i][3] = (word[i] ^ newkey[i][2]);
    }

    return newkey;
}

unsigned char** AddRoundKey(unsigned char block1[4][4], unsigned char block2[4][4])
{
    unsigned char** xored = 0;
    xored = new unsigned char*[4];
    for(int i = 0; i < 4; i++)
    {
        xored[i] = new unsigned char[4];
        for(int j = 0; j < 4; j++)
        {
            xored[i][j] = bitxor(block1[i][j], block2[i][j]);
        }
    }
    return xored;
}

void SubBytes(unsigned char** block)
{
    for(int i = 0; i < 4; i++)
    {
        for(int j = 0; j < 4; j++)
        {
            std::bitset<8> x(block[i][j]);
            std::bitset<4> first(0);
            std::bitset<4> second(0);
            for(int i = 0; i < 8; i++)
            {
                if(i < 4)
                    second[i] = x[i];
                else
                    first[i-4] = x[i];
            }
            block[i][j] = sbox[(int)(first.to_ulong())][(int)(second.to_ulong())];
        }
    }
}

void ShiftRows(unsigned char** block)
{
    int row[4];
    for(int i = 0; i < 4; i++)
    {
        row[i] = block[1][i];
    }
    for(int i = 0; i < 3; i++)
    {
        block[1][i] = row[i+1];
    }
    block[1][3] = row[0];

    for(int i = 0; i < 4; i++)
    {
        row[i] = block[2][i];
    }
    for(int i = 0; i < 2; i++)
    {
        block[2][i] = row[i+2];
    }
    block[2][2] = row[0];
    block[2][3] = row[1];

    for(int i = 0; i < 4; i++)
    {
        row[i] = block[3][i];
    }
    block[3][0] = row[3];
    for(int i = 1; i < 4; i++)
    {
        block[3][i] = row[i-1];
    }
}

void gmix_column(unsigned char *r) {
    unsigned char a[4];
    unsigned char b[4];
    unsigned char c;
    unsigned char h;
    /* The array 'a' is simply a copy of the input array 'r'
     * The array 'b' is each element of the array 'a' multiplied by 2
     * in Rijndael's Galois field
     * a[n] ^ b[n] is element n multiplied by 3 in Rijndael's Galois field */
    for (c = 0; c < 4; c++) {
        a[c] = r[c];
        /* h is 0xff if the high bit of r[c] is set, 0 otherwise */
        h = (unsigned char)((signed char)r[c] >> 7); /* arithmetic right shift, thus shifting in either zeros or ones */
        b[c] = r[c] << 1; /* implicitly removes high bit because b[c] is an 8-bit char, so we xor by 0x1b and not 0x11b in the next line */
        b[c] ^= 0x1B & h; /* Rijndael's Galois field */
    }
    r[0] = b[0] ^ a[3] ^ a[2] ^ b[1] ^ a[1]; /* 2 * a0 + a3 + a2 + 3 * a1 */
    r[1] = b[1] ^ a[0] ^ a[3] ^ b[2] ^ a[2]; /* 2 * a1 + a0 + a3 + 3 * a2 */
    r[2] = b[2] ^ a[1] ^ a[0] ^ b[3] ^ a[3]; /* 2 * a2 + a1 + a0 + 3 * a3 */
    r[3] = b[3] ^ a[2] ^ a[1] ^ b[0] ^ a[0]; /* 2 * a3 + a2 + a1 + 3 * a0 */
}

void MixColumns(unsigned char** block)
{
    unsigned char column[4];
    for(int j = 0; j < 4; j++)
    {
        for(int i = 0; i < 4; i++)
        {
            column[i] = block[i][j];
        }
        gmix_column(column);
        for(int i = 0; i < 4; i++)
        {
            block[i][j] = column[i];
        }
    }
}

std::string AES(int key, std::string text)
{
    unsigned char block[4][4] = {
       {(unsigned char)50, (unsigned char)136, (unsigned char)49, (unsigned char)224},
       {(unsigned char)67, (unsigned char)90, (unsigned char)49, (unsigned char)55},
       {(unsigned char)246, (unsigned char)48, (unsigned char)152, (unsigned char)7},
       {(unsigned char)168, (unsigned char)141, (unsigned char)162, (unsigned char)52}
    };

    unsigned char cipherkey[4][4] = {
       {(unsigned char)43, (unsigned char)40, (unsigned char)171, (unsigned char)9},
       {(unsigned char)126, (unsigned char)174, (unsigned char)247, (unsigned char)207},
       {(unsigned char)21, (unsigned char)210, (unsigned char)21, (unsigned char)79},
       {(unsigned char)22, (unsigned char)166, (unsigned char)136, (unsigned char)60}
    };

    //std::cout << "Basic block:" << std::endl;
    //DrawMatrix(block);

    std::cout << "Cypher key:" << std::endl;
    DrawMatrix(cipherkey);

    unsigned char** addedRoundKey = AddRoundKey(block, cipherkey);

    //std::cout << "After addrounds:" << std::endl;
    //DrawRoundMatrix(addedRoundKey);

    //std::cout << "After subbytes:" << std::endl;
    SubBytes(addedRoundKey);
    //DrawRoundMatrix(addedRoundKey);

    //std::cout << "After shiftrows:" << std::endl;
    ShiftRows(addedRoundKey);
    //DrawRoundMatrix(addedRoundKey);

    //std::cout << "After mixcolumns:" << std::endl;
    MixColumns(addedRoundKey);
    //DrawRoundMatrix(addedRoundKey);

    unsigned char** newkey = KeySchedule(cipherkey, 0);
    //DrawRoundMatrix(newkey);

    unsigned char** newerblock = addedRoundKey(addedRoundKey, newkey);
    DrawRoundMatrix(newerblock);
    //std::cout << (int)(first.to_ulong()) << " " << (int)(second.to_ulong()) << std::endl;
    return "Encrypted data";
}

int main()
{
  /*long long alice;// = 4;
  long long bob;// = 3;
  long long prime = 31;
  long long generated = 5; //a primitive root modulo 'prime'
  alice = rand() % prime; //must be lower than generated
  bob = rand() % prime;
  std::cout << "Alice's private key is " << alice << std::endl;
  std::cout << "Bob's private key is " << bob << std::endl;

  long long ag = (long long)pow(generated, alice) % prime;
  std::cout << "Alice's ag is " << ag << " and sends it to Bob." << std::endl;

  long long bg = (long long)pow(generated, bob) % prime;
  std::cout << "Bob's bg is " << bg << " and sends it to Alice." << std::endl;

  int aliceSecretKey = (long long)pow(bg, alice) % prime;
  std::cout << "Alice's computed key is " << aliceSecretKey << std::endl;

  int bobSecretKey = (long long)pow(ag, bob) % prime;
  std::cout << "Bob's computed key is " << bobSecretKey << std::endl;

  std::cout << "Public values are: P=23, generated value=5, ag=" << ag << " and bg=" << bg << std::endl;
  */

  std::cout << AES(25, "text");
  return 0;
}