#include <Arduino.h>
#include "check_firmware.h"
#include "monocypher.h"
#include "parameters.h"
#include <string.h>
#include "util.h"
/*
该函数基于加密上下文（crypto_check_ctx），先初始化签名校验环境（绑定固件签名和公钥）→ 
分批次传入固件数据（含可选的引导字节）→ 最终验证固件哈希值是否与描述符中的签名匹配，返回校验结果；
*/
bool CheckFirmware::check_partition(const uint8_t *flash, uint32_t flash_len,
                                    const uint8_t *lead_bytes, uint32_t lead_length,
                                    const app_descriptor_t *ad, const uint8_t public_key[32])
{
    crypto_check_ctx ctx {};  // 1. 初始化加密上下文（清空内存）
    crypto_check_ctx_abstract *actx = (crypto_check_ctx_abstract*)&ctx; // 2. 转换为抽象上下文指针（适配通用加密接口）
    crypto_check_init(actx, ad->sign_signature, public_key); // 3. 绑定固件签名和公钥，初始化校验环境
    if (lead_length > 0) { // 1. 若有引导字节，先传入引导字节数据
        crypto_check_update(actx, lead_bytes, lead_length);
    }
    crypto_check_update(actx, &flash[lead_length], flash_len-lead_length); // 2. 传入剩余固件数据（总长度 - 引导字节长度）
    return crypto_check_final(actx) == 0; // 验证最终哈希值是否与描述符中的签名匹配（公钥解密签名后对比）
}
/*
OTA 固件合法性校验的核心实现，完整覆盖了 “分区映射→固件描述符校验→板型匹配→公钥签名验证” 
全流程 —— 只有通过所有校验的固件，才会被判定为合法并允许运行。
*/
bool CheckFirmware::check_OTA_partition(const esp_partition_t *part, const uint8_t *lead_bytes, uint32_t lead_length, uint32_t &board_id)
{
    Serial.printf("Checking partition %s\n", part->label);
    spi_flash_mmap_handle_t handle;
    const void *ptr = nullptr;
    auto ret = esp_partition_mmap(part, 0, part->size, SPI_FLASH_MMAP_DATA, &ptr, &handle);
    if (ret != ESP_OK) {
        Serial.printf("mmap failed\n");
        return false;
    }
    // 1. 反转签名字节序（APP_DESCRIPTOR_REV 是RemoteID厂商 预定义的8字节签名）可自行修改
    const uint8_t sig_rev[] = APP_DESCRIPTOR_REV;
    uint8_t sig[8];
    for (uint8_t i=0; i<8; i++) {
        sig[i] = sig_rev[7-i];
    }
    // 2. 在映射内存中查找签名（定位固件描述符）  必须用脚本签名后才能找到这个app描述符，不然无效
    const app_descriptor_t *ad = (app_descriptor_t *)memmem(ptr, part->size, sig, sizeof(sig));
    if (ad == nullptr) {
        Serial.printf("app_descriptor not found\n");
        spi_flash_munmap(handle); // 必须释放映射，否则内存泄漏
        return false;
    }

    Serial.printf("app descriptor at 0x%x size=%u id=%u (own id %u)\n", unsigned(ad)-unsigned(ptr), ad->image_size, ad->board_id,BOARD_ID);
    // 1. 计算固件实际长度（从分区起始到描述符的偏移）
    const uint32_t img_len = uint32_t(uintptr_t(ad) - uintptr_t(ptr)); // 固件描述符紧跟在固件的后面
    // 2. 校验描述符中记录的固件大小是否匹配实际长度
    if (ad->image_size != img_len) {
        Serial.printf("app_descriptor bad size %u\n", ad->image_size);
        spi_flash_munmap(handle);
        return false;
    }
    board_id = ad->board_id; // 3. 提取板型ID（输出参数，供上层函数使用）
    // 无公钥时跳过签名校验（降级策略）
    if (g.no_public_keys()) {
        Serial.printf("No public keys - accepting firmware\n");
        spi_flash_munmap(handle);
        return true;
    }
    // 遍历公钥校验固件签名（核心安全校验）
    for (uint8_t i=0; i<MAX_PUBLIC_KEYS; i++) {
        uint8_t key[32];
        if (!g.get_public_key(i, key)) { // 读取预存的第i个公钥
            continue;
        }
        // 用当前公钥校验固件签名
        if (check_partition((const uint8_t *)ptr, img_len, lead_bytes, lead_length, ad, key)) {
            Serial.printf("check firmware good for key %u\n", i);
            spi_flash_munmap(handle);
            return true;
        }
        Serial.printf("check failed key %u\n", i);
    }
    // 校验失败处理 
    // 所有公钥校验失败 → 判定固件非法，返回 false；
    // 最终释放内存映射，避免泄漏。
    spi_flash_munmap(handle);
    Serial.printf("firmware failed checks\n");
    return false;
}

bool CheckFirmware::check_OTA_next(const esp_partition_t *part, const uint8_t *lead_bytes, uint32_t lead_length)
{
    Serial.printf("Running partition %s\n", esp_ota_get_running_partition()->label); // 打印当前运行分区

    uint32_t board_id = 0;
    bool sig_ok = check_OTA_partition(part, lead_bytes, lead_length, board_id); // 校验目标 OTA 分区固件

    if (g.lock_level == -1) { // 降级策略（lock_level=-1 放行所有固件）
        // only if lock_level is -1 then accept any firmware
        return true;
    }

    // if app descriptor has a board ID and the ID is wrong then reject
    if (board_id != 0 && board_id != BOARD_ID) { // 板型 ID 校验（防止错配固件）
        return false;
    }

    return sig_ok;
}
/*
OTA 固件运行合法性校验的核心入口，作用是获取当前正在运行的 OTA 分区信息，
并调用 check_OTA_partition() 校验该分区固件的合法性（比如是否为合规固件、是否匹配硬件板型等）。
*/
bool CheckFirmware::check_OTA_running(void)
{
    const auto *running_part = esp_ota_get_running_partition(); // 获取当前运行的 OTA 分区
    if (running_part == nullptr) {
        Serial.printf("No running OTA partition\n");
        return false;
    }
    uint32_t board_id=0;
    return check_OTA_partition(running_part, nullptr, 0, board_id); // 定义板型 ID 变量并调用校验函数
}
        
esp_err_t esp_partition_read_raw(const esp_partition_t* partition,
                                 size_t src_offset, void* dst, size_t size);