#include "subghz_keystore.h"

#include <furi.h>
#include <furi-hal.h>

#include <storage/storage.h>
#include <toolbox/hex.h>
#include <flipper_file/flipper_file.h>

#define TAG "SubGhzKeystore"

#define FILE_BUFFER_SIZE 64

#define SUBGHZ_KEYSTORE_FILE_TYPE "Flipper SubGhz Keystore File"
#define SUBGHZ_KEYSTORE_FILE_RAW_TYPE "Flipper SubGhz Keystore RAW File"
#define SUBGHZ_KEYSTORE_FILE_VERSION 0

#define SUBGHZ_KEYSTORE_FILE_ENCRYPTION_KEY_SLOT 1
#define SUBGHZ_KEYSTORE_FILE_DECRYPTED_LINE_SIZE 512
#define SUBGHZ_KEYSTORE_FILE_ENCRYPTED_LINE_SIZE (SUBGHZ_KEYSTORE_FILE_DECRYPTED_LINE_SIZE * 2)

typedef enum {
    SubGhzKeystoreEncryptionNone,
    SubGhzKeystoreEncryptionAES256,
} SubGhzKeystoreEncryption;

struct SubGhzKeystore {
    SubGhzKeyArray_t data;
};

SubGhzKeystore* subghz_keystore_alloc() {
    SubGhzKeystore* instance = furi_alloc(sizeof(SubGhzKeystore));

    SubGhzKeyArray_init(instance->data);

    return instance;
}

void subghz_keystore_free(SubGhzKeystore* instance) {
    furi_assert(instance);

    for
        M_EACH(manufacture_code, instance->data, SubGhzKeyArray_t) {
            string_clear(manufacture_code->name);
            manufacture_code->key = 0;
        }
    SubGhzKeyArray_clear(instance->data);

    free(instance);
}

static void subghz_keystore_add_key(
    SubGhzKeystore* instance,
    const char* name,
    uint64_t key,
    uint16_t type) {
    SubGhzKey* manufacture_code = SubGhzKeyArray_push_raw(instance->data);
    string_init_set_str(manufacture_code->name, name);
    manufacture_code->key = key;
    manufacture_code->type = type;
}

static bool subghz_keystore_process_line(SubGhzKeystore* instance, char* line) {
    uint64_t key = 0;
    uint16_t type = 0;
    char skey[17] = {0};
    char name[65] = {0};
    int ret = sscanf(line, "%16s:%hu:%64s", skey, &type, name);
    key = strtoull(skey, NULL, 16);
    if(ret == 3) {
        subghz_keystore_add_key(instance, name, key, type);
        return true;
    } else {
        FURI_LOG_E(TAG, "Failed to load line: %s\r\n", line);
        return false;
    }
}

static void subghz_keystore_mess_with_iv(uint8_t* iv) {
    // Please do not share decrypted manufacture keys
    // Sharing them will bring some discomfort to legal owners
    // And potential legal action against you
    // While you reading this code think about your own personal responsibility
    asm volatile("movs   r0, #0x0    \n"
                 "movs   r1, #0x0    \n"
                 "movs   r2, #0x0    \n"
                 "movs   r3, #0x0    \n"
                 "nani:              \n"
                 "ldrb   r1, [r0, %0]\n"
                 "mov    r2, r1      \n"
                 "add    r1, r3      \n"
                 "mov    r3, r2      \n"
                 "strb   r1, [r0, %0]\n"
                 "adds   r0, #0x1    \n"
                 "cmp    r0, #0xF    \n"
                 "bls    nani        \n"
                 :
                 : "r"(iv)
                 : "r0", "r1", "r2", "r3", "memory");
}

static bool subghz_keystore_read_file(SubGhzKeystore* instance, File* file, uint8_t* iv) {
    bool result = true;
    char buffer[FILE_BUFFER_SIZE];

    char* decrypted_line = furi_alloc(SUBGHZ_KEYSTORE_FILE_DECRYPTED_LINE_SIZE);
    char* encrypted_line = furi_alloc(SUBGHZ_KEYSTORE_FILE_ENCRYPTED_LINE_SIZE);
    size_t encrypted_line_cursor = 0;

    if(iv) furi_hal_crypto_store_load_key(SUBGHZ_KEYSTORE_FILE_ENCRYPTION_KEY_SLOT, iv);

    size_t ret = 0;
    do {
        ret = storage_file_read(file, buffer, FILE_BUFFER_SIZE);
        for(uint16_t i = 0; i < ret; i++) {
            if(buffer[i] == '\n' && encrypted_line_cursor > 0) {
                // Process line
                if(iv) {
                    // Data alignment check, 32 instead of 16 because of hex encoding
                    size_t len = strlen(encrypted_line);
                    if(len % 32 == 0) {
                        // Inplace hex to bin conversion
                        for(size_t i = 0; i < len; i += 2) {
                            uint8_t hi_nibble = 0;
                            uint8_t lo_nibble = 0;
                            hex_char_to_hex_nibble(encrypted_line[i], &hi_nibble);
                            hex_char_to_hex_nibble(encrypted_line[i + 1], &lo_nibble);
                            encrypted_line[i / 2] = (hi_nibble << 4) | lo_nibble;
                        }
                        len /= 2;

                        if(furi_hal_crypto_decrypt(
                               (uint8_t*)encrypted_line, (uint8_t*)decrypted_line, len)) {
                            subghz_keystore_process_line(instance, decrypted_line);
                        } else {
                            FURI_LOG_E(TAG, "Decryption failed");
                            result = false;
                            break;
                        }
                    } else {
                        FURI_LOG_E(
                            TAG, "Invalid encrypted data: %s", encrypted_line);
                    }
                } else {
                    subghz_keystore_process_line(instance, encrypted_line);
                }
                // reset line buffer
                memset(decrypted_line, 0, SUBGHZ_KEYSTORE_FILE_DECRYPTED_LINE_SIZE);
                memset(encrypted_line, 0, SUBGHZ_KEYSTORE_FILE_ENCRYPTED_LINE_SIZE);
                encrypted_line_cursor = 0;
            } else if(buffer[i] == '\r' || buffer[i] == '\n') {
                // do not add line endings to the buffer
            } else {
                if(encrypted_line_cursor < SUBGHZ_KEYSTORE_FILE_ENCRYPTED_LINE_SIZE) {
                    encrypted_line[encrypted_line_cursor] = buffer[i];
                    encrypted_line_cursor++;
                } else {
                    FURI_LOG_E(TAG, "Malformed file");
                    result = false;
                    break;
                }
            }
        }
    } while(ret > 0 && result);

    if(iv) furi_hal_crypto_store_unload_key(SUBGHZ_KEYSTORE_FILE_ENCRYPTION_KEY_SLOT);

    free(encrypted_line);
    free(decrypted_line);

    return result;
}

bool subghz_keystore_load(SubGhzKeystore* instance, const char* file_name) {
    furi_assert(instance);
    bool result = false;
    uint8_t iv[16];
    uint32_t version;
    SubGhzKeystoreEncryption encryption;

    string_t filetype;
    string_init(filetype);

    Storage* storage = furi_record_open("storage");

    FlipperFile* flipper_file = flipper_file_alloc(storage);
    do {
        if(!flipper_file_open_existing(flipper_file, file_name)) {
            FURI_LOG_E(TAG, "Unable to open file for read: %s", file_name);
            break;
        }
        if(!flipper_file_read_header(flipper_file, filetype, &version)) {
            FURI_LOG_E(TAG, "Missing or incorrect header");
            break;
        }
        if(!flipper_file_read_uint32(flipper_file, "Encryption", (uint32_t*)&encryption, 1)) {
            FURI_LOG_E(TAG, "Missing encryption type");
            break;
        }

        if(strcmp(string_get_cstr(filetype), SUBGHZ_KEYSTORE_FILE_TYPE) != 0 ||
           version != SUBGHZ_KEYSTORE_FILE_VERSION) {
            FURI_LOG_E(TAG, "Type or version mismatch");
            break;
        }

        File* file = flipper_file_get_file(flipper_file);
        if(encryption == SubGhzKeystoreEncryptionNone) {
            result = subghz_keystore_read_file(instance, file, NULL);
        } else if(encryption == SubGhzKeystoreEncryptionAES256) {
            if(!flipper_file_read_hex(flipper_file, "IV", iv, 16)) {
                FURI_LOG_E(TAG, "Missing IV");
                break;
            }
            subghz_keystore_mess_with_iv(iv);
            result = subghz_keystore_read_file(instance, file, iv);
        } else {
            FURI_LOG_E(TAG, "Unknown encryption");
            break;
        }
    } while(0);
    flipper_file_close(flipper_file);
    flipper_file_free(flipper_file);

    furi_record_close("storage");

    string_clear(filetype);

    return result;
}

bool subghz_keystore_save(SubGhzKeystore* instance, const char* file_name, uint8_t* iv) {
    furi_assert(instance);
    bool result = false;

    Storage* storage = furi_record_open("storage");
    char* decrypted_line = furi_alloc(SUBGHZ_KEYSTORE_FILE_DECRYPTED_LINE_SIZE);
    char* encrypted_line = furi_alloc(SUBGHZ_KEYSTORE_FILE_ENCRYPTED_LINE_SIZE);

    FlipperFile* flipper_file = flipper_file_alloc(storage);
    do {
        if(!flipper_file_open_always(flipper_file, file_name)) {
            FURI_LOG_E(TAG, "Unable to open file for write: %s", file_name);
            break;
        }
        if(!flipper_file_write_header_cstr(
               flipper_file, SUBGHZ_KEYSTORE_FILE_TYPE, SUBGHZ_KEYSTORE_FILE_VERSION)) {
            FURI_LOG_E(TAG, "Unable to add header");
            break;
        }
        uint32_t encryption = SubGhzKeystoreEncryptionAES256;
        if(!flipper_file_write_uint32(flipper_file, "Encryption", &encryption, 1)) {
            FURI_LOG_E(TAG, "Unable to add Encryption");
            break;
        }
        if(!flipper_file_write_hex(flipper_file, "IV", iv, 16)) {
            FURI_LOG_E(TAG, "Unable to add IV");
            break;
        }

        subghz_keystore_mess_with_iv(iv);

        if(!furi_hal_crypto_store_load_key(SUBGHZ_KEYSTORE_FILE_ENCRYPTION_KEY_SLOT, iv)) {
            FURI_LOG_E(TAG, "Unable to load encryption key");
            break;
        }

        File* file = flipper_file_get_file(flipper_file);
        size_t encrypted_line_count = 0;
        for
            M_EACH(key, instance->data, SubGhzKeyArray_t) {
                // Wipe buffer before packing
                memset(decrypted_line, 0, SUBGHZ_KEYSTORE_FILE_DECRYPTED_LINE_SIZE);
                memset(encrypted_line, 0, SUBGHZ_KEYSTORE_FILE_ENCRYPTED_LINE_SIZE);
                // Form unecreypted line
                int len = snprintf(
                    decrypted_line,
                    SUBGHZ_KEYSTORE_FILE_DECRYPTED_LINE_SIZE,
                    "%08lX%08lX:%hu:%s",
                    (uint32_t)(key->key >> 32),
                    (uint32_t)key->key,
                    key->type,
                    string_get_cstr(key->name));
                // Verify length and align
                furi_assert(len > 0);
                if(len % 16 != 0) {
                    len += (16 - len % 16);
                }
                furi_assert(len % 16 == 0);
                furi_assert(len <= SUBGHZ_KEYSTORE_FILE_DECRYPTED_LINE_SIZE);
                // Form encrypted line
                if(!furi_hal_crypto_encrypt(
                       (uint8_t*)decrypted_line, (uint8_t*)encrypted_line, len)) {
                    FURI_LOG_E(TAG, "Encryption failed");
                    break;
                }
                // HEX Encode encrypted line
                const char xx[] = "0123456789ABCDEF";
                for(size_t i = 0; i < len; i++) {
                    size_t cursor = len - i - 1;
                    size_t hex_cursor = len * 2 - i * 2 - 1;
                    encrypted_line[hex_cursor] = xx[encrypted_line[cursor] & 0xF];
                    encrypted_line[hex_cursor - 1] = xx[(encrypted_line[cursor] >> 4) & 0xF];
                }
                storage_file_write(file, encrypted_line, strlen(encrypted_line));
                storage_file_write(file, "\n", 1);
                encrypted_line_count++;

                FURI_LOG_I(
                    TAG, "Encrypted: `%s` -> `%s`", decrypted_line, encrypted_line);
            }
        furi_hal_crypto_store_unload_key(SUBGHZ_KEYSTORE_FILE_ENCRYPTION_KEY_SLOT);
        result = encrypted_line_count == SubGhzKeyArray_size(instance->data);
    } while(0);
    flipper_file_close(flipper_file);
    flipper_file_free(flipper_file);

    free(encrypted_line);
    free(decrypted_line);
    furi_record_close("storage");

    return result;
}

SubGhzKeyArray_t* subghz_keystore_get_data(SubGhzKeystore* instance) {
    furi_assert(instance);
    return &instance->data;
}

bool subghz_keystore_raw_encrypted_save(
    const char* input_file_name,
    const char* output_file_name,
    uint8_t* iv) {
    bool encrypted = false;
    uint32_t version;
    string_t filetype;
    string_init(filetype);
    SubGhzKeystoreEncryption encryption;

    Storage* storage = furi_record_open("storage");

    char* encrypted_line = furi_alloc(SUBGHZ_KEYSTORE_FILE_ENCRYPTED_LINE_SIZE);

    FlipperFile* input_flipper_file = flipper_file_alloc(storage);
    do {
        if(!flipper_file_open_existing(input_flipper_file, input_file_name)) {
            FURI_LOG_E(TAG, "Unable to open file for read: %s", input_file_name);
            break;
        }
        if(!flipper_file_read_header(input_flipper_file, filetype, &version)) {
            FURI_LOG_E(TAG, "Missing or incorrect header");
            break;
        }
        if(!flipper_file_read_uint32(input_flipper_file, "Encryption", (uint32_t*)&encryption, 1)) {
            FURI_LOG_E(TAG, "Missing encryption type");
            break;
        }

        if(strcmp(string_get_cstr(filetype), SUBGHZ_KEYSTORE_FILE_RAW_TYPE) != 0 ||
           version != SUBGHZ_KEYSTORE_FILE_VERSION) {
            FURI_LOG_E(TAG, "Type or version mismatch");
            break;
        }

        if(encryption != SubGhzKeystoreEncryptionNone) {
            FURI_LOG_E(TAG, "Already encryption");
            break;
        }
        File* input_file = flipper_file_get_file(input_flipper_file);

        FlipperFile* output_flipper_file = flipper_file_alloc(storage);

        if(!flipper_file_open_always(output_flipper_file, output_file_name)) {
            FURI_LOG_E(TAG, "Unable to open file for write: %s", output_file_name);
            break;
        }
        if(!flipper_file_write_header_cstr(
               output_flipper_file, string_get_cstr(filetype), SUBGHZ_KEYSTORE_FILE_VERSION)) {
            FURI_LOG_E(TAG, "Unable to add header");
            break;
        }
        uint32_t encryption = SubGhzKeystoreEncryptionAES256;
        if(!flipper_file_write_uint32(
               output_flipper_file, "Encryption", &encryption, 1)) {
            FURI_LOG_E(TAG, "Unable to add Encryption");
            break;
        }
        if(!flipper_file_write_hex(output_flipper_file, "IV", iv, 16)) {
            FURI_LOG_E(TAG, "Unable to add IV");
            break;
        }

        if(!flipper_file_write_string_cstr(output_flipper_file, "Encrypt_data", "RAW")) {
            FURI_LOG_E(TAG, "Unable to add Encrypt_data");
            break;
        }

        subghz_keystore_mess_with_iv(iv);

        if(!furi_hal_crypto_store_load_key(SUBGHZ_KEYSTORE_FILE_ENCRYPTION_KEY_SLOT, iv)) {
            FURI_LOG_E(TAG, "Unable to load encryption key");
            break;
        }

        File* output_file = flipper_file_get_file(output_flipper_file);
        char buffer[FILE_BUFFER_SIZE];
        bool result = true;

        size_t ret = 0;
        furi_assert(FILE_BUFFER_SIZE % 16 == 0);

        //skip the end of the previous line "\n"
        storage_file_read(input_file, buffer, 1);

        do {
            memset(buffer, 0, FILE_BUFFER_SIZE);
            ret = storage_file_read(input_file, buffer, FILE_BUFFER_SIZE);
            if(ret == 0) {
                break;
            }

            for(uint16_t i = 0; i < FILE_BUFFER_SIZE - 1; i += 2) {
                uint8_t hi_nibble = 0;
                uint8_t lo_nibble = 0;
                hex_char_to_hex_nibble(buffer[i], &hi_nibble);
                hex_char_to_hex_nibble(buffer[i + 1], &lo_nibble);
                buffer[i / 2] = (hi_nibble << 4) | lo_nibble;
            }

            memset(encrypted_line, 0, SUBGHZ_KEYSTORE_FILE_ENCRYPTED_LINE_SIZE);
            // Form encrypted line
            if(!furi_hal_crypto_encrypt(
                   (uint8_t*)buffer, (uint8_t*)encrypted_line, FILE_BUFFER_SIZE / 2)) {
                FURI_LOG_E(TAG, "Encryption failed");
                result = false;
                break;
            }

            // HEX Encode encrypted line
            const char xx[] = "0123456789ABCDEF";
            for(size_t i = 0; i < FILE_BUFFER_SIZE / 2; i++) {
                size_t cursor = FILE_BUFFER_SIZE / 2 - i - 1;
                size_t hex_cursor = FILE_BUFFER_SIZE - i * 2 - 1;
                encrypted_line[hex_cursor] = xx[encrypted_line[cursor] & 0xF];
                encrypted_line[hex_cursor - 1] = xx[(encrypted_line[cursor] >> 4) & 0xF];
            }
            storage_file_write(output_file, encrypted_line, strlen(encrypted_line));

        } while(ret > 0 && result);

        flipper_file_close(output_flipper_file);
        flipper_file_free(output_flipper_file);

        furi_hal_crypto_store_unload_key(SUBGHZ_KEYSTORE_FILE_ENCRYPTION_KEY_SLOT);

        if(!result) break;

        encrypted = true;
    } while(0);

    flipper_file_close(input_flipper_file);
    flipper_file_free(input_flipper_file);

    free(encrypted_line);

    furi_record_close("storage");

    return encrypted;
}

bool subghz_keystore_raw_get_data(const char* file_name, size_t offset, uint8_t* data, size_t len) {
    bool result = false;
    uint8_t iv[16];
    uint32_t version;
    SubGhzKeystoreEncryption encryption;

    string_t str_temp;
    string_init(str_temp);

    Storage* storage = furi_record_open("storage");
    char* decrypted_line = furi_alloc(SUBGHZ_KEYSTORE_FILE_DECRYPTED_LINE_SIZE);

    FlipperFile* flipper_file = flipper_file_alloc(storage);
    do {
        if(!flipper_file_open_existing(flipper_file, file_name)) {
            FURI_LOG_E(TAG, "Unable to open file for read: %s", file_name);
            break;
        }
        if(!flipper_file_read_header(flipper_file, str_temp, &version)) {
            FURI_LOG_E(TAG, "Missing or incorrect header");
            break;
        }
        if(!flipper_file_read_uint32(flipper_file, "Encryption", (uint32_t*)&encryption, 1)) {
            FURI_LOG_E(TAG, "Missing encryption type");
            break;
        }

        if(strcmp(string_get_cstr(str_temp), SUBGHZ_KEYSTORE_FILE_RAW_TYPE) != 0 ||
           version != SUBGHZ_KEYSTORE_FILE_VERSION) {
            FURI_LOG_E(TAG, "Type or version mismatch");
            break;
        }

        File* file = flipper_file_get_file(flipper_file);
        if(encryption != SubGhzKeystoreEncryptionAES256) {
            FURI_LOG_E(TAG, "Unknown encryption");
            break;
        }

        if(offset < 16) {
            if(!flipper_file_read_hex(flipper_file, "IV", iv, 16)) {
                FURI_LOG_E(TAG, "Missing IV");
                break;
            }
            subghz_keystore_mess_with_iv(iv);
        }

        if(!flipper_file_read_string(flipper_file, "Encrypt_data", str_temp)) {
            FURI_LOG_E(TAG, "Missing Encrypt_data");
            break;
        }

        size_t bufer_size;
        if(len <= (16 - offset % 16)) {
            bufer_size = 32;
        } else {
            bufer_size = (((len) / 16) + 2) * 32;
        }
        furi_assert(SUBGHZ_KEYSTORE_FILE_DECRYPTED_LINE_SIZE >= bufer_size / 2);

        char buffer[bufer_size];
        size_t ret = 0;
        bool decrypted = true;
        //skip the end of the previous line "\n"
        storage_file_read(file, buffer, 1);

        size_t size = storage_file_size(file);
        size -= storage_file_tell(file);
        if(size < (offset * 2 + len * 2)) {
            FURI_LOG_E(TAG, "Seek position exceeds file size");
            break;
        }

        if(offset >= 16) {
            storage_file_seek(file, ((offset / 16) - 1) * 32, false);
            ret = storage_file_read(file, buffer, 32);
            furi_assert(ret == 32);
            for(uint16_t i = 0; i < ret - 1; i += 2) {
                uint8_t hi_nibble = 0;
                uint8_t lo_nibble = 0;
                hex_char_to_hex_nibble(buffer[i], &hi_nibble);
                hex_char_to_hex_nibble(buffer[i + 1], &lo_nibble);
                iv[i / 2] = (hi_nibble << 4) | lo_nibble;
            }
        }

        if(!furi_hal_crypto_store_load_key(SUBGHZ_KEYSTORE_FILE_ENCRYPTION_KEY_SLOT, iv)) {
            FURI_LOG_E(TAG, "Unable to load encryption key");
            break;
        }

        do {
            memset(buffer, 0, bufer_size);
            ret = storage_file_read(file, buffer, bufer_size);
            furi_assert(ret == bufer_size);
            for(uint16_t i = 0; i < ret - 1; i += 2) {
                uint8_t hi_nibble = 0;
                uint8_t lo_nibble = 0;
                hex_char_to_hex_nibble(buffer[i], &hi_nibble);
                hex_char_to_hex_nibble(buffer[i + 1], &lo_nibble);
                buffer[i / 2] = (hi_nibble << 4) | lo_nibble;
            }

            memset(decrypted_line, 0, SUBGHZ_KEYSTORE_FILE_DECRYPTED_LINE_SIZE);

            if(!furi_hal_crypto_decrypt(
                   (uint8_t*)buffer, (uint8_t*)decrypted_line, bufer_size / 2)) {
                decrypted = false;
                FURI_LOG_E(TAG, "Decryption failed");
                break;
            }
            memcpy(data, (uint8_t*)decrypted_line + (offset - (offset / 16) * 16), len);

        } while(0);
        furi_hal_crypto_store_unload_key(SUBGHZ_KEYSTORE_FILE_ENCRYPTION_KEY_SLOT);
        if(decrypted) result = true;
    } while(0);
    flipper_file_close(flipper_file);
    flipper_file_free(flipper_file);

    furi_record_close("storage");

    free(decrypted_line);

    string_clear(str_temp);

    return result;
}