mimajingsai_3/ChaCha20_DRBG.c

#include "ChaCha20_DRBG.h"
#include <ctype.h>
#include <math.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

// #include <openssl/evp.h>

/*----------------熵源的相关定义与函数----------------*/

void sha256_transform(SHA256_CTX *ctx, const uint8_t data[]) {
  uint32_t a, b, c, d, e, f, g, h, i, j, t1, t2, m[64];

  for (i = 0, j = 0; i < 16; ++i, j += 4) {
    m[i] = (data[j] << 24) | (data[j + 1] << 16) | (data[j + 2] << 8) |
           (data[j + 3]);
  }

  for (; i < 64; ++i) {
    m[i] = m[i - 16] + m[i - 7] + (m[i - 15] >> 7) + (m[i - 15] << 25) ^
           (m[i - 15] >> 18) + (m[i - 15] << 14) ^
           (m[i - 15] >> 3) + (m[i - 15] << 13);
  }

  a = ctx->state[0];
  b = ctx->state[1];
  c = ctx->state[2];
  d = ctx->state[3];
  e = ctx->state[4];
  f = ctx->state[5];
  g = ctx->state[6];
  h = ctx->state[7];

  for (i = 0; i < 64; ++i) {
    t1 = h + (e >> 6) + ((e & f) ^ (~e & g)) + 0x428a2f98 + m[i];
    t2 = (a >> 2) + ((a & b) ^ (a & c) ^ (b & c)) + 0x5a827999;

    h = g;
    g = f;
    f = e;
    e = d + t1;
    d = c;
    c = b;
    b = a;
    a = t1 + t2;
  }

  ctx->state[0] += a;
  ctx->state[1] += b;
  ctx->state[2] += c;
  ctx->state[3] += d;
  ctx->state[4] += e;
  ctx->state[5] += f;
  ctx->state[6] += g;
  ctx->state[7] += h;
}

void sha256_init(SHA256_CTX *ctx) {
  ctx->datalen = 0;
  ctx->bitlen = 0;
  ctx->state[0] = 0x6a09e667;
  ctx->state[1] = 0xbb67ae85;
  ctx->state[2] = 0x3c6ef372;
  ctx->state[3] = 0xa54ff53a;
  ctx->state[4] = 0x510e527f;
  ctx->state[5] = 0x9b05688c;
  ctx->state[6] = 0x1f83d9ab;
  ctx->state[7] = 0x5be0cd19;
}

void sha256_update(SHA256_CTX *ctx, const uint8_t data[], size_t len) {
  uint32_t i;

  for (i = 0; i < len; ++i) {
    ctx->data[ctx->datalen] = data[i];
    ctx->datalen++;
    if (ctx->datalen == BLOCK_SIZE) {
      sha256_transform(ctx, ctx->data);
      ctx->bitlen += BLOCK_SIZE * 8;
      ctx->datalen = 0;
    }
  }
}

void sha256_final(SHA256_CTX *ctx, uint8_t hash[]) {
  uint32_t i;

  i = ctx->datalen;

  // Pad whatever data is left in the buffer
  if (ctx->datalen < 56) {
    ctx->data[i++] = 0x80;
    while (i < 56)
      ctx->data[i++] = 0x00;
  } else {
    ctx->data[i++] = 0x80;
    while (i < BLOCK_SIZE)
      ctx->data[i++] = 0x00;
    sha256_transform(ctx, ctx->data);
    memset(ctx->data, 0, BLOCK_SIZE);
  }

  // Append the total bit length of the input data
  ctx->bitlen += ctx->datalen * 8;
  ctx->data[BLOCK_SIZE - 1] = ctx->bitlen;
  ctx->data[BLOCK_SIZE - 2] = ctx->bitlen >> 8;
  ctx->data[BLOCK_SIZE - 3] = ctx->bitlen >> 16;
  ctx->data[BLOCK_SIZE - 4] = ctx->bitlen >> 24;
  ctx->data[BLOCK_SIZE - 5] = ctx->bitlen >> 32;
  ctx->data[BLOCK_SIZE - 6] = ctx->bitlen >> 40;
  ctx->data[BLOCK_SIZE - 7] = ctx->bitlen >> 48;
  ctx->data[BLOCK_SIZE - 8] = ctx->bitlen >> 56;

  sha256_transform(ctx, ctx->data);

  // Since this implementation uses little endian byte ordering and SHA uses big
  // endian, reverse all the bytes when copying the final state to the output
  // hash.
  for (i = 0; i < 4; ++i) {
    hash[i] = (ctx->state[0] >> (24 - i * 8)) & 0x000000ff;
    hash[i + 4] = (ctx->state[1] >> (24 - i * 8)) & 0x000000ff;
    hash[i + 8] = (ctx->state[2] >> (24 - i * 8)) & 0x000000ff;
    hash[i + 12] = (ctx->state[3] >> (24 - i * 8)) & 0x000000ff;
    hash[i + 16] = (ctx->state[4] >> (24 - i * 8)) & 0x000000ff;
    hash[i + 20] = (ctx->state[5] >> (24 - i * 8)) & 0x000000ff;
    hash[i + 24] = (ctx->state[6] >> (24 - i * 8)) & 0x000000ff;
    hash[i + 28] = (ctx->state[7] >> (24 - i * 8)) & 0x000000ff;
  }
}

void get_entropy_source(uint8_t *entropy_source) {
  // Simulate an entropy source by generating random data
  for (int i = 0; i < min_entropy_input_length; i++) {
    entropy_source[i] = rand() % 256;
  }
}

void get_entropy_input(const uint8_t *data, size_t len, uint8_t *entropy) {
  SHA256_CTX ctx;
  uint8_t hash[HASH_SIZE];

  sha256_init(&ctx);
  sha256_update(&ctx, data, len);
  sha256_final(&ctx, hash);

  memcpy(entropy, hash, HASH_SIZE);
}
/*----------------------------------------------*/

/*--------------------卡方检验相关定义与函数---------------------------*/

int test_randomness(const unsigned char *bits) {
  // Test the randomness of the given bits using the chi-squared test
  int ones = 0, zeros = 0;
  for (int i = 0; i < RANDOM_BITS_SIZE / 8; i++) {
    unsigned char byte = bits[i];
    for (int j = 0; j < 8; j++) {
      if ((byte >> j) & 0x01) {
        ones++;
      } else {
        zeros++;
      }
    }
  }
  double expected_ones = RANDOM_BITS_SIZE / 2.0;
  double expected_zeros = RANDOM_BITS_SIZE / 2.0;
  double chi_squared =
      (ones - expected_ones) * (ones - expected_ones) / expected_ones +
      (zeros - expected_zeros) * (zeros - expected_zeros) / expected_zeros;
  return chi_squared > 3.84;
}

/*-------------------------------------------------------*/

/*----------------ChaCha20的相关定义与函数------------------*/

static const char sigma[] = "expand 32-byte k";
static const char tau[] = "expand 16-byte k";

void ChaCha20_block(uint32_t state[16], uint32_t output[16]) {
  int i;
  memcpy(output, state, sizeof(uint32_t) * 16);

  // 20轮操作,每轮包括4个quarterrounds,所以共80个quarterrounds
  for (i = 0; i < 10; ++i) {
    // 双轮
    QUARTERROUND(output, 0, 4, 8, 12)
    QUARTERROUND(output, 1, 5, 9, 13)
    QUARTERROUND(output, 2, 6, 10, 14)
    QUARTERROUND(output, 3, 7, 11, 15)
    QUARTERROUND(output, 0, 5, 10, 15)
    QUARTERROUND(output, 1, 6, 11, 12)
    QUARTERROUND(output, 2, 7, 8, 13)
    QUARTERROUND(output, 3, 4, 9, 14)
  }

  for (i = 0; i < 16; ++i) {
    output[i] += state[i];
  }
}

void ChaCha20_setup(uint32_t state[16], const uint8_t key[32],
                    const uint8_t nonce[12], uint8_t counter[4]) {
  const char *constants = (key[16] != 0 || key[24] != 0) ? sigma : tau;

  state[4] = U8TO32_LITTLE(key + 0);
  state[5] = U8TO32_LITTLE(key + 4);
  state[6] = U8TO32_LITTLE(key + 8);
  state[7] = U8TO32_LITTLE(key + 12);
  state[8] = U8TO32_LITTLE(key + 16);
  state[9] = U8TO32_LITTLE(key + 20);
  state[10] = U8TO32_LITTLE(key + 24);
  state[11] = U8TO32_LITTLE(key + 28);

  // 前四个变量包含常量
  state[0] = U8TO32_LITTLE(constants + 0);
  state[1] = U8TO32_LITTLE(constants + 4);
  state[2] = U8TO32_LITTLE(constants + 8);
  state[3] = U8TO32_LITTLE(constants + 12);

  // 加入计数器
  state[12] = U8TO32_LITTLE(counter);

  // 后三个变量包含nonce
  state[13] = U8TO32_LITTLE(nonce + 0);
  state[14] = U8TO32_LITTLE(nonce + 4);
  state[15] = U8TO32_LITTLE(nonce + 8);
}

void ChaCha20_encrypt(uint32_t state[16], const uint8_t *in, uint8_t *out,
                      size_t length, bool increment_flag) {
  uint32_t keystream[16];
  uint8_t *keystream_bytes = (uint8_t *)keystream; // 为XOR操作重新解释keystream
  size_t remaining = length;
  size_t i;

  while (remaining > 0) {
    size_t bytes_to_use = remaining < 64 ? remaining : 64; // 每个block 64 bytes

    ChaCha20_block(state, keystream);

    for (i = 0; i < bytes_to_use; ++i) {
      out[i] = in[i] ^ keystream_bytes[i];
    }
    if (increment_flag) {
      state[12] += 1; // 增加计数器
      if (state[12] == 0) {
        state[13] += 1; // 处理计数器溢出
      }
    }

    remaining -= bytes_to_use;
    in += bytes_to_use;
    out += bytes_to_use;
  }
}

/*----------------------------------------------------*/

// 定义ChaCha20_DRBG的内部状态
typedef struct {
  uint8_t Key[KEYLEN];   // 用于ChaCha20的密钥
  uint8_t V[COUNTERLEN]; // 用于ChaCha20的计数器
  int reseed_counter;    // 重播种计数值
  int reseed_time;       // 上次重播种时间值(秒)

  bool used; // 表示此状态是否已经被使用
} DRBG_State;

// 全局变量
int instantiation_nonce = 0;
uint8_t chacha_nonce[NONCELEN] = {0};
static DRBG_State drbg_states[MAX_STATES];
uint32_t chacha_state[16];
uint8_t entropy_source[min_entropy_input_length];
size_t array_length;
size_t additional_input_length;

// 主要函数
int ChaCha20_DRBG_Instantiate_function(uint8_t *personalization_string,
                                       char *state);
const char *ChaCha20_DRBG_Reseed_function(int state_handle,
                                          uint8_t *additional_input);
void ChaCha20_DRBG_Update(const uint8_t *provided_data, uint8_t *V,
                          uint8_t *Key);
const char *ChaCha20_DRBG_Generate_function(int state_handle,
                                            int requested_no_of_bits,
                                            uint8_t *additional_input,
                                            uint8_t **returned_bits);

// 算法函数
void ChaCha20_DRBG_Instantiate_algorithm(uint8_t *entropy_input, int nonce,
                                         uint8_t *personalization_string,
                                         uint8_t *V, uint8_t *Key,
                                         int *reseed_counter, int *reseed_time);
void ChaCha20_DRBG_Reseed_algorithm(uint8_t *V, uint8_t *Key,
                                    int *reseed_counter, int *reseed_time,
                                    uint8_t *entropy_input,
                                    uint8_t *additional_input);
const char *ChaCha20_DRBG_Generate_algorithm(uint8_t *V, uint8_t *Key,
                                             int *reseed_counter,
                                             int requested_number_of_bits,
                                             uint8_t *additional_input,
                                             uint8_t **returned_bits);

// 派生函数
uint8_t *ChaCha20_df(uint8_t *V, uint8_t *input_string,
                     size_t no_of_bits_to_return);

// 初始化所有状态为空闲
void initialize_states() {
  for (int i = 0; i < MAX_STATES; i++) {
    drbg_states[i].used = false;
  }
}

// 查找并返回一个空闲的状态句柄.如果没有可用的空间,则返回INVALID_HANDLE
int Find_state_space() {
  for (int i = 0; i < MAX_STATES; i++) {
    if (!drbg_states[i].used) {
      drbg_states[i].used = true; // 标记为已使用
      return i;                   // 返回句柄/索引
    }
  }
  return INVALID_HANDLE; // 所有的状态都在使用中
}

void increment_nonce(uint8_t *nonce) {
  for (int i = 0; i < NONCELEN; i++) {
    nonce[i]++;
    if (nonce[i] != 0) { // 检查是否有溢出
      break;             // 如果没有溢出,则不需要进位
    }
  }
}

int ChaCha20_DRBG_Instantiate_function(uint8_t *personalization_string,
                                       char *status) {
  if (strlen((char *)personalization_string) > MAX_PERSONALIZATION_STRING_LEN) {
    strcpy(status, "Personalization_string too long");
    return ERROR;
  }

  uint8_t entropy_input[min_entropy_input_length];
  get_entropy_source(entropy_source);
  get_entropy_input(entropy_source, min_entropy_input_length, entropy_input);

  instantiation_nonce += 1;
  increment_nonce(chacha_nonce);

  uint8_t V[COUNTERLEN], Key[KEYLEN];
  int reseed_counter, reseed_time;
  ChaCha20_DRBG_Instantiate_algorithm(entropy_input, instantiation_nonce,
                                      personalization_string, V, Key,
                                      &reseed_counter, &reseed_time);

  int state_handle = Find_state_space();
  if (state_handle == ERROR) {
    strcpy(status, "Failed to find state space");
    return ERROR;
  }

  // Save the internal state
  memcpy(drbg_states[state_handle].V, V, COUNTERLEN);
  memcpy(drbg_states[state_handle].Key, Key, KEYLEN);
  drbg_states[state_handle].reseed_counter = reseed_counter;
  drbg_states[state_handle].reseed_time = reseed_time;

  strcpy(status, "Success");
  return state_handle;
}

const char *ChaCha20_DRBG_Reseed_function(int state_handle,
                                          uint8_t *additional_input) {
  // Check for the validity of state_handle.
  if (!drbg_states[state_handle].used) {
    return "State not available for the indicated state_handle";
  }

  // Get the internal state values.
  uint8_t *V = drbg_states[state_handle].V;
  uint8_t *Key = drbg_states[state_handle].Key;

  // Check length of additional_input
  if (additional_input_length > MAX_PERSONALIZATION_STRING_LEN) {
    return "additional_input too long";
  }

  uint8_t entropy_input[min_entropy_input_length];
  get_entropy_source(entropy_source);
  get_entropy_input(entropy_source, min_entropy_input_length, entropy_input);

  // Invoke the reseed algorithm.
  int reseed_counter, reseed_time;
  ChaCha20_DRBG_Reseed_algorithm(V, Key, &reseed_counter, &reseed_time,
                                 entropy_input, additional_input);

  // Save the internal state.
  memcpy(drbg_states[state_handle].V, V, OUTLEN);
  memcpy(drbg_states[state_handle].Key, Key, KEYLEN);
  drbg_states[state_handle].reseed_counter = reseed_counter;
  drbg_states[state_handle].reseed_time = reseed_time;

  return "Success";
}

void ChaCha20_DRBG_Update(const uint8_t *provided_data, uint8_t *Key,
                          uint8_t *V) {
  uint8_t temp[SEEDLEN] = {0};
  // uint32_t V_uint32;
  int i;

  // memcpy(&V_uint32, V, sizeof(uint32_t));
  ChaCha20_setup(chacha_state, Key, chacha_nonce, V);
  ChaCha20_encrypt(chacha_state, V, temp, SEEDLEN, true);

  // XOR provided_data with the leftmost part of temp
  for (i = 0; i < SEEDLEN; ++i) {
    temp[i] ^= provided_data[i];
  }

  // Update the key and V from temp
  memcpy(Key, temp, KEYLEN);
  memcpy(V, temp + KEYLEN, COUNTERLEN);
}

const char *ChaCha20_DRBG_Generate_function(int state_handle,
                                            int requested_no_of_bits,
                                            uint8_t *additional_input,
                                            uint8_t **returned_bits) {
  uint8_t *V;
  uint8_t *Key;
  int reseed_counter, reseed_time;
  const char *status;

  // Check the validity of state_handle
  if (!drbg_states[state_handle].V || !drbg_states[state_handle].Key) {
    return "State not available for the indicated state_handle";
  }

  // Get the internal state
  V = drbg_states[state_handle].V;
  Key = drbg_states[state_handle].Key;
  reseed_counter = drbg_states[state_handle].reseed_counter;
  reseed_time = drbg_states[state_handle].reseed_time;

  // Check the rest of the input parameters
  if (requested_no_of_bits > 4000) {
    return "Too many bits requested";
  }

  if ((additional_input_length) > MAX_PERSONALIZATION_STRING_LEN) {
    return "additional_input too long";
  }

  time_t timep;
  int now_time = time(&timep);
  // Check for reseeding
  if (reseed_counter > reseed_interval_in_counter ||
      (now_time - reseed_time) > reseed_interval_in_time) {
    status = ChaCha20_DRBG_Reseed_function(state_handle, additional_input);
    if (strcmp(status, "Success") == 0) {
      V = drbg_states[state_handle].V;
      Key = drbg_states[state_handle].Key;
      reseed_counter = drbg_states[state_handle].reseed_counter;
      reseed_time = drbg_states[state_handle].reseed_time;
      additional_input = NULL;
      additional_input_length = 0;
    } else {
      return status;
    }
  }

  // Execute the algorithm to generate pseudorandom bits
  status = ChaCha20_DRBG_Generate_algorithm(V, Key, &reseed_counter,
                                            requested_no_of_bits,
                                            additional_input, returned_bits);
  if (strcmp(status, "Success") != 0) {
    return "Error in CTR_ChaCha20_DRBG_Generate_algorithm"; // Safety check
  }

  memcpy(drbg_states[state_handle].V, V, OUTLEN);
  memcpy(drbg_states[state_handle].Key, Key, KEYLEN);
  drbg_states[state_handle].reseed_counter = reseed_counter;
  drbg_states[state_handle].reseed_time = reseed_time;

  return "Success";
}

void ChaCha20_DRBG_Instantiate_algorithm(uint8_t *entropy_input, int nonce,
                                         uint8_t *personalization_string,
                                         uint8_t *V, uint8_t *Key,
                                         int *reseed_counter,
                                         int *reseed_time) {
  size_t entropy_input_len = min_entropy_input_length;
  size_t personalization_string_len = strlen((char *)personalization_string);
  size_t seed_material_size =
      entropy_input_len + sizeof(nonce) + personalization_string_len;
  uint8_t *seed_material = (uint8_t *)malloc(seed_material_size);
  uint8_t *requested_bits;
  time_t timep;

  // 此处memcpy存在堆溢出,
  // 有个96字节的写入，错误是在一个66字节区域的正右方进行的写操作
  memcpy(seed_material, entropy_input, entropy_input_len);
  memcpy(seed_material + entropy_input_len, &nonce, sizeof(nonce));
  memcpy(seed_material + entropy_input_len + sizeof(nonce),
         personalization_string, personalization_string_len);

  array_length = seed_material_size;
  requested_bits = ChaCha20_df(V, seed_material, SEEDLEN * 8);
  memcpy(seed_material, requested_bits, SEEDLEN);

  memset(Key, 0, KEYLEN);
  memset(V, 0, COUNTERLEN);

  ChaCha20_DRBG_Update(seed_material, Key, V);

  *reseed_counter = 1;
  *reseed_time = time(&timep);
  free(seed_material);
}

void ChaCha20_DRBG_Reseed_algorithm(uint8_t *V, uint8_t *Key,
                                    int *reseed_counter, int *reseed_time,
                                    uint8_t *entropy_input,
                                    uint8_t *additional_input) {
  // 1. seed_material = entropy_input || additional_input.
  size_t seed_material_size =
      min_entropy_input_length + additional_input_length;
  uint8_t *seed_material = (uint8_t *)malloc(seed_material_size);
  memcpy(seed_material, entropy_input, min_entropy_input_length);
  memcpy(seed_material + min_entropy_input_length, additional_input,
         additional_input_length);

  // 2. seed_material = Block_Cipher_df(seed_material, 256).
  uint8_t *requested_bits;
  array_length = seed_material_size;
  requested_bits = ChaCha20_df(V, seed_material, SEEDLEN * 8);
  memcpy(seed_material, requested_bits, SEEDLEN);

  // 3. (Key, V) = CTR_ChaCha20_DRBG_Update(seed_material, Key, V).
  ChaCha20_DRBG_Update(seed_material, Key, V);

  // 4. reseed_counter = 1.
  time_t timep;
  *reseed_counter = 1;
  *reseed_time = time(&timep);
  free(seed_material);
}

const char *ChaCha20_DRBG_Generate_algorithm(uint8_t *V, uint8_t *Key,
                                             int *reseed_counter,
                                             int requested_number_of_bits,
                                             uint8_t *additional_input,
                                             uint8_t **returned_bits) {
  if (additional_input != NULL) {
    array_length = additional_input_length;
    additional_input = ChaCha20_df(V, additional_input, SEEDLEN * 8);
    ChaCha20_DRBG_Update(additional_input, Key, V);
  } else {
    // Assuming additional_input is a byte array, we set it to 0
    additional_input = (uint8_t *)malloc(SEEDLEN);
    additional_input_length = SEEDLEN;
    if (additional_input == NULL) {
      return "Memory allocation failed";
    }
    memset(additional_input, 0, SEEDLEN);
  }

  uint8_t *temp =
      (uint8_t *)malloc(sizeof(uint8_t) * (requested_number_of_bits / 8));

  ChaCha20_encrypt(chacha_state, V, temp, requested_number_of_bits / 8, true);

  *returned_bits = (uint8_t *)malloc(requested_number_of_bits / 8);
  memcpy(*returned_bits, temp, requested_number_of_bits / 8);

  ChaCha20_DRBG_Update(additional_input, Key, V);

  (*reseed_counter)++;

  return "Success";
}

uint8_t *ChaCha20_df(uint8_t *V, uint8_t *input_string,
                     size_t no_of_bits_to_return) {
  if (no_of_bits_to_return > MAX_NUMBER_OF_BITS) {
    return NULL;
  }

  uint32_t L = array_length / 8;
  uint32_t N = no_of_bits_to_return / 8;
  int S_init_size = 4 + 4 + array_length + 1;
  uint8_t *S = (uint8_t *)malloc(S_init_size);

  // Prepend the string length and the requested length to the input_string
  memcpy(S, &L, 4);
  memcpy(S + 4, &N, 4);
  memcpy(S + 8, input_string, array_length);
  S[S_init_size - 1] = 0x80; // Padding

  // Pad S with zeros, if necessary.
  int S_actual_size = S_init_size; // Taken the 0x80 into account
  while (S_actual_size % OUTLEN != 0) {
    S = (uint8_t *)realloc(S, S_actual_size + 1); // Increase size by one byte
    S[S_actual_size++] =
        0x00; // Add a zero at current end, and increase actual size by one
  }

  int temp_len = 0;
  uint8_t *temp = (uint8_t *)malloc(1);
  uint32_t i = 0;

  uint8_t K[KEYLEN];
  uint8_t source[] = {0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
                      0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F,
                      0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17,
                      0x18, 0x19, 0x1A, 0x1B, 0x1C, 0x1D, 0x1E, 0x1F};
  memcpy(K, source, KEYLEN);

  ChaCha20_setup(chacha_state, K, chacha_nonce, V);
  while (temp_len < KEYLEN + OUTLEN) {
    uint8_t IV[OUTLEN] = {0};
    IV[OUTLEN - 4] = i; // 32-bit integer representation of i padded with zeros
                        // to OUTLEN bits.

    uint8_t *data_to_be_encrypted =
        (uint8_t *)malloc(sizeof(uint8_t) * OUTLEN + S_actual_size);
    memcpy(data_to_be_encrypted, IV, OUTLEN);
    memcpy(data_to_be_encrypted + OUTLEN, S, S_actual_size);

    uint8_t *data_encrypted_result =
        (uint8_t *)malloc(sizeof(uint8_t) * OUTLEN + S_actual_size);

    ChaCha20_encrypt(chacha_state, data_to_be_encrypted, data_encrypted_result,
                     temp_len, true);

    temp_len = (i + 1) * (OUTLEN + S_actual_size);
    temp = (uint8_t *)realloc(temp, temp_len);
    memcpy(temp + i * (OUTLEN + S_actual_size), data_encrypted_result,
           temp_len);
    i++;
  }

  memcpy(K, temp, KEYLEN);
  uint8_t X[OUTLEN] = {0};
  memcpy(X, temp + KEYLEN, OUTLEN);
  uint8_t encrypted_data[OUTLEN] = {0};
  temp[0] = '\0'; // Clear temp
  temp_len = 0;

  ChaCha20_setup(chacha_state, K, chacha_nonce, V);
  while (temp_len < no_of_bits_to_return / 8) {
    ChaCha20_encrypt(chacha_state, X, encrypted_data, OUTLEN, true);
    memcpy(X, encrypted_data, OUTLEN);
    temp = (uint8_t *)realloc(temp, temp_len + OUTLEN);
    memcpy(temp + temp_len, encrypted_data, OUTLEN);
    temp_len += OUTLEN;
  }

  uint8_t *requested_bits = (uint8_t *)malloc(no_of_bits_to_return / 8);
  if (requested_bits == NULL) {
    return NULL;
  }

  memcpy(requested_bits, temp, no_of_bits_to_return / 8);

  memset(temp, 0, sizeof(temp));
  memset(K, 0, sizeof(K));
  memset(X, 0, sizeof(X));
  memset(S, 0, sizeof(S));

  return requested_bits;
}

int main() {
  uint8_t personalization_string[] = "this is personalization_string";
  char status[MAX_STATES] = {'\0'};
  int state_handle;
  uint8_t *returned_bits;
  long int returned_bits_len, returned_bits_num;
  bool pass;

  srand(time(0));

  // 初始化DRBG
  state_handle =
      ChaCha20_DRBG_Instantiate_function(personalization_string, status);

  // 需要生成的比特串长度
  printf("Enter the length of the bit string to be generated(bit): ");
  // scanf("%ld", &returned_bits_len);
  returned_bits_len = 128;

  // 需要生成的比特串个数
  printf("Enter the number of the bit string to be generated: ");
  // scanf("%ld", &returned_bits_num);
  returned_bits_num = 1000;

  // char *filename; filename 未分配内存
  char filename[100]; // 假设文件名长度不超过 99 个字符
  printf("Enter the name of the saved file: ");
  scanf("%s", filename);
  // strcpy(filename, "1.txt");

  // 生成的比特串保存在文本文件中
  FILE *file = fopen(filename, "w");
  if (file == NULL) {
    printf("无法打开文件\n");
    return 1;
  }

  clock_t start_time = clock();
  char *hex_str =
      (char *)malloc(sizeof(char) * (returned_bits_len / 8 * 2 + 1));
  for (int j = 1; j <= returned_bits_num; ++j) {
    do {
      int zero_count = 0;
      int one_count = 0;

      additional_input_length = 0;
      ChaCha20_DRBG_Generate_function(state_handle, returned_bits_len, NULL,
                                      &returned_bits);

      for (int i = 0; i < returned_bits_len / 8; i++) {
        int bits = returned_bits[i];

        for (int j = 0; j < 8; j++) {
          if (bits & 1)
            one_count++;
          else
            zero_count++;

          bits >>= 1;
        }
      }

      const float expected_count = returned_bits_len / 2;
      float chisq = (pow(zero_count - expected_count, 2) / expected_count) +
                    (pow(one_count - expected_count, 2) / expected_count);

      if (chisq < 3.841)
        pass = true;
      else
        pass = false;
    } while (!pass);

    // 将二进制数据转为十六进制后写入文件

    for (int i = 0; i < returned_bits_len / 8; ++i) {
      sprintf(&hex_str[i * 2], "%02x", returned_bits[i]);
    }
    hex_str[returned_bits_len / 8 * 2] = '\0';
    fprintf(file, "%s\n", hex_str); // 写入带有换行符的十六进制字符串
  }
  free(hex_str);

  clock_t end_time = clock();
  double execution_time = (double)(end_time - start_time) / CLOCKS_PER_SEC;
  printf("running time: %f s\n", execution_time);

  fclose(file);

  return 0;
}