/**
 * HSS - Hierarchical Signatures System (RFC 8554)
 * (C) 2023 Jack Lloyd
 *     2023 Fabian Albert, Philippe Lieser - Rohde & Schwarz Cybersecurity GmbH
 *
 * Botan is released under the Simplified BSD License (see license.txt)
 */

#include <botan/internal/hss.h>

#include <botan/internal/fmt.h>
#include <botan/internal/hss_lms_utils.h>
#include <botan/internal/int_utils.h>
#include <botan/internal/scan_name.h>
#include <botan/internal/stl_util.h>

#include <algorithm>
#include <limits>

namespace Botan {

namespace {

/**
 * @brief The maximum number of levels in a HSS-LMS tree.
 *
 * RFC 8554 Section 6:
 *   The number of levels is denoted as L and is between one
 *   and eight, inclusive.
 */
constexpr HSS_Level HSS_MAX_LEVELS(8);

/**
 * @brief Domain-separation parameter for generation the seed of a child LMS tree.
 *
 * This comes from https://github.com/cisco/hash-sigs.
 */
constexpr uint16_t SEED_CHILD_SEED = 0xfffe;

/**
 * @brief Domain-separation parameter for generation the identifier of a child LMS tree.
 *
 * This comes from https://github.com/cisco/hash-sigs.
 */
constexpr uint16_t SEED_CHILD_I = 0xffff;

/**
 * @brief Check that the given @p hash_name is one of the supported hash functions for HSS-LMS.
 */
constexpr bool is_supported_hash_function(std::string_view hash_name) {
   return hash_name == "SHA-256" || hash_name == "Truncated(SHA-256,192)" || hash_name == "SHAKE-256(256)" ||
          hash_name == "SHAKE-256(192)";
}

/**
 * Given an HSS index, i.e. the number of already created HSS signatures, return the lms leaf indices for
 * the different LMS layers from root layer to bottom layer.
 */
std::vector<LMS_Tree_Node_Idx> derive_lms_leaf_indices_from_hss_index(HSS_Sig_Idx hss_idx,
                                                                      const HSS_LMS_Params& hss_params) {
   std::vector<LMS_Tree_Node_Idx> q(hss_params.L().get());
   for(int32_t layer_ctr = hss_params.L().get() - 1; layer_ctr >= 0; --layer_ctr) {
      HSS_Level layer(layer_ctr);
      const HSS_LMS_Params::LMS_LMOTS_Params_Pair& layer_params = hss_params.params_at_level(layer);
      size_t layer_h = layer_params.lms_params().h();
      q.at(layer.get()) =
         checked_cast_to<LMS_Tree_Node_Idx>(hss_idx.get() % checked_cast_to<uint64_t>(1ULL << layer_h));
      hss_idx = hss_idx >> layer_h;
   }
   BOTAN_ARG_CHECK(hss_idx == HSS_Sig_Idx(0), "HSS Tree is exhausted");

   return q;
}

}  // namespace

HSS_LMS_Params::HSS_LMS_Params(std::vector<LMS_LMOTS_Params_Pair> lm_lmots_params) :
      m_lms_lmots_params(std::move(lm_lmots_params)), m_max_sig_count(calc_max_sig_count()) {
   BOTAN_ARG_CHECK(!m_lms_lmots_params.empty() && m_lms_lmots_params.size() <= HSS_MAX_LEVELS,
                   "Invalid number of levels");
}

HSS_LMS_Params::HSS_LMS_Params(std::string_view algo_params) {
   const auto wrap_in_hss_lms = [&]() {
      if(algo_params.starts_with("HSS-LMS(")) {
         return std::string(algo_params);
      } else {
         return fmt("HSS-LMS({})", algo_params);
      }
   }();
   SCAN_Name scan(wrap_in_hss_lms);

   BOTAN_ARG_CHECK(scan.arg_count() >= 2 && scan.arg_count() <= HSS_MAX_LEVELS + 1, "Invalid number of arguments");
   std::string hash = scan.arg(0);
   BOTAN_ARG_CHECK(is_supported_hash_function(hash), "Supported HSS-LMS hash function");

   for(size_t i = 1; i < scan.arg_count(); ++i) {
      SCAN_Name scan_layer(scan.arg(i));
      BOTAN_ARG_CHECK(scan_layer.algo_name() == "HW", "Invalid name for layer parameters");
      BOTAN_ARG_CHECK(scan_layer.arg_count() == 2, "Invalid number of layer parameters");
      const auto h =
         checked_cast_to_or_throw<uint8_t, Invalid_Argument>(scan_layer.arg_as_integer(0), "Invalid tree height");
      const auto w = checked_cast_to_or_throw<uint8_t, Invalid_Argument>(scan_layer.arg_as_integer(1),
                                                                         "Invalid Winternitz parameter");
      m_lms_lmots_params.push_back({LMS_Params::create_or_throw(hash, h), LMOTS_Params::create_or_throw(hash, w)});
   }
   m_max_sig_count = calc_max_sig_count();
}

HSS_Sig_Idx HSS_LMS_Params::calc_max_sig_count() const {
   uint32_t total_height_counter = 0;
   for(HSS_Level level(0); level < L(); level++) {
      total_height_counter += params_at_level(level).lms_params().h();
   }
   if(total_height_counter >= sizeof(HSS_Sig_Idx) * 8) {
      return HSS_Sig_Idx(std::numeric_limits<HSS_Sig_Idx::wrapped_type>::max());
   }
   return HSS_Sig_Idx(1) << total_height_counter;
}

HSS_LMS_PrivateKeyInternal::HSS_LMS_PrivateKeyInternal(const HSS_LMS_Params& hss_params, RandomNumberGenerator& rng) :
      m_hss_params(hss_params), m_current_idx(0), m_sig_size(HSS_Signature::size(m_hss_params)) {
   m_hss_seed = rng.random_vec<LMS_Seed>(m_hss_params.params_at_level(HSS_Level(0)).lms_params().m());
   m_identifier = rng.random_vec<LMS_Identifier>(LMS_IDENTIFIER_LEN);
}

std::shared_ptr<HSS_LMS_PrivateKeyInternal> HSS_LMS_PrivateKeyInternal::from_bytes_or_throw(
   std::span<const uint8_t> key_bytes) {
   if(key_bytes.size() < sizeof(HSS_Level) + sizeof(HSS_Sig_Idx)) {
      throw Decoding_Error("Too few private key bytes.");
   }
   BufferSlicer slicer(key_bytes);

   const auto L = load_be<HSS_Level>(slicer.take<sizeof(HSS_Level)>());
   if(L == 0U || L > HSS_MAX_LEVELS) {
      throw Decoding_Error("Invalid number of HSS layers in private HSS-LMS key.");
   }

   const auto sig_idx = load_be<HSS_Sig_Idx>(slicer.take<sizeof(HSS_Sig_Idx)>());

   std::vector<HSS_LMS_Params::LMS_LMOTS_Params_Pair> params;
   for(size_t layer = 1; layer <= L; ++layer) {
      if(slicer.remaining() < sizeof(LMS_Algorithm_Type) + sizeof(LMOTS_Algorithm_Type)) {
         throw Decoding_Error("Out of bytes while parsing private HSS-LMS key.");
      }
      const auto lms_type = load_be<LMS_Algorithm_Type>(slicer.take<sizeof(LMS_Algorithm_Type)>());
      const auto lmots_type = load_be<LMOTS_Algorithm_Type>(slicer.take<sizeof(LMOTS_Algorithm_Type)>());
      params.push_back({LMS_Params::create_or_throw(lms_type), LMOTS_Params::create_or_throw(lmots_type)});
   }
   std::string hash_name = params.at(0).lms_params().hash_name();
   if(std::any_of(params.begin(), params.end(), [&hash_name](HSS_LMS_Params::LMS_LMOTS_Params_Pair& lms_lmots_params) {
         bool invalid_lmots_hash = lms_lmots_params.lmots_params().hash_name() != hash_name;
         bool invalid_lms_hash = lms_lmots_params.lms_params().hash_name() != hash_name;
         return invalid_lmots_hash || invalid_lms_hash;
      })) {
      throw Decoding_Error("Inconsistent hash functions are not allowed.");
   }

   if(slicer.remaining() < params.at(0).lms_params().m() + LMS_IDENTIFIER_LEN) {
      throw Decoding_Error("Out of bytes while parsing private HSS-LMS key.");
   }
   auto hss_seed = slicer.copy<LMS_Seed>(params.at(0).lms_params().m());
   auto identifier = slicer.copy<LMS_Identifier>(LMS_IDENTIFIER_LEN);

   if(!slicer.empty()) {
      throw Decoding_Error("Private HSS-LMS key contains more bytes than expected.");
   }
   auto sk = std::shared_ptr<HSS_LMS_PrivateKeyInternal>(
      new HSS_LMS_PrivateKeyInternal(HSS_LMS_Params(std::move(params)), std::move(hss_seed), std::move(identifier)));

   sk->set_idx(sig_idx);
   return sk;
}

secure_vector<uint8_t> HSS_LMS_PrivateKeyInternal::to_bytes() const {
   secure_vector<uint8_t> sk_bytes(size());
   BufferStuffer stuffer(sk_bytes);

   stuffer.append(store_be(hss_params().L()));
   stuffer.append(store_be(get_idx()));

   for(HSS_Level layer(1); layer <= hss_params().L(); ++layer) {
      const auto& params = hss_params().params_at_level(layer - 1);
      stuffer.append(store_be(params.lms_params().algorithm_type()));
      stuffer.append(store_be(params.lmots_params().algorithm_type()));
   }
   stuffer.append(m_hss_seed);
   stuffer.append(m_identifier);
   BOTAN_ASSERT_NOMSG(stuffer.full());

   return sk_bytes;
}

void HSS_LMS_PrivateKeyInternal::set_idx(HSS_Sig_Idx idx) {
   m_current_idx = idx;
}

HSS_Sig_Idx HSS_LMS_PrivateKeyInternal::reserve_next_idx() {
   HSS_Sig_Idx next_idx = m_current_idx;
   if(next_idx >= m_hss_params.max_sig_count()) {
      throw Decoding_Error("HSS private key is exhausted");
   }
   set_idx(m_current_idx + 1);
   return next_idx;
}

size_t HSS_LMS_PrivateKeyInternal::size() const {
   size_t sk_size = sizeof(HSS_Level) + sizeof(HSS_Sig_Idx);
   // The concatenated algorithm types for all layers
   sk_size += hss_params().L().get() * (sizeof(LMS_Algorithm_Type) + sizeof(LMOTS_Algorithm_Type));
   sk_size += m_hss_seed.size() + m_identifier.size();
   return sk_size;
}

HSS_LMS_PrivateKeyInternal::HSS_LMS_PrivateKeyInternal(HSS_LMS_Params hss_params,
                                                       LMS_Seed hss_seed,
                                                       LMS_Identifier identifier) :
      m_hss_params(std::move(hss_params)),
      m_hss_seed(std::move(hss_seed)),
      m_identifier(std::move(identifier)),
      m_current_idx(0),
      m_sig_size(HSS_Signature::size(m_hss_params)) {
   BOTAN_ARG_CHECK(m_hss_seed.size() == m_hss_params.params_at_level(HSS_Level(0)).lms_params().m(),
                   "Invalid seed size");
   BOTAN_ARG_CHECK(m_identifier.size() == LMS_IDENTIFIER_LEN, "Invalid identifier size");
}

std::vector<uint8_t> HSS_LMS_PrivateKeyInternal::sign(std::span<const uint8_t> msg) {
   std::vector<uint8_t> sig(HSS_Signature::size(hss_params()));
   BufferStuffer sig_stuffer(sig);
   sig_stuffer.append(store_be(hss_params().L() - 1));

   std::vector<LMS_Tree_Node_Idx> q = derive_lms_leaf_indices_from_hss_index(reserve_next_idx(), hss_params());

   // Derive LMS private keys and compute buffers
   std::vector<LMS_PrivateKey> lms_key_at_layer;
   std::vector<StrongSpan<LMS_Signature_Bytes>> out_lms_sig_buffer_at_layer;
   std::vector<std::span<uint8_t>> out_child_pk_buffer_at_layer;
   for(HSS_Level layer(0); layer < hss_params().L(); ++layer) {
      // Generate key for current layer
      const HSS_LMS_Params::LMS_LMOTS_Params_Pair& layer_params = hss_params().params_at_level(layer);
      if(layer == HSS_Level(0)) {
         lms_key_at_layer.push_back(hss_derive_root_lms_private_key());
      } else {
         lms_key_at_layer.push_back(
            hss_derive_child_lms_private_key(layer_params, lms_key_at_layer.back(), q.at(layer.get() - 1)));
         out_child_pk_buffer_at_layer.push_back(sig_stuffer.next(LMS_PublicKey::size(layer_params.lms_params())));
      }
      out_lms_sig_buffer_at_layer.push_back(sig_stuffer.next<LMS_Signature_Bytes>(
         LMS_Signature::size(layer_params.lms_params(), layer_params.lmots_params())));
   }
   BOTAN_ASSERT_NOMSG(sig_stuffer.full());

   // Sign and write the signature from bottom layer to root layer
   std::vector<uint8_t> current_pk;
   for(int32_t layer_it = hss_params().L().get() - 1; layer_it >= 0; --layer_it) {
      HSS_Level layer(layer_it);
      if(layer == hss_params().L() - 1) {
         current_pk =
            lms_key_at_layer.at(layer.get())
               .sign_and_get_pk(out_lms_sig_buffer_at_layer.at(layer.get()), q.at(layer.get()), LMS_Message(msg))
               .to_bytes();
      } else {
         copy_mem(out_child_pk_buffer_at_layer.at(layer.get()), current_pk);
         current_pk =
            lms_key_at_layer.at(layer.get())
               .sign_and_get_pk(out_lms_sig_buffer_at_layer.at(layer.get()), q.at(layer.get()), LMS_Message(current_pk))
               .to_bytes();
      }
   }

   return sig;
}

LMS_PrivateKey HSS_LMS_PrivateKeyInternal::hss_derive_root_lms_private_key() const {
   auto& top_params = hss_params().params_at_level(HSS_Level(0));
   return LMS_PrivateKey(top_params.lms_params(), top_params.lmots_params(), m_identifier, m_hss_seed);
}

LMS_PrivateKey HSS_LMS_PrivateKeyInternal::hss_derive_child_lms_private_key(
   const HSS_LMS_Params::LMS_LMOTS_Params_Pair& child_lms_lmots_params,
   const LMS_PrivateKey& parent_sk,
   LMS_Tree_Node_Idx parent_q) {
   const auto hash = HashFunction::create_or_throw(child_lms_lmots_params.lms_params().hash_name());

   // CHILD_SEED = H( PARENT_I || PARENT_Q || SEED_CHILD_SEED || 0xff || PARENT_SEED )
   PseudorandomKeyGeneration seed_generator(parent_sk.identifier());
   seed_generator.set_q(parent_q.get());
   seed_generator.set_i(SEED_CHILD_SEED);
   seed_generator.set_j(0xff);
   auto child_seed = seed_generator.gen<LMS_Seed>(*hash, parent_sk.seed());

   // CHILD_I = H( PARENT_I || PARENT_Q || SEED_CHILD_I || 0xff || PARENT_SEED )
   seed_generator.set_i(SEED_CHILD_I);
   auto child_identifier = seed_generator.gen<LMS_Identifier>(*hash, parent_sk.seed());
   child_identifier.resize(LMS_IDENTIFIER_LEN);
   CT::unpoison(child_identifier);  // identifiers are part of the signature

   return LMS_PrivateKey(child_lms_lmots_params.lms_params(),
                         child_lms_lmots_params.lmots_params(),
                         std::move(child_identifier),
                         std::move(child_seed));
}

HSS_LMS_PublicKeyInternal HSS_LMS_PublicKeyInternal::create(const HSS_LMS_PrivateKeyInternal& hss_sk) {
   auto& hss_params = hss_sk.hss_params();

   const auto root_sk = hss_sk.hss_derive_root_lms_private_key();
   LMS_PublicKey top_pub_key = LMS_PublicKey(root_sk);

   return HSS_LMS_PublicKeyInternal(hss_params.L(), std::move(top_pub_key));
}

std::shared_ptr<HSS_LMS_PublicKeyInternal> HSS_LMS_PublicKeyInternal::from_bytes_or_throw(
   std::span<const uint8_t> key_bytes) {
   if(key_bytes.size() < sizeof(HSS_Level)) {
      throw Decoding_Error("Too few public key bytes.");
   }
   BufferSlicer slicer(key_bytes);

   const auto L = load_be<HSS_Level>(slicer.take<sizeof(HSS_Level)>());
   if(L > HSS_MAX_LEVELS) {
      throw Decoding_Error("Invalid number of HSS layers in public HSS-LMS key.");
   }

   LMS_PublicKey lms_pub_key = LMS_PublicKey::from_bytes_or_throw(slicer);

   if(!slicer.empty()) {
      throw Decoding_Error("Public HSS-LMS key contains more bytes than expected.");
   }
   return std::make_shared<HSS_LMS_PublicKeyInternal>(L, std::move(lms_pub_key));
}

std::vector<uint8_t> HSS_LMS_PublicKeyInternal::to_bytes() const {
   return concat<std::vector<uint8_t>>(store_be(m_L), m_top_lms_pub_key.to_bytes());
}

AlgorithmIdentifier HSS_LMS_PublicKeyInternal::algorithm_identifier() const {
   return AlgorithmIdentifier(object_identifier(), AlgorithmIdentifier::USE_EMPTY_PARAM);
}

OID HSS_LMS_PublicKeyInternal::object_identifier() const {
   return OID::from_string(algo_name());
}

size_t HSS_LMS_PublicKeyInternal::size() const {
   return sizeof(m_L) + LMS_PublicKey::size(m_top_lms_pub_key.lms_params());
}

bool HSS_LMS_PublicKeyInternal::verify_signature(std::span<const uint8_t> msg, const HSS_Signature& sig) const {
   if(checked_cast_to<HSS_Level>(sig.Nspk()) + 1 != m_L) {
      // HSS levels in the public key does not match with the signature's
      return false;
   }

   const LMS_PublicKey* lms_pk = &lms_pub_key();
   const auto hash_name = lms_pk->lms_params().hash_name();

   // Verify the signature by the above layer over the LMS public keys for layer 1 to Nspk.
   for(HSS_Level layer(0); layer < sig.Nspk(); ++layer) {
      const HSS_Signature::Signed_Pub_Key& signed_pub_key = sig.signed_pub_key(layer);
      if(signed_pub_key.public_key().lms_params().hash_name() != hash_name ||
         signed_pub_key.public_key().lmots_params().hash_name() != hash_name) {
         // We do not allow HSS-LMS instances with multiple different hash functions.
         return false;
      }
      if(!lms_pk->verify_signature(LMS_Message(signed_pub_key.public_key().to_bytes()), signed_pub_key.signature())) {
         return false;
      }
      lms_pk = &signed_pub_key.public_key();
   }

   // Verify the signature by the bottom layer over the message.
   return lms_pk->verify_signature(LMS_Message(msg), sig.bottom_sig());
}

HSS_Signature::Signed_Pub_Key::Signed_Pub_Key(LMS_Signature sig, LMS_PublicKey pub) :
      m_sig(std::move(sig)), m_pub(std::move(pub)) {}

HSS_Signature HSS_Signature::from_bytes_or_throw(std::span<const uint8_t> sig_bytes) {
   if(sig_bytes.size() < sizeof(uint32_t)) {
      throw Decoding_Error("Too few HSS signature bytes.");
   }
   BufferSlicer slicer(sig_bytes);

   const auto Nspk = load_be(slicer.take<sizeof(uint32_t)>());
   if(Nspk >= HSS_MAX_LEVELS) {
      throw Decoding_Error("Invalid number of HSS layers in signature.");
   }

   std::vector<Signed_Pub_Key> signed_pub_keys;
   for(size_t i = 0; i < Nspk; ++i) {
      LMS_Signature sig = LMS_Signature::from_bytes_or_throw(slicer);
      LMS_PublicKey pub_key = LMS_PublicKey::from_bytes_or_throw(slicer);
      signed_pub_keys.push_back(Signed_Pub_Key(std::move(sig), std::move(pub_key)));
   }

   auto sig = LMS_Signature::from_bytes_or_throw(slicer);

   if(!slicer.empty()) {
      throw Decoding_Error("HSS-LMS signature contains more bytes than expected.");
   }
   return HSS_Signature(std::move(signed_pub_keys), std::move(sig));
}

size_t HSS_Signature::size(const HSS_LMS_Params& params) {
   size_t size = sizeof(uint32_t);
   size += LMS_Signature::size(params.params_at_level(HSS_Level(0)).lms_params(),
                               params.params_at_level(HSS_Level(0)).lmots_params());
   for(HSS_Level layer(1); layer < params.L(); ++layer) {
      const auto& param = params.params_at_level(layer);
      size += LMS_PublicKey::size(param.lms_params());
      size += LMS_Signature::size(param.lms_params(), param.lmots_params());
   }
   return size;
}

}  // namespace Botan
