#include <stdlib.h>
#include <time.h>
#include <cstring>
#include <stdio.h>
#include <cstdint>
#include <fcntl.h>
#include <unistd.h>
#include <string>

#include "lzo/minilzo.h"
#include "zstd/zstd.h"
#include "ucl/ucl.h"
#include "lz4/lz4.h"

#include "cipher/aes.h"
#include "cipher/des.h"
#include "cipher/des3.h"
#include "cipher/rc4.h"
#include "cipher/rsa.h"
#include "cipher/ecc.h"

using namespace std;

#ifdef __cplusplus
extern "C" {
#endif
#include "lzma/FileIO.h"
#include "lzma/TinyLzmaCompress.h"
#include "lzma/TinyLzmaDecompress.h"
#ifdef __cplusplus
}
#endif

#define HEAP_ALLOC(var, size) \
    lzo_align_t __LZO_MMODEL var [ ((size) + (sizeof(lzo_align_t) - 1)) / sizeof(lzo_align_t) ]

// 定义各个字段的偏移量
#define SYMMETRIC_ENC_TYP_OFFSET    0
#define ASYMMETRIC_ENC_TYP_OFFSET   1
#define COMP_TYP_OFFSET             2
#define INITIAL_FILE_SIZE_OFFSET    3
#define COMPRESSED_FILE_SIZE_OFFSET 7
#define ENCRYPTED_FILE_SIZE_OFFSET  11
#define ENCRYPTED_KEY_SIZE_OFFSET   15
#define PROTECT_TYPE_OFFSET         19
#define MAC_NUMS_OFFSET             20
#define RUN_ADD_OFFSET              21
/**
 *
 * 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
 */

static HEAP_ALLOC(wrkmem, LZO1X_1_MEM_COMPRESS);

const int MAGIC_LENGTH = 32;

static int get_random_bytes(unsigned char *buf, size_t len) {
//  FILE *f;
//  f = fopen("/dev/urandom", "r");
//  fread(buf, len, 1, f);
//  fclose(f);
  srand((unsigned)time(NULL));
  for(int i = 0; i < len; i++) {
        buf[i] = rand() % 255;
  }
  return 0;
}

int hexStringToByteArray(const char *hexString, unsigned char *byteArray, int byteArraySize) {
  int count = 0;
  const char *pos = hexString;

  // 转换16进制字符串为字节数组
  while (count < byteArraySize && sscanf(pos, "%2hhx", &byteArray[count]) == 1) {
    pos += 2;
    count++;
  }
  return count; // 返回转换的字节数
}

void replaceWithMiddleEight(unsigned char arr[], int length) {
    if (length < 8) {
        printf("数组长度不足8位\n");
        return;
    }

    // 计算中间8位的起始索引
    int start = (length - 8) / 2;

    // 提取中间的8个元素
    int middleEight[8];
    for (int i = 0; i < 8; i++) {
        middleEight[i] = arr[start + i];
    }

    // 用中间的8位替换源数组的内容
    for (int i = 0; i < length; i++) {
        arr[i] = middleEight[i % 8]; // 循环使用中间的8位
    }
}

int main(int argc, char *argv[]) {
  FILE *input_file;

  char *input_file_path = argv[1];
  uint8_t symmetric_enc_typ = (uint8_t) atoi(argv[2]);
  uint8_t asymmetric_enc_typ = (uint8_t) atoi(argv[3]);
  uint8_t comp_typ = (uint8_t) atoi(argv[4]);
  
  uint8_t cp_symmetric_enc_typ = symmetric_enc_typ;
  uint8_t cp_asymmetric_enc_typ = asymmetric_enc_typ;
  uint8_t cp_comp_typ = comp_typ;
  
  char *output_file_path = argv[5];
  // 保护模式
  uint8_t protect_type = (uint8_t) atoi(argv[6]);
  char *puf_path = argv[7];

  uint8_t mac_enable = (uint8_t) atoi(argv[8]);
  char *mac_path = argv[9];

  // 运行地址 0xb0000000
  char *str_run_add = argv[10];
//  printf("str_run_add %s\n", str_run_add);
  unsigned int run_add = strtoull(str_run_add, NULL, 16);

  printf("input_file_path %s\n", input_file_path);
  printf("symmetric_enc_typ %u\n", symmetric_enc_typ);
  printf("asymmetric_enc_typ %u\n", asymmetric_enc_typ);
  printf("comp_typ %u\n", comp_typ);
  printf("output_file_path %s\n", output_file_path);
  printf("protect type %u\n", protect_type);
  printf("puf_path %s\n", puf_path);
  printf("mac_enable %u\n", mac_enable);
  printf("mac_path %s\n", mac_path);


  input_file = fopen(input_file_path, "rb");
  if (input_file == NULL) {
    printf("unable to open file.\n");
    return 1;
  }

  // 获取文件大小
  fseek(input_file, 0, SEEK_END);
  uint32_t initial_file_size = ftell(input_file);
  printf("initial file size:%u\n", initial_file_size);
  rewind(input_file);

  /**
   * DES: 64bit  1 key_length 8
   * AES: 128bit 2 key_length 32
   * 3DES: 64bit 3 key_length 24
   * RC4: 256bit 4 key_length 32
   */
   
  
  int symmetric_map[5] = {0, 4, 1, 3, 2};
  symmetric_enc_typ = symmetric_map[symmetric_enc_typ];
  
  uint8_t input_length, key_length;
  if (symmetric_enc_typ == 1) input_length = 8, key_length = 8;
  else if (symmetric_enc_typ == 2) input_length = 16, key_length = 32;
  else if (symmetric_enc_typ == 3) input_length = 8, key_length = 24;
  else if (symmetric_enc_typ == 4) input_length = 1, key_length = 32;
  else input_length = 0, key_length = 0;

  // 分配足够的内存来存储文件内容
  uint8_t *byte_array = (uint8_t *) malloc(initial_file_size * sizeof(uint8_t));

  if (byte_array == NULL) {
    printf("unable to allocate memory.\n");
    return 0;
  }

  // 读取文件内容到字节数组
  int num_elements = fread(byte_array, sizeof(uint8_t), initial_file_size, input_file);
  if (num_elements != initial_file_size) {
    printf("an error occurred while reading the file.\n");
    return 0;
  }
  fclose(input_file);

  const int serial_length = 39;
  unsigned char puf_key[serial_length];
  memset(puf_key, 0, serial_length);
  unsigned char puf_value[serial_length+2];
  memset(puf_value, 0, serial_length);
  if (protect_type == 1) {
    printf("[STATE] node:1 ; message:PUF交互\n");
    FILE *file = fopen(puf_path, "r");
    if (file == NULL) {
      perror("Failed to open file");
      return EXIT_FAILURE;
    }
    char line1[256];
    char line2[256];
    char line3[256];

    // 读取第一行
    if (fgets(line1, sizeof(line1), file) == NULL) {
      perror("Failed to read line 1");
      fclose(file);
      return EXIT_FAILURE;
    }

    // 读取第二行
    if (fgets(line2, sizeof(line2), file) == NULL) {
      perror("Failed to read line 2");
      fclose(file);
      return EXIT_FAILURE;
    }

    // 读取第三行
    if (fgets(line3, sizeof(line3), file) == NULL) {
      perror("Failed to read line 3");
      fclose(file);
      return EXIT_FAILURE;
    }
    // 打印读取的内容
    printf("Line 1: %s", line1);
    printf("Line 2: %s", line2);
    printf("Line 3: %s\n", line3);

    // 关闭文件
    fclose(file);
    // 转换字符串为字节数组
    int convertedCount2 = hexStringToByteArray(line2, puf_key, serial_length);
    int convertedCount3 = hexStringToByteArray(line3, puf_value, serial_length);

    for (int i = 0; i < 39; i++) {
      printf("0x%02x,", puf_value[i]);
    }
    printf("\n");

    // 检查转换的结果
    if (convertedCount2 != serial_length || convertedCount3 != serial_length) {
      printf("Conversion error\n");
      return EXIT_FAILURE;
    }
  }

  printf("[STATE] node:2 ; message:获取密钥\n");
  uint8_t module_key[32] = {0};
  if (protect_type == 1){
      replaceWithMiddleEight(puf_value, 41);
      memcpy(module_key, puf_value + 4, 32);
  }
  else {
      get_random_bytes(module_key, 32);
  }
  printf("module key:\n");
  for (int i = 0; i < 32; i++) {
    printf("%02x ", module_key[i]);
  }
  printf("\n");

  uint8_t mac_nums = 0;
  unsigned char mac_array[10][18];
  memset(mac_array, 0, 180);
  if (mac_enable == 1) {
    // 从本地文件中读取MAC地址
    int mac_fd = open(mac_path, O_RDONLY);
    if (mac_fd <= 0) {
      printf("mac_fd : %d\n", mac_fd);
      printf("本地未找到MAC地址列表文件\n");
      return -1;
    }
    while (mac_nums < 10 && (read(mac_fd, mac_array[mac_nums], 18) > 0)) {
      mac_nums++;
    }
    close(mac_fd);
  } else {
    mac_array[0][0] = 'P';
  }

  int comp_map[6] = {0, 3, 1, 2, 4, 5};
  comp_typ = comp_map[comp_typ];

  unsigned long compressed_file_size;
  uint8_t *compress_out = (uint8_t *) calloc(initial_file_size, sizeof(uint8_t));
  int r = -1;
  if (comp_typ == 1) {
      printf("compress with lzo...\n");
    r = lzo1x_1_compress(byte_array, initial_file_size, compress_out, reinterpret_cast<lzo_uint *>(&compressed_file_size), wrkmem);
  } else if (comp_typ == 2) {
      printf("compress with ucl...\n");
      int level = 2;
    r = ucl_nrv2b_99_compress(byte_array, initial_file_size, compress_out,
                              reinterpret_cast<ucl_uint *>(&compressed_file_size), NULL, level, NULL, NULL);
    compressed_file_size = reinterpret_cast<unsigned long>(compressed_file_size);
  } else if (comp_typ == 3) {
      printf("compress with lzma...\n");
      compressed_file_size = initial_file_size;
    r = tinyLzmaCompress(byte_array, initial_file_size, compress_out,
                         reinterpret_cast<size_t *>(&compressed_file_size));
    compressed_file_size = reinterpret_cast<unsigned long>(compressed_file_size);

  } else if (comp_typ == 4) {
      printf("compress with zstd...\n");
      compressed_file_size = ZSTD_compressBound(initial_file_size);
    printf("%d\n", compressed_file_size);
    compressed_file_size = ZSTD_compress(compress_out, compressed_file_size, byte_array, initial_file_size, 2);
  } else if (comp_typ == 5) {
      printf("compress with lz4...\n");
      int max_dst_size = LZ4_compressBound(initial_file_size);
    compressed_file_size = LZ4_compress_default(byte_array, compress_out, initial_file_size, max_dst_size);
  }
  if (r == 0 || r == -1) {
    printf("compressed %u bytes into %u bytes\n", initial_file_size, compressed_file_size);
    printf("[STATE] node: 6 ; message:整体压缩\n");
  } else {
    /* this should NEVER happen */
    printf("internal error - compression failed: %d\n", r);
    return 2;
  }
  /* check for an incompressible block */
  if ((unsigned int) compressed_file_size >= initial_file_size) {
    printf("This block contains incompressible data.\n");
    return 0;
  }

  uint8_t *new_compress_out = (uint8_t *) realloc(compress_out, compressed_file_size);
  if (new_compress_out != NULL) {
    compress_out = new_compress_out;
  }


  uint8_t key[key_length];
  for (int i = 0; i < key_length; i++) key[i] = module_key[i];

  printf("key:\n");
  for (int i = 0; i < key_length; i++) {
    printf("%02x ", key[i]);
  }
  printf("\n");

  uint32_t encrypted_file_size = (compressed_file_size + input_length - 1) / input_length * input_length;
  printf("encrypted_file_size %u\n", encrypted_file_size);

  printf("compress bytes:\n");
  for (int i = 0; i < 8; i++) {
    printf("%02x ", compress_out[i]);
  }
  printf("\n");

  uint8_t *buff = (uint8_t *) malloc(encrypted_file_size * sizeof(uint8_t));

  if (symmetric_enc_typ == 1) {
    printf("enc_typ = 1\n");
    DesContext des_context;
    desInit(&des_context, key, key_length);
    for (int i = 0; i < encrypted_file_size; i += input_length)
      desEncryptBlock(&des_context, compress_out + i, buff + i);
  }
  if (symmetric_enc_typ == 2) {
    printf("enc_typ = 2\n");
    AesContext aes_context;
    aesInit(&aes_context, key, key_length);
    for (int i = 0; i < encrypted_file_size; i += input_length)
      aesEncryptBlock(&aes_context, compress_out + i, buff + i);
  }
  if (symmetric_enc_typ == 3) {
    printf("enc_typ = 3\n");
    Des3Context des3_context;
    des3Init(&des3_context, key, key_length);
    for (int i = 0; i < encrypted_file_size; i += input_length)
      des3EncryptBlock(&des3_context, compress_out + i, buff + i);
  }
  if (symmetric_enc_typ == 4) {
    printf("enc_typ = 4\n");
    rc4_encryption(compress_out, encrypted_file_size, key, key_length);
    memcpy(buff, compress_out, encrypted_file_size);
  }

  printf("final bytes:\n");
  for (int i = 0; i < 8; i++) {
    printf("%02x ", buff[i]);
  }
  printf("\n");

  for (int i = encrypted_file_size - 8; i < encrypted_file_size; i++) {
    printf("%02x ", buff[i]);
  }
  printf("\n");

  printf("[STATE] node: 4 ; message:整体加密\n");


//  uint8_t *tmp_encrypted = NULL;
  uint8_t *encrypted = NULL;
//  uint32_t encrypted_key_size = sizeof key * sizeof(long long);
  uint32_t encrypted_key_size = key_length;
  // 如果是软硬结合，加密的key为PUF激励长度39字节
  if (protect_type == 1) encrypted_key_size = 39;
  /**
   * rsa key
   */
  struct public_key_class pub[1];
  struct private_key_class priv[1];
  /**
   * ecc key
   */
  uint8_t pub_key[ECC_KEYSIZE];
  uint8_t prv_key[ECC_KEYSIZE];
  if (asymmetric_enc_typ != 0) {
    if (asymmetric_enc_typ == 5) {
      rsa_gen_keys(pub, priv);
      encrypted_key_size *= 8;
      printf("Private Key:\n Modulus: %lld\n Exponent: %lld\n", (long long) priv->modulus, (long long) priv->exponent);
      printf("Public Key:\n Modulus: %lld\n Exponent: %lld\n", (long long) pub->modulus, (long long) pub->exponent);


      // 1是软硬结合，加密的是激励
      // 0是软加密，加密的是秘钥
      if (protect_type == 1) encrypted = rsa_encrypt(puf_key, 39, pub);
      else if (protect_type == 0) encrypted = rsa_encrypt(key, key_length, pub);

//      encrypted = (uint8_t *) tmp_encrypted;
      if (encrypted == NULL) {
        fprintf(stderr, "error in encryption!\n");
        return 1;
      }

    } else if (asymmetric_enc_typ == 6) {
      printf("enc with ecc\n");
      ecc_init_keys(pub_key, prv_key);
      printf("ecc init end\n");

      // 1是软硬结合，加密的是激励，申请的空间是39
      // 0是软加密，加密的是秘钥，申请的空间是sizeof key
      if (protect_type == 0) {
        encrypted = (uint8_t *) realloc(encrypted, sizeof key);
        ecc_encrypt(pub_key, key, key_length, encrypted, &encrypted_key_size);
      } else if (protect_type == 1) {
        encrypted = (uint8_t *) realloc(encrypted, 39);
        ecc_encrypt(pub_key, puf_key, 39, encrypted, &encrypted_key_size);
      }
      printf("encrypted_key_size is %u with ecc\n", encrypted_key_size);
    }
  }

  printf("enc key\n");
  for (int i = 0; i < encrypted_key_size; i++) {
    printf("0x%02x, ", encrypted[i]);
  }
  printf("\n");


  /**
   * enc_typ              1B
   * comp_typ             1B
   * initial_file_size    4B
   * compressed_file_size 4B
   * encrypted_file_size  4B
   */
  printf("uint8_t symmetric_enc_typ: %u\n", symmetric_enc_typ);
  printf("uint8_t asymmetric_enc_typ: %u\n", asymmetric_enc_typ);
  printf("uint8_t comp_typ: %u\n", comp_typ);
  printf("uint32_t file_size: %u\n", initial_file_size);
  printf("uint32_t compressed_file_size: %u\n", compressed_file_size);
  printf("uint32_t encrypted_file_size: %u\n", encrypted_file_size);
  printf("uint8_t input_length: %u\n", input_length);
  printf("uint8_t key_length: %u\n", key_length);

  uint8_t magic[MAGIC_LENGTH];
  // 将各个字段的值存入 magic 数组
  magic[SYMMETRIC_ENC_TYP_OFFSET] = symmetric_enc_typ;
  magic[ASYMMETRIC_ENC_TYP_OFFSET] = asymmetric_enc_typ;
  magic[COMP_TYP_OFFSET] = comp_typ;
  memcpy(magic + INITIAL_FILE_SIZE_OFFSET, &initial_file_size, sizeof(uint32_t));
  memcpy(magic + COMPRESSED_FILE_SIZE_OFFSET, &compressed_file_size, sizeof(uint32_t));
  memcpy(magic + ENCRYPTED_FILE_SIZE_OFFSET, &encrypted_file_size, sizeof(uint32_t));
  memcpy(magic + ENCRYPTED_KEY_SIZE_OFFSET, &encrypted_key_size, sizeof(uint32_t));
  // 保护类型
  memcpy(magic + PROTECT_TYPE_OFFSET, &protect_type, sizeof(uint8_t));
  // MAC地址绑定的数量
  memcpy(magic + MAC_NUMS_OFFSET, &mac_nums, sizeof(uint8_t));
  // 运行地址
  printf("str_run_add %x\n", run_add);
  memcpy(magic + RUN_ADD_OFFSET, &run_add, 4);

  FILE *p = fopen(output_file_path, "wb");
  if (p == NULL) {
    printf("open filed\n");
  }

  fwrite(magic, sizeof(uint8_t), MAGIC_LENGTH, p);
  fwrite(buff, sizeof(uint8_t), encrypted_file_size, p);

  fwrite(encrypted, sizeof(uint8_t), encrypted_key_size, p);
  if (asymmetric_enc_typ == 5) {
    fwrite(pub, sizeof(struct public_key_class), 1, p);
    fwrite(priv, sizeof(struct private_key_class), 1, p);
  } else if (asymmetric_enc_typ == 6) {
    fwrite(pub_key, sizeof(uint8_t), ECC_KEYSIZE, p);
    fwrite(prv_key, sizeof(uint8_t), ECC_KEYSIZE, p);
  }

  for (int i = 0; i < mac_nums; i++) {
    for (int j = 0; j < 18; j++) {
      printf("%02x ", mac_array[i][j]);
    }
    printf("\n");
  }


  for (int i = 0; i < mac_nums; i++) fwrite(mac_array[i], sizeof(uint8_t), 18, p);

  printf("[STATE] node: 7 ; message:加壳\n");
  fclose(p);
  free(byte_array);
  free(buff);
  printf("[STATE] node: 8 ; message:成功\n");
  printf("成功\n");
  return 0;
}


/**
 * magic
 * encrypted_file
 * enc_key/enc serial_challenge
 * pub_key
 * prv_key
 * mac_info[0-10]
 */