
#include "src/lib/util/hybrid_tink_encrypted_message_maker.h"

#include "src/logging.h"
#include "src/tracing.h"
#include "third_party/tink/cc/hybrid/hybrid_config.h"
#include "third_party/tink/cc/hybrid_encrypt.h"
#include "third_party/tink/cc/keyset_handle.h"
#include "third_party/tink/cc/util/status.h"

namespace cobalt::util {

Status StatusFromTinkStatus(const ::crypto::tink::util::Status& tink_status) {
  #ifndef TINK_USE_ABSL_STATUS
  return Status(static_cast<StatusCode>(tink_status.error_code()), tink_status.error_message());
  #else
  return Status(static_cast<StatusCode>(tink_status.code()), std::string(tink_status.message()));
  #endif
}

// Make a HybridTinkEncryptedMessageMaker from a serialized encoded keyset.
lib::statusor::StatusOr<std::unique_ptr<EncryptedMessageMaker>> MakeHybridTinkEncryptedMessageMaker(
    // TODO(fxbug.dev/85571): NOLINTNEXTLINE(bugprone-easily-swappable-parameters)
    const std::string& public_keyset_bytes, const std::string& context_info, uint32_t key_index) {
  auto status = ::crypto::tink::HybridConfig::Register();
  if (!status.ok()) {
    return StatusFromTinkStatus(status);
  }

  auto read_result = ::crypto::tink::KeysetHandle::ReadNoSecret(public_keyset_bytes);
  if (!read_result.ok()) {
    return StatusFromTinkStatus(read_result.status());
  }
  auto keyset_handle = std::move(read_result.ValueOrDie());

  auto primitive_result = keyset_handle->GetPrimitive<::crypto::tink::HybridEncrypt>();
  if (!primitive_result.ok()) {
    return StatusFromTinkStatus(primitive_result.status());
  }

  std::unique_ptr<EncryptedMessageMaker> maker(new HybridTinkEncryptedMessageMaker(
      std::move(primitive_result.ValueOrDie()), context_info, key_index));

  return maker;
}

HybridTinkEncryptedMessageMaker::HybridTinkEncryptedMessageMaker(
    std::unique_ptr<::crypto::tink::HybridEncrypt> encrypter, std::string context_info,
    uint32_t key_index)
    : encrypter_(std::move(encrypter)),
      context_info_(std::move(context_info)),
      key_index_(key_index) {}

bool HybridTinkEncryptedMessageMaker::Encrypt(const google::protobuf::MessageLite& message,
                                              EncryptedMessage* encrypted_message) const {
  TRACE_DURATION("cobalt_core", "HybridTinkEncryptedMessageMaker::Encrypt");
  if (!encrypted_message) {
    return false;
  }

  std::string serialized_message;
  message.SerializeToString(&serialized_message);

  VLOG(5) << "EncryptedMessage: encryption_scheme=HYBRID_TINK.";

  auto encrypted_result = encrypter_->Encrypt(serialized_message, context_info_);
  if (!encrypted_result.ok()) {
    VLOG(5) << "EncryptedMessage: Tink could not encrypt message: "
            << encrypted_result.status().error_message();
    return false;
  }
  encrypted_message->set_ciphertext(encrypted_result.ValueOrDie());
  if (key_index_ == 0) {
    encrypted_message->set_scheme(EncryptedMessage::HYBRID_TINK);
  } else {
    encrypted_message->set_key_index(key_index_);
  }

  return true;
}

}  // namespace cobalt::util