// Copyright 2017 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

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

#include <utility>
#include <vector>

#include "google/protobuf/message_lite.h"
#include "src/lib/util/not_null.h"
#include "src/logging.h"
#include "src/pb/encrypted_message.pb.h"
#include "src/pb/key.pb.h"
#include "src/public/lib/status.h"
#include "src/public/lib/status_codes.h"
#include "src/tracing.h"

#if TQ_COBALT_USE_TINK
#include "src/lib/util/hybrid_tink_encrypted_message_maker.h"
#endif

namespace cobalt::util {

namespace {

#if TQ_COBALT_USE_TINK
constexpr char kShufflerContextInfo[] = "cobalt-1.0-shuffler";
constexpr char kAnalyzerContextInfo[] = "cobalt-1.0-analyzer";
#endif

// Parse and validate a serialized CobaltEncryptionKey.
lib::statusor::StatusOr<cobalt::CobaltEncryptionKey> ParseCobaltEncryptionKey(
    const std::string& cobalt_encryption_key_bytes) {
  cobalt::CobaltEncryptionKey cobalt_encryption_key;
  if (!cobalt_encryption_key.ParseFromString(cobalt_encryption_key_bytes)) {
    return Status(StatusCode::INVALID_ARGUMENT,
                  "EncryptedMessageMaker: Invalid Cobalt encryption key.");
  }

  if (cobalt_encryption_key.key_index() == 0) {
    return Status(StatusCode::INVALID_ARGUMENT,
                  "EncryptedMessageMaker: Invalid Cobalt key_index: 0.");
  }

  if (cobalt_encryption_key.purpose() != cobalt::CobaltEncryptionKey::SHUFFLER &&
      cobalt_encryption_key.purpose() != cobalt::CobaltEncryptionKey::ANALYZER) {
    return Status(StatusCode::INVALID_ARGUMENT,
                  "EncryptedMessageMaker: Invalid Cobalt key purpose.");
  }

  return cobalt_encryption_key;
}
}  // namespace

util::NotNullUniquePtr<EncryptedMessageMaker> EncryptedMessageMaker::MakeUnencrypted() {
  VLOG(5) << "WARNING: Cobalt data will not be encrypted!";
  return util::MakeNotNullUniquePtr<UnencryptedMessageMaker>();
}

lib::statusor::StatusOr<util::NotNullUniquePtr<EncryptedMessageMaker>>
EncryptedMessageMaker::MakeForEnvelopes(const std::string& cobalt_encryption_key_bytes) {
  auto cobalt_encryption_key_or_result = ParseCobaltEncryptionKey(cobalt_encryption_key_bytes);
  if (!cobalt_encryption_key_or_result.ok()) {
    return cobalt_encryption_key_or_result.status();
  }
  auto cobalt_encryption_key = cobalt_encryption_key_or_result.value();

  if (cobalt_encryption_key.purpose() != cobalt::CobaltEncryptionKey::SHUFFLER) {
    return Status(StatusCode::INVALID_ARGUMENT, "EncryptedMessageMaker: Expected a shuffler key.");
  }

#if TQ_COBALT_USE_TINK
  return MakeHybridTinkEncryptedMessageMaker(cobalt_encryption_key.serialized_key(),
                                             kShufflerContextInfo,
                                             cobalt_encryption_key.key_index());
#else
  return Status(StatusCode::UNIMPLEMENTED, "EncryptedMessageMaker: Cobalt was built without Tink.");
#endif
}

lib::statusor::StatusOr<util::NotNullUniquePtr<EncryptedMessageMaker>>
EncryptedMessageMaker::MakeForObservations(const std::string& cobalt_encryption_key_bytes) {
  auto cobalt_encryption_key_or_result = ParseCobaltEncryptionKey(cobalt_encryption_key_bytes);
  if (!cobalt_encryption_key_or_result.ok()) {
    return cobalt_encryption_key_or_result.status();
  }
  auto cobalt_encryption_key = cobalt_encryption_key_or_result.value();

  if (cobalt_encryption_key.purpose() != cobalt::CobaltEncryptionKey::ANALYZER) {
    return Status(StatusCode::INVALID_ARGUMENT, "EncryptedMessageMaker: Expected an analyzer key.");
  }

#if TQ_COBALT_USE_TINK
  return MakeHybridTinkEncryptedMessageMaker(cobalt_encryption_key.serialized_key(),
                                             kAnalyzerContextInfo,
                                             cobalt_encryption_key.key_index());
#else
  return Status(StatusCode::UNIMPLEMENTED, "EncryptedMessageMaker: Cobalt was built without Tink.");
#endif
}

bool UnencryptedMessageMaker::Encrypt(const google::protobuf::MessageLite& message,
                                      EncryptedMessage* encrypted_message) const {
  TRACE_DURATION("cobalt_core", "UnencryptedMessageMaker::Encrypt");
  if (!encrypted_message) {
    return false;
  }
  std::string serialized_message;
  message.SerializeToString(&serialized_message);
  encrypted_message->set_ciphertext(serialized_message);
  VLOG(5) << "EncryptedMessage: encryption_scheme=NONE.";
  return true;
}

}  // namespace cobalt::util