//
// Created by root on 24-4-17.
//

#include "aes_cbc_iv_128_append.h"

//test only
void print_hex(const char *title, const unsigned char *s, int len) {
    std::cout << title << ": ";
    for (int i = 0; i < len; ++i) {
        printf("%02x", (unsigned int) s[i]);
    }
    std::cout << std::endl;
}

// PKCS#7 padding implementation
int apply_pkcs7_padding(const unsigned char *input, int input_len, unsigned char *padded_input, int block_size) {
    int padding_len = block_size - (input_len % block_size);
    memcpy(padded_input, input, input_len);
    for (int i = 0; i < padding_len; ++i) {
        padded_input[input_len + i] = padding_len;
    }
    return input_len + padding_len;
}

// PKCS#7 padding removal
int remove_pkcs7_padding(unsigned char *data, int input_len) {
    int padding_len = data[input_len - 1];
    // 确认所有填充字节均符合预期
    for (int i = 0; i < padding_len; i++) {
        if (data[input_len - 1 - i] != padding_len) {
            std::cerr << "Padding error: padding byte mismatch." << std::endl;
            return -1;  // 错误的填充
        }
    }
    return input_len - padding_len;
}

//创建iv
inline void iv_create(unsigned char *iv) {
    if (!RAND_bytes(iv, AES_BLOCK_SIZE)) {
        std::cerr << "Failed to generate random IV." << std::endl;
        return;
    }
}

//将iv从[ciphertext|iv]格式中提取出来并删除iv
void iv_extract(unsigned char *iv, const unsigned char *encrypted, int encrypted_len) {
    //TODO 合理性检查

    memcpy(iv, encrypted + encrypted_len, AES_BLOCK_SIZE);
}


//将iv添加到密文后
void iv_append(unsigned char *iv, unsigned char *encrypted, int encrypted_len, unsigned char *output) {
    // 首先复制原始的加密数据到输出缓冲区
    memcpy(output, encrypted, encrypted_len);

    // 然后将IV追加到加密数据的后面
    memcpy(output + encrypted_len, iv, AES_BLOCK_SIZE);
}

//这里是用iv的cbc模式
int
aes_encrypt(const unsigned char *plaintext, int plaintext_len, const unsigned char *key, unsigned char *encrypted) {

    //检查输入合法性
    if (plaintext == NULL || key == NULL || encrypted == NULL) {
        return -1;
    }

    //设置AES key
    AES_KEY aes_key;
    if (AES_set_encrypt_key(key, 128, &aes_key) < 0) {
        return -2;
    }

    //设置iv
    unsigned char iv_enc[AES_BLOCK_SIZE];
    iv_create(iv_enc);
    unsigned char iv_tem[AES_BLOCK_SIZE];

    //加密操作会改变IV，所以备份一下
    memcpy(iv_tem, iv_enc, AES_BLOCK_SIZE);

    //test only
    // print_hex("IV", iv_enc, AES_BLOCK_SIZE);
    // std::cout <<std::endl;

    //对明文进行PKCS#7 padding
    unsigned char padded_plaintext[BUFFER_SIZE];
    int padded_plaintext_len = apply_pkcs7_padding((unsigned char *) plaintext, plaintext_len, padded_plaintext,
                                                   AES_BLOCK_SIZE);

    //加密
    AES_cbc_encrypt(padded_plaintext, encrypted, padded_plaintext_len, &aes_key, iv_tem, AES_ENCRYPT);

    //test only
    // print_hex("Encrypted", encrypted, padded_plaintext_len);
    // std::cout <<std::endl;

    //将iv添加到密文结尾
    unsigned char out[padded_plaintext_len + AES_BLOCK_SIZE];
    iv_append(iv_enc, encrypted, padded_plaintext_len, out);

    //TODO:可能会有内存泄漏的问题，要改
    memcpy(encrypted, out, padded_plaintext_len + AES_BLOCK_SIZE);

    //test only
    // print_hex("Appended ", out, padded_plaintext_len + AES_BLOCK_SIZE);
    // std::cout <<std::endl;

    return padded_plaintext_len;
}

//解密
int aes_decrypt(const unsigned char *ciphertext, int ciphertext_len, const unsigned char *key, unsigned char *decrypted,
                int padded_plaintext_len) {

    //检查输入合法性
    if (ciphertext == NULL || key == NULL || decrypted == NULL) {
        return -1;
    }

    //设置AES key
    AES_KEY aes_key;
    if (AES_set_decrypt_key(key, 128, &aes_key) < 0) {
        return -2;
    }

    //获取iv
    unsigned char iv_dec[AES_BLOCK_SIZE];
    unsigned char out[ciphertext_len];
    iv_extract(iv_dec, ciphertext, ciphertext_len);

    //test only
    // print_hex("Encrypted_dec ", ciphertext, ciphertext_len + AES_BLOCK_SIZE);
    // std::cout <<std::endl;
    // print_hex("IV_from_decrypt ", iv_dec, AES_BLOCK_SIZE);
    // std::cout <<std::endl;

    //恢复密文长度
    memcpy(out, ciphertext, ciphertext_len);

    // print_hex("ciphertext ", out, ciphertext_len);
    // std::cout <<std::endl;

    //解密
    AES_cbc_encrypt(out, decrypted, ciphertext_len, &aes_key, iv_dec, AES_DECRYPT);

    //去除padding
    int final_len = remove_pkcs7_padding(decrypted, padded_plaintext_len);
    if (final_len < 0) {
        std::cout << "Invalid padding found!" << std::endl;
        return -3;
    }

    //没看懂要返回什么，先返回decrypted长度
    return final_len;
}