/*
 * 加密模块实现
 * 使用AES-128-CBC + Base64编码
 */

#include "Crypto.h"
#include "aes.h"
#include "base64.h"
#include <string.h>

// ⚠️ 重要：密钥和IV必须与服务器端完全一致！
// AES-128密钥（16字节）
static const uint8_t AES_KEY[16] = {
    0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6,
    0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf, 0x4f, 0x3c
};

// 初始化向量（16字节）
static const uint8_t AES_IV[16] = {
    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
    0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f
};

static struct AES_ctx aes_ctx;

void Crypto_Init(void)
{
    // 初始化AES上下文（使用CBC模式）
    AES_init_ctx_iv(&aes_ctx, AES_KEY, AES_IV);
}

/**
 * @brief PKCS#7填充
 * 将数据填充到16字节的倍数
 */
static uint32_t pkcs7_pad(uint8_t* data, uint32_t data_len, uint32_t buffer_size)
{
    // 计算需要填充的字节数
    uint32_t padding_len = 16 - (data_len % 16);

    // 检查缓冲区是否足够
    if (data_len + padding_len > buffer_size) {
        return 0; // 缓冲区不足
    }

    // 填充（PKCS#7标准）
    uint32_t i;
    for (i = 0; i < padding_len; i++) {
        data[data_len + i] = (uint8_t)padding_len;
    }

    return data_len + padding_len;
}

static uint32_t pkcs7_unpad(uint8_t* data, uint32_t data_len)
{
    uint8_t padding_len;
    uint32_t i;

    if (data_len == 0) {
        return 0;
    }

    padding_len = data[data_len - 1];
    if (padding_len < 1 || padding_len > 16 || padding_len > data_len) {
        return data_len;
    }

    for (i = data_len - padding_len; i < data_len; i++) {
        if (data[i] != padding_len) {
            return data_len;
        }
    }

    return data_len - padding_len;
}

uint32_t Crypto_EncryptJSON(const char* json_data, char* encrypted_output)
{
    uint32_t json_len = strlen(json_data);
    uint8_t buffer[512]; // 临时缓冲区

    // 1. 复制JSON数据到缓冲区
    if (json_len > sizeof(buffer) - 16) {
        return 0; // 数据太长
    }
    memcpy(buffer, json_data, json_len);

    // 2. PKCS#7填充到16字节的倍数
    uint32_t padded_len = pkcs7_pad(buffer, json_len, sizeof(buffer));
    if (padded_len == 0) {
        return 0; // 填充失败
    }

    // 3. 手动实现标准CBC加密（绕过aes.c的bug）
    struct AES_ctx ctx;
    uint8_t iv[16];
    uint32_t i, j;

    // 初始化密钥和IV
    AES_init_ctx(&ctx, AES_KEY);
    memcpy(iv, AES_IV, 16);

    // CBC加密：每个块先XOR IV，再加密，然后用密文更新IV
    for (i = 0; i < padded_len; i += 16) {
        // XOR with IV
        for (j = 0; j < 16; j++) {
            buffer[i + j] ^= iv[j];
        }

        // AES加密单个块
        AES_ECB_encrypt(&ctx, &buffer[i]);

        // 更新IV为当前密文块
        memcpy(iv, &buffer[i], 16);
    }

    // 4. Base64编码
    uint32_t encoded_len = base64_encode(buffer, padded_len, encrypted_output);

    return encoded_len;
}

uint32_t Crypto_DecryptJSON(const char* encrypted_base64, char* plain_output, uint32_t output_size)
{
    uint32_t encrypted_len;
    uint32_t binary_len;
    uint32_t plain_len;
    uint8_t buffer[768];
    struct AES_ctx ctx;

    if (encrypted_base64 == NULL || plain_output == NULL || output_size == 0) {
        return 0;
    }

    encrypted_len = strlen(encrypted_base64);
    if (encrypted_len == 0 || encrypted_len > BASE64_ENCODE_SIZE(sizeof(buffer))) {
        return 0;
    }

    binary_len = base64_decode(encrypted_base64, encrypted_len, buffer);
    if (binary_len == 0 || (binary_len % 16) != 0 || binary_len >= sizeof(buffer)) {
        return 0;
    }

    AES_init_ctx_iv(&ctx, AES_KEY, AES_IV);
    AES_CBC_decrypt_buffer(&ctx, buffer, binary_len);
    plain_len = pkcs7_unpad(buffer, binary_len);
    if (plain_len >= output_size) {
        return 0;
    }

    memcpy(plain_output, buffer, plain_len);
    plain_output[plain_len] = '\0';
    return plain_len;
}