/*
 * Copyright 2020 Software Radio Systems Limited
 * Author: Vadim Yanitskiy <axilirator@gmail.com>
 * Sponsored by Positive Technologies
 *
 * This file is part of srsLTE.
 *
 * srsLTE is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of
 * the License, or (at your option) any later version.
 *
 * srsLTE is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * A copy of the GNU Affero General Public License can be found in
 * the LICENSE file in the top-level directory of this distribution
 * and at http://www.gnu.org/licenses/.
 *
 */

#include "srslte/common/buffer_pool.h"
#include "srslte/common/common.h"
#include "srslte/common/log.h"
#include "srslte/common/nas_pcap.h"
#include "srslte/common/security.h"
#include "srslte/common/stack_procedure.h"
#include "srslte/interfaces/ue_interfaces.h"

#include "srsue/hdr/stack/upper/nas.h"
#include "srsue/hdr/stack/upper/nas_common.h"
#include "srsue/hdr/stack/upper/nas_ext.h"
#include "srsue/hdr/stack/upper/nas_extif_unix.h"
#include "srsue/hdr/stack/upper/nas_metrics.h"

using namespace srslte;

namespace srsue {

void nas_ext::init(usim_interface_nas* usim_, rrc_interface_nas* rrc_, gw_interface_nas* gw_)
{
  usim = usim_;
  rrc  = rrc_;
  gw   = gw_;

  // RRCTL PDU handler
  auto rrctl_rx_cb = [this](const srslte::byte_buffer_t& pdu) {
    rrctl::proto::msg_type type;
    rrctl::proto::msg_disc disc;
    const uint8_t* payload;
    std::string desc;
    uint16_t length;

    // Parse the message header
    try {
      payload = rrctl::codec::dec_hdr(pdu, type, disc, length);
      desc = rrctl::proto::msg_hdr_desc(type, disc, length);
      nas_log->info("Got RRCTL message: %s\n", desc.c_str());
    } catch (const rrctl::codec::error& e) {
      nas_log->warning("Got malformed RRCTL message: %s\n", e.what());
      return;
    }

    // Call the corresponding handler
    switch (type) {
    case rrctl::proto::RRCTL_RESET:
      handle_rrctl_reset(disc, payload, length);
      break;
    case rrctl::proto::RRCTL_PLMN_SEARCH:
      handle_rrctl_plmn_search(disc, payload, length);
      break;
    case rrctl::proto::RRCTL_PLMN_SELECT:
      handle_rrctl_plmn_select(disc, payload, length);
      break;
    case rrctl::proto::RRCTL_CONN_ESTABLISH:
      handle_rrctl_conn_establish(disc, payload, length);
      break;
    case rrctl::proto::RRCTL_DATA:
      handle_rrctl_data(disc, payload, length);
      break;
    case rrctl::proto::RRCTL_PARAM:
      handle_rrctl_param(disc, payload, length);
      break;
    case rrctl::proto::RRCTL_SEC_MODE:
      handle_rrctl_sec_mode(disc, payload, length);
      break;
    case rrctl::proto::RRCTL_EXT_USIM:
      handle_rrctl_ext_usim(payload, length);
      break;
    case rrctl::proto::RRCTL_CONN_RELEASE:
    default:
      nas_log->warning("%s is not handled\n", desc.c_str());
    }
  };

  std::unique_ptr<nas_extif_unix> iface_(new nas_extif_unix(rrctl_rx_cb, cfg.sock_path));
  iface = std::move(iface_);
}

void nas_ext::rrctl_send_confirm(rrctl::proto::msg_type type)
{
  srslte::byte_buffer_t pdu;

  rrctl::codec::enc_hdr(pdu, type, rrctl::proto::RRCTL_CNF);
  iface->write(pdu);
}

void nas_ext::rrctl_send_error(rrctl::proto::msg_type type)
{
  srslte::byte_buffer_t pdu;

  rrctl::codec::enc_hdr(pdu, type, rrctl::proto::RRCTL_ERR);
  iface->write(pdu);
}

void nas_ext::handle_rrctl_reset(rrctl::proto::msg_disc disc, const uint8_t* msg, size_t len)
{
  /* Reset security state (no EEA, no EIA) */
  memset(&ctxt, 0x00, sizeof(ctxt));

  rrctl_send_confirm(rrctl::proto::RRCTL_RESET);
}

void nas_ext::handle_rrctl_plmn_search(rrctl::proto::msg_disc disc, const uint8_t* msg, size_t len)
{
  // Response to be sent when RRC responds
  rrc->plmn_search();
}

void nas_ext::handle_rrctl_plmn_select(rrctl::proto::msg_disc disc, const uint8_t* msg, size_t len)
{
  std::pair<uint16_t, uint16_t> mcc_mnc;
  srslte::plmn_id_t plmn_id;

  // Parse PLMN ID
  try {
    rrctl::codec::dec_plmn_select_req(mcc_mnc, msg, len);
  } catch (const rrctl::codec::error& e) {
    nas_log->warning("Failed to parse RRCTL message: %s\n", e.what());
    rrctl_send_error(rrctl::proto::RRCTL_PLMN_SELECT);
    return;
  }

  if (plmn_id.from_number(mcc_mnc.first, mcc_mnc.second) != SRSLTE_SUCCESS) {
    nas_log->warning("Failed to parse PLMN ID from PLMN Select Request\n");
    rrctl_send_error(rrctl::proto::RRCTL_PLMN_SELECT);
  }

  rrc->plmn_select(plmn_id);
  rrctl_send_confirm(rrctl::proto::RRCTL_PLMN_SELECT);
}

void nas_ext::handle_rrctl_conn_establish(rrctl::proto::msg_disc disc, const uint8_t* msg, size_t len)
{
  srslte::establishment_cause_t cause;
  const uint8_t* pdu;
  size_t pdu_len;

  try {
    rrctl::codec::dec_conn_establish_req(cause, pdu, pdu_len, msg, len);
  } catch (const rrctl::codec::error& e) {
    nas_log->warning("Failed to parse RRCTL message: %s\n", e.what());
    rrctl_send_error(rrctl::proto::RRCTL_CONN_ESTABLISH);
    return;
  }

  set_k_enb_count(ctxt.tx_count);
  ctxt.tx_count++;

  // Allocate a new NAS PDU on heap
  unique_byte_buffer_t nas_pdu = srslte::allocate_unique_buffer(*pool, true);
  nas_pdu->append_bytes(pdu, pdu_len);

  if (pcap != nullptr)
    pcap->write_nas(nas_pdu->msg, nas_pdu->N_bytes);

  rrc->connection_request(cause, std::move(nas_pdu));
}

void nas_ext::handle_rrctl_data(rrctl::proto::msg_disc disc, const uint8_t* msg, size_t len)
{
  if (not rrc->is_connected()) {
    nas_log->warning("Received DATA.req, but there is no active connection\n");
    rrctl_send_error(rrctl::proto::RRCTL_DATA);
    return;
  }

  // Allocate a new NAS PDU on heap
  unique_byte_buffer_t nas_pdu = srslte::allocate_unique_buffer(*pool, true);
  nas_pdu->append_bytes(msg, len);

  if (pcap != nullptr)
    pcap->write_nas(nas_pdu->msg, nas_pdu->N_bytes);

  // Apply pre-configured EEA algorythm (if enabled)
  cipher_encrypt(nas_pdu.get());
  // Apply pre-configured EIA algorythm (if enabled)
  integrity_generate(&k_nas_int[16], ctxt.tx_count, SECURITY_DIRECTION_UPLINK,
                     &nas_pdu->msg[5], nas_pdu->N_bytes - 5, &nas_pdu->msg[1]);

  rrc->write_sdu(std::move(nas_pdu));
  ctxt.tx_count++;

  rrctl_send_confirm(rrctl::proto::RRCTL_DATA);
}

void nas_ext::handle_rrctl_param(rrctl::proto::msg_disc disc, const uint8_t* msg, size_t len)
{
  const struct rrctl::proto::msg_param_req* param;
  srslte::s_tmsi_t ue_identity;

  param = reinterpret_cast<const struct rrctl::proto::msg_param_req*> (msg);
  nas_log->warning("Rx PARAM.req (type=%02x, len=%u)\n", param->type, param->len);

  if (param->len != (len - 2)) { /* XXX: type + length */
    nas_log->info("Received malformed PARAM.req (len=%u)\n", param->len);
    rrctl_send_error(rrctl::proto::RRCTL_PARAM);
    return;
  }

  switch (param->type) {
  case rrctl::proto::RRCTL_PARAM_UEID:
    ue_identity.m_tmsi = ntohl(param->u.ueid.m_tmsi);
    ue_identity.mmec = param->u.ueid.mmec;

    nas_log->info("Setting UEID: mmec=%x m_tmsi=%x\n",
                  ue_identity.mmec, ue_identity.m_tmsi);
    rrc->set_ue_identity(ue_identity);
    rrctl_send_confirm(rrctl::proto::RRCTL_PARAM);
    break;
  default:
    nas_log->warning("Unhandled PARAM.req type (0x%02x)\n", param->type);
    rrctl_send_error(rrctl::proto::RRCTL_PARAM);
  }
}

void nas_ext::handle_rrctl_sec_mode(rrctl::proto::msg_disc disc, const uint8_t* msg, size_t len)
{
  const struct rrctl::proto::msg_sec_mode_req* req;

  req = reinterpret_cast<const struct rrctl::proto::msg_sec_mode_req*> (msg);

  if (len < sizeof(*req)) { /* XXX: mandatory fields only */
    nas_log->error("Received malformed PARAM.req (len=%zu)\n", len);
    rrctl_send_error(rrctl::proto::RRCTL_SEC_MODE);
    return;
  }

  /* Skip the header */
  len -= sizeof(*req);

  nas_log->info("Rx SecurityMode.req (EEA%u, EIA%u)\n", req->eea, req->eia);

  ctxt.cipher_algo = static_cast<srslte::CIPHERING_ALGORITHM_ID_ENUM> (req->eea);
  ctxt.integ_algo = static_cast<srslte::INTEGRITY_ALGORITHM_ID_ENUM> (req->eia);

  if (req->flags & SEC_MODE_F_RESET_RX_CTR)
    ctxt.rx_count = 0;
  if (req->flags & SEC_MODE_F_RESET_TX_CTR)
    ctxt.tx_count = 0;

  if (req->eea != 0x00 or req->eia != 0x00) {
    /* Ensure that Kasme is present */
    if (len != sizeof(ctxt.k_asme)) {
      nas_log->error("Kasme is expected, but not present");
      rrctl_send_error(rrctl::proto::RRCTL_SEC_MODE);
      return;
    }

    memcpy(&ctxt.k_asme[0], &req->k_asme[0], sizeof(ctxt.k_asme));

    /* Derive both Knas_enc and Knas_int from Kasme */
    /* NOTE: how is this related to (U)SIM at all?!? */
    usim->generate_nas_keys(&ctxt.k_asme[0],  // in
                            &k_nas_enc[0],    // out
                            &k_nas_int[0],    // out
                            ctxt.cipher_algo,
                            ctxt.integ_algo);
  }

  rrctl_send_confirm(rrctl::proto::RRCTL_SEC_MODE);
}

void nas_ext::handle_rrctl_ext_usim(const uint8_t* _msg, size_t len)
{
  enum rrctl::proto::ext_usim_msg_type msg_type;
  const struct rrctl::proto::ext_usim_msg* msg;
  std::string msg_desc;

  if (len < 4) { /* XXX: type + padding */
    nas_log->error("Received too short (U)SIM specific message (len=%zu)\n", len);
    rrctl_send_error(rrctl::proto::RRCTL_EXT_USIM);
    return;
  }

  /* Skip the header */
  len -= 4;

  msg = reinterpret_cast<const struct rrctl::proto::ext_usim_msg*> (_msg);
  msg_type = static_cast<rrctl::proto::ext_usim_msg_type> (msg->type);

  switch (msg_type) {
  case rrctl::proto::EXT_USIM_GEN_AUTH_RESP:
    handle_usim_gen_auth_resp_req(msg, len);
    break;
  case rrctl::proto::EXT_USIM_RAW_APDU:
  case rrctl::proto::EXT_USIM_READ_FILE:
  case rrctl::proto::EXT_USIM_UPDATE_FILE:
  case rrctl::proto::EXT_USIM_RESERVED:
  default:
      nas_log->warning("(U)SIM specific message 0x%02x is not handled\n", msg->type);
  }
}

void nas_ext::handle_usim_gen_auth_resp_req(const struct rrctl::proto::ext_usim_msg* msg, size_t len)
{
  const struct rrctl::proto::ext_usim_gen_auth_resp_req* req;
  struct rrctl::proto::ext_usim_gen_auth_resp_rsp* rsp;
  struct rrctl::proto::ext_usim_msg* rsp_msg;
  std::pair<uint16_t, uint16_t> mcc_mnc;
  struct rrctl::proto::msg_hdr* hdr;
  srslte::byte_buffer_t pdu;

  /* Allocate the response message in advance */
  hdr = rrctl::codec::enc_hdr(pdu, rrctl::proto::RRCTL_EXT_USIM, rrctl::proto::RRCTL_CNF, 4);
  rsp_msg = reinterpret_cast<struct rrctl::proto::ext_usim_msg*> (&pdu.msg[pdu.N_bytes]);
  rsp_msg->type = msg->type;
  pdu.N_bytes += 4;

  if (len < sizeof(*req)) {
    nas_log->error("Received too short (U)SIM GenAuthResp.req (len=%zu)\n", len);
    hdr->disc = (uint8_t) rrctl::proto::RRCTL_ERR;
    iface->write(pdu);
    return;
  }

  rsp = &rsp_msg->u.gen_auth_resp_rsp;
  req = &msg->u.gen_auth_resp_req;

  // Parse PLMN ID
  try {
    /* HACK: generalize this function */
    rrctl::codec::dec_plmn_select_req(mcc_mnc, &req->mcc[0], len);
  } catch (const rrctl::codec::error& e) {
    nas_log->warning("Failed to parse MCC/MNC: %s\n", e.what());
    hdr->disc = (uint8_t) rrctl::proto::RRCTL_ERR;
    iface->write(pdu);
    return;
  }

  uint8_t k_asme[32];
  uint8_t res[16];
  int res_len;

  /* NOTE: this is a blocking call => no other RRCTL messages can be processed in parallel?
   * FIXME: the authors of srsUE apparently are not aware of 'const', so we have to cast. */
  auth_result_t auth_res = usim->generate_authentication_response((uint8_t*) &req->rand[0],
                                                                  (uint8_t*) &req->autn[0],
                                                                  mcc_mnc.first, mcc_mnc.second,
                                                                  &res[0], &res_len, &k_asme[0]);
  switch (auth_res) {
  case AUTH_OK:
    nas_log->info("Authentication vector has been generated successfully\n");

    set_k_enb_count(0);

    pdu.N_bytes += sizeof(*rsp);
    pdu.append_bytes(&k_asme[0], sizeof(k_asme));
    pdu.append_bytes(&res[0], res_len);
    rsp->res_len = (uint8_t) res_len;
    break;
  case AUTH_SYNCH_FAILURE:
    nas_log->warning("Synchronization is required to generate an authentication vector\n");

    pdu.N_bytes += sizeof(*rsp);
    pdu.append_bytes(&res[0], res_len);
    rsp->res_len = (uint8_t) res_len;
    rsp->out_of_sync = true;
    break;
  case AUTH_FAILED:
  default:
    nas_log->warning("Could not generate an authentication vector\n");
    hdr->disc = (uint8_t) rrctl::proto::RRCTL_ERR;
  }

  hdr->len = htons(pdu.N_bytes - sizeof(*hdr));
  iface->write(pdu);
}

void nas_ext::get_metrics(nas_metrics_t* m)
{
  nas_metrics_t metrics = {};
  // FIXME: is there anything we could fill in?
  *m = metrics;
}

void nas_ext::stop()
{
  // Close the UNIX domain socket connection
  iface->close();
  iface.release();
}

/*******************************************************************************
 * UE interface (dummy)
 ******************************************************************************/

void nas_ext::start_attach_proc(srslte::proc_state_t* result, srslte::establishment_cause_t cause_)
{
  nas_log->info("The UE has requested us to perform Attach Request, however we ignore it\n");
  if (result != nullptr) {
    result->set_val();
  }
}

bool nas_ext::detach_request(const bool switch_off)
{
  nas_log->info("The UE has requested us to perform Detach Request, however we ignore it\n");
  return false;
}

void nas_ext::timer_expired(uint32_t timeout_id)
{
  nas_log->info("Timer id=%u is expired, however we ignore it\n", timeout_id);
}

/*******************************************************************************
 * RRC interface
 ******************************************************************************/

// TODO: investigate the meaning of these signals
void nas_ext::set_barring(barring_t barring) {}
void nas_ext::left_rrc_connected() {}

bool nas_ext::paging(srslte::s_tmsi_t* ue_identity)
{
  srslte::byte_buffer_t msg;

  nas_log->info("Received paging from RRC\n");

  rrctl::codec::enc_paging_ind(msg, ue_identity);
  iface->write(msg);

  // TODO: what are we supposed to return?
  return false;
}

void nas_ext::write_pdu(uint32_t lcid, srslte::unique_byte_buffer_t pdu)
{
  srslte::byte_buffer_t msg;
  uint8 pd, sec_hdr_type;

  nas_log->info_hex(pdu->msg, pdu->N_bytes, "Received DL %s PDU from RRC\n", rrc->get_rb_name(lcid).c_str());

  // Parse the message security header
  liblte_mme_parse_msg_sec_header((LIBLTE_BYTE_MSG_STRUCT*) pdu.get(), &pd, &sec_hdr_type);
  switch (sec_hdr_type) {
  case LIBLTE_MME_SECURITY_HDR_TYPE_INTEGRITY_AND_CIPHERED_WITH_NEW_EPS_SECURITY_CONTEXT:
  case LIBLTE_MME_SECURITY_HDR_TYPE_INTEGRITY_WITH_NEW_EPS_SECURITY_CONTEXT:
  case LIBLTE_MME_SECURITY_HDR_TYPE_SERVICE_REQUEST:
  case LIBLTE_MME_SECURITY_HDR_TYPE_PLAIN_NAS:
  case LIBLTE_MME_SECURITY_HDR_TYPE_INTEGRITY:
    break;
  case LIBLTE_MME_SECURITY_HDR_TYPE_INTEGRITY_AND_CIPHERED:
    // Apply pre-configured EEA algorythm (if enabled)
    cipher_decrypt(pdu.get());
    break;
  default:
    nas_log->error("Received DL NAS PDU with unknown sec_hdr=%02x\n", sec_hdr_type);
  }

  if (pcap != nullptr)
    pcap->write_nas(pdu->msg, pdu->N_bytes);

  rrctl::codec::enc_data_ind(msg, pdu->msg, pdu->N_bytes, lcid);
  iface->write(msg);
}

bool nas_ext::is_attached()
{
  // FIXME: we probably need to maintain the state
  return false; // return a dummy value for now
}

bool nas_ext::get_k_asme(uint8_t* k_asme_, uint32_t n)
{
  // FIXME: we probably need to maintain a security context
  return false; // return a dummy value for now
}

uint32_t nas_ext::get_ipv4_addr()
{
  // FIXME: where can we get it? maybe from GW?
  return 0x00000000;
}

bool nas_ext::get_ipv6_addr(uint8_t* ipv6_addr)
{
  // FIXME: where can we get it? maybe from GW?
  return false;
}

void nas_ext::plmn_search_completed(
    const rrc_interface_nas::found_plmn_t found_plmns[rrc_interface_nas::MAX_FOUND_PLMNS],
    int                                   nof_plmns)
{
  srslte::byte_buffer_t pdu;

  nas_log->info("RRC has completed PLMN search, %d carriers found\n", nof_plmns);

  // Send PLMN_SEARCH.res to an external entity
  if (nof_plmns >= 0) {
    rrctl::codec::enc_plmn_search_res(pdu, found_plmns, nof_plmns);
    iface->write(pdu);
  } else {
    nas_log->warning("PLMN search completed with an error\n");
    rrctl_send_error(rrctl::proto::RRCTL_PLMN_SEARCH);
  }
}

bool nas_ext::connection_request_completed(bool outcome)
{
  nas_log->info("RRC has %s connection establisment\n", outcome ? "completed" : "failed");

  if (outcome)
    rrctl_send_confirm(rrctl::proto::RRCTL_CONN_ESTABLISH);
  else
    rrctl_send_error(rrctl::proto::RRCTL_CONN_ESTABLISH);

  return false; // FIXME: what should we return here?
}

} // namespace srsue