round2-full-3-known-plaintext.c

/***
 * Differential Attack on full 2-round AES (with last MixColumns) with only 3 known-plaintext
 * Author: Merricx
 * AES implementation is from https://github.com/kokke/tiny-AES-c
 * 
 * Compile:
 * gcc aes.c round2-full-3-known-plaintext.c -o round2-full-3-known-plaintext -lpthread
*/
#include <stdio.h>
#include <string.h>
#include <stdint.h>
#include <stdlib.h>
#include <pthread.h>

#include "aes.h"

#define MAXBUFLEN 16
#define NUM_THREAD 4 // Number of thread in one execution

unsigned char *plaintext1;
unsigned char *plaintext2;
unsigned char *plaintext3;
unsigned char *ciphertext1;
unsigned char *ciphertext2;
unsigned char *ciphertext3;

const uint8_t SBOX[256] = {
  //0     1    2      3     4    5     6     7      8    9     A      B    C     D     E     F
  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 };

const uint8_t INV_SBOX[256] = {
  0x52, 0x09, 0x6a, 0xd5, 0x30, 0x36, 0xa5, 0x38, 0xbf, 0x40, 0xa3, 0x9e, 0x81, 0xf3, 0xd7, 0xfb,
  0x7c, 0xe3, 0x39, 0x82, 0x9b, 0x2f, 0xff, 0x87, 0x34, 0x8e, 0x43, 0x44, 0xc4, 0xde, 0xe9, 0xcb,
  0x54, 0x7b, 0x94, 0x32, 0xa6, 0xc2, 0x23, 0x3d, 0xee, 0x4c, 0x95, 0x0b, 0x42, 0xfa, 0xc3, 0x4e,
  0x08, 0x2e, 0xa1, 0x66, 0x28, 0xd9, 0x24, 0xb2, 0x76, 0x5b, 0xa2, 0x49, 0x6d, 0x8b, 0xd1, 0x25,
  0x72, 0xf8, 0xf6, 0x64, 0x86, 0x68, 0x98, 0x16, 0xd4, 0xa4, 0x5c, 0xcc, 0x5d, 0x65, 0xb6, 0x92,
  0x6c, 0x70, 0x48, 0x50, 0xfd, 0xed, 0xb9, 0xda, 0x5e, 0x15, 0x46, 0x57, 0xa7, 0x8d, 0x9d, 0x84,
  0x90, 0xd8, 0xab, 0x00, 0x8c, 0xbc, 0xd3, 0x0a, 0xf7, 0xe4, 0x58, 0x05, 0xb8, 0xb3, 0x45, 0x06,
  0xd0, 0x2c, 0x1e, 0x8f, 0xca, 0x3f, 0x0f, 0x02, 0xc1, 0xaf, 0xbd, 0x03, 0x01, 0x13, 0x8a, 0x6b,
  0x3a, 0x91, 0x11, 0x41, 0x4f, 0x67, 0xdc, 0xea, 0x97, 0xf2, 0xcf, 0xce, 0xf0, 0xb4, 0xe6, 0x73,
  0x96, 0xac, 0x74, 0x22, 0xe7, 0xad, 0x35, 0x85, 0xe2, 0xf9, 0x37, 0xe8, 0x1c, 0x75, 0xdf, 0x6e,
  0x47, 0xf1, 0x1a, 0x71, 0x1d, 0x29, 0xc5, 0x89, 0x6f, 0xb7, 0x62, 0x0e, 0xaa, 0x18, 0xbe, 0x1b,
  0xfc, 0x56, 0x3e, 0x4b, 0xc6, 0xd2, 0x79, 0x20, 0x9a, 0xdb, 0xc0, 0xfe, 0x78, 0xcd, 0x5a, 0xf4,
  0x1f, 0xdd, 0xa8, 0x33, 0x88, 0x07, 0xc7, 0x31, 0xb1, 0x12, 0x10, 0x59, 0x27, 0x80, 0xec, 0x5f,
  0x60, 0x51, 0x7f, 0xa9, 0x19, 0xb5, 0x4a, 0x0d, 0x2d, 0xe5, 0x7a, 0x9f, 0x93, 0xc9, 0x9c, 0xef,
  0xa0, 0xe0, 0x3b, 0x4d, 0xae, 0x2a, 0xf5, 0xb0, 0xc8, 0xeb, 0xbb, 0x3c, 0x83, 0x53, 0x99, 0x61,
  0x17, 0x2b, 0x04, 0x7e, 0xba, 0x77, 0xd6, 0x26, 0xe1, 0x69, 0x14, 0x63, 0x55, 0x21, 0x0c, 0x7d };

const uint8_t RCON[11] = {
  0x8d, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1b, 0x36 };

const uint8_t gmul_by_2[256] = {
    0x00,0x02,0x04,0x06,0x08,0x0a,0x0c,0x0e,0x10,0x12,0x14,0x16,0x18,0x1a,0x1c,0x1e,
    0x20,0x22,0x24,0x26,0x28,0x2a,0x2c,0x2e,0x30,0x32,0x34,0x36,0x38,0x3a,0x3c,0x3e,
    0x40,0x42,0x44,0x46,0x48,0x4a,0x4c,0x4e,0x50,0x52,0x54,0x56,0x58,0x5a,0x5c,0x5e,
    0x60,0x62,0x64,0x66,0x68,0x6a,0x6c,0x6e,0x70,0x72,0x74,0x76,0x78,0x7a,0x7c,0x7e,
    0x80,0x82,0x84,0x86,0x88,0x8a,0x8c,0x8e,0x90,0x92,0x94,0x96,0x98,0x9a,0x9c,0x9e,
    0xa0,0xa2,0xa4,0xa6,0xa8,0xaa,0xac,0xae,0xb0,0xb2,0xb4,0xb6,0xb8,0xba,0xbc,0xbe,
    0xc0,0xc2,0xc4,0xc6,0xc8,0xca,0xcc,0xce,0xd0,0xd2,0xd4,0xd6,0xd8,0xda,0xdc,0xde,
    0xe0,0xe2,0xe4,0xe6,0xe8,0xea,0xec,0xee,0xf0,0xf2,0xf4,0xf6,0xf8,0xfa,0xfc,0xfe,
    0x1b,0x19,0x1f,0x1d,0x13,0x11,0x17,0x15,0x0b,0x09,0x0f,0x0d,0x03,0x01,0x07,0x05,
    0x3b,0x39,0x3f,0x3d,0x33,0x31,0x37,0x35,0x2b,0x29,0x2f,0x2d,0x23,0x21,0x27,0x25,
    0x5b,0x59,0x5f,0x5d,0x53,0x51,0x57,0x55,0x4b,0x49,0x4f,0x4d,0x43,0x41,0x47,0x45,
    0x7b,0x79,0x7f,0x7d,0x73,0x71,0x77,0x75,0x6b,0x69,0x6f,0x6d,0x63,0x61,0x67,0x65,
    0x9b,0x99,0x9f,0x9d,0x93,0x91,0x97,0x95,0x8b,0x89,0x8f,0x8d,0x83,0x81,0x87,0x85,
    0xbb,0xb9,0xbf,0xbd,0xb3,0xb1,0xb7,0xb5,0xab,0xa9,0xaf,0xad,0xa3,0xa1,0xa7,0xa5,
    0xdb,0xd9,0xdf,0xdd,0xd3,0xd1,0xd7,0xd5,0xcb,0xc9,0xcf,0xcd,0xc3,0xc1,0xc7,0xc5,
    0xfb,0xf9,0xff,0xfd,0xf3,0xf1,0xf7,0xf5,0xeb,0xe9,0xef,0xed,0xe3,0xe1,0xe7,0xe5
};

const uint8_t gmul_by_3[256] = {
    0x00,0x03,0x06,0x05,0x0c,0x0f,0x0a,0x09,0x18,0x1b,0x1e,0x1d,0x14,0x17,0x12,0x11,
    0x30,0x33,0x36,0x35,0x3c,0x3f,0x3a,0x39,0x28,0x2b,0x2e,0x2d,0x24,0x27,0x22,0x21,
    0x60,0x63,0x66,0x65,0x6c,0x6f,0x6a,0x69,0x78,0x7b,0x7e,0x7d,0x74,0x77,0x72,0x71,
    0x50,0x53,0x56,0x55,0x5c,0x5f,0x5a,0x59,0x48,0x4b,0x4e,0x4d,0x44,0x47,0x42,0x41,
    0xc0,0xc3,0xc6,0xc5,0xcc,0xcf,0xca,0xc9,0xd8,0xdb,0xde,0xdd,0xd4,0xd7,0xd2,0xd1,
    0xf0,0xf3,0xf6,0xf5,0xfc,0xff,0xfa,0xf9,0xe8,0xeb,0xee,0xed,0xe4,0xe7,0xe2,0xe1,
    0xa0,0xa3,0xa6,0xa5,0xac,0xaf,0xaa,0xa9,0xb8,0xbb,0xbe,0xbd,0xb4,0xb7,0xb2,0xb1,
    0x90,0x93,0x96,0x95,0x9c,0x9f,0x9a,0x99,0x88,0x8b,0x8e,0x8d,0x84,0x87,0x82,0x81,
    0x9b,0x98,0x9d,0x9e,0x97,0x94,0x91,0x92,0x83,0x80,0x85,0x86,0x8f,0x8c,0x89,0x8a,
    0xab,0xa8,0xad,0xae,0xa7,0xa4,0xa1,0xa2,0xb3,0xb0,0xb5,0xb6,0xbf,0xbc,0xb9,0xba,
    0xfb,0xf8,0xfd,0xfe,0xf7,0xf4,0xf1,0xf2,0xe3,0xe0,0xe5,0xe6,0xef,0xec,0xe9,0xea,
    0xcb,0xc8,0xcd,0xce,0xc7,0xc4,0xc1,0xc2,0xd3,0xd0,0xd5,0xd6,0xdf,0xdc,0xd9,0xda,
    0x5b,0x58,0x5d,0x5e,0x57,0x54,0x51,0x52,0x43,0x40,0x45,0x46,0x4f,0x4c,0x49,0x4a,
    0x6b,0x68,0x6d,0x6e,0x67,0x64,0x61,0x62,0x73,0x70,0x75,0x76,0x7f,0x7c,0x79,0x7a,
    0x3b,0x38,0x3d,0x3e,0x37,0x34,0x31,0x32,0x23,0x20,0x25,0x26,0x2f,0x2c,0x29,0x2a,
    0x0b,0x08,0x0d,0x0e,0x07,0x04,0x01,0x02,0x13,0x10,0x15,0x16,0x1f,0x1c,0x19,0x1a
};

void xor(uint8_t* diff, const char* c1, const char* c2, int size){
    for(int i=0; i<size; i++){
        diff[i] = c1[i] ^ c2[i];
    }
}

void readFile(char* buf, const char* name){
    FILE *fp = fopen(name, "r");
    if (fp != NULL){
        size_t newLen = fread(buf, sizeof(char), MAXBUFLEN, fp);
        if ( ferror( fp ) != 0 ) {
            fputs("Error reading file", stderr);
        } else {
            buf[newLen++] = '\0';
        }

        fclose(fp);
    }
}

void writeFile(char* buf, const char* name){
    FILE *fp = fopen(name, "w");
    if(fp != NULL){
        fprintf(fp, "%.16s", buf);
        fclose(fp);
    }
    
}

/**
 * Attack Phase 2
 * 
 * After we recover Rk0 of bytes (0,5,10,15), we start iterate over Rk0 of bytes (7, 8)
 */
int attack_phase2(uint8_t round_key[][16], const uint8_t* diff_c1_c2, const uint8_t* diff_c1_c3){

    for(int i=0; i < 256; i++){
        for(int j=0; j < 256; j++){

            /**
             * We construct following Rk0
             * where K is known value
             * and x is value we brute-force
             * 
             * [
             *   K, 0, x, 0,
             *   0, K, 0, K,
             *   K, 0, K, 0,
             *   0, x, 0, K
             * ]
             * 
             */
            round_key[0][7] = j;
            round_key[0][8] = i;

            uint8_t guess_key[16];

            memcpy(guess_key, round_key[0], 16);

            int possible_key1[4][256] = {{-1,-1,-1,-1}, {-1,-1,-1,-1}, {-1,-1,-1,-1}, {-1,-1,-1,-1}};
            int possible_key2[4][256] = {{-1,-1,-1,-1}, {-1,-1,-1,-1}, {-1,-1,-1,-1}, {-1,-1,-1,-1}};

            unsigned char p1[16];
            unsigned char p2[16];
            unsigned char p3[16];

            strncpy(p1, plaintext1, 16);
            strncpy(p2, plaintext2, 16);
            strncpy(p3, plaintext3, 16);

            uint8_t diff_p1_p2[16];
            uint8_t diff_p2_p3[16];
            uint8_t diff_p1_p3[16];

            // Plaintext 1 to Round 2 before SubBytes
            AddRoundKey(0, (state_t*)p1, guess_key);
            SubBytes((state_t*)p1);
            ShiftRows((state_t*)p1);
            MixColumns((state_t*)p1);

            // Plaintext 2 to Round 2 before SubBytes
            AddRoundKey(0, (state_t*)p2, guess_key);
            SubBytes((state_t*)p2);
            ShiftRows((state_t*)p2);
            MixColumns((state_t*)p2);

            // Plaintext 3 to Round 2 before SubBytes
            AddRoundKey(0, (state_t*)p3, guess_key);
            SubBytes((state_t*)p3);
            ShiftRows((state_t*)p3);
            MixColumns((state_t*)p3);

            // Calculate differential of P1 and P2
            xor(diff_p1_p2, p1, p2, 16);

            // Calculate differential of P1 and P3
            xor(diff_p1_p3, p1, p3, 16);

            int is_correct = 1;
            for(int bi=0; bi<4;bi++){
                int current_index1 = 0;
                int current_index2 = 0;
                for(unsigned int m=0; m<256;m++){
                    for(unsigned int n=0; n<256;n++){
                        if(m == n){
                            continue;
                        }

                        unsigned int input_diff = m ^ n;
                                    
                        if(input_diff == diff_p1_p2[8+bi]){
                            unsigned int sbox_out_diff = SBOX[m] ^ SBOX[n];

                            if(sbox_out_diff == diff_c1_c2[8+bi]){
                                possible_key1[bi][current_index1] = m ^ p1[8+bi];
                                current_index1++;
                            }
                        }

                        if(input_diff == diff_p1_p3[8+bi]){
                            unsigned int sbox_out_diff = SBOX[m] ^ SBOX[n];

                            if(sbox_out_diff == diff_c1_c3[8+bi]){
                                possible_key2[bi][current_index2] = m ^ p3[8+bi];
                                current_index2++;
                            }
                        }
                                    
                    }
                }

                if(current_index1 == 0 || current_index2 == 0){
                    is_correct = 0;
                    break;
                }
            }

            if(is_correct==0){
                continue;
            }

            int correct_rk1_col2[4];
            int count = 0;
            for(int bi=0; bi<4;bi++){
                for(int x=0; x<4;x++){
                    if(possible_key1[bi][x] < 0){
                        break;
                    }
                    for(int y=0; y<4;y++){
                        if(possible_key1[bi][x] < 256 && possible_key1[bi][x] >= 0 && 
                            possible_key1[bi][x] == possible_key2[bi][y]){
                            correct_rk1_col2[bi] = possible_key1[bi][x];
                            count++;
                        }
                    }
                }
            }

            if(count>=4){
                printf("FOUND ROUND KEY -> [%d, %d, %d, %d]\n", correct_rk1_col2[0], correct_rk1_col2[1], correct_rk1_col2[2], correct_rk1_col2[3]);

                round_key[1][8]  = correct_rk1_col2[0];
                round_key[1][9]  = correct_rk1_col2[1];
                round_key[1][10] = correct_rk1_col2[2];
                round_key[1][11] = correct_rk1_col2[3];

                // Get value of Rk1[7], Rk1[13], and Rk1[15] from Key Schedule relation
                round_key[1][7] = round_key[0][7] ^ round_key[1][3];
                round_key[1][13] = round_key[0][13] ^ round_key[1][9];
                round_key[1][15] = round_key[0][15] ^ round_key[1][11];

                // Get value of Rk1[4] and Rk1[6] from key schedule relation
                round_key[1][4] = round_key[1][8] ^ round_key[0][8];
                round_key[1][6] = round_key[1][10] ^ round_key[0][10];

                // Get value of Rk2[0] and Rk2[2] using Key schedule relation
                round_key[2][0] = SBOX[round_key[1][13]] ^ round_key[1][0] ^ RCON[2];
                round_key[2][2] = SBOX[round_key[1][15]] ^ round_key[1][2];

                // Get value of InvMixColumns(Rk2[0]) and InvMixColumns(Rk2[2])
                // by comparing p1 after ShiftRows round 2 with InvMixColumns(c1)
                unsigned char c1[16];
                uint8_t u2[16]; // InvMixColumns(Rk2)
                strncpy(c1, ciphertext1, 16);
                InvMixColumns((state_t*)c1);

                AddRoundKey(0, (state_t*)p1, round_key[1]);
                SubBytes((state_t*)p1);
                ShiftRows((state_t*)p1);

                xor(u2, p1, c1, 16);

                // Get value of Rk2[1] and Rk2[3] using relation of InvMixColumns(Rk2) and Rk2
                // Basically we solve following two linear equation:
                // Rk2[0] = 2*u2[0] ^ 3*u2[1] ^ u2[2] ^ u2[3]
                // Rk2[2] = u2[0] ^ u2[1] ^ 2*u2[2] ^ 3*u2[3]
                // where we know value of Rk2[0], Rk[2], u2[0], and u2[2]
                // note: u2 = InvMixColumns(Rk2)
                for(unsigned int x=0; x < 256; x++){
                    unsigned int x3 = gmul_by_3[x];

                    unsigned int y = gmul_by_2[u2[0]] ^ u2[2] ^ round_key[2][0] ^ x3;
                    unsigned int check = u2[0] ^ x ^ gmul_by_2[u2[2]] ^ round_key[2][2];
                    unsigned int y3 = gmul_by_3[y];

                    if(check == y3){
                        round_key[2][1] = u2[0] ^ gmul_by_2[x] ^ gmul_by_3[u2[2]] ^ y;
                        round_key[2][3] = gmul_by_3[u2[0]] ^ x ^ u2[2] ^ gmul_by_2[y];
                        break;
                    }
                }

                // Get Rk1[12] and Rk1[14] from Key schedule relation
                round_key[1][12] = INV_SBOX[round_key[1][3] ^ round_key[2][3]];
                round_key[1][14] = INV_SBOX[round_key[1][1] ^ round_key[2][1]];
                
                // We've got all bytes of Rk1 now!!!
                // Write Round Key 1 to data/output.bin and let python handle the rest
                writeFile((char *)round_key[1], "data/output.bin");
                printf("ROUND KEY 1: [");
                for(int bi=0;bi<15;bi++){
                    printf("%d,", round_key[1][bi]);
                }
                printf("%d]\n", round_key[1][15]);

                return 1;
            }
        }
    }
    return 0;
}

/**
 * Attack Phase 1
 * 
 * Iterate over Rk0 of bytes (0,5,10,15) and recover Rk1 of bytes (0,1,2,3)
 * Use Multi-threading for faster execution
 */
void *attack(void *vargp){

    uint8_t round_key[3][16];
    
    uint8_t diff_c1_c2[16];
    uint8_t diff_c1_c3[16];
    uint8_t diff_c2_c3[16];

    int start = *((int *) vargp);
    int end = start + (256 / NUM_THREAD);

    xor(diff_c1_c2, ciphertext1, ciphertext2, 16);
    xor(diff_c1_c3, ciphertext1, ciphertext3, 16);
    xor(diff_c2_c3, ciphertext2, ciphertext3, 16);

    InvMixColumns((state_t*)diff_c1_c2);
    InvShiftRows((state_t*)diff_c1_c2);

    InvMixColumns((state_t*)diff_c1_c3);
    InvShiftRows((state_t*)diff_c1_c3);

    InvMixColumns((state_t*)diff_c2_c3);
    InvShiftRows((state_t*)diff_c2_c3);

    // Iterate over Rk0 of index byte (0,5,10,15)
    for(int i=start; i<end; i++){
        for(int j=0; j<256; j++){
            printf("Enumerate Rk0 (%d,%d,x,x)...\n", i, j);
            for(int k=0; k<256; k++){
                for(int l=0; l<256; l++){

                    uint8_t guess_key[16] = {
                        i,0,0,0,
                        0,j,0,0,
                        0,0,k,0,
                        0,0,0,l
                    };

                    int possible_key1[4][256] = {{-1,-1,-1,-1}, {-1,-1,-1,-1}, {-1,-1,-1,-1}, {-1,-1,-1,-1}};
                    int possible_key2[4][256] = {{-1,-1,-1,-1}, {-1,-1,-1,-1}, {-1,-1,-1,-1}, {-1,-1,-1,-1}};

                    unsigned char p1[16];
                    unsigned char p2[16];
                    unsigned char p3[16];

                    strncpy(p1, plaintext1, 16);
                    strncpy(p2, plaintext2, 16);
                    strncpy(p3, plaintext3, 16);

                    uint8_t diff_p1_p2[16];
                    uint8_t diff_p2_p3[16];
                    uint8_t diff_p1_p3[16];

                    // Plaintext 1 to Round 2 before SubBytes
                    AddRoundKey(0, (state_t*)p1, guess_key);
                    SubBytes((state_t*)p1);
                    ShiftRows((state_t*)p1);
                    MixColumns((state_t*)p1);

                    // Plaintext 2 to Round 2 before SubBytes
                    AddRoundKey(0, (state_t*)p2, guess_key);
                    SubBytes((state_t*)p2);
                    ShiftRows((state_t*)p2);
                    MixColumns((state_t*)p2);

                    // Plaintext 3 to Round 2 before SubBytes
                    AddRoundKey(0, (state_t*)p3, guess_key);
                    SubBytes((state_t*)p3);
                    ShiftRows((state_t*)p3);
                    MixColumns((state_t*)p3);

                    // Calculate differential of P1 and P2
                    xor(diff_p1_p2, p1, p2, 16);

                    // Calculate differential of P1 and P3
                    xor(diff_p1_p3, p1, p3, 16);

                    // Calculate differential of P2 and P3
                    //xor(diff_p2_p3, p2, p3, 16);
                        
                    // Check value on Rk1_col(0) by comparing input and output of SBOX
                    int is_correct = 1;
                    for(int bi=0; bi<4;bi++){
                        int current_index1 = 0;
                        int current_index2 = 0;
                        for(unsigned int m=0; m<256;m++){
                            for(unsigned int n=0; n<256;n++){
                                if(m == n){
                                    continue;
                                }

                                unsigned int input_diff = m ^ n;
                                    
                                if(input_diff == diff_p1_p2[bi]){
                                    unsigned int sbox_out_diff = SBOX[m] ^ SBOX[n];

                                    if(sbox_out_diff == diff_c1_c2[bi]){
                                        possible_key1[bi][current_index1] = m ^ p1[bi];
                                        current_index1++;
                                    }
                                }

                                if(input_diff == diff_p1_p3[bi]){
                                    unsigned int sbox_out_diff = SBOX[m] ^ SBOX[n];

                                    if(sbox_out_diff == diff_c1_c3[bi]){
                                        possible_key2[bi][current_index2] = m ^ p3[bi];
                                        current_index2++;
                                    }
                                }
                                    
                            }
                        }

                        if(current_index1 == 0 || current_index2 == 0){
                            is_correct = 0;
                            break;
                        }
                    }

                    if(is_correct==0){
                        continue;
                    }

                    // Check duplicate possible key
                    // If there's duplicate on each bi of the key, then we get the correct key
                    int correct_rk1_col0[4] = {-1,-1,-1,-1};
                    int count = 0;
                    for(int bi=0; bi<4;bi++){
                        for(int x=0; x<4;x++){
                            if(possible_key1[bi][x] < 0){
                                break;
                            }
                            for(int y=0; y<4;y++){
                                if(possible_key1[bi][x] < 256 && possible_key1[bi][x] >= 0 && 
                                    possible_key1[bi][x] == possible_key2[bi][y]){
                                    correct_rk1_col0[bi] = possible_key1[bi][x];
                                    count++;
                                }
                            }
                        }
                    }

                    // We found correct key!!! 
                    // There's still possibility of false positive tho with prob. 2**-32
                    // Until here we know actual value of Rk0(0,5,10,15) and Rk1(0,1,2,3)
                    if(correct_rk1_col0[0]>=0 && correct_rk1_col0[1]>=0 && correct_rk1_col0[2]>=0 && correct_rk1_col0[3]>=0){
                        if(count != 4){
                            printf("There's multiple possible solution");
                        }
                        printf("FOUND ROUND KEY -> [%d,%d,%d,%d] -> ", i,j,k,l);
                        printf("[%d,%d,%d,%d]\n", correct_rk1_col0[0],correct_rk1_col0[1],correct_rk1_col0[2],correct_rk1_col0[3]);

                        round_key[0][0] = i;
                        round_key[0][5] = j;
                        round_key[0][10] = k;
                        round_key[0][15] = l;

                        round_key[1][0] = correct_rk1_col0[0];
                        round_key[1][1] = correct_rk1_col0[1];
                        round_key[1][2] = correct_rk1_col0[2];
                        round_key[1][3] = correct_rk1_col0[3];

                        // Get value of Rk0[2] and Rk0[13] from key schedule relation
                        round_key[0][2] = round_key[1][2] ^ SBOX[round_key[0][15]];
                        round_key[0][13] = INV_SBOX[round_key[0][0] ^ round_key[1][0] ^ RCON[1]];

                        // Get value of Rk1[5] where Rk1[5] = Rk0[5] ^ Rk1[2]
                        round_key[1][5] = round_key[0][5] ^ round_key[1][1];
        
                        // Start Attack Phase 2                       
                        if(!attack_phase2(round_key, diff_c1_c2, diff_c1_c3)){
                            continue;
                        }
                        free(vargp);
                        exit(0); // Too lazy to stop other thread lol
                    }
                        
                }
            }
        }
    }
    free(vargp);
}

int main(int argc, char *argv[]){

    if(argc == 7){
        plaintext1 = malloc(18);
        plaintext2 = malloc(18);
        plaintext3 = malloc(18);
        ciphertext1 = malloc(18);
        ciphertext2 = malloc(18);
        ciphertext3 = malloc(18);

        readFile(plaintext1, argv[1]);
        readFile(plaintext2, argv[2]);
        readFile(plaintext3, argv[3]);
        readFile(ciphertext1, argv[4]);
        readFile(ciphertext2, argv[5]);
        readFile(ciphertext3, argv[6]);

        
        pthread_t t_id;

        for(int i=0; i<256;i+=256/NUM_THREAD){
            int *arg = malloc(sizeof(*arg));
            if ( arg == NULL ) {
                fprintf(stderr, "Couldn't allocate memory for thread arg.\n");
                exit(EXIT_FAILURE);
            }

            *arg = i;
            pthread_create(&t_id, NULL, attack, arg);
        }

        pthread_exit(NULL);
        free(plaintext1);
        free(plaintext2);
        free(plaintext3);
        free(ciphertext1);
        free(ciphertext2);
        free(ciphertext3);
        
    } else {
        printf("%s <p1> <p2> <p3> <c1> <c2> <c3>\n", argv[0]);
    }
    
    
    return 0;
}