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

#define ECB 1

#include "aes.h"


int isPrime(unsigned int number) {
    if (number <= 1) return 0;
    unsigned int i;

    for (i=2; i*i<=number; i++) {
        if (number % i == 0) return 0;
    }

    return 1;
}


int getPrimeIndex(int n) {
    int primeCounter = 0;
    if (isPrime(n)) {
        for (int i=1; i<=n; i++) {
            primeCounter += isPrime(i);
        }
        return primeCounter;
    } else {
        return 0;
    }
}

int main(void)
{
    int number;
    char str[7];

    char key1[7];
    char key2[7] = "r|xzzr";
    char key3[17];
    char enc[17] = "\xce\xe7\xcd\x07\x9a\xf8\x40\x2f\xb1\xe0\xeb\xe0\x3c\x5b\x77\xbc";

    printf("Enter key: ");

    scanf("%d-%6s", &number, &str);

    int ind;
    if ((ind = getPrimeIndex(number)) != 33) {
        printf("Bad key\r\n");
        return -1;
    }

    sprintf(key1, "%d%d", number, number);

    //printf("key1: %s\r\n", key1);

    for (int i=0; i<6; i++) {
        key2[i] ^= key1[i];
    }

    // printf("key2: %s\r\n", key2);
    
    if (strcmp(key2, str) == 0) {
        printf("Decoding...\r\n");
        sprintf(key3, "%s-%s-OK", key1, key2);
        //printf("key3: %s\r\n", key3);

        struct AES_ctx ctx;
    
        AES_init_ctx(&ctx, key3);
        AES_ECB_decrypt(&ctx, enc);
        printf("Decoded: IGCTF{%s}\r\n", enc);

    } else {
        printf("Bad key\r\n");
        return -2;
    }

    return 0;
}