/*
 * 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

int aes_128_cbc_pkcs7(rt_uint8_t *data, rt_uint16_t len, rt_uint8_t *dout)
{

//  工作模式 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_kprintf("%d -- %d %d new len is %d.\n", len,m,n,cnt);

    rt_memcpy(pbuff, data, len);    //data -> buff
    rt_memset(pbuff + len, (16 - n), 16 - n);    //按PKCS7规则填充

//    for (size_t i = 0; i < (m + 1) * 16; i++)
//    {
//        rt_kprintf("0x%02X ", p_key[i]);
//    }

    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, (m + 1) * 16, iv, pbuff, edata);    //加密
//    rt_kprintf("--%d\n", rst);
    mbedtls_aes_free(&ctx);    //释放结构体

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

//    for (rt_uint16_t i = 0; i < cnt; i++)
//    {
//        rt_kprintf("0x%02X ", edata[i]);
//    }

//    rt_kprintf("\n---DONE.---\n");

    return cnt;

}
void aes_string_test(int argc, char **argv)
{
    rt_uint8_t data[128];
    size_t len;
    if (argc == 1)
    {
        len=aes_128_cbc_pkcs7("Test String", strlen("Test String"), data);
    }
    else if (argc == 2)
    {
        len=aes_128_cbc_pkcs7(argv[1], strlen(argv[1]), data);
    }
    for (size_t var = 0; var < len; ++var) {
        rt_kprintf("0x%02X ", data[var]);
    }
    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;
    }

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

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

    size_t len = aes_128_cbc_pkcs7(buffer, file_size, outbuffer);
    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 aes_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");

    }

}

#include <finsh.h>
/* 导出到自动初始化 */
MSH_CMD_EXPORT(aes_string_test, 使用AES-256加密字符串。crypt string using AES-256.);
MSH_CMD_EXPORT(aes_file_test, 使用AES-256加密文件。crypt file using AES-256.);