/*
 * (c) 2015-2017 Marcos Del Sol Vives
 * (c) 2016      javiMaD
 *
 * SPDX-License-Identifier: MIT
 */

#include "amiibo.h"
#include "md.h"
#include "aes.h"
#include "commonutil.h"
#include "../src/fileutils.h"

#define HMAC_POS_DATA 0x008
#define HMAC_POS_TAG 0x1B4
#define AMIBOO_KEY_FN "key_retail.bin"

static void nfc3d_amiibo_calc_seed(const uint8_t *dump, uint8_t *key) {
    memcpy(key + 0x00, dump + 0x029, 0x02);
    memset(key + 0x02, 0x00, 0x0E);
    memcpy(key + 0x10, dump + 0x1D4, 0x08);
    memcpy(key + 0x18, dump + 0x1D4, 0x08);
    memcpy(key + 0x20, dump + 0x1E8, 0x20);
}

static void nfc3d_amiibo_keygen(const nfc3d_keygen_masterkeys_t *masterKeys, const uint8_t *dump, nfc3d_keygen_derivedkeys_t *derivedKeys) {
    uint8_t seed[NFC3D_KEYGEN_SEED_SIZE] = {0};
    nfc3d_amiibo_calc_seed(dump, seed);
    nfc3d_keygen(masterKeys, seed, derivedKeys);
}

static void nfc3d_amiibo_cipher(const nfc3d_keygen_derivedkeys_t *keys, const uint8_t *in, uint8_t *out) {
    mbedtls_aes_context aes;
    size_t nc_off = 0;
    unsigned char nonce_counter[16];
    unsigned char stream_block[16];

    mbedtls_aes_setkey_enc(&aes, keys->aesKey, 128);
    memset(nonce_counter, 0, sizeof(nonce_counter));
    memset(stream_block, 0, sizeof(stream_block));
    memcpy(nonce_counter, keys->aesIV, sizeof(nonce_counter));
    mbedtls_aes_crypt_ctr(&aes, 0x188, &nc_off, nonce_counter, stream_block, in + 0x02C, out + 0x02C);

    memcpy(out + 0x000, in + 0x000, 0x008);
    // Data signature NOT copied
    memcpy(out + 0x028, in + 0x028, 0x004);
    // Tag signature NOT copied
    memcpy(out + 0x1D4, in + 0x1D4, 0x034);
}

static void nfc3d_amiibo_tag_to_internal(const uint8_t *tag, uint8_t *intl) {
    memcpy(intl + 0x000, tag + 0x008, 0x008);
    memcpy(intl + 0x008, tag + 0x080, 0x020);
    memcpy(intl + 0x028, tag + 0x010, 0x024);
    memcpy(intl + 0x04C, tag + 0x0A0, 0x168);
    memcpy(intl + 0x1B4, tag + 0x034, 0x020);
    memcpy(intl + 0x1D4, tag + 0x000, 0x008);
    memcpy(intl + 0x1DC, tag + 0x054, 0x02C);
}

static void nfc3d_amiibo_internal_to_tag(const uint8_t *intl, uint8_t *tag) {
    memcpy(tag + 0x008, intl + 0x000, 0x008);
    memcpy(tag + 0x080, intl + 0x008, 0x020);
    memcpy(tag + 0x010, intl + 0x028, 0x024);
    memcpy(tag + 0x0A0, intl + 0x04C, 0x168);
    memcpy(tag + 0x034, intl + 0x1B4, 0x020);
    memcpy(tag + 0x000, intl + 0x1D4, 0x008);
    memcpy(tag + 0x054, intl + 0x1DC, 0x02C);
}

bool nfc3d_amiibo_unpack(const nfc3d_amiibo_keys_t *amiiboKeys, const uint8_t *tag, uint8_t *plain) {

    uint8_t internal[NFC3D_AMIIBO_SIZE] = {0};

    nfc3d_keygen_derivedkeys_t dataKeys;
    nfc3d_keygen_derivedkeys_t tagKeys;

    // Convert format
    nfc3d_amiibo_tag_to_internal(tag, internal);

    // Generate keys
    nfc3d_amiibo_keygen(&amiiboKeys->data, internal, &dataKeys);
    nfc3d_amiibo_keygen(&amiiboKeys->tag, internal, &tagKeys);

    // Decrypt
    nfc3d_amiibo_cipher(&dataKeys, internal, plain);

    // Regenerate tag HMAC. Note: order matters, data HMAC depends on tag HMAC!
    mbedtls_md_hmac(mbedtls_md_info_from_type(MBEDTLS_MD_SHA256)
                    , tagKeys.hmacKey
                    , sizeof(tagKeys.hmacKey)
                    , plain + 0x1D4
                    , 0x34
                    , plain + HMAC_POS_TAG
                   );

    // Regenerate data HMAC
    mbedtls_md_hmac(mbedtls_md_info_from_type(MBEDTLS_MD_SHA256)
                    , dataKeys.hmacKey
                    , sizeof(dataKeys.hmacKey)
                    , plain + 0x029
                    , 0x1DF
                    , plain + HMAC_POS_DATA
                   );

    return ((memcmp(plain + HMAC_POS_DATA, internal + HMAC_POS_DATA, 32) == 0) &&
            (memcmp(plain + HMAC_POS_TAG, internal + HMAC_POS_TAG, 32) == 0));
}

void nfc3d_amiibo_pack(const nfc3d_amiibo_keys_t *amiiboKeys, const uint8_t *plain, uint8_t *tag) {
    uint8_t cipher[NFC3D_AMIIBO_SIZE] = {0};
    nfc3d_keygen_derivedkeys_t tagKeys;
    nfc3d_keygen_derivedkeys_t dataKeys;

    // Generate keys
    nfc3d_amiibo_keygen(&amiiboKeys->tag, plain, &tagKeys);
    nfc3d_amiibo_keygen(&amiiboKeys->data, plain, &dataKeys);

    // Generate tag HMAC
    mbedtls_md_hmac(mbedtls_md_info_from_type(MBEDTLS_MD_SHA256)
                    , tagKeys.hmacKey
                    , sizeof(tagKeys.hmacKey)
                    , plain + 0x1D4
                    , 0x34
                    , cipher + HMAC_POS_TAG
                   );

    // Init mbedtls HMAC context
    mbedtls_md_context_t ctx;
    mbedtls_md_init(&ctx);
    mbedtls_md_setup(&ctx, mbedtls_md_info_from_type(MBEDTLS_MD_SHA256), 1);

    // Generate data HMAC
    mbedtls_md_hmac_starts(&ctx, dataKeys.hmacKey, sizeof(dataKeys.hmacKey));
    mbedtls_md_hmac_update(&ctx, plain + 0x029, 0x18B);   // Data
    mbedtls_md_hmac_update(&ctx, cipher + HMAC_POS_TAG, 0x20);   // Tag HMAC
    mbedtls_md_hmac_update(&ctx, plain + 0x1D4, 0x34);   // Here be dragons

    mbedtls_md_hmac_finish(&ctx, cipher + HMAC_POS_DATA);

    // HMAC cleanup
    mbedtls_md_free(&ctx);

    // Encrypt
    nfc3d_amiibo_cipher(&dataKeys, plain, cipher);

    // Convert back to hardware
    nfc3d_amiibo_internal_to_tag(cipher, tag);
}

bool nfc3d_amiibo_load_keys(nfc3d_amiibo_keys_t *amiiboKeys) {

    uint8_t *dump = NULL;
    size_t bytes_read = 0;
    if (loadFile_safe(AMIBOO_KEY_FN, "", (void **)&dump, &bytes_read) != PM3_SUCCESS) {
        return false;
    }

    if (bytes_read != sizeof(*amiiboKeys)) {
        free(dump);
        return false;
    }

    memcpy(amiiboKeys, dump, bytes_read);
    free(dump);

    if ((amiiboKeys->data.magicBytesSize > 16) || (amiiboKeys->tag.magicBytesSize > 16)) {
        return false;
    }

    return true;
}

void nfc3d_amiibo_copy_app_data(const uint8_t *src, uint8_t *dst) {


    //uint16_t *ami_nb_wr = (uint16_t*)(dst + 0x29);
    //uint16_t *cfg_nb_wr = (uint16_t*)(dst + 0xB4);

    /* increment write counters */
    //*ami_nb_wr = htobe16(be16toh(*ami_nb_wr) + 1);
    //*cfg_nb_wr = htobe16(be16toh(*cfg_nb_wr) + 1);

    uint16_t ami_nb_wr = ((uint16_t)bytes_to_num(dst + 0x29, 2)) + 1;
    uint16_t cfg_nb_wr = ((uint16_t)bytes_to_num(dst + 0xB4, 2)) + 1;

    num_to_bytes(ami_nb_wr, 2, dst + 0x29);
    num_to_bytes(cfg_nb_wr, 2, dst + 0xB4);

    /* copy flags */
    dst[0x2C] = src[0x2C];
    /* copy programID */
    memcpy(dst + 0xAC, src + 0xAC, 8);
    /* copy AppID */
    memcpy(dst + 0xB6, src + 0xB6, 4);
    /* copy AppData */
    memcpy(dst + 0xDC, src + 0xDC, 216);
}