#include "Aes_Hmac.h"

// 辅助函数：根据密钥长度选择AES算法
const EVP_CIPHER* get_aes_cipher(size_t key_len) {
    switch (key_len) {
        case 16: return EVP_aes_128_ecb();
        case 24: return EVP_aes_192_ecb();
        case 32: return EVP_aes_256_ecb();
        default: return nullptr;
    }
}

void encrypt_aeshmack(const uint8_t *in, const size_t inlen, 
                     uint8_t *&out, size_t &outlen, uint8_t *ckey) {
    // 初始化OpenSSL错误
    ERR_clear_error();
    out = nullptr;
    outlen = 0;

    // 获取密钥长度（需要额外参数）
    size_t ckey_len = 16; // 这里假设密钥长度为16字节，实际需要从参数获取

    // Step 1: AES加密
    const EVP_CIPHER* cipher = get_aes_cipher(ckey_len);
    if (!cipher) {
        return;
    }

    EVP_CIPHER_CTX* ctx = EVP_CIPHER_CTX_new();
    if (!ctx) return;

    int rc = EVP_EncryptInit_ex(ctx, cipher, nullptr, ckey, nullptr);
    if (rc != 1) EVP_CIPHER_CTX_free(ctx);

    EVP_CIPHER_CTX_set_padding(ctx, 1);

    size_t block_size = EVP_CIPHER_block_size(cipher);
    size_t buf_len = inlen + block_size;
    uint8_t* encrypted = new uint8_t[buf_len];
    
    int tmp_len = 0;
    rc = EVP_EncryptUpdate(ctx, encrypted, &tmp_len, in, inlen);
    if (rc != 1) {
        delete[] encrypted;
        EVP_CIPHER_CTX_free(ctx);
    }
    int encrypted_len = tmp_len;

    rc = EVP_EncryptFinal_ex(ctx, encrypted + encrypted_len, &tmp_len);
    if (rc != 1) {
        delete[] encrypted;
        EVP_CIPHER_CTX_free(ctx);
    }
    encrypted_len += tmp_len;

    // Step 2: 计算HMAC
    constexpr size_t HMAC_LEN = 32; // SHA256
    uint8_t hmac[HMAC_LEN];
    unsigned int hmac_len = 0;
    
    HMAC(EVP_sha256(), ckey, ckey_len, encrypted, encrypted_len, hmac, &hmac_len);

    // Step 3: 拼接结果
    out = new uint8_t[encrypted_len + HMAC_LEN];
    memcpy(out, encrypted, encrypted_len);
    memcpy(out + encrypted_len, hmac, HMAC_LEN);
    outlen = encrypted_len + HMAC_LEN;

    delete[] encrypted;
    
}

int decrypt_aeshmack(const uint8_t *in, const size_t inlen,
                    uint8_t *&out, size_t &outlen, uint8_t *ckey) {
    ERR_clear_error();
    out = nullptr;
    outlen = 0;

    if(in == nullptr || inlen == 0 || inlen < 32) {
        return -3; // 输入不合法
    }

    // Step 1: 拆分数据
    constexpr size_t HMAC_LEN = 32;

    size_t data_len = inlen - HMAC_LEN;
    const uint8_t* encrypted = in;
    const uint8_t* received_hmac = in + data_len;

    // Step 2: 验证HMAC
    size_t ckey_len = 16; // 假设同上
    uint8_t calculated_hmac[HMAC_LEN];
    unsigned int calculated_len = 0;
    
    HMAC(EVP_sha256(), ckey, ckey_len, encrypted, data_len, 
        calculated_hmac, &calculated_len);

    if (CRYPTO_memcmp(received_hmac, calculated_hmac, HMAC_LEN) != 0) {
        return -2;
    }

    // Step 3: AES解密
    const EVP_CIPHER* cipher = get_aes_cipher(ckey_len);
    if (!cipher) return -1;

    EVP_CIPHER_CTX* ctx = EVP_CIPHER_CTX_new();
    if (!ctx) return -1;

    int rc = EVP_DecryptInit_ex(ctx, cipher, nullptr, ckey, nullptr);
    if (rc != 1) {
        EVP_CIPHER_CTX_free(ctx);
        return -1;
    }

    EVP_CIPHER_CTX_set_padding(ctx, 1);

    uint8_t* decrypted = new uint8_t[data_len];
    memset(decrypted, 0, data_len+1);
    int tmp_len = 0;
    int decrypted_len = 0;

    rc = EVP_DecryptUpdate(ctx, decrypted, &tmp_len, encrypted, data_len);
    if (rc != 1) {
        delete[] decrypted;
        EVP_CIPHER_CTX_free(ctx);
        return -1;
    }
    decrypted_len = tmp_len;

    rc = EVP_DecryptFinal_ex(ctx, decrypted + decrypted_len, &tmp_len);
    if (rc != 1) {
        delete[] decrypted;
        EVP_CIPHER_CTX_free(ctx);
        return -1;
    }
    decrypted_len += tmp_len;

    out = decrypted;
    outlen = decrypted_len;
    out[outlen] = '\0';

    EVP_CIPHER_CTX_free(ctx);
    return 1;
}
