#include "sm4.h"
#include <new>        // 必须包含这个


sm4::sm4(uint8_t *ckey){
    size_t keylen = L_SM4_GCM_KEY_SIZE;


    //INIT CRYPT_KEY
    uint8_t *key = new uint8_t[keylen];
    memcpy(key, ckey, keylen);
    sm4_set_encrypt_key(&encrypt_key, key);
    sm4_set_decrypt_key(&decrypt_key, key);

    this->hmac_keylen = 16;  // SM3_HMAC推荐最小16字节
    this->hmac_key = new uint8_t[hmac_keylen];
    memcpy(this->hmac_key, key, hmac_keylen);
    delete[] key;

}

sm4::~sm4(){
    delete[] hmac_key;
}

int sm4::encrypt_gcm(const uint8_t *in, const size_t inlen, uint8_t *&out, size_t &outlen){
    out = nullptr;
    outlen = 0;

    size_t ivlen = L_SM4_GCM_IV_SIZE;
    size_t taglen = L_SM4_GCM_TAG_SIZE;
    size_t cipherlen = inlen;

    uint8_t iv[ivlen];
    uint8_t tag[taglen];
    uint8_t cipher[cipherlen];
    int ret;

    //随即生成iv，并且给iv填充，这里使用pkcs7_pad填充
    rand_bytes(iv,ivlen);

    //加密input
    ret = sm4_gcm_encrypt(&encrypt_key,iv,ivlen,NULL,0,in,inlen,cipher,taglen,tag);

    //拼接 cipher tag 和 iv
    outlen = cipherlen + taglen + ivlen;
    out = new uint8_t[outlen];

    memcpy(out,cipher,cipherlen);
    memcpy(out+cipherlen,tag,taglen);
    memcpy(out+cipherlen+taglen,iv,ivlen);

    return ret;
}

int sm4::decrypt_gcm(const uint8_t *in, const size_t inlen, uint8_t *&out, size_t &outlen){
    out = nullptr;
    outlen = 0;


    size_t ivlen = L_SM4_GCM_IV_SIZE;
    size_t taglen = L_SM4_GCM_TAG_SIZE;
    size_t cipherlen = inlen - ivlen -taglen;

    if(in == nullptr || inlen == 0 || inlen < ivlen + taglen ){
        return -1; // 输入长度不足
    }

    uint8_t iv[ivlen];
    uint8_t tag[taglen];
    uint8_t cipher[cipherlen];
    int ret;

    // 从输入中提取 IV、TAG 和 密文
    memcpy(iv, in + cipherlen + taglen, ivlen);
    memcpy(tag, in + cipherlen, taglen);                   
    memcpy(cipher, in, cipherlen);   

    //解密cipher得到out
    outlen = cipherlen;
    out = new uint8_t[outlen];
    memset(out,0,outlen+1);
    ret = sm4_gcm_decrypt(&encrypt_key, iv, ivlen, NULL, 0, cipher, cipherlen, tag, taglen, out);

    return ret;
}

int sm4::encrypt_ecb(const uint8_t *in, const size_t inlen, uint8_t *&out, size_t &outlen){
    out = nullptr;
    outlen = 0;

    // 检查输入参数
    if (in == nullptr || inlen == 0) {
        return -1;
    }

    uint8_t *pad;
    int ret = pkcs7_pad(in,inlen,pad,outlen);
    if (ret != 1) {
        return -1; // 填充失败
    }

    try {
    size_t nblocks  = outlen/L_PCKS7_PAD_SIZE;
    out = new uint8_t[outlen];
    sm4_encrypt_blocks(&encrypt_key,pad,nblocks,out);
    } catch (const std::bad_alloc&) {
        delete[] pad;
        return -1; // 内存分配失败
    } catch (...) {
        delete[] pad;
        delete[] out;
        out = nullptr;
        return -1; // 加密失败
    }

    delete[] pad;
    return 1;
}

int sm4::decrypt_ecb(const uint8_t *in, const size_t inlen, uint8_t *&out, size_t &outlen){
    out = nullptr;
    outlen = 0;

    // 检查输入参数
    if (in == nullptr || inlen == 0 || inlen % L_PCKS7_PAD_SIZE != 0) {
        return -3;
    }

    uint8_t *pad = nullptr;

    try {
    size_t nblocks = inlen/L_PCKS7_PAD_SIZE;
    pad = new uint8_t[inlen];
    sm4_encrypt_blocks(&decrypt_key,in,nblocks,pad);
    } catch (const std::bad_alloc&) {
        return -1; // 内存分配失败
    } catch (...) {
        delete[] pad;
        return -1; // 其他异常
    }

    int ret = pkcs7_unpad(pad,inlen,out,outlen);
    if(ret != 1){
        delete[] pad;
        return -1; // 去除填充失败
    }

    delete[] pad;
    return 1;
}

/**
 * in:data, datalen
 * out:data + padding , data_len + padding_len
 */
int sm4::pkcs7_pad(const uint8_t *in, const size_t inlen, uint8_t *&out, size_t &outlen){
    out = nullptr;
    outlen = 0;

    // 检查输入参数
    if (in == nullptr && inlen > 0) {
        return -1; // 输入不合法
    }

    //补充 pad_byte 个字节，如果是16的倍数则补充16个字节
    size_t pcks7_size = L_PCKS7_PAD_SIZE;
    size_t pad_byte = pcks7_size - inlen % pcks7_size;
    if(pad_byte == 0) pad_byte = pcks7_size;

    //分配内存并复制原始数据
    outlen = inlen + pad_byte;

    try {
    out = new uint8_t[outlen];
    memcpy(out,in,inlen);

    //填充
    for(size_t i = inlen; i< outlen; i++){
        out[i] = static_cast<uint8_t>(pad_byte);;
    }
    } catch (const std::bad_alloc&) {
        outlen = 0; // 内存分配失败
        return -1;
    } catch (...) {
        delete[] out;
        out = nullptr;
        outlen = 0;
        return -1; // 其他异常
    }

    return 1;
}

/**
 * in:data + padding, data_len + padding_len
 * out:data, padding_len
 */
int sm4::pkcs7_unpad(const uint8_t *in, const size_t inlen, uint8_t *&out, size_t &outlen){
    out = nullptr;
    outlen = 0;
    // 检查输入参数
    if (in == nullptr || inlen == 0 || inlen % L_PCKS7_PAD_SIZE != 0) {
        return -1; // 输入不合法
    }

    size_t pcks7_size = L_PCKS7_PAD_SIZE;
    //获取填充字节
    uint8_t pad_byte = in[inlen - 1];
    if (pad_byte < 1 || pad_byte > pcks7_size) {
        return -1; // 填充字节不合法
    }

    // 验证所有填充字节是否正确
    size_t pad_start = inlen - pad_byte;
    for (size_t i = pad_start; i < inlen; ++i) {
        if (in[i] != pad_byte) {
            return -1; // 填充字节不合法
        }
    }

    // 分配内存并复制原始数据
    outlen = pad_start;

    try {
    out = new uint8_t[outlen];
    memset(out,0,outlen+1);
    memcpy(out, in, outlen);
    } catch (const std::bad_alloc&) {
        outlen = 0;
        return -1; // 内存分配失败
    } catch (...) {
        delete[] out;
        out = nullptr;
        outlen = 0;
        return -1; // 其他异常
    }

    return 1;
}


int sm4::encrypt_ecb_with_checksum(const uint8_t *in, const size_t inlen, uint8_t *&out, size_t &outlen) {
    out = nullptr;
    outlen = 0;

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

    // 1. 执行ECB加密
    uint8_t *cipher = nullptr;
    size_t cipher_len = 0;
    int ret = encrypt_ecb(in, inlen, cipher, cipher_len);
    if(ret != 1){
        return -1; // 加密失败
    }

    // 2. 计算SM3_HMAC校验码
    uint8_t hmac[SM3_HMAC_SIZE];
    SM3_DIGEST_CTX ctx;

    if (sm3_digest_init(&ctx, this->hmac_key, this->hmac_keylen) != 1) {
        delete[] cipher;
        return -2; // 初始化失败
    }
    if (sm3_digest_update(&ctx, cipher, cipher_len) != 1) {
        delete[] cipher;
        return -2; // 更新失败
    }
    if (sm3_digest_finish(&ctx, hmac) != 1) {
        delete[] cipher;
        return -2; // 计算失败
    }

    // 3. 拼接密文和HMAC
    outlen = cipher_len + SM3_HMAC_SIZE;
    out = new uint8_t[outlen];
    if (!out) {
        delete[] cipher;
        return -2; // 内存分配失败
    }

    memcpy(out, cipher, cipher_len);
    memcpy(out + cipher_len, hmac, SM3_HMAC_SIZE);

    delete[] cipher;
    return 1; // 成功
}

int sm4::decrypt_ecb_with_checksum(const uint8_t *in, const size_t inlen, uint8_t *&out, size_t &outlen) {
    out = nullptr;
    outlen = 0;

    // 1. 分离密文和HMAC
    const size_t hmac_size = SM3_HMAC_SIZE;
    if (inlen < hmac_size) {
        return -1; // 输入长度不足
    }

    const size_t cipher_len = inlen - hmac_size;
    const uint8_t* cipher = in;
    const uint8_t* received_hmac = in + cipher_len;

    // 2. 计算SM3_HMAC校验
    uint8_t computed_hmac[SM3_HMAC_SIZE];
    SM3_DIGEST_CTX ctx;

    if (sm3_digest_init(&ctx, this->hmac_key, this->hmac_keylen) != 1) {
        return -2; // 初始化失败
    }
    if (sm3_digest_update(&ctx, cipher, cipher_len) != 1) {
        return -2; // 更新失败
    }
    if (sm3_digest_finish(&ctx, computed_hmac) != 1) {
        return -2; // 计算失败
    }

    // 3. 校验比对
    if (memcmp(received_hmac, computed_hmac, hmac_size) != 0) {
        return -2; // 校验失败
    }

    // 4. 执行ECB解密
    int ret = decrypt_ecb(cipher, cipher_len, out, outlen);
    if(ret != 1) 
        return -1; // 解密失败
        
    return 1; // 成功
}