/*
 * Classic McEliece key generation with Internal Private and Public Key classes
 * (C) 2023 Jack Lloyd
 *     2023,2024 Fabian Albert, Amos Treiber - Rohde & Schwarz Cybersecurity
 *
 * Botan is released under the Simplified BSD License (see license.txt)
 **/

#include <botan/internal/cmce_keys_internal.h>

#include <botan/internal/buffer_slicer.h>
#include <botan/internal/concat_util.h>

namespace Botan {

namespace {

/**
 * @brief Try to generate a Classic McEliece keypair for a given seed.
 *
 * @param[out] out_next_seed The next seed to use for key generation, if this iteration fails
 * @param params Classic McEliece parameters
 * @param seed The seed to used for this key generation iteration
 * @return a keypair on success, std::nullopt otherwise
 */
std::optional<Classic_McEliece_KeyPair_Internal> try_generate_keypair(std::span<uint8_t> out_next_seed,
                                                                      const Classic_McEliece_Parameters& params,
                                                                      CmceKeyGenSeed seed) {
   BOTAN_ASSERT_EQUAL(seed.size(), 32, "Valid seed length");
   BOTAN_ASSERT_EQUAL(out_next_seed.size(), 32, "Valid output seed length");

   auto big_e_xof = params.prg(seed);

   auto s = big_e_xof->output<CmceRejectionSeed>(params.n() / 8);
   auto ordering_bits = big_e_xof->output<CmceOrderingBits>((params.sigma2() * params.q()) / 8);
   auto irreducible_bits = big_e_xof->output<CmceIrreducibleBits>((params.sigma1() * params.t()) / 8);
   big_e_xof->output(out_next_seed);

   // Field-ordering generation - Classic McEliece ISO 8.2
   auto field_ordering = Classic_McEliece_Field_Ordering::create_field_ordering(params, ordering_bits);
   if(!field_ordering) {
      return std::nullopt;
   }

   // Irreducible-polynomial generation - Classic McEliece ISO 8.1
   auto g = params.poly_ring().compute_minimal_polynomial(irreducible_bits);
   if(!g) {
      return std::nullopt;
   }

   // Matrix generation for Goppa codes - Classic McEliece ISO 7.2
   auto pk_matrix_and_pivots =
      Classic_McEliece_Matrix::create_matrix_and_apply_pivots(params, field_ordering.value(), g.value());
   if(!pk_matrix_and_pivots) {
      return std::nullopt;
   }
   auto& [pk_matrix, pivots] = pk_matrix_and_pivots.value();

   // Key generation was successful - Create and return keys
   return Classic_McEliece_KeyPair_Internal{
      .private_key = std::make_shared<Classic_McEliece_PrivateKeyInternal>(
         params, std::move(seed), pivots, std::move(g.value()), std::move(field_ordering.value()), std::move(s)),
      .public_key = std::make_shared<Classic_McEliece_PublicKeyInternal>(params, std::move(pk_matrix))};
}

}  // namespace

Classic_McEliece_PrivateKeyInternal Classic_McEliece_PrivateKeyInternal::from_bytes(
   const Classic_McEliece_Parameters& params, std::span<const uint8_t> sk_bytes) {
   BOTAN_ASSERT(sk_bytes.size() == params.sk_size_bytes(), "Valid private key size");
   BufferSlicer sk_slicer(sk_bytes);

   auto delta = sk_slicer.copy<CmceKeyGenSeed>(params.seed_len());
   auto c = CmceColumnSelection(sk_slicer.take(params.sk_c_bytes()));
   auto g = Classic_McEliece_Minimal_Polynomial::from_bytes(sk_slicer.take(params.sk_poly_g_bytes()), params.poly_f());
   auto field_ordering = Classic_McEliece_Field_Ordering::create_from_control_bits(
      params, secure_bitvector(sk_slicer.take(params.sk_alpha_control_bytes())));
   auto s = sk_slicer.copy<CmceRejectionSeed>(params.sk_s_bytes());
   BOTAN_ASSERT_NOMSG(sk_slicer.empty());

   return Classic_McEliece_PrivateKeyInternal(
      params, std::move(delta), std::move(c), std::move(g), std::move(field_ordering), std::move(s));
}

secure_vector<uint8_t> Classic_McEliece_PrivateKeyInternal::serialize() const {
   auto control_bits = m_field_ordering.alphas_control_bits();

   /* NIST Impl. guide 6.1 Control-Bit Gen:
    *     As low-cost protection against faults in the control-bit computation, implementors are advised
    *     to check after the computation that applying the Benes network produces pi, and to
    *     restart key generation if this test fails; applying the Benes network is very fast.
    *
    * Here, we just assert that applying the Benes network produces pi.
    */
   BOTAN_ASSERT(Classic_McEliece_Field_Ordering::create_from_control_bits(m_params, control_bits)
                   .ct_is_equal(m_field_ordering)
                   .as_bool(),
                "Control Bit Computation Check");

   return concat(m_delta.get(), m_c.get().to_bytes(), m_g.serialize(), control_bits.to_bytes(), m_s);
}

bool Classic_McEliece_PrivateKeyInternal::check_key() const {
   auto prg = m_params.prg(m_delta);

   const auto s = prg->output<CmceRejectionSeed>(m_params.n() / 8);
   const auto ordering_bits = prg->output<CmceOrderingBits>((m_params.sigma2() * m_params.q()) / 8);
   const auto irreducible_bits = prg->output<CmceIrreducibleBits>((m_params.sigma1() * m_params.t()) / 8);

   // Recomputing s as hash of delta
   auto ret = CT::Mask<size_t>::expand(CT::is_equal<uint8_t>(s.data(), m_s.data(), m_params.n() / 8));

   // Checking weight of c
   ret &= CT::Mask<size_t>::is_equal(c().hamming_weight(), 32);

   if(auto g = m_params.poly_ring().compute_minimal_polynomial(irreducible_bits)) {
      for(size_t i = 0; i < g->degree() - 1; ++i) {
         ret &= CT::Mask<size_t>::expand(GF_Mask::is_equal(g->coef_at(i), m_g.coef_at(i)).elem_mask());
      }
   } else {
      ret = CT::Mask<size_t>::cleared();
   }

   // Check alpha control bits
   if(auto field_ord_from_seed = Classic_McEliece_Field_Ordering::create_field_ordering(m_params, ordering_bits)) {
      field_ord_from_seed->permute_with_pivots(m_params, c());
      ret &= CT::Mask<size_t>::expand(field_ord_from_seed->ct_is_equal(field_ordering()));
   } else {
      ret = CT::Mask<size_t>::cleared();
   }

   return ret.as_bool();
}

std::shared_ptr<Classic_McEliece_PublicKeyInternal> Classic_McEliece_PublicKeyInternal::create_from_private_key(
   const Classic_McEliece_PrivateKeyInternal& sk) {
   auto pk_matrix_and_pivot = Classic_McEliece_Matrix::create_matrix(sk.params(), sk.field_ordering(), sk.g());
   if(!pk_matrix_and_pivot.has_value()) {
      throw Decoding_Error("Cannot create public key from private key. Private key is invalid.");
   }
   auto& [pk_matrix, pivot] = pk_matrix_and_pivot.value();

   // There should not be a pivot of any other form. Otherwise the gauss
   // algorithm failed effectively.
   if(!CT::driveby_unpoison(pivot.equals(bitvector{0xff, 0xff, 0xff, 0xff, 0x00, 0x00, 0x00, 0x00}))) {
      throw Decoding_Error("Cannot create public key from private key. Private key is invalid.");
   }

   return std::make_shared<Classic_McEliece_PublicKeyInternal>(sk.params(), std::move(pk_matrix));
}

Classic_McEliece_KeyPair_Internal Classic_McEliece_KeyPair_Internal::generate(const Classic_McEliece_Parameters& params,
                                                                              StrongSpan<const CmceInitialSeed> seed) {
   BOTAN_ASSERT_EQUAL(seed.size(), params.seed_len(), "Valid seed length");

   CmceKeyGenSeed next_seed(seed.size());
   CmceKeyGenSeed current_seed(seed.begin(), seed.end());

   while(true) {
      if(auto keypair = try_generate_keypair(next_seed, params, std::move(current_seed))) {
         return keypair.value();
      }
      current_seed = next_seed;
   }
}

}  // namespace Botan
