//
// Created by 甜甜 on 2020/12/2.
//


#include "LRUCache.hpp"

#ifdef ENCLAVE

#include "encryption.h"
extern "C" {
#include "aes/taes.h"
}
#include <sgx_tcrypto.h>
// #if !defined(ENCLAVE)
//     #include "LogBase.h"
// #endif

#include <string.h>
#include <memory>
#include "t.h"
#include <assert.h>
using namespace std;
// using namespace util;
// 根据数据字节数 获取加密字节数

int encryption::get_crypt_size(int data_len)
{
    return ((data_len / AES_BLOCKLEN) + ((data_len % AES_BLOCKLEN) == 0 ? 0 : 1)) * AES_BLOCKLEN;
}

// 加密
int encryption::aes_encrypt(char *in, char *key, char *out, size_t inlen) //, int olen)
{
    if (NULL == in || NULL == key || NULL == out)
        return -1;


    // if(inlen % AES_BLOCKLEN != 0) {
    //     Log("aes encrypt: in size % AES_BLOCK != 0", log::error)
    // }

    // assert(inlen % AES_BLOCKLEN == 0);

    struct AES_ctx ctx;
    AES_init_ctx(&ctx, (const uint8_t *)key);
    int en_len = 0;

    memcpy_s(out, inlen, in, inlen);
    while (en_len < inlen) //输入输出字符串够长。而且是AES_BLOCK_SIZE的整数倍，须要严格限制
    {
        AES_ECB_encrypt(&ctx, (unsigned char *)out);
        out += AES_BLOCKLEN;
        en_len += AES_BLOCKLEN;
    }
    return get_crypt_size(inlen);
}

// 解密
int encryption::aes_decrypt(char *in, int inlen, char *key, char *out)
{
    if (NULL == in || NULL == key || NULL == out)
        return -1;
    // if(inlen % AES_BLOCKLEN != 0) {
    //     Log("aes decrypt: in size % AES_BLOCK != 0", log::error)
    // }
    // assert(inlen % AES_BLOCKLEN == 0);
    struct AES_ctx ctx;
    AES_init_ctx(&ctx, (const uint8_t *)key);
    int en_len = 0;

    memcpy_s(out, inlen, in, inlen);
    while (en_len < inlen) //输入输出字符串够长。而且是AES_BLOCK_SIZE的整数倍，须要严格限制
    {
        AES_ECB_decrypt(&ctx, (unsigned char *)out);
        out += AES_BLOCKLEN;
        en_len += AES_BLOCKLEN;
    }
    return get_crypt_size(inlen);
}


int encryption::aeshmac_encrypt(char *in, int inlen, char *key, char *out) {
    if (NULL == in || NULL == key || NULL == out)
        return -1;
    // 1. AES加密
    int cipher_len = aes_encrypt(in, key, out, inlen);
    // 2.HMAC
    int HMAC_LEN = 32;
    int keylen = 16;
    uint8_t mac[HMAC_LEN];
    sgx_status_t ret = sgx_hmac_sha256_msg(
        (uint8_t*)out, cipher_len,
        (uint8_t*)key, keylen,
        mac, HMAC_LEN
    );
    memcpy(out + cipher_len, mac, HMAC_LEN);
    return cipher_len + HMAC_LEN;
}

int encryption::aeshmac_decrypt(char *in, int inlen, char *key, char *out) {
    if (NULL == in || NULL == key || NULL == out || inlen < 32)
        return -1;
    int HMAC_LEN = 32;
    int cipher_len = inlen - HMAC_LEN;
    int keylen = 16;
    // 1. 校验HMAC
    uint8_t mac[HMAC_LEN];
    sgx_status_t ret = sgx_hmac_sha256_msg(
        (uint8_t*)in, cipher_len,
        (uint8_t*)key, keylen,
        mac, HMAC_LEN
    );
    if (ret != SGX_SUCCESS) return -1;
    if (memcmp(mac, in + cipher_len, 32) != 0) {
        return -1; // HMAC校验失败
    }
    // 2. AES解密
    return aes_decrypt(in, cipher_len, key, out);
}

long long encryption::charToLongLong(char b[])
{

    return *((long long *)b);
}

//! not used
char *encryption::DET_encrypt(char *m, size_t *cipher_length)
{
}

char *encryption::DET_decrypt(char *ciphertext, std::size_t *cipher_len)
{
}

#else

#include "encryption.h"
#include <stdio.h>
#include <fstream>
#include <ostream>
#include <string.h>
#include <memory>
#include <iostream>
#include <cstring>
#include <openssl/rand.h>
#include <openssl/aes.h>
#include <openssl/hmac.h>
#include <openssl/evp.h>
#include <stdlib.h>
#include "t.h"

using namespace std;

//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);
}

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);

    //对明文进行PKCS#7 padding
    unsigned char *padded_plaintext = (unsigned char *) malloc(plaintext_len + AES_BLOCK_SIZE);
    if (padded_plaintext == NULL) {
        abort();
        // return -3; // 无法分配内存
    }
    int padded_plaintext_len = apply_pkcs7_padding((unsigned char *) plaintext, plaintext_len, padded_plaintext,
                                                   AES_BLOCK_SIZE);

//	print_hex("test ", plaintext, plaintext_len);

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

    //将iv添加到密文结尾
    unsigned char *out = (unsigned char *) malloc(padded_plaintext_len + AES_BLOCK_SIZE);
    if (out == NULL) {
        // free(padded_plaintext);
        abort();
        // return -4; // 无法分配内存
    }
    iv_append(iv_enc, encrypted, padded_plaintext_len, out);
    memcpy(encrypted, out, padded_plaintext_len + AES_BLOCK_SIZE);

    //释放动态分配的内存
    free(padded_plaintext);
    free(out);

    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) {
        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 - AES_BLOCK_SIZE);

    //test only
//    print_hex("Encrypted_dec ", ciphertext, ciphertext_len);
//    print_hex("IV_from_decrypt ", iv_dec, AES_BLOCK_SIZE);

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

//    print_hex("ciphertext ", out, ciphertext_len);

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

//	print_hex("decrypted ", decrypted, ciphertext_len - AES_BLOCK_SIZE);

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

//	print_hex("decrypted 1", decrypted, final_len);

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

#ifdef USE_LRU
LRUCache<long, std::string> AESCache;
#endif

// 根据数据字节数 获取加密字节数
int encryption::get_crypt_size(int data_len) {
    return ((data_len / AES_BLOCK_SIZE) + ((data_len % AES_BLOCK_SIZE) == 0 ? 0 : 1)) * AES_BLOCK_SIZE;
}

// 加密
int encryption::aes_encrypt(char *in, char *key, char *out, size_t inlen) //, int olen)
{
    if (NULL == in || NULL == key || NULL == out)
        return -1;

    AES_KEY aes;
    if (AES_set_encrypt_key((unsigned char *) key, 128, &aes) < 0) {
        return -2;
    }

    //int len=strlen(in);
    int en_len = 0;
    while (en_len < inlen) //输入输出字符串够长。而且是AES_BLOCK_SIZE的整数倍，须要严格限制
    {
        AES_encrypt((unsigned char *) in, (unsigned char *) out, &aes);
        in += AES_BLOCK_SIZE;
        out += AES_BLOCK_SIZE;
        en_len += AES_BLOCK_SIZE;
    }

    return get_crypt_size(inlen);
}

// AES+HMAC加密
int encryption::aeshmac_encrypt(char *in, int inlen, char *key, char *out) {
    if (NULL == in || NULL == key || NULL == out)
        return -1;
    // 1. AES加密
    int cipher_len = aes_encrypt(in, key, out, inlen);
    if (cipher_len < 0) return -1;
    // 2. HMAC
    int keylen = 16;
    int HMAC_LEN = 32;
    unsigned char mac[HMAC_LEN];
    unsigned int mac_len = 0;
    HMAC(EVP_sha256(), key, keylen, (unsigned char*)out, cipher_len, mac, &mac_len);
    // 3. 拼接
    memcpy(out + cipher_len, mac, HMAC_LEN);
    return cipher_len + HMAC_LEN;
}

// 解密
int encryption::aes_decrypt(char *in, int in_len, char *key, char *out) {
#ifdef USE_LRU
    string instr(in, in_len);
    auto hashv = hash<string>()(instr);
    auto tmpstr = AESCache.get(hashv, "");

    if(!tmpstr.empty()) {
        memcpy(out, tmpstr.c_str(), tmpstr.size());
        return tmpstr.size();
    }

    char* outbase = out;
#endif
    if (NULL == in || NULL == key || NULL == out)
        return -1;

    AES_KEY aes;
    if (AES_set_decrypt_key((unsigned char *) key, 128, &aes) < 0) {
        return -2;
    }

    int en_len = 0;
    while (en_len < in_len) {
        AES_decrypt((unsigned char *) in, (unsigned char *) out, &aes);
        in += AES_BLOCK_SIZE;
        out += AES_BLOCK_SIZE;
        en_len += AES_BLOCK_SIZE;
    }
	//fixed by qxy
	memset(out,0,AES_BLOCK_SIZE);
	en_len += AES_BLOCK_SIZE;

#ifdef USE_LRU
    AESCache.put(hashv, string(outbase, en_len));
#endif
    return en_len;
}

// AES+HMAC解密
int encryption::aeshmac_decrypt(char *in, int inlen, char *key, char *out) {
    if (NULL == in || NULL == key || NULL == out || inlen < 32)
        return -1;
    int HMAC_LEN = 32;
    int cipher_len = inlen - 32;
    // 1. 校验HMAC
    int keylen = 16;
    unsigned char mac[HMAC_LEN];
    unsigned int mac_len = 0;
    HMAC(EVP_sha256(), key, keylen, (unsigned char*)in, cipher_len, mac, &mac_len);
    if (mac_len != 32) return -1;
    if (memcmp(mac, in + cipher_len, HMAC_LEN) != 0) {
        return -1; // HMAC校验失败
    }
    // 2. AES解密
    return aes_decrypt(in, cipher_len, key, out);
}

long long encryption::charToLongLong(char b[]) {
    // long long result=0L;
    // for(int i=0;i<sizeof(long long);i++)
    // {
    //     result<<=
    // }
    return *((long long *) b);
}

//! not used
char *encryption::DET_encrypt(char *m, size_t *cipher_length) {
    //encryption encrypt;
    t change;
    char buf[100];
    memset(buf, 1, sizeof(buf));
    strcpy((char *) buf, m);
    size_t inlen;
    //std::// plain value is :" << buf << std::endl;

    inlen = strlen(buf);

    char aes_keybuf[32];
    memset(aes_keybuf, 0, sizeof(aes_keybuf));

    ifstream infile;
    infile.open("SAHE_key.txt");
    if (!infile.is_open())
        abort();
    while (!infile.eof()) {
        infile.getline(aes_keybuf, 32);
    }
    infile.close();

    char *aes_en = new char[BUF_SIZE]();
    //char aes_en[BUF_SIZE] = {0};

    int aes_en_len = aes_encrypt(buf, aes_keybuf, aes_en, inlen);
    *cipher_length = aes_en_len;

    // aes-128 加密
    //    if (aes_en_len < 0)
    //    {
    //        printf("aes_encrypt error ret: %d\n", aes_en_len);
    //        return -1;
    //    }

    // size_t bufsz1 = strlen(aes_en) + 1;
    // size_t hexsz1 = 2 * (strlen(aes_en) + 1);
    // char* h = new char[hexsz1 * 2];
    // change.buf2hex((uint8_t*)aes_en, strlen(aes_en), h, &hexsz1);

    return aes_en;
}

char *encryption::DET_decrypt(char *ciphertext, std::size_t *cipher_len) {
    encryption encrypt;
    t change;
    // size_t bufsz2 = sizeof(ciphertext) + 1;
    // size_t hexsz2 = 2 * (strlen(ciphertext) + 1);
    // char* nc2 = new char[hexsz2];
    // change.hex2buf((uint8_t*)nc2, &hexsz2, ciphertext, hexsz2);
    // cout<<"cipher1_parse: "<<nc2<<endl;
    char *aes_de = new char[BUF_SIZE]();
    //char aes_de[BUF_SIZE] = {0};
    char aes_keybuf[32];
    memset(aes_keybuf, 0, sizeof(aes_keybuf));
    //memset(aes_keybuf, 0, sizeof(aes_keybuf));

    ifstream infile2;
    infile2.open("SAHE_key.txt");
    if (!infile2.is_open())
        // SAHE_KEY file failure" << endl;
        while (!infile2.eof()) {
            infile2.getline(aes_keybuf, 32);
        }
    infile2.close();

    //strcpy((char*)aes_keybuf, CRYPT_USER_KEY);

    //const   char*   p2   =   (const   char*)(char*)aes_de;
    int aes_de_len = encrypt.aes_decrypt(ciphertext, strlen(ciphertext), aes_keybuf, aes_de); // aes-128 解密
    *cipher_len = aes_de_len;

    // printf("aes_de_len: %d aes_de: %s\n", strlen(aes_de), aes_de);
    // " << strlen(aes_de) << endl;

    return aes_de;
}

#endif
