/*
 * Copyright (c) 2021 scandal_uk
 *
 * This program is free software; you can redistribute it and/or modify it
 * under the terms and conditions of the GNU General Public License,
 * version 2, as published by the Free Software Foundation.
 *
 * This program is distributed in the hope 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 <http://www.gnu.org/licenses/>.
 */

#include "incognito.h"

#include "../config.h"
#include <gfx/di.h>
#include <gfx_utils.h>
#include "../gfx/tui.h"
#include "../hos/pkg1.h"
#include "../hos/pkg2.h"
#include "../hos/sept.h"
#include <libs/fatfs/ff.h>
#include <mem/heap.h>
#include <mem/mc.h>
#include <mem/sdram.h>
#include <sec/se.h>
#include <sec/se_t210.h>
#include <sec/tsec.h>
#include <soc/fuse.h>
#include <mem/smmu.h>
#include <soc/t210.h>
#include "../storage/emummc.h"
#include "../storage/nx_emmc.h"
#include "../storage/nx_emmc_bis.h"
#include <storage/sdmmc.h>
#include <utils/btn.h>
#include <utils/list.h>
#include <utils/sprintf.h>
#include <utils/util.h>

#include "key_sources.inl"

#include "io/io.h"
#include <string.h>

#define RETRY_COUNT 5
#define RETRY(exp)                                                                              \
    ({                                                                                          \
        u8 _attemptc_ = RETRY_COUNT;                                                            \
        bool _resultb_ = false;                                                                 \
        while (_attemptc_--)                                                                    \
        {                                                                                       \
            if ((_resultb_ = exp))                                                              \
                break;                                                                          \
            gfx_printf("%kretry %d/%d...\n", COLOR_RED, RETRY_COUNT - _attemptc_, RETRY_COUNT); \
        }                                                                                       \
        _resultb_;                                                                              \
    })

extern bool sd_mount();
extern void sd_unmount();
extern int sd_save_to_file(void *buf, u32 size, const char *filename);

extern hekate_config h_cfg;

sdmmc_storage_t storage;
sdmmc_t sdmmc;
emmc_part_t *system_part;
emmc_part_t *prodinfo_part;

#define SECTORS_IN_CLUSTER 32
#define PRODINFO_SIZE 0x3FBC00

#define BACKUP_NAME_EMUNAND "sd:/prodinfo_emunand.bin"
#define BACKUP_NAME_SYSNAND "sd:/prodinfo_sysnand.bin"

static u8 temp_key[0x10],
    bis_key[4][0x20] = {0},
    device_key[0x10] = {0},
    device_key_4x[0x10] = {0},
    // FS-related keys
    header_key[0x20] = {0},
    // keyblob-derived families
    keyblob_key[KB_FIRMWARE_VERSION_600 + 1][0x10] = {0},
    keyblob_mac_key[KB_FIRMWARE_VERSION_600 + 1][0x10] = {0},
    package1_key[KB_FIRMWARE_VERSION_600 + 1][0x10] = {0},
    // master key-derived families
    master_kek[KB_FIRMWARE_VERSION_MAX + 1][0x10] = {0},
    master_key[KB_FIRMWARE_VERSION_MAX + 1][0x10] = {0};

keyblob_t keyblob[KB_FIRMWARE_VERSION_600 + 1] = {0};

LIST_INIT(gpt);

// key functions
static int   _key_exists(const void *data) { return memcmp(data, zeros, 0x10) != 0; };
static void  _generate_kek(u32 ks, const void *key_source, void *master_key, const void *kek_seed, const void *key_seed);
static void  _get_device_key(u32 ks, void *out_device_key, u32 revision, const void *device_key, const void *master_key);

static ALWAYS_INLINE u8 *_find_tsec_fw(const u8 *pkg1) {
    const u32 tsec_fw_align = 0x100;
    const u32 tsec_fw_first_instruction = 0xCF42004D;

    for (const u32 *pos = (const u32 *)pkg1; (u8 *)pos < pkg1 + PKG1_MAX_SIZE; pos += tsec_fw_align / sizeof(u32))
        if (*pos == tsec_fw_first_instruction)
            return (u8 *)pos;

    return NULL;
}

static ALWAYS_INLINE u32 _get_tsec_fw_size(tsec_key_data_t *key_data) {
    return key_data->blob0_size + sizeof(tsec_key_data_t) + key_data->blob1_size + key_data->blob2_size + key_data->blob3_size + key_data->blob4_size;
}

unsigned int crc_16_table[16] = {
    0x0000, 0xCC01, 0xD801, 0x1400, 0xF001, 0x3C00, 0x2800, 0xE401,
    0xA001, 0x6C00, 0x7800, 0xB401, 0x5000, 0x9C01, 0x8801, 0x4400 };

unsigned short int get_crc_16 (const char *p, int n) {
    unsigned short int crc = 0x55AA;
    int r;

    while (n-- > 0) {
        r = crc_16_table[crc & 0xF];
        crc = (crc >> 4) & 0x0FFF;
        crc = crc ^ r ^ crc_16_table[*p & 0xF];

        r = crc_16_table[crc & 0xF];
        crc = (crc >> 4) & 0x0FFF;
        crc = crc ^ r ^ crc_16_table[(*p >> 4) & 0xF];

        p++;
    }

    return(crc);
}

u16 calculateCrc(u32 offset, u32 size, u8 *blob)
{
    const char buffer[size + 1];
    if (blob == NULL)
        readData((u8 *)buffer, offset, size, NULL);
    else
        memcpy((u8 *)buffer, blob + offset, size);

    return get_crc_16(buffer, size);
}

u16 readCrc(u32 offset, u8 *blob)
{
    u16 buffer;
    if (blob == NULL)
        readData((u8 *)&buffer, offset, sizeof(u16), NULL);
    else
        memcpy((u8 *)&buffer, blob + offset, sizeof(u16));
    
    return buffer;
}

bool validateCrc(u32 offset, u32 size, u8 *blob)
{
    return calculateCrc(offset, size, blob) == readCrc(offset + size, blob);
}

bool calculateAndWriteCrc(u32 offset, u32 size)
{
    u16 crcValue = calculateCrc(offset, size, NULL);
    u8 crc[2] = { crcValue & 0xff, crcValue >> 8 }; // bytes of u16
    return writeData(crc, offset + size, sizeof(u16), NULL);
}

void validateChecksums(u8 *blob)
{
    if (!validateCrc(0x0250, 0x1E, blob))
        gfx_printf("%kWarning - invalid serial crc\n", COLOR_RED);

    if (!validateCrc(0x0480, 0x18E, blob))
        gfx_printf("%kWarning - invalid ECC-B233 crc...\n", COLOR_RED);

    if (!validateCrc(0x3AE0, 0x13E, blob))
        gfx_printf("%kWarning - invalid ext SSL key crc...\n", COLOR_RED);

    if (!validateCrc(0x35A0, 0x07E, blob))
        gfx_printf("%kWarning - invalid ECDSA cert crc...\n", COLOR_RED);

    if (!validateCrc(0x36A0, 0x09E, blob))
        gfx_printf("%kWarning - invalid ECQV-BLS cert crc...\n", COLOR_RED);
}

bool dump_keys()
{
    display_backlight_brightness(100, 1000);
    gfx_clear_partial_grey(0x1B, 0, 1256);
    gfx_con_setpos(0, 0);

    gfx_print_header();

    gfx_printf("%kGetting bis_keys...\n", COLOR_YELLOW);

    u32 retries = 0;

    tsec_ctxt_t tsec_ctxt;

    if (emummc_storage_init_mmc(&storage, &sdmmc) == 2)
    {
        EPRINTF("Unable to init MMC.");
        return false;
    }

    // Read package1.
    u8 *pkg1 = (u8 *)malloc(PKG1_MAX_SIZE);
    if (!emummc_storage_set_mmc_partition(&storage, EMMC_BOOT0)) {
        EPRINTF("Unable to set partition.");
        return false;
    }
    if (!emummc_storage_read(&emmc_storage, PKG1_OFFSET / NX_EMMC_BLOCKSIZE, PKG1_MAX_SIZE / NX_EMMC_BLOCKSIZE, pkg1)) {
        EPRINTF("Unable to read pkg1.");
        return false;
    }
    const pkg1_id_t *pkg1_id = pkg1_identify(pkg1);
    if (!pkg1_id)
    {
        EPRINTF("Unknown pkg1 version.\n Make sure you have the latest Incognito_RCM.\n If a new firmware version just came out,\n Incognito_RCM must be updated.\n Check Github for new release.");
        free(pkg1);
        return false;
    }

    tsec_ctxt.fw = _find_tsec_fw(pkg1);
    if (!tsec_ctxt.fw)
    {
        EPRINTF("Failed to locate TSEC firmware.");
        free(pkg1);
        return false;
    }

    tsec_ctxt.pkg1 = pkg1;
    tsec_ctxt.size = _get_tsec_fw_size((tsec_key_data_t *)(tsec_ctxt.fw + TSEC_KEY_DATA_OFFSET));
    if (tsec_ctxt.size > PKG1_MAX_SIZE) {
        EPRINTF("Unexpected TSEC firmware size.");
        return false;
    }

    if (pkg1_id->kb >= KB_FIRMWARE_VERSION_700)
    {
        se_aes_key_read(se_key_acc_ctrl_get(12) == 0x6A ? 13 : 12, master_key[KB_FIRMWARE_VERSION_MAX], 0x10);
    }

    //get_tsec: ;
    u8 tsec_keys[0x10 * 2] = {0};

    if (pkg1_id->kb == KB_FIRMWARE_VERSION_620)
    {
        u8 *tsec_paged = (u8 *)page_alloc(3);
        memcpy(tsec_paged, (void *)tsec_ctxt.fw, tsec_ctxt.size);
        tsec_ctxt.fw = tsec_paged;
    }

    int res = 0;

    mc_disable_ahb_redirect();

    while (tsec_query(tsec_keys, pkg1_id->kb, &tsec_ctxt) < 0)
    {
        memset(tsec_keys, 0x00, 0x20);
        retries++;
        if (retries > 15)
        {
            res = -1;
            break;
        }
    }
    free(pkg1);

    mc_enable_ahb_redirect();

    if (res < 0)
    {
        EPRINTFARGS("ERROR %x dumping TSEC.\n", res);
        return false;
    }

    // Master key derivation

    // on firmware 6.2.0 only, tsec_query delivers the tsec_root_key
    if (pkg1_id->kb == KB_FIRMWARE_VERSION_620 && _key_exists(tsec_keys + 0x10)) {
        se_aes_key_set(8, tsec_keys + 0x10, 0x10); // mkek6 = unwrap(mkeks6, tsecroot)
        se_aes_crypt_block_ecb(8, 0, master_kek[6], master_kek_sources[0]);
        se_aes_key_set(8, master_kek[6], 0x10); // mkey = unwrap(mkek, mks)
        se_aes_crypt_block_ecb(8, 0, master_key[6], master_key_source);
    }

    u8 *keyblob_block = (u8 *)calloc(KB_FIRMWARE_VERSION_600 + 1, NX_EMMC_BLOCKSIZE);
    encrypted_keyblob_t *current_keyblob = (encrypted_keyblob_t *)keyblob_block;
    u8 keyblob_mac[0x10] = {0};
    u32 sbk[4] = {FUSE(FUSE_PRIVATE_KEY0), FUSE(FUSE_PRIVATE_KEY1),
                  FUSE(FUSE_PRIVATE_KEY2), FUSE(FUSE_PRIVATE_KEY3)};
    se_aes_key_set(8, tsec_keys, sizeof(tsec_keys) / 2);
    se_aes_key_set(9, sbk, sizeof(sbk));

    if (!emummc_storage_read(&storage, KEYBLOB_OFFSET / NX_EMMC_BLOCKSIZE, KB_FIRMWARE_VERSION_600 + 1, keyblob_block)) {
        EPRINTF("Unable to read keyblob.");
    }

    for (u32 i = 0; i <= KB_FIRMWARE_VERSION_600; i++, current_keyblob++)
    {
        se_aes_crypt_block_ecb(8, 0, keyblob_key[i], keyblob_key_source[i]); // temp = unwrap(kbks, tsec)
        se_aes_crypt_block_ecb(9, 0, keyblob_key[i], keyblob_key[i]);        // kbk = unwrap(temp, sbk)
        se_aes_key_set(7, keyblob_key[i], sizeof(keyblob_key[i]));
        se_aes_crypt_block_ecb(7, 0, keyblob_mac_key[i], keyblob_mac_key_source); // kbm = unwrap(kbms, kbk)
        if (i == 0)
        {
            se_aes_crypt_block_ecb(7, 0, device_key, per_console_key_source); // devkey = unwrap(pcks, kbk0)
            se_aes_crypt_block_ecb(7, 0, device_key_4x, device_master_key_source_kek_source);
        }

        // verify keyblob is not corrupt
        se_aes_key_set(10, keyblob_mac_key[i], sizeof(keyblob_mac_key[i]));
        se_aes_cmac(10, keyblob_mac, sizeof(keyblob_mac), current_keyblob->iv, sizeof(current_keyblob->iv) + sizeof(keyblob_t));
        if (memcmp(current_keyblob, keyblob_mac, sizeof(keyblob_mac)) != 0) {
            EPRINTFARGS("Keyblob %x corrupt.", i);
            continue;
        }

        // decrypt keyblobs
        se_aes_key_set(6, keyblob_key[i], sizeof(keyblob_key[i]));
        se_aes_crypt_ctr(6, &keyblob[i], sizeof(keyblob_t), &current_keyblob->key_data, sizeof(keyblob_t), current_keyblob->iv);

        memcpy(package1_key[i], keyblob[i].package1_key, sizeof(package1_key[i]));
        memcpy(master_kek[i], keyblob[i].master_kek, sizeof(master_kek[i]));
        se_aes_key_set(7, master_kek[i], sizeof(master_kek[i]));
        se_aes_crypt_block_ecb(7, 0, master_key[i], master_key_source);
    }
    free(keyblob_block);

    /*  key = unwrap(source, wrapped_key):
        key_set(ks, wrapped_key), block_ecb(ks, 0, key, source) -> final key in key
    */
    u32 key_generation = fuse_read_odm_keygen_rev();
    if (key_generation)
        key_generation--;

    if (_key_exists(device_key))
    {
        if (key_generation)
        {
            _get_device_key(8, temp_key, key_generation, device_key_4x, master_key[0]);
        }
        else
            memcpy(temp_key, device_key, 0x10);
        se_aes_key_set(8, temp_key, 0x10);
        se_aes_unwrap_key(8, 8, retail_specific_aes_key_source);                   // kek = unwrap(rsaks, devkey)
        se_aes_crypt_block_ecb(8, 0, bis_key[0] + 0x00, bis_key_source[0] + 0x00); // bkey = unwrap(bkeys, kek)
        se_aes_crypt_block_ecb(8, 0, bis_key[0] + 0x10, bis_key_source[0] + 0x10);
        // kek = generate_kek(bkeks, devkey, aeskek, aeskey)
        _generate_kek(8, bis_kek_source, temp_key, aes_kek_generation_source, aes_key_generation_source);
        se_aes_crypt_block_ecb(8, 0, bis_key[1] + 0x00, bis_key_source[1] + 0x00); // bkey = unwrap(bkeys, kek)
        se_aes_crypt_block_ecb(8, 0, bis_key[1] + 0x10, bis_key_source[1] + 0x10);
        se_aes_crypt_block_ecb(8, 0, bis_key[2] + 0x00, bis_key_source[2] + 0x00);
        se_aes_crypt_block_ecb(8, 0, bis_key[2] + 0x10, bis_key_source[2] + 0x10);
        memcpy(bis_key[3], bis_key[2], 0x20);
    }

    // Set BIS keys.

    se_aes_key_set(8, bis_key[0] + 0x00, AES_128_KEY_SIZE);
    se_aes_key_set(9, bis_key[0] + 0x10, AES_128_KEY_SIZE);

    if (!emummc_storage_set_mmc_partition(&storage, EMMC_GPP)) {
        EPRINTF("Unable to set partition.");
        return false;
    }
    // Parse eMMC GPT.
    LIST_INIT(gpt);
    nx_emmc_gpt_parse(&gpt, &storage);

    // Find PRODINFO partition.
    prodinfo_part = nx_emmc_part_find(&gpt, "PRODINFO");
    if (!prodinfo_part)
    {
        EPRINTF("Unable to locate PRODINFO partition.");
        return false;
    }

    gfx_printf("%kGot keys!\n%kValidate...", COLOR_GREEN, COLOR_YELLOW);

    const char magic[4] = "CAL0";
    char buffer[4];
    readData((u8 *)buffer, 0, 4, NULL);
    if (memcmp(magic, buffer, 4) == 0)
    {
        gfx_printf("%kOK!\n", COLOR_GREEN);
    }
    else
    {
        gfx_printf("%kError!\n", COLOR_RED);
        return false;
    }

    char serial[15];
    readData((u8 *)serial, 0x250, 14, NULL);

    gfx_printf("%kCurrent serial: [%s]\n\n", COLOR_BLUE, serial);

    return true;
}

bool erase(u32 offset, u32 length)
{
    u8 *tmp = (u8 *)calloc(length, sizeof(u8));
    bool result = writeData(tmp, offset, length, NULL);
    free(tmp);
    return result;
}

bool writeSerial()
{
    const char *junkSerial;
    if (isSysNAND())
    {
        junkSerial = "XAW00000000000";
    }
    else
    {
        junkSerial = "XAW00000000001";
    }

    const u32 serialOffset = 0x250;
    if (!writeData((u8 *)junkSerial, serialOffset, 14, NULL))
        return false;

    return calculateAndWriteCrc(serialOffset, 0x1E);
}

bool incognito()
{
    gfx_printf("%kChecking if backup exists...\n", COLOR_YELLOW);
    if (!checkBackupExists())
    {
        gfx_printf("%kI'm sorry Dave, I'm afraid I can't do that..\n%kWill make a backup first...\n", COLOR_RED, COLOR_YELLOW);
        if (!backupProdinfo())
            return false;
    }

    validateChecksums(NULL);

    gfx_printf("%kWriting fake serial...\n", COLOR_YELLOW);
    if (!writeSerial())
        return false;
/*
    gfx_printf("%kErasing ECC-B233 device cert...\n", COLOR_YELLOW);
    if (!erase(0x0480, 0x180))
        return false;

    if (!calculateAndWriteCrc(0x0480, 0x18E)) // whatever I do here, it crashes Atmos..?
        return false;
*/
    gfx_printf("%kErasing SSL cert...\n", COLOR_YELLOW);
    if (!erase(0x0AE0, 0x800))
        return false;

    gfx_printf("%kErasing extended SSL key...\n", COLOR_YELLOW);
    if (!erase(0x3AE0, 0x130))
        return false;

    gfx_printf("%kWriting checksum...\n", COLOR_YELLOW);
    if (!calculateAndWriteCrc(0x3AE0, 0x13E))
        return false;

    gfx_printf("%kErasing Amiibo ECDSA cert...\n", COLOR_YELLOW);
    if (!erase(0x35A0, 0x070))
        return false;

    gfx_printf("%kWriting checksum...\n", COLOR_YELLOW);
    if (!calculateAndWriteCrc(0x35A0, 0x07E))
        return false;

    gfx_printf("%kErasing Amiibo ECQV-BLS root cert...\n", COLOR_YELLOW);
    if (!erase(0x36A0, 0x090))
        return false;

    gfx_printf("%kWriting checksum...\n", COLOR_YELLOW);
    if (!calculateAndWriteCrc(0x36A0, 0x09E))
        return false;

    gfx_printf("%kErasing RSA-2048 extended device key...\n", COLOR_YELLOW);
    if (!erase(0x3D70, 0x240)) // seems empty & unused!
        return false;

    gfx_printf("%kErasing RSA-2048 device certificate...\n", COLOR_YELLOW);
    if (!erase(0x3FC0, 0x240)) // seems empty & unused!
        return false;

    gfx_printf("%kWriting SSL cert hash...\n", COLOR_YELLOW);
    if (!writeClientCertHash())
        return false;

    gfx_printf("%kWriting body hash...\n", COLOR_YELLOW);
    if (!writeCal0Hash())
        return false;

    gfx_printf("\n%kIncognito done!\n", COLOR_GREEN);
    return true;
}

u32 divideCeil(u32 x, u32 y)
{
    return 1 + ((x - 1) / y);
}

void cleanUp()
{
	h_cfg.emummc_force_disable = emu_cfg.sector == 0 && !emu_cfg.path;
    //nx_emmc_gpt_free(&gpt);
    //emummc_storage_end(&storage);
}

static void _generate_kek(u32 ks, const void *key_source, void *master_key, const void *kek_seed, const void *key_seed) {
    if (!_key_exists(key_source) || !_key_exists(master_key) || !_key_exists(kek_seed))
        return;

    se_aes_key_set(ks, master_key, 0x10);
    se_aes_unwrap_key(ks, ks, kek_seed);
    se_aes_unwrap_key(ks, ks, key_source);
    if (key_seed && _key_exists(key_seed))
        se_aes_unwrap_key(ks, ks, key_seed);
}

static void _get_device_key(u32 ks, void *out_device_key, u32 revision, const void *device_key, const void *master_key) {
    if (revision < KB_FIRMWARE_VERSION_400)
        memcpy(out_device_key, device_key, 0x10);

    revision -= KB_FIRMWARE_VERSION_400;
    u8 temp_key[0x10] = {0};
    se_aes_key_set(ks, device_key, 0x10);
    se_aes_crypt_ecb(ks, 0, temp_key, 0x10, device_master_key_source_sources[revision], 0x10);
    se_aes_key_set(ks, master_key, 0x10);
    se_aes_unwrap_key(ks, ks, device_master_kek_sources[revision]);
    se_aes_crypt_ecb(ks, 0, out_device_key, 0x10, temp_key, 0x10);
}

bool readData(u8 *buffer, u32 offset, u32 length, void (*progress_callback)(u32, u32))
{
    if (progress_callback != NULL)
    {
        (*progress_callback)(0, length);
    }
    bool result = false;
    u32 sector = (offset / NX_EMMC_BLOCKSIZE);
    u32 newOffset = (offset % NX_EMMC_BLOCKSIZE);

    u32 sectorCount = divideCeil(newOffset + length, NX_EMMC_BLOCKSIZE);

    u8 *tmp = (u8 *)malloc(sectorCount * NX_EMMC_BLOCKSIZE);

    u32 clusterOffset = sector % SECTORS_IN_CLUSTER;
    u32 sectorOffset = 0;
    while (clusterOffset + sectorCount > SECTORS_IN_CLUSTER)
    {
        u32 sectorsToRead = SECTORS_IN_CLUSTER - clusterOffset;
        if (!RETRY(prodinfo_read(tmp + (sectorOffset * NX_EMMC_BLOCKSIZE), sector, sectorsToRead)))
            goto out;

        sector += sectorsToRead;
        sectorCount -= sectorsToRead;
        clusterOffset = 0;
        sectorOffset += sectorsToRead;
        if (progress_callback != NULL)
        {
            (*progress_callback)(sectorOffset * NX_EMMC_BLOCKSIZE, length);
        }
    }
    if (sectorCount == 0)
        goto done;

    if (!RETRY(prodinfo_read(tmp + (sectorOffset * NX_EMMC_BLOCKSIZE), sector, sectorCount)))
        goto out;

    memcpy(buffer, tmp + newOffset, length);
done:
    result = true;
    if (progress_callback != NULL)
    {
        (*progress_callback)(length, length);
    }
out:
    free(tmp);
    return result;
}

bool writeData(u8 *buffer, u32 offset, u32 length, void (*progress_callback)(u32, u32))
{
    if (progress_callback != NULL)
    {
        (*progress_callback)(0, length);
    }
    bool result = false;

    u32 initialLength = length;

    u8 *tmp_sec = (u8 *)malloc(NX_EMMC_BLOCKSIZE);
    u8 *tmp = NULL;

    u32 sector = (offset / NX_EMMC_BLOCKSIZE);
    u32 newOffset = (offset % NX_EMMC_BLOCKSIZE);

    // if there is a sector offset, read involved sector, write data to it with offset and write back whole sector to be sector aligned
    if (newOffset > 0)
    {
        u32 bytesToRead = NX_EMMC_BLOCKSIZE - newOffset;
        u32 bytesToWrite;
        if (length >= bytesToRead)
        {
            bytesToWrite = bytesToRead;
        }
        else
        {
            bytesToWrite = length;
        }
        if (!RETRY(prodinfo_read(tmp_sec, sector, 1)))
            goto out;

        memcpy(tmp_sec + newOffset, buffer, bytesToWrite);
        if (!RETRY(prodinfo_write(tmp_sec, sector, 1)))
            goto out;

        sector++;
        length -= bytesToWrite;
        newOffset = bytesToWrite;

        if (progress_callback != NULL)
        {
            (*progress_callback)(initialLength - length, initialLength);
        }
        // are we done?
        if (length == 0)
            goto done;
    }

    // write whole sectors in chunks while being cluster aligned
    u32 sectorCount = length / NX_EMMC_BLOCKSIZE;
    tmp = (u8 *)malloc(sectorCount * NX_EMMC_BLOCKSIZE);

    u32 clusterOffset = sector % SECTORS_IN_CLUSTER;
    u32 sectorOffset = 0;
    while (clusterOffset + sectorCount >= SECTORS_IN_CLUSTER)
    {
        u32 sectorsToRead = SECTORS_IN_CLUSTER - clusterOffset;
        if (!RETRY(prodinfo_write(buffer + newOffset + (sectorOffset * NX_EMMC_BLOCKSIZE), sector, sectorsToRead)))
            goto out;

        sector += sectorsToRead;
        sectorOffset += sectorsToRead;
        sectorCount -= sectorsToRead;
        clusterOffset = 0;
        length -= sectorsToRead * NX_EMMC_BLOCKSIZE;

        if (progress_callback != NULL)
        {
            (*progress_callback)(initialLength - length, initialLength);
        }
    }

    // write remaining sectors
    if (sectorCount > 0)
    {
        if (!RETRY(prodinfo_write(buffer + newOffset + (sectorOffset * NX_EMMC_BLOCKSIZE), sector, sectorCount)))
            goto out;

        length -= sectorCount * NX_EMMC_BLOCKSIZE;
        sector += sectorCount;
        sectorOffset += sectorCount;

        if (progress_callback != NULL)
        {
            (*progress_callback)(initialLength - length, initialLength);
        }
    }

    // if there is data remaining that is smaller than a sector, read that sector, write remaining data to it and write back whole sector
    if (length == 0)
        goto done;

    if (length > NX_EMMC_BLOCKSIZE)
    {
        gfx_printf("%kERROR, ERRO! Length is %d!\n", COLOR_RED, length);
        goto out;
    }

    if (!RETRY(prodinfo_read(tmp_sec, sector, 1)))
        goto out;

    memcpy(tmp_sec, buffer + newOffset + (sectorOffset * NX_EMMC_BLOCKSIZE), length);
    if (!RETRY(prodinfo_write(tmp_sec, sector, 1)))
        goto out;

done:
    result = true;
    if (progress_callback != NULL)
    {
        (*progress_callback)(initialLength, initialLength);
    }
out:
    free(tmp_sec);
    free(tmp);
    return result;
}

bool writeHash(u32 hashOffset, u32 offset, u32 sz)
{
    bool result = false;
    u8 *buffer = (u8 *)malloc(sz);
    if (!readData(buffer, offset, sz, NULL))
    {
        goto out;
    }
    u8 hash[0x20];
    se_calc_sha256(hash, buffer, sz);

    if (!writeData(hash, hashOffset, 0x20, NULL))
    {
        goto out;
    }
    result = true;
out:
    free(buffer);
    return result;
}

#ifdef DEBUG
void screenshot(const char *filename)
{
    sd_mount();

    FIL fp;
    if (f_open(&fp, filename, FA_CREATE_ALWAYS | FA_WRITE) != FR_OK)
    {
        gfx_printf("\n%kCannot write image!\n", COLOR_RED);
        return;
    }
    u32 size;
    u8 *buffer = gfx_bmp_screenshot(&size);

    f_write(&fp, buffer, size, NULL);
    f_close(&fp);
    free(buffer);
}
#endif

bool verifyHash(u32 hashOffset, u32 offset, u32 sz, u8 *blob)
{
    bool result = false;
    u8 *buffer = (u8 *)malloc(sz);
    if (blob == NULL)
    {
        if (!readData(buffer, offset, sz, NULL))
            goto out;
    }
    else
    {
        memcpy(buffer, blob + offset, sz);
    }
    u8 hash1[0x20];
    se_calc_sha256(hash1, buffer, sz);

    u8 hash2[0x20];

    if (blob == NULL)
    {
        if (!readData(hash2, hashOffset, 0x20, NULL))
            goto out;
    }
    else
    {
        memcpy(hash2, blob + hashOffset, 0x20);
    }

    if (memcmp(hash1, hash2, 0x20) != 0)
    {
        EPRINTF("error: hash verification failed\n");
        // gfx_hexdump(0, hash1, 0x20);
        // gfx_hexdump(0, hash2, 0x20);
        goto out;
    }

    result = true;
out:
    free(buffer);
    return result;
}

s32 getClientCertSize(u8 *blob)
{
    s32 buffer;
    if (blob == NULL)
    {
        if (!RETRY(readData((u8 *)&buffer, 0x0AD0, sizeof(buffer), NULL)))
        {
            return -1;
        }
    }
    else
    {
        memcpy(&buffer, blob + 0x0AD0, sizeof(buffer));
    }
    return buffer;
}

s32 getCalibrationDataSize(u8 *blob)
{
    s32 buffer;
    if (blob == NULL)
    {
        if (!RETRY(readData((u8 *)&buffer, 0x08, sizeof(buffer), NULL)))
        {
            return -1;
        }
    }
    else
    {
        memcpy(&buffer, blob + 0x08, sizeof(buffer));
    }

    return buffer;
}

bool writeCal0Hash()
{
    s32 calibrationSize = getCalibrationDataSize(NULL);
    if (calibrationSize == -1)
        return false;

    return writeHash(0x20, 0x40, calibrationSize);
}

bool writeClientCertHash()
{
    s32 certSize = getClientCertSize(NULL);
    if (certSize == -1)
        return false;

    return writeHash(0x12E0, 0xAE0, certSize);
}

bool verifyCal0Hash(u8 *blob)
{
    s32 calibrationSize = getCalibrationDataSize(blob);
    if (calibrationSize == -1)
        return false;

    return verifyHash(0x20, 0x40, calibrationSize, blob);
}

bool verifyClientCertHash(u8 *blob)
{
    s32 certSize = getClientCertSize(blob);
    if (certSize == -1)
        return false;

    return verifyHash(0x12E0, 0xAE0, certSize, blob);
}

bool verifyProdinfo(u8 *blob)
{
    gfx_printf("%kVerifying client cert hash and CAL0 hash%s...\n", COLOR_YELLOW, blob != NULL ? "\nfrom backup" : "");

    if (verifyClientCertHash(blob) && verifyCal0Hash(blob))
    {
        validateChecksums(blob);

        char serial[15] = "";
        if (blob == NULL)
        {
            readData((u8 *)serial, 0x250, 14, NULL);
        }
        else
        {
            memcpy(serial, blob + 0x250, 14);
        }

        gfx_printf("%kVerification successful!\n%kSerial:%s\n", COLOR_GREEN, COLOR_BLUE, serial);
        return true;
    }
    gfx_printf("%kVerification not successful!\n", COLOR_RED);
    return false;
}

void print_progress(u32 count, u32 max)
{
    tui_pbar(0, gfx_con.y + 1, (int)(count * 100 / (float)max), COLOR_BLUE, COLOR_ORANGE);
}

bool isSysNAND()
{
    return (!emu_cfg.enabled || h_cfg.emummc_force_disable);
}

bool checkBackupExists()
{
    char *name;
    if (isSysNAND())
    {
        name = BACKUP_NAME_SYSNAND;
    }
    else
    {
        name = BACKUP_NAME_EMUNAND;
    }
    return f_stat(name, NULL) == FR_OK;
}

bool backupProdinfo()
{
    bool result = false;
    char *name;
    if (isSysNAND())
    {
        name = BACKUP_NAME_SYSNAND;
    }
    else
    {
        name = BACKUP_NAME_EMUNAND;
    }

    gfx_printf("%kBacking up %s...\n", COLOR_YELLOW, name);
    if (checkBackupExists())
    {
        gfx_printf("%kBackup already exists!\nWill rename old backup.\n", COLOR_ORANGE);
        u32 filenameSuffix = 0;
        char newName[255];
        do
        {
            sprintf(newName, "%s.%d", name, filenameSuffix);
            filenameSuffix++;
        } while (f_stat(newName, NULL) == FR_OK);
        f_rename(name, newName);
        gfx_printf("%kOld backup renamed to:\n%s\n", COLOR_YELLOW, newName);
    }

    FIL fp;
    if (f_open(&fp, name, FA_CREATE_ALWAYS | FA_WRITE) != FR_OK)
    {
        gfx_printf("\n%kCannot write to %s!\n", COLOR_RED, name);
        return false;
    }

    u8 *bufferNX = (u8 *)malloc(PRODINFO_SIZE);
    gfx_printf("%kReading from NAND...\n", COLOR_YELLOW);
    if (!readData(bufferNX, 0, PRODINFO_SIZE, print_progress))
    {
        gfx_printf("\n%kError reading from NAND!\n", COLOR_RED);
        goto out;
    }
    gfx_putc('\n');
    if (!verifyProdinfo(bufferNX))
    {
        goto out;
    }
    gfx_printf("%k\nWriting to file...\n", COLOR_YELLOW);
    u32 bytesWritten;
    if (f_write(&fp, bufferNX, PRODINFO_SIZE, &bytesWritten) != FR_OK || bytesWritten != PRODINFO_SIZE)
    {
        gfx_printf("\n%kError writing to file!\nPlease try again. If this doesn't work, you don't have a working backup!\n", COLOR_RED);
        goto out;
    }
    f_sync(&fp);

    result = true;
    gfx_printf("\n%kBackup to %s done!\n", COLOR_GREEN, name);

out:
    f_close(&fp);
    free(bufferNX);

    return result;
}

bool restoreProdinfo()
{
    bool result = false;
    sd_mount();

    const char *name;
    if (isSysNAND())
    {
        name = BACKUP_NAME_SYSNAND;
    }
    else
    {
        name = BACKUP_NAME_EMUNAND;
    }

    gfx_printf("%kRestoring from %s...\n", COLOR_YELLOW, name);

    FIL fp;
    if (f_open(&fp, name, FA_READ) != FR_OK)
    {
        gfx_printf("\n%kCannot open %s!\n", COLOR_RED, name);
        return false;
    }

    u8 *bufferNX = (u8 *)malloc(PRODINFO_SIZE);
    u32 bytesRead;
    gfx_printf("%kReading from file...\n", COLOR_YELLOW);
    if (f_read(&fp, bufferNX, PRODINFO_SIZE, &bytesRead) != FR_OK || bytesRead != PRODINFO_SIZE)
    {
        gfx_printf("\n%kError reading from file!\n", COLOR_RED);
        goto out;
    }
    if (!verifyProdinfo(bufferNX))
    {
        goto out;
    }
    gfx_printf("%kWriting to NAND...\n", COLOR_YELLOW);
    if (!writeData(bufferNX, 0, PRODINFO_SIZE, print_progress))
    {
        gfx_printf("\n%kError writing to NAND!\nThis is bad. Try again, because your switch probably won't boot.\n"
                   "If you see this error again, you should restore via NAND backup in hekate.\n",
                   COLOR_RED);
        goto out;
    }

    result = true;
    gfx_printf("\n%kRestore from %s done!\n\n", COLOR_GREEN, name);
out:
    f_close(&fp);
    free(bufferNX);

    return result;
}