/*
 * Copyright (c) 2006-2021, RT-Thread Development Team
 *
 * SPDX-License-Identifier: Apache-2.0
 *
 * Change Logs:
 * Date           Author       Notes
 * 2023-04-24     murmur       the first version
 */

#include <aes/aes.h>
#include <rtthread.h>
//#include <dfs_posix.h>
#include <dfs_file.h>

#define LOG_TAG "cryp"
#define LOG_LVL  LOG_LVL_DBG
#include <ulog.h>

#ifndef KEY_LEN
#define KEY_LEN 32//密钥长度
#endif

size_t aes_256_cbc_pkcs7(rt_uint8_t *data, size_t len, rt_uint8_t *dout)
{
//  AES-256
//  工作模式 CBC ，填充模式 PKCS7，不同语言要保持一致。
//   PKCS7规则是长度不够时缺几位补几个几，长度够时也要补
//   key长度由KEY_LEN定义，默认32位
//   iv取密钥前16位
//  已多平台验证

    unsigned char key[KEY_LEN];
    rt_memset(key, 0x00, KEY_LEN);    //初始化密钥，填充0x00
    rt_strcpy(key, "Cssc722.Cssc722.Cssc722..");    //写入密钥

    unsigned char iv[16];
    rt_memcpy(iv, key, 16);    //取密钥前16为偏移量

    rt_uint16_t m = len / 16;
    rt_uint8_t n = len % 16;
    rt_uint16_t cnt = (m + 1) * 16;

    rt_uint8_t *pbuff = rt_malloc(cnt);    //分配内存

    rt_memcpy(pbuff, data, len);    //data -> buff
    rt_memset(pbuff + len, (16 - n), 16 - n);    //按PKCS7规则填充
    unsigned char edata[cnt];//输出buff
    mbedtls_aes_context ctx;    //创建结构体
    mbedtls_aes_init(&ctx);    //初始化
    mbedtls_aes_setkey_enc(&ctx, key, KEY_LEN * 8);    //设置密钥
//    rt_kprintf("--%d\n", rst);
    mbedtls_aes_crypt_cbc(&ctx, MBEDTLS_AES_ENCRYPT, cnt, iv, pbuff, edata);    //加密
//    rt_kprintf("--%d\n", rst);
    mbedtls_aes_free(&ctx);    //释放结构体

    rt_memcpy(dout, edata, cnt);
    rt_free(pbuff);    //释放内存

    return cnt;

}

size_t aes_256_cbc_pkcs7_de(rt_uint8_t *data, size_t len, rt_uint8_t *dout)
{
//  AES-256
//  工作模式 CBC ，填充模式 PKCS7，不同语言要保持一致。
//   PKCS7规则是长度不够时缺几位补几个几，长度够时也要补
//   key长度由KEY_LEN定义，默认32位
//   iv取密钥前16位
//  已多平台验证

    unsigned char key[KEY_LEN];
    rt_memset(key, 0x00, KEY_LEN);    //初始化密钥，填充0x00
    rt_strcpy(key, "Cssc722.Cssc722.Cssc722..");    //写入密钥

    unsigned char iv[16];
    rt_memcpy(iv, key, 16);    //取密钥前16为偏移量

    unsigned char edata[len];//输出buff
    rt_memset(edata,0xff,len);

    mbedtls_aes_context ctx;    //创建结构体
    mbedtls_aes_init(&ctx);    //初始化
    mbedtls_aes_setkey_dec(&ctx, key, KEY_LEN * 8);    //设置密钥
    mbedtls_aes_crypt_cbc(&ctx, MBEDTLS_AES_DECRYPT, len, iv, data, edata);    //加密
    mbedtls_aes_free(&ctx);    //释放结构体

//    LOG_HEX("de",16,edata,len);
    uint8_t v= edata[len-1];
    if (v==0 || !(v<len)) {//判断解密是否成功，依据PKCS7填充模式，解密数据最后按
        return 0;
    }

    uint8_t t[v];
    size_t rst = len-v;
    rt_memset(t,v,v);
    if (memcmp(edata+rst,t,v) != 0) {//依据PKCS7填充模式，正确解密的数据
        return 0;
    }

    edata[rst]='\0';
    rt_memcpy(dout, edata, rst+1);

    return rst;

}

/**
 * 解密数据
 * @param din 待解密数据
 * @param len 待解密数据长度
 * @param dout 解密后数据
 * @return 解密后数据长度
 */
size_t decryp_data(const uint8_t *din, size_t len, uint8_t *dout)
{
    return aes_256_cbc_pkcs7_de(din, len, dout);
}

/**
 * 加密数据
 * @param din 待加密数据
 * @param len 待加密数据长度
 * @param dout 加密后数据
 * @return 加密后数据长度，是16的整数倍
 */
size_t cryp_data(const uint8_t *din, int len, uint8_t *dout)
{
    return aes_256_cbc_pkcs7(din, len, dout);
}
void aes_string_test(int argc, char **argv)
{

    rt_uint8_t data[128];
    size_t len;
    rt_uint8_t str[100];
    if (argc == 1)
    {
        len=cryp_data("Test String", strlen("Test String"), data);
    }
    else if (argc == 2)
    {
        len=cryp_data(argv[1], strlen(argv[1]), data);
    }
    LOG_HEX("cry",16,data,len);
    len = decryp_data(data, len, str);
    rt_kprintf("decrypt str is '%s'\n", str,len);
//    rt_kprintf("\n---DONE.---\n");

}
/**
 *  将文件内容采用AES-256方式加密
 * @param fin 待加密的文件名称字符串
 * @param fout 加密后文件名称字符串
 */
void aes_file(const char *fin, const char *fout)
{
    rt_uint8_t *buffer = RT_NULL, *outbuffer = RT_NULL;
    int fd_in = -1, fd_out = -1;

    fd_in = open(fin, O_RDONLY, 0);
    if (fd_in < 0)
    {
        LOG_E("open the input file : %s error!\n", fin);
        goto _exit;
    }

    rt_uint16_t file_size = lseek(fd_in, 0, SEEK_END);
    lseek(fd_in, 0, SEEK_SET);

    buffer = (rt_uint8_t *) malloc(file_size);
    outbuffer = (rt_uint8_t *) malloc(file_size + 16);    //加密后最大大16字节
    if (!buffer || !outbuffer)
    {
        LOG_E("No memory for AES!\n");
        goto _exit;
    }

    read(fd_in, buffer, file_size);
    close(fd_in);//读取完毕后及时关闭，使支持fin=fout覆盖写入

    size_t len = aes_256_cbc_pkcs7(buffer, file_size, outbuffer);
    fd_out = open(fout, O_WRONLY | O_CREAT | O_TRUNC, 0);
    if (fd_out < 0)
    {
        LOG_E("open the output file : %s error!\n", fout);
        goto _exit;
    }
    write(fd_out, outbuffer, len);

//    LOG_I("AESed to %s done. File size from %d bytes to %d.\n", fout,file_size, len);

//    goto _exit;

    _exit: if (buffer)
    {
        rt_free(buffer);
    }
    if (outbuffer)
    {
        rt_free(outbuffer);
    }
    if (fd_in >= 0)
    {
        close(fd_in);
    }

    if (fd_out >= 0)
    {
        close(fd_out);
    }

//    return RT_EOK;
}

void crypt_file_test(int argc, char **argv)
{

    if (argc == 3)
    {
        aes_file(argv[1], argv[2]);
    }
    else
    {
        rt_kprintf("Usage:\n");
        rt_kprintf("aes_file_test  [input_file] [output_file]        \"input_file\" to \"output_file\" \n");

    }

}

/**
 *  将文件内容采用AES-256方式加密
 * @param fin 待加密的文件名称字符串
 * @param fout 加密后文件名称字符串
 */
void deaes_file(const char *fin, const char *fout)
{
    rt_uint8_t *buffer = RT_NULL, *outbuffer = RT_NULL;
    int fd_in = -1, fd_out = -1;

    fd_in = open(fin, O_RDONLY, 0);
    if (fd_in < 0)
    {
        LOG_E("open the input file : %s error!\n", fin);
        goto _exit;
    }

    rt_uint16_t file_size = lseek(fd_in, 0, SEEK_END);
    lseek(fd_in, 0, SEEK_SET);

    buffer = (rt_uint8_t *) malloc(file_size);
    outbuffer = (rt_uint8_t *) malloc(file_size);    //加密后最大大16字节
    if (!buffer || !outbuffer)
    {
        LOG_E("No memory for AES!\n");
        goto _exit;
    }

    read(fd_in, buffer, file_size);
    close(fd_in);//读取完毕后及时关闭，使支持fin=fout覆盖写入

    size_t len = aes_256_cbc_pkcs7_de(buffer, file_size, outbuffer);
    fd_out = open(fout, O_WRONLY | O_CREAT | O_TRUNC, 0);
    if (fd_out < 0)
    {
        LOG_E("open the output file : %s error!\n", fout);
        goto _exit;
    }
    write(fd_out, outbuffer, len);

    LOG_I("deAESed to %s done. File size from %d bytes to %d.\n", fout,file_size, len);

//    goto _exit;

    _exit: if (buffer)
    {
        rt_free(buffer);
    }
    if (outbuffer)
    {
        rt_free(outbuffer);
    }
    if (fd_in >= 0)
    {
        close(fd_in);
    }

    if (fd_out >= 0)
    {
        close(fd_out);
    }

//    return RT_EOK;
}
void decrypt_file_test(int argc, char **argv)
{

    if (argc == 3)
    {
        deaes_file(argv[1], argv[2]);
    }
    else
    {
        rt_kprintf("Usage:\n");
        rt_kprintf("deaes_file_test  [input_file] [output_file]        \"input_file\" to \"output_file\" \n");

    }

}

#include <finsh.h>
/* 导出到自动初始化 */
MSH_CMD_EXPORT(aes_string_test, crypt string using AES-256.);
MSH_CMD_EXPORT(crypt_file_test, crypt file using AES-256.);
MSH_CMD_EXPORT(decrypt_file_test, decrypt file using AES-256.);