// Copyright 2021 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "ash/quick_pair/fast_pair_handshake/fast_pair_encryption.h"

#include <algorithm>
#include <array>
#include <optional>

#include "ash/quick_pair/fast_pair_handshake/fast_pair_key_pair.h"
#include "base/check.h"
#include "base/containers/span.h"
#include "base/types/fixed_array.h"
#include "chromeos/ash/services/quick_pair/public/cpp/fast_pair_message_type.h"
#include "components/cross_device/logging/logging.h"
#include "third_party/boringssl/src/include/openssl/aes.h"
#include "third_party/boringssl/src/include/openssl/base.h"
#include "third_party/boringssl/src/include/openssl/ec.h"
#include "third_party/boringssl/src/include/openssl/ec_key.h"
#include "third_party/boringssl/src/include/openssl/ecdh.h"
#include "third_party/boringssl/src/include/openssl/evp.h"
#include "third_party/boringssl/src/include/openssl/hmac.h"
#include "third_party/boringssl/src/include/openssl/nid.h"
#include "third_party/boringssl/src/include/openssl/sha.h"

namespace {

using ash::quick_pair::fast_pair_encryption::kBlockSizeBytes;

// Converts the public anti-spoofing key into an EC_Point.
bssl::UniquePtr<EC_POINT> GetEcPointFromPublicAntiSpoofingKey(
    const bssl::UniquePtr<EC_GROUP>& ec_group,
    const std::string& decoded_public_anti_spoofing) {
  std::array<uint8_t, kPublicKeyByteSize + 1> buffer;
  buffer[0] = POINT_CONVERSION_UNCOMPRESSED;
  std::ranges::copy(decoded_public_anti_spoofing, buffer.begin() + 1);

  bssl::UniquePtr<EC_POINT> new_ec_point(EC_POINT_new(ec_group.get()));

  if (!EC_POINT_oct2point(ec_group.get(), new_ec_point.get(), buffer.data(),
                          buffer.size(), nullptr)) {
    return nullptr;
  }

  return new_ec_point;
}

// Key derivation function to be used in hashing the generated secret key.
void* KDF(const void* in, size_t inlen, void* out, size_t* outlen) {
  // Set this to 16 since that's the amount of bytes we want to use
  // for the key, even though more will be written by SHA256 below.
  *outlen = kPrivateKeyByteSize;
  return SHA256(static_cast<const uint8_t*>(in), inlen,
                static_cast<uint8_t*>(out));
}

}  // namespace

namespace ash {
namespace quick_pair {
namespace fast_pair_encryption {

std::optional<KeyPair> GenerateKeysWithEcdhKeyAgreement(
    const std::string& decoded_public_anti_spoofing) {
  if (decoded_public_anti_spoofing.size() != kPublicKeyByteSize) {
    CD_LOG(WARNING, Feature::FP) << "Expected " << kPublicKeyByteSize
                                 << " byte value for anti-spoofing key. Got:"
                                 << decoded_public_anti_spoofing.size();
    return std::nullopt;
  }

  // Generate the secp256r1 key-pair.
  bssl::UniquePtr<EC_GROUP> ec_group(
      EC_GROUP_new_by_curve_name(NID_X9_62_prime256v1));
  bssl::UniquePtr<EC_KEY> ec_key(
      EC_KEY_new_by_curve_name(NID_X9_62_prime256v1));

  if (!EC_KEY_generate_key(ec_key.get())) {
    CD_LOG(WARNING, Feature::FP) << __func__ << ": Failed to generate ec key";
    return std::nullopt;
  }

  // The ultimate goal here is to get a 64-byte public key. We accomplish this
  // by converting the generated public key into the uncompressed X9.62 format,
  // which is 0x04 followed by padded x and y coordinates.
  std::array<uint8_t, kPublicKeyByteSize + 1> uncompressed_private_key;
  int point_bytes_written = EC_POINT_point2oct(
      ec_group.get(), EC_KEY_get0_public_key(ec_key.get()),
      POINT_CONVERSION_UNCOMPRESSED, uncompressed_private_key.data(),
      uncompressed_private_key.size(), nullptr);

  if (point_bytes_written != uncompressed_private_key.size()) {
    CD_LOG(WARNING, Feature::FP)
        << __func__
        << ": EC_POINT_point2oct failed to convert public key to "
           "uncompressed x9.62 format.";
    return std::nullopt;
  }

  bssl::UniquePtr<EC_POINT> public_anti_spoofing_point =
      GetEcPointFromPublicAntiSpoofingKey(ec_group,
                                          decoded_public_anti_spoofing);

  if (!public_anti_spoofing_point) {
    CD_LOG(WARNING, Feature::FP)
        << __func__
        << ": Failed to convert Public Anti-Spoofing key to EC_POINT";
    return std::nullopt;
  }

  std::array<uint8_t, SHA256_DIGEST_LENGTH> secret;
  int computed_key_size =
      ECDH_compute_key(secret.data(), secret.size(),
                       public_anti_spoofing_point.get(), ec_key.get(), &KDF);

  if (computed_key_size != kPrivateKeyByteSize) {
    CD_LOG(WARNING, Feature::FP) << __func__ << ": ECDH_compute_key failed.";
    return std::nullopt;
  }

  // Take first 16 bytes from secret as the private key.
  std::array<uint8_t, kPrivateKeyByteSize> private_key;
  base::span(private_key)
      .copy_from(base::span(secret).first<kPrivateKeyByteSize>());

  // Ignore the first byte since it is 0x04, from the above uncompressed X9 .62
  // format.
  std::array<uint8_t, kPublicKeyByteSize> public_key;
  std::copy(uncompressed_private_key.begin() + 1,
            uncompressed_private_key.end(), public_key.begin());

  return KeyPair(private_key, public_key);
}

const std::array<uint8_t, kHmacSizeBytes> GenerateHmacSha256(
    const std::array<uint8_t, kSecretKeySizeBytes>& secret_key,
    std::array<uint8_t, kNonceSizeBytes> nonce,
    const std::vector<uint8_t>& data) {
  const size_t nonce_data_concat_size = kNonceSizeBytes + data.size();
  base::FixedArray<uint8_t> nonce_data_concat(nonce_data_concat_size);
  base::span<uint8_t> nonce_data_concat_span(nonce_data_concat);
  nonce_data_concat_span.first<kNonceSizeBytes>().copy_from(nonce);
  nonce_data_concat_span.subspan<kNonceSizeBytes>().copy_from(data);

  std::array<uint8_t, kHmacKeySizeBytes> K = {};
  base::span(K).first<kSecretKeySizeBytes>().copy_from(secret_key);

  std::array<uint8_t, kHmacSizeBytes> output;
  unsigned int output_size;
  HMAC(/*evp_md=*/EVP_sha256(), /*key=*/K.data(),
       /*key_len=*/kHmacKeySizeBytes, /*data=*/nonce_data_concat.data(),
       /*data_len=*/nonce_data_concat.size(),
       /*out=*/output.data(), /*out_len*/ &output_size);
  return output;
}

const std::array<uint8_t, kBlockSizeBytes> EncryptBytes(
    const std::array<uint8_t, kBlockSizeBytes>& aes_key_bytes,
    const std::array<uint8_t, kBlockSizeBytes>& bytes_to_encrypt) {
  AES_KEY aes_key;
  int aes_key_was_set = AES_set_encrypt_key(aes_key_bytes.data(),
                                            aes_key_bytes.size() * 8, &aes_key);
  DCHECK(aes_key_was_set == 0) << "Invalid AES key size.";
  std::array<uint8_t, kBlockSizeBytes> encrypted_bytes;
  AES_encrypt(bytes_to_encrypt.data(), encrypted_bytes.data(), &aes_key);
  return encrypted_bytes;
}

const std::vector<uint8_t> EncryptAdditionalData(
    const std::array<uint8_t, kSecretKeySizeBytes>& secret_key,
    std::array<uint8_t, kNonceSizeBytes> nonce,
    const std::vector<uint8_t>& data) {
  if (data.empty()) {
    return {};
  }

  AES_KEY aes_key;
  int aes_key_was_set =
      AES_set_encrypt_key(secret_key.data(), secret_key.size() * 8, &aes_key);
  DCHECK(aes_key_was_set == 0) << "Invalid AES key size.";

  unsigned int bytes_read = 0;
  std::array<uint8_t, AES_BLOCK_SIZE> ivec = {};
  std::array<uint8_t, AES_BLOCK_SIZE> ecount = {};
  constexpr size_t kNonceOffsetBytes = 8;

  base::FixedArray<uint8_t> encrypted_data(data.size());

  // The Fast Pair Spec AES-CTR version increments the first byte of the
  // initialization vector; the typical AES-CTR algorithm increments the
  // last byte. So, instead of calling AES_ctr128_encrypt() once on all of
  // `data`, it is called on each 128-bit block of `data` and the counter is
  // incremented manually.
  size_t bytes_to_encrypt = data.size();
  size_t i = 0;
  while (bytes_to_encrypt > 0) {
    const size_t block_size =
        std::min(bytes_to_encrypt, size_t{AES_BLOCK_SIZE});
    ivec.fill(0);
    base::span(ivec).subspan<kNonceOffsetBytes, kNonceSizeBytes>().copy_from(
        nonce);
    ivec.front() = static_cast<uint8_t>(i);
    const size_t offset = data.size() - bytes_to_encrypt;
    const base::span<const uint8_t> input =
        base::span(data).subspan(offset, block_size);
    base::span<uint8_t> output =
        base::span(encrypted_data).subspan(offset, block_size);
    AES_ctr128_encrypt(/*in=*/input.data(), /*out=*/output.data(),
                       /*len=*/input.size(), &aes_key, /*ivec=*/ivec.data(),
                       /*ecount_buf=*/ecount.data(), &bytes_read);

    bytes_to_encrypt -= block_size;
    i++;
  }

  CHECK(!bytes_to_encrypt);

  return std::vector<uint8_t>(encrypted_data.begin(), encrypted_data.end());
}

}  // namespace fast_pair_encryption
}  // namespace quick_pair
}  // namespace ash