/*
* TLS Handshaking
* (C) 2004-2006,2011,2012,2015,2016 Jack Lloyd
*     2017 Harry Reimann, Rohde & Schwarz Cybersecurity
*
* Botan is released under the Simplified BSD License (see license.txt)
*/

#include <botan/internal/tls_handshake_state.h>

#include <botan/kdf.h>
#include <botan/pkix_types.h>
#include <botan/tls_callbacks.h>
#include <botan/tls_messages_12.h>
#include <botan/tls_policy.h>
#include <botan/tls_signature_scheme.h>
#include <botan/internal/tls_record.h>

namespace Botan::TLS {

std::string Handshake_Message::type_string() const {
   return handshake_type_to_string(type());
}

/*
* Initialize the SSL/TLS Handshake State
*/
Handshake_State::~Handshake_State() = default;

Handshake_State::Handshake_State(std::unique_ptr<Handshake_IO> io, Callbacks& cb) :
      m_callbacks(cb), m_handshake_io(std::move(io)), m_version(m_handshake_io->initial_record_version()) {}

void Handshake_State::note_message(const Handshake_Message& msg) {
   m_callbacks.tls_inspect_handshake_msg(msg);
}

void Handshake_State::hello_verify_request(const Hello_Verify_Request& hello_verify) {
   note_message(hello_verify);

   BOTAN_ASSERT_NONNULL(m_client_hello);
   m_client_hello->update_hello_cookie(hello_verify);
   hash().reset();
   hash().update(handshake_io().send(*m_client_hello));
   note_message(*m_client_hello);
}

void Handshake_State::client_hello(std::unique_ptr<Client_Hello_12> client_hello) {
   // Legacy behavior (exception to the rule): Allow client_hello to be nullptr to reset state.
   if(client_hello == nullptr) {
      m_client_hello.reset();
      hash().reset();
   } else {
      m_client_hello = std::move(client_hello);
      note_message(*m_client_hello);
   }
}

void Handshake_State::server_hello(std::unique_ptr<Server_Hello_12> server_hello) {
   BOTAN_ASSERT_NONNULL(server_hello);
   m_server_hello = std::move(server_hello);
   m_ciphersuite = Ciphersuite::by_id(m_server_hello->ciphersuite());
   note_message(*m_server_hello);
}

void Handshake_State::server_certs(std::unique_ptr<Certificate_12> server_certs) {
   BOTAN_ASSERT_NONNULL(server_certs);
   m_server_certs = std::move(server_certs);
   note_message(*m_server_certs);
}

void Handshake_State::server_cert_status(std::unique_ptr<Certificate_Status> server_cert_status) {
   BOTAN_ASSERT_NONNULL(server_cert_status);
   m_server_cert_status = std::move(server_cert_status);
   note_message(*m_server_cert_status);
}

void Handshake_State::server_kex(std::unique_ptr<Server_Key_Exchange> server_kex) {
   BOTAN_ASSERT_NONNULL(server_kex);
   m_server_kex = std::move(server_kex);
   note_message(*m_server_kex);
}

void Handshake_State::cert_req(std::unique_ptr<Certificate_Request_12> cert_req) {
   BOTAN_ASSERT_NONNULL(cert_req);
   m_cert_req = std::move(cert_req);
   note_message(*m_cert_req);
}

void Handshake_State::server_hello_done(std::unique_ptr<Server_Hello_Done> server_hello_done) {
   BOTAN_ASSERT_NONNULL(server_hello_done);
   m_server_hello_done = std::move(server_hello_done);
   note_message(*m_server_hello_done);
}

void Handshake_State::client_certs(std::unique_ptr<Certificate_12> client_certs) {
   BOTAN_ASSERT_NONNULL(client_certs);
   m_client_certs = std::move(client_certs);
   note_message(*m_client_certs);
}

void Handshake_State::client_kex(std::unique_ptr<Client_Key_Exchange> client_kex) {
   BOTAN_ASSERT_NONNULL(client_kex);
   m_client_kex = std::move(client_kex);
   note_message(*m_client_kex);
}

void Handshake_State::client_verify(std::unique_ptr<Certificate_Verify_12> client_verify) {
   BOTAN_ASSERT_NONNULL(client_verify);
   m_client_verify = std::move(client_verify);
   note_message(*m_client_verify);
}

void Handshake_State::server_verify(std::unique_ptr<Certificate_Verify_12> server_verify) {
   BOTAN_ASSERT_NONNULL(server_verify);
   m_server_verify = std::move(server_verify);
   note_message(*m_server_verify);
}

void Handshake_State::new_session_ticket(std::unique_ptr<New_Session_Ticket_12> new_session_ticket) {
   BOTAN_ASSERT_NONNULL(new_session_ticket);
   m_new_session_ticket = std::move(new_session_ticket);
   note_message(*m_new_session_ticket);
}

void Handshake_State::server_finished(std::unique_ptr<Finished_12> server_finished) {
   BOTAN_ASSERT_NONNULL(server_finished);
   m_server_finished = std::move(server_finished);
   note_message(*m_server_finished);
}

void Handshake_State::client_finished(std::unique_ptr<Finished_12> client_finished) {
   BOTAN_ASSERT_NONNULL(client_finished);
   m_client_finished = std::move(client_finished);
   note_message(*m_client_finished);
}

const Ciphersuite& Handshake_State::ciphersuite() const {
   if(!m_ciphersuite.has_value()) {
      throw Invalid_State("Cipher suite is not set");
   }
   return m_ciphersuite.value();
}

std::optional<std::string> Handshake_State::psk_identity() const {
   if(!m_client_kex) {
      return std::nullopt;
   }
   return m_client_kex->psk_identity();
}

void Handshake_State::set_version(const Protocol_Version& version) {
   m_version = version;
}

void Handshake_State::compute_session_keys() {
   BOTAN_ASSERT_NONNULL(client_kex());
   m_session_keys = Session_Keys(this, client_kex()->pre_master_secret(), false);
}

void Handshake_State::compute_session_keys(const secure_vector<uint8_t>& resume_master_secret) {
   m_session_keys = Session_Keys(this, resume_master_secret, true);
}

void Handshake_State::confirm_transition_to(Handshake_Type handshake_msg) {
   m_transitions.confirm_transition_to(handshake_msg);
}

void Handshake_State::set_expected_next(Handshake_Type handshake_msg) {
   m_transitions.set_expected_next(handshake_msg);
}

bool Handshake_State::received_handshake_msg(Handshake_Type handshake_msg) const {
   return m_transitions.received_handshake_msg(handshake_msg);
}

std::pair<Handshake_Type, std::vector<uint8_t>> Handshake_State::get_next_handshake_msg(size_t max_handshake_msg_size) {
   return m_handshake_io->get_next_record(m_transitions.change_cipher_spec_expected(), max_handshake_msg_size);
}

Session_Ticket Handshake_State::session_ticket() const {
   if(const auto* nst = new_session_ticket()) {
      const auto& ticket = nst->ticket();
      if(!ticket.empty()) {
         return ticket;
      }
   }

   BOTAN_ASSERT_NONNULL(client_hello());
   return client_hello()->session_ticket();
}

std::unique_ptr<KDF> Handshake_State::protocol_specific_prf() const {
   return m_callbacks.tls12_protocol_specific_kdf(ciphersuite().prf_algo());
}

std::pair<std::string, Signature_Format> Handshake_State::choose_sig_format(const Private_Key& key,
                                                                            Signature_Scheme& chosen_scheme,
                                                                            bool for_client_auth,
                                                                            const Policy& policy) const {
   const std::string sig_algo = key.algo_name();

   const std::vector<Signature_Scheme> allowed = policy.allowed_signature_schemes();

   if(for_client_auth) {
      BOTAN_ASSERT_NONNULL(cert_req());
   } else {
      BOTAN_ASSERT_NONNULL(client_hello());
   }

   std::vector<Signature_Scheme> requested =
      (for_client_auth) ? cert_req()->signature_schemes() : client_hello()->signature_schemes();

   for(const Signature_Scheme scheme : allowed) {
      if(!scheme.is_available()) {
         continue;
      }

      if(scheme.algorithm_name() == sig_algo) {
         if(std::find(requested.begin(), requested.end(), scheme) != requested.end()) {
            chosen_scheme = scheme;
            break;
         }
      }
   }

   if(!chosen_scheme.is_set()) {
      throw TLS_Exception(Alert::HandshakeFailure,
                          "Could not agree on a signature scheme with peer for " + sig_algo + " key");
   }

   const std::string hash = chosen_scheme.hash_function_name();

   if(!policy.allowed_signature_hash(hash)) {
      throw TLS_Exception(Alert::HandshakeFailure, "Policy refuses to accept signing with any hash supported by peer");
   }

   if(!chosen_scheme.format().has_value()) {
      throw Invalid_Argument(sig_algo + " is invalid/unknown for TLS signatures");
   }

   return std::make_pair(chosen_scheme.padding_string(), chosen_scheme.format().value());
}

namespace {

bool supported_algos_include(const std::vector<Signature_Scheme>& schemes, Signature_Scheme received_scheme) {
   for(const Signature_Scheme scheme : schemes) {
      if(scheme == received_scheme && scheme.is_available()) {
         return true;
      }
   }

   return false;
}

}  // namespace

std::pair<std::string, Signature_Format> Handshake_State::parse_sig_format(
   const Public_Key& key,
   Signature_Scheme scheme,
   const std::vector<Signature_Scheme>& offered_schemes,
   bool for_client_auth,
   const Policy& policy) const {
   const std::string key_type = key.algo_name();

   if(!policy.allowed_signature_method(key_type)) {
      throw TLS_Exception(Alert::HandshakeFailure, "Rejecting " + key_type + " signature");
   }

   if(!scheme.is_available()) {
      throw TLS_Exception(Alert::IllegalParameter, "Peer sent unknown signature scheme");
   }

   if(key_type != scheme.algorithm_name()) {
      throw Decoding_Error("Counterparty sent inconsistent key and sig types");
   }

   if(for_client_auth && cert_req() == nullptr) {
      throw TLS_Exception(Alert::HandshakeFailure, "No CertificateVerify message received");
   }

   /*
   Confirm the signature type we just received against the
   supported_algos list that we sent; it better be there.
   */

   const std::vector<Signature_Scheme> supported_algos =
      for_client_auth ? cert_req()->signature_schemes() : offered_schemes;

   if(!scheme.is_compatible_with(Protocol_Version::TLS_V12)) {
      throw TLS_Exception(Alert::IllegalParameter, "Peer sent unacceptable signature scheme");
   }

   if(!supported_algos_include(supported_algos, scheme)) {
      throw TLS_Exception(Alert::IllegalParameter,
                          "TLS signature extension did not allow for " + scheme.to_string() + " signature");
   }

   if(!scheme.format().has_value()) {
      throw Invalid_Argument(key_type + " is invalid/unknown for TLS signatures");
   }

   return std::make_pair(scheme.padding_string(), scheme.format().value());
}

}  // namespace Botan::TLS
