/**
 *
 * \section COPYRIGHT
 *
 * Copyright 2013-2021 Software Radio Systems Limited
 *
 * By using this file, you agree to the terms and conditions set
 * forth in the LICENSE file which can be found at the top level of
 * the distribution.
 *
 */

#include "srsran/common/security.h"
#include "srsran/common/liblte_security.h"
#include "srsran/common/s3g.h"
#include "srsran/common/ssl.h"
#include "srsran/config.h"

#include <arpa/inet.h>

#ifdef HAVE_MBEDTLS
#include "mbedtls/md5.h"
#endif
#ifdef HAVE_POLARSSL
#include "polarssl/md5.h"
#endif

#define FC_EPS_K_ASME_DERIVATION 0x10
#define FC_EPS_K_ENB_DERIVATION 0x11
#define FC_EPS_NH_DERIVATION 0x12
#define FC_EPS_K_ENB_STAR_DERIVATION 0x13
#define FC_EPS_ALGORITHM_KEY_DERIVATION 0x15

#define ALGO_EPS_DISTINGUISHER_NAS_ENC_ALG 0x01
#define ALGO_EPS_DISTINGUISHER_NAS_INT_ALG 0x02
#define ALGO_EPS_DISTINGUISHER_RRC_ENC_ALG 0x03
#define ALGO_EPS_DISTINGUISHER_RRC_INT_ALG 0x04
#define ALGO_EPS_DISTINGUISHER_UP_ENC_ALG 0x05
#define ALGO_EPS_DISTINGUISHER_UP_INT_ALG 0x06

#define FC_5G_ALGORITHM_KEY_DERIVATION 0x69
#define FC_5G_KAUSF_DERIVATION 0x6A
#define FC_5G_RES_STAR_DERIVATION 0x6B
#define FC_5G_KSEAF_DERIVATION 0x6C
#define FC_5G_KAMF_DERIVATION 0x6D
#define FC_5G_KGNB_KN3IWF_DERIVATION 0x6E
#define FC_5G_NH_GNB_DERIVATION 0x6F

#define ALGO_5G_DISTINGUISHER_NAS_ENC_ALG 0x01
#define ALGO_5G_DISTINGUISHER_NAS_INT_ALG 0x02
#define ALGO_5G_DISTINGUISHER_RRC_ENC_ALG 0x03
#define ALGO_5G_DISTINGUISHER_RRC_INT_ALG 0x04
#define ALGO_5G_DISTINGUISHER_UP_ENC_ALG 0x05
#define ALGO_5G_DISTINGUISHER_UP_INT_ALG 0x06
namespace srsran {

/******************************************************************************
 * Key Generation
 *****************************************************************************/

uint8_t
security_generate_k_asme(uint8_t* ck, uint8_t* ik, uint8_t* ak_xor_sqn_, uint16_t mcc, uint16_t mnc, uint8_t* k_asme)
{
  if (ck == NULL || ik == NULL || ak_xor_sqn_ == NULL || k_asme == NULL) {
    log_error("Invalid inputs");
    return SRSRAN_ERROR;
  }
  std::array<uint8_t, 32> key;

  // The input key Key shall be equal to the concatenation CK || IK of CK and IK.
  memcpy(key.data(), ck, 16);
  memcpy(key.data() + 16, ik, 16);

  // Serving Network id
  std::vector<uint8_t> sn_id;
  sn_id.resize(3);
  sn_id[0] = (mcc & 0x00F0) | ((mcc & 0x0F00) >> 8); // First byte of P0
  if ((mnc & 0xFF00) == 0xFF00) {
    // 2-digit MNC
    sn_id[1] = 0xF0 | (mcc & 0x000F);                         // Second byte of P0
    sn_id[2] = ((mnc & 0x000F) << 4) | ((mnc & 0x00F0) >> 4); // Third byte of P0
  } else {
    // 3-digit MNC
    sn_id[1] = ((mnc & 0x000F) << 4) | (mcc & 0x000F);   // Second byte of P0
    sn_id[2] = ((mnc & 0x00F0)) | ((mnc & 0x0F00) >> 8); // Third byte of P0
  }

  // AK XOR SQN
  std::vector<uint8_t> ak_xor_sqn;
  ak_xor_sqn.resize(AK_LEN);
  memcpy(ak_xor_sqn.data(), ak_xor_sqn_, ak_xor_sqn.size());

  uint8_t output[32];
  if (kdf_common(FC_EPS_K_ASME_DERIVATION, key, sn_id, ak_xor_sqn, output) != SRSRAN_SUCCESS) {
    log_error("Failed to run kdf_common");
    return SRSRAN_ERROR;
  }

  memcpy(k_asme, output, 32);
  return SRSRAN_SUCCESS;
}

uint8_t security_generate_k_ausf(uint8_t*    ck,
                                 uint8_t*    ik,
                                 uint8_t*    ak_xor_sqn_,
                                 const char* serving_network_name,
                                 uint8_t*    k_ausf)
{
  if (ck == NULL || ik == NULL || ak_xor_sqn_ == NULL || serving_network_name == NULL || k_ausf == NULL) {
    log_error("Invalid inputs");
    return SRSRAN_ERROR;
  }

  std::array<uint8_t, 32> key;

  // The input key Key shall be equal to the concatenation CK || IK of CK and IK.
  memcpy(key.data(), ck, 16);
  memcpy(key.data() + 16, ik, 16);

  // Serving Network Name
  std::vector<uint8_t> ssn;
  ssn.resize(strlen(serving_network_name));
  memcpy(ssn.data(), serving_network_name, ssn.size());

  // AK XOR SQN
  std::vector<uint8_t> ak_xor_sqn;
  ak_xor_sqn.resize(AK_LEN);
  memcpy(ak_xor_sqn.data(), ak_xor_sqn_, ak_xor_sqn.size());

  uint8_t output[32];
  if (kdf_common(FC_5G_RES_STAR_DERIVATION, key, ssn, ak_xor_sqn, output) != SRSRAN_SUCCESS) {
    log_error("Failed to run kdf_common");
    return SRSRAN_ERROR;
  }
  memcpy(k_ausf, output, 32);

  return SRSRAN_SUCCESS;
}

uint8_t security_generate_k_enb(uint8_t* k_asme, uint32_t nas_count_, uint8_t* k_enb)
{
  if (k_asme == NULL || k_enb == NULL) {
    log_error("Invalid inputs");
    return SRSRAN_ERROR;
  }
  std::array<uint8_t, 32> key;

  memcpy(key.data(), k_asme, 32);

  // NAS Count
  std::vector<uint8_t> nas_count;
  nas_count.resize(4);
  nas_count[0] = (nas_count_ >> 24) & 0xFF;
  nas_count[1] = (nas_count_ >> 16) & 0xFF;
  nas_count[2] = (nas_count_ >> 8) & 0xFF;
  nas_count[3] = nas_count_ & 0xFF;

  if (kdf_common(FC_EPS_K_ENB_DERIVATION, key, nas_count, k_enb) != SRSRAN_SUCCESS) {
    log_error("Failed to run kdf_common");
    return SRSRAN_ERROR;
  }
  return SRSRAN_SUCCESS;
}

uint8_t security_generate_k_enb_star(uint8_t* k_enb, uint32_t pci_, uint32_t earfcn_, uint8_t* k_enb_star)
{
  if (k_enb == NULL || k_enb_star == NULL) {
    log_error("Invalid inputs");
    return SRSRAN_ERROR;
  }
  std::array<uint8_t, 32> key;
  memcpy(key.data(), k_enb, 32);

  // PCI
  std::vector<uint8_t> pci;
  pci.resize(2);
  pci[0] = (pci_ >> 8) & 0xFF;
  pci[1] = pci_ & 0xFF;

  // EARFCN
  std::vector<uint8_t> earfcn;
  earfcn.resize(2);
  earfcn[0] = (earfcn_ >> 8) & 0xFF;
  earfcn[1] = earfcn_ & 0xFF;

  if (kdf_common(FC_EPS_K_ENB_STAR_DERIVATION, key, pci, earfcn, k_enb_star) != SRSRAN_SUCCESS) {
    log_error("Failed to run kdf_common");
    return SRSRAN_ERROR;
  }
  return SRSRAN_SUCCESS;
}

uint8_t security_generate_nh(uint8_t* k_asme, uint8_t* sync_, uint8_t* nh)
{
  if (k_asme == NULL || sync_ == NULL || nh == NULL) {
    log_error("Invalid inputs");
    return SRSRAN_ERROR;
  }
  std::array<uint8_t, 32> key;
  memcpy(key.data(), k_asme, 32);

  // PCI
  std::vector<uint8_t> sync;
  sync.resize(32);
  memcpy(sync.data(), sync_, 32);

  if (kdf_common(FC_EPS_NH_DERIVATION, key, sync, nh) != SRSRAN_SUCCESS) {
    log_error("Failed to run kdf_common");
    return SRSRAN_ERROR;
  }
  return SRSRAN_SUCCESS;
}

uint8_t security_generate_k_nas(uint8_t*                    k_asme,
                                CIPHERING_ALGORITHM_ID_ENUM enc_alg_id,
                                INTEGRITY_ALGORITHM_ID_ENUM int_alg_id,
                                uint8_t*                    k_nas_enc,
                                uint8_t*                    k_nas_int)
{
  if (k_asme == NULL || k_nas_enc == NULL || k_nas_int == NULL) {
    log_error("Invalid inputs");
    return SRSRAN_ERROR;
  }
  std::array<uint8_t, 32> key;

  memcpy(key.data(), k_asme, 32);

  // Derive NAS ENC
  // algorithm type distinguisher
  std::vector<uint8_t> algo_distinguisher;
  algo_distinguisher.resize(1);
  algo_distinguisher[0] = ALGO_EPS_DISTINGUISHER_NAS_ENC_ALG;

  // algorithm type distinguisher
  std::vector<uint8_t> algorithm_identity;
  algorithm_identity.resize(1);
  algorithm_identity[0] = enc_alg_id;

  if (kdf_common(FC_EPS_ALGORITHM_KEY_DERIVATION, key, algo_distinguisher, algorithm_identity, k_nas_enc) !=
      SRSRAN_SUCCESS) {
    log_error("Failed to run kdf_common");
    return SRSRAN_ERROR;
  }

  // Derive NAS INT
  // algorithm type distinguisher
  algo_distinguisher.resize(1);
  algo_distinguisher[0] = ALGO_EPS_DISTINGUISHER_NAS_INT_ALG;

  // algorithm type distinguisher
  algorithm_identity.resize(1);
  algorithm_identity[0] = int_alg_id;

  // Derive NAS int
  if (kdf_common(FC_EPS_ALGORITHM_KEY_DERIVATION, key, algo_distinguisher, algorithm_identity, k_nas_int) !=
      SRSRAN_SUCCESS) {
    log_error("Failed to run kdf_common");
    return SRSRAN_ERROR;
  }
  return SRSRAN_SUCCESS;
}

uint8_t security_generate_k_rrc(uint8_t*                    k_enb,
                                CIPHERING_ALGORITHM_ID_ENUM enc_alg_id,
                                INTEGRITY_ALGORITHM_ID_ENUM int_alg_id,
                                uint8_t*                    k_rrc_enc,
                                uint8_t*                    k_rrc_int)
{
  if (k_enb == NULL || k_rrc_enc == NULL || k_rrc_int == NULL) {
    log_error("Invalid inputs");
    return SRSRAN_ERROR;
  }
  std::array<uint8_t, 32> key;

  memcpy(key.data(), k_enb, 32);

  // Derive RRC ENC
  // algorithm type distinguisher
  std::vector<uint8_t> algo_distinguisher;
  algo_distinguisher.resize(1);
  algo_distinguisher[0] = ALGO_EPS_DISTINGUISHER_RRC_ENC_ALG;

  // algorithm type distinguisher
  std::vector<uint8_t> algorithm_identity;
  algorithm_identity.resize(1);
  algorithm_identity[0] = enc_alg_id;

  if (kdf_common(FC_EPS_ALGORITHM_KEY_DERIVATION, key, algo_distinguisher, algorithm_identity, k_rrc_enc) !=
      SRSRAN_SUCCESS) {
    log_error("Failed to run kdf_common");
    return SRSRAN_ERROR;
  }

  // Derive RRC INT
  // algorithm type distinguisher
  algo_distinguisher.resize(1);
  algo_distinguisher[0] = ALGO_EPS_DISTINGUISHER_RRC_INT_ALG;

  // algorithm type distinguisher
  algorithm_identity.resize(1);
  algorithm_identity[0] = int_alg_id;

  // Derive RRC int
  if (kdf_common(FC_EPS_ALGORITHM_KEY_DERIVATION, key, algo_distinguisher, algorithm_identity, k_rrc_int) !=
      SRSRAN_SUCCESS) {
    log_error("Failed to run kdf_common");
    return SRSRAN_ERROR;
  }
  return SRSRAN_SUCCESS;
}

uint8_t security_generate_k_up(uint8_t*                    k_enb,
                               CIPHERING_ALGORITHM_ID_ENUM enc_alg_id,
                               INTEGRITY_ALGORITHM_ID_ENUM int_alg_id,
                               uint8_t*                    k_up_enc,
                               uint8_t*                    k_up_int)
{
  if (k_enb == NULL || k_up_enc == NULL || k_up_int == NULL) {
    log_error("Invalid inputs");
    return SRSRAN_ERROR;
  }
  std::array<uint8_t, 32> key;

  memcpy(key.data(), k_enb, 32);

  // Derive UP ENC
  // algorithm type distinguisher
  std::vector<uint8_t> algo_distinguisher;
  algo_distinguisher.resize(1);
  algo_distinguisher[0] = ALGO_EPS_DISTINGUISHER_UP_ENC_ALG;

  // algorithm type distinguisher
  std::vector<uint8_t> algorithm_identity;
  algorithm_identity.resize(1);
  algorithm_identity[0] = enc_alg_id;

  if (kdf_common(FC_EPS_ALGORITHM_KEY_DERIVATION, key, algo_distinguisher, algorithm_identity, k_up_enc) !=
      SRSRAN_SUCCESS) {
    log_error("Failed to run kdf_common");
    return SRSRAN_ERROR;
  }

  // Derive UP INT
  // algorithm type distinguisher
  algo_distinguisher.resize(1);
  algo_distinguisher[0] = ALGO_EPS_DISTINGUISHER_UP_INT_ALG;

  // algorithm type distinguisher
  algorithm_identity.resize(1);
  algorithm_identity[0] = int_alg_id;

  // Derive UP int
  if (kdf_common(FC_EPS_ALGORITHM_KEY_DERIVATION, key, algo_distinguisher, algorithm_identity, k_up_int) !=
      SRSRAN_SUCCESS) {
    log_error("Failed to run kdf_common");
    return SRSRAN_ERROR;
  }
  return SRSRAN_SUCCESS;
}

uint8_t security_generate_sk_gnb(uint8_t* k_enb, uint8_t* sk_gnb, uint16_t scg_count_)
{
  if (k_enb == NULL || sk_gnb == NULL) {
    log_error("Invalid inputs");
    return SRSRAN_ERROR;
  }
  std::array<uint8_t, 32> key;

  memcpy(key.data(), k_enb, 32);

  // SCG Count
  std::vector<uint8_t> scg_count;
  scg_count.resize(2);
  scg_count[0] = (scg_count_ >> 8) & 0xFF; // first byte of P0
  scg_count[1] = scg_count_ & 0xFF;        // second byte of P0

  // Derive sk_gnb
  uint8_t output[32];
  if (kdf_common(0x1C, key, scg_count, output) != SRSRAN_SUCCESS) {
    log_error("Failed to run kdf_common");
    return SRSRAN_ERROR;
  }
  memcpy(sk_gnb, output, 32);

  return SRSRAN_SUCCESS;
}

uint8_t security_generate_k_nr_rrc(uint8_t*                    k_gnb,
                                   CIPHERING_ALGORITHM_ID_ENUM enc_alg_id,
                                   INTEGRITY_ALGORITHM_ID_ENUM int_alg_id,
                                   uint8_t*                    k_rrc_enc,
                                   uint8_t*                    k_rrc_int)
{
  if (k_gnb == NULL || k_rrc_enc == NULL || k_rrc_int == NULL) {
    log_error("Invalid inputs");
    return SRSRAN_ERROR;
  }
  std::array<uint8_t, 32> key;

  memcpy(key.data(), k_gnb, 32);

  // Derive RRC ENC
  // algorithm type distinguisher
  std::vector<uint8_t> algo_distinguisher;
  algo_distinguisher.resize(1);
  algo_distinguisher[0] = ALGO_5G_DISTINGUISHER_RRC_ENC_ALG;

  // algorithm type distinguisher
  std::vector<uint8_t> algorithm_identity;
  algorithm_identity.resize(1);
  algorithm_identity[0] = enc_alg_id;

  if (kdf_common(FC_5G_ALGORITHM_KEY_DERIVATION, key, algo_distinguisher, algorithm_identity, k_rrc_enc) !=
      SRSRAN_SUCCESS) {
    log_error("Failed to run kdf_common");
    return SRSRAN_ERROR;
  }

  // Derive RRC INT
  // algorithm type distinguisher
  algo_distinguisher.resize(1);
  algo_distinguisher[0] = ALGO_5G_DISTINGUISHER_RRC_INT_ALG;

  // algorithm type distinguisher
  algorithm_identity.resize(1);
  algorithm_identity[0] = int_alg_id;

  // Derive RRC int
  if (kdf_common(FC_5G_ALGORITHM_KEY_DERIVATION, key, algo_distinguisher, algorithm_identity, k_rrc_int) !=
      SRSRAN_SUCCESS) {
    log_error("Failed to run kdf_common");
    return SRSRAN_ERROR;
  }
  return SRSRAN_SUCCESS;
}

uint8_t security_generate_k_nr_up(uint8_t*                    k_gnb,
                                  CIPHERING_ALGORITHM_ID_ENUM enc_alg_id,
                                  INTEGRITY_ALGORITHM_ID_ENUM int_alg_id,
                                  uint8_t*                    k_up_enc,
                                  uint8_t*                    k_up_int)
{
  if (k_gnb == NULL || k_up_enc == NULL || k_up_int == NULL) {
    log_error("Invalid inputs");
    return SRSRAN_ERROR;
  }
  std::array<uint8_t, 32> key;

  memcpy(key.data(), k_gnb, 32);

  // Derive UP ENC
  // algorithm type distinguisher
  std::vector<uint8_t> algo_distinguisher;
  algo_distinguisher.resize(1);
  algo_distinguisher[0] = ALGO_5G_DISTINGUISHER_UP_ENC_ALG;

  // algorithm type distinguisher
  std::vector<uint8_t> algorithm_identity;
  algorithm_identity.resize(1);
  algorithm_identity[0] = enc_alg_id;

  if (kdf_common(FC_5G_ALGORITHM_KEY_DERIVATION, key, algo_distinguisher, algorithm_identity, k_up_enc) !=
      SRSRAN_SUCCESS) {
    log_error("Failed to run kdf_common");
    return SRSRAN_ERROR;
  }

  // Derive UP INT
  // algorithm type distinguisher
  algo_distinguisher.resize(1);
  algo_distinguisher[0] = ALGO_5G_DISTINGUISHER_UP_INT_ALG;

  // algorithm type distinguisher
  algorithm_identity.resize(1);
  algorithm_identity[0] = int_alg_id;

  // Derive UP int
  if (kdf_common(FC_5G_ALGORITHM_KEY_DERIVATION, key, algo_distinguisher, algorithm_identity, k_up_int) !=
      SRSRAN_SUCCESS) {
    log_error("Failed to run kdf_common");
    return SRSRAN_ERROR;
  }
  return SRSRAN_SUCCESS;
}

uint8_t security_generate_res_star(uint8_t*    ck,
                                   uint8_t*    ik,
                                   const char* serving_network_name,
                                   uint8_t*    rand_,
                                   uint8_t*    res_,
                                   size_t      res_len_,
                                   uint8_t*    res_star)
{
  if (ck == NULL || ik == NULL || serving_network_name == NULL || rand_ == NULL || res_ == NULL || res_star == NULL) {
    log_error("Invalid inputs");
    return SRSRAN_ERROR;
  }
  std::array<uint8_t, 32> key;

  // The input key Key shall be equal to the concatenation CK || IK of CK and IK.
  memcpy(key.data(), ck, 16);
  memcpy(key.data() + 16, ik, 16);

  // Serving Network Name
  std::vector<uint8_t> ssn;
  ssn.resize(strlen(serving_network_name));
  memcpy(ssn.data(), serving_network_name, strlen(serving_network_name));

  // RAND
  std::vector<uint8_t> rand;
  rand.resize(AKA_RAND_LEN);
  memcpy(rand.data(), rand_, rand.size());

  // RES
  std::vector<uint8_t> res;
  res.resize(res_len_);
  memcpy(res.data(), res_, res.size());

  uint8_t output[32];
  if (kdf_common(FC_5G_RES_STAR_DERIVATION, key, ssn, rand, res, output) != SRSRAN_SUCCESS) {
    log_error("Failed to run kdf_common");
    return SRSRAN_ERROR;
  }
  memcpy(res_star, output + 16, 16);

  return SRSRAN_SUCCESS;
}

int kdf_common(const uint8_t fc, const std::array<uint8_t, 32>& key, const std::vector<uint8_t>& P0, uint8_t* output)
{
  uint8_t* s;
  uint32_t s_len = 1 + P0.size() + 2;

  s = (uint8_t*)calloc(s_len, sizeof(uint8_t));

  if (s == nullptr) {
    log_error("Unable to allocate memory in %s()", __FUNCTION__);
    return SRSRAN_ERROR;
  }

  uint32_t i = 0;
  s[i]       = fc; // FC
  i++;

  // P0
  memcpy(&s[i], P0.data(), P0.size());
  i += P0.size();
  uint16_t p0_length_value = htons(P0.size());
  memcpy(&s[i], &p0_length_value, sizeof(p0_length_value));
  i += sizeof(p0_length_value);

  sha256(key.data(), key.size(), s, i, output, 0);
  free(s);

  return SRSRAN_SUCCESS;
}

int kdf_common(const uint8_t                  fc,
               const std::array<uint8_t, 32>& key,
               const std::vector<uint8_t>&    P0,
               const std::vector<uint8_t>&    P1,
               uint8_t*                       output)
{
  uint8_t* s;
  uint32_t s_len = 1 + P0.size() + 2 + P1.size() + 2;

  s = (uint8_t*)calloc(s_len, sizeof(uint8_t));

  if (s == nullptr) {
    log_error("Unable to allocate memory in %s()", __FUNCTION__);
    return SRSRAN_ERROR;
  }

  uint32_t i = 0;
  s[i]       = fc; // FC
  i++;

  // P0
  memcpy(&s[i], P0.data(), P0.size());
  i += P0.size();
  uint16_t p0_length_value = htons(P0.size());
  memcpy(&s[i], &p0_length_value, sizeof(p0_length_value));
  i += sizeof(p0_length_value);

  // P1
  memcpy(&s[i], P1.data(), P1.size());
  i += P1.size();
  uint16_t p1_length_value = htons(P1.size());
  memcpy(&s[i], &p1_length_value, sizeof(p1_length_value));
  i += sizeof(p1_length_value);

  sha256(key.data(), key.size(), s, i, output, 0);
  free(s);

  return SRSRAN_SUCCESS;
}

int kdf_common(const uint8_t                  fc,
               const std::array<uint8_t, 32>& key,
               const std::vector<uint8_t>&    P0,
               const std::vector<uint8_t>&    P1,
               const std::vector<uint8_t>&    P2,
               uint8_t*                       output)
{
  uint8_t* s;
  uint32_t s_len = 1 + P0.size() + 2 + P1.size() + 2 + P2.size() + 2;

  s = (uint8_t*)calloc(s_len, sizeof(uint8_t));

  if (s == nullptr) {
    log_error("Unable to allocate memory in %s()", __FUNCTION__);
    return SRSRAN_ERROR;
  }

  uint32_t i = 0;
  s[i]       = fc; // FC
  i++;

  // P0
  memcpy(&s[i], P0.data(), P0.size());
  i += P0.size();
  uint16_t p0_length_value = htons(P0.size());
  memcpy(&s[i], &p0_length_value, sizeof(p0_length_value));
  i += sizeof(p0_length_value);

  // P1
  memcpy(&s[i], P1.data(), P1.size());
  i += P1.size();
  uint16_t p1_length_value = htons(P1.size());
  memcpy(&s[i], &p1_length_value, sizeof(p1_length_value));
  i += sizeof(p1_length_value);

  // P2
  memcpy(&s[i], P2.data(), P2.size());
  i += P2.size();
  uint16_t p2_length_value = htons(P2.size());
  memcpy(&s[i], &p2_length_value, sizeof(p2_length_value));
  i += sizeof(p2_length_value);

  sha256(key.data(), key.size(), s, i, output, 0);
  free(s);

  return SRSRAN_SUCCESS;
}
/******************************************************************************
 * Integrity Protection
 *****************************************************************************/

uint8_t security_128_eia1(const uint8_t* key,
                          uint32_t       count,
                          uint32_t       bearer,
                          uint8_t        direction,
                          uint8_t*       msg,
                          uint32_t       msg_len,
                          uint8_t*       mac)
{
  return liblte_security_128_eia1(key, count, bearer, direction, msg, msg_len, mac);
}

uint8_t security_128_eia2(const uint8_t* key,
                          uint32_t       count,
                          uint32_t       bearer,
                          uint8_t        direction,
                          uint8_t*       msg,
                          uint32_t       msg_len,
                          uint8_t*       mac)
{
  return liblte_security_128_eia2(key, count, bearer, direction, msg, msg_len, mac);
}

uint8_t security_128_eia3(const uint8_t* key,
                          uint32_t       count,
                          uint32_t       bearer,
                          uint8_t        direction,
                          uint8_t*       msg,
                          uint32_t       msg_len,
                          uint8_t*       mac)
{
  return liblte_security_128_eia3(key, count, bearer, direction, msg, msg_len * 8, mac);
}

uint8_t security_md5(const uint8_t* input, size_t len, uint8_t* output)
{
  memset(output, 0x00, 16);
#ifdef HAVE_MBEDTLS
  mbedtls_md5(input, len, output);
#endif // HAVE_MBEDTLS
#ifdef HAVE_POLARSSL
  md5(input, len, output);
#endif
  return SRSRAN_SUCCESS;
}

/******************************************************************************
 * Encryption / Decryption
 *****************************************************************************/

uint8_t security_128_eea1(uint8_t* key,
                          uint32_t count,
                          uint8_t  bearer,
                          uint8_t  direction,
                          uint8_t* msg,
                          uint32_t msg_len,
                          uint8_t* msg_out)
{
  return liblte_security_encryption_eea1(key, count, bearer, direction, msg, msg_len * 8, msg_out);
}

uint8_t security_128_eea2(uint8_t* key,
                          uint32_t count,
                          uint8_t  bearer,
                          uint8_t  direction,
                          uint8_t* msg,
                          uint32_t msg_len,
                          uint8_t* msg_out)
{
  return liblte_security_encryption_eea2(key, count, bearer, direction, msg, msg_len * 8, msg_out);
}

uint8_t security_128_eea3(uint8_t* key,
                          uint32_t count,
                          uint8_t  bearer,
                          uint8_t  direction,
                          uint8_t* msg,
                          uint32_t msg_len,
                          uint8_t* msg_out)
{
  return liblte_security_encryption_eea3(key, count, bearer, direction, msg, msg_len * 8, msg_out);
}

/******************************************************************************
 * Authentication
 *****************************************************************************/
uint8_t compute_opc(uint8_t* k, uint8_t* op, uint8_t* opc)
{
  return liblte_compute_opc(k, op, opc);
}

uint8_t security_milenage_f1(uint8_t* k, uint8_t* op, uint8_t* rand, uint8_t* sqn, uint8_t* amf, uint8_t* mac_a)
{
  return liblte_security_milenage_f1(k, op, rand, sqn, amf, mac_a);
}

uint8_t security_milenage_f1_star(uint8_t* k, uint8_t* op, uint8_t* rand, uint8_t* sqn, uint8_t* amf, uint8_t* mac_s)
{
  return liblte_security_milenage_f1_star(k, op, rand, sqn, amf, mac_s);
}

uint8_t
security_milenage_f2345(uint8_t* k, uint8_t* op, uint8_t* rand, uint8_t* res, uint8_t* ck, uint8_t* ik, uint8_t* ak)
{
  return liblte_security_milenage_f2345(k, op, rand, res, ck, ik, ak);
}

uint8_t security_milenage_f5_star(uint8_t* k, uint8_t* op, uint8_t* rand, uint8_t* ak)
{
  return liblte_security_milenage_f5_star(k, op, rand, ak);
}

int security_xor_f2345(uint8_t* k, uint8_t* rand, uint8_t* res, uint8_t* ck, uint8_t* ik, uint8_t* ak)
{
  uint8_t xdout[16];
  uint8_t cdout[8];
  // Use RAND and K to compute RES, CK, IK and AK
  for (uint32_t i = 0; i < 16; i++) {
    xdout[i] = k[i] ^ rand[i];
  }
  for (uint32_t i = 0; i < 16; i++) {
    res[i] = xdout[i];
    ck[i]  = xdout[(i + 1) % 16];
    ik[i]  = xdout[(i + 2) % 16];
  }
  for (uint32_t i = 0; i < 6; i++) {
    ak[i] = xdout[i + 3];
  }
  return SRSRAN_SUCCESS;
}

int security_xor_f1(uint8_t* k, uint8_t* rand, uint8_t* sqn, uint8_t* amf, uint8_t* mac_a)
{
  uint8_t xdout[16];
  uint8_t cdout[8];
  // Use RAND and K to compute RES, CK, IK and AK
  for (uint32_t i = 0; i < 16; i++) {
    xdout[i] = k[i] ^ rand[i];
  }
  // Generate cdout
  for (uint32_t i = 0; i < 6; i++) {
    cdout[i] = sqn[i];
  }
  for (uint32_t i = 0; i < 2; i++) {
    cdout[6 + i] = amf[i];
  }

  // Generate MAC
  for (uint32_t i = 0; i < 8; i++) {
    mac_a[i] = xdout[i] ^ cdout[i];
  }
  return SRSRAN_SUCCESS;
}

} // namespace srsran