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

#include "components/trusted_vault/icloud_keychain_recovery_factor.h"

#include <algorithm>
#include <optional>

#include "base/functional/callback_helpers.h"
#include "base/strings/strcat.h"
#include "base/task/bind_post_task.h"
#include "components/trusted_vault/icloud_recovery_key_mac.h"
#include "components/trusted_vault/proto/vault.pb.h"
#include "components/trusted_vault/securebox.h"
#include "components/trusted_vault/trusted_vault_connection.h"
#include "components/trusted_vault/trusted_vault_crypto.h"
#include "components/trusted_vault/trusted_vault_histograms.h"
#include "components/trusted_vault/trusted_vault_server_constants.h"

namespace trusted_vault {

namespace {

constexpr char kICloudKeychainRecoveryKeyAccessGroupSuffix[] =
    ".com.google.common.folsom";

std::optional<std::vector<std::vector<uint8_t>>> DecryptTrustedVaultWrappedKeys(
    const SecureBoxPrivateKey& private_key,
    const std::vector<MemberKeys>& member_keys) {
  std::vector<std::vector<uint8_t>> decrypted_keys;

  for (const auto& member_key : member_keys) {
    std::optional<std::vector<uint8_t>> decrypted_key =
        DecryptTrustedVaultWrappedKey(private_key, member_key.wrapped_key);
    if (!decrypted_key) {
      return std::nullopt;
    }

    decrypted_keys.emplace_back(*decrypted_key);
  }

  return decrypted_keys;
}

}  // namespace

ICloudKeychainRecoveryFactor::ICloudKeychainRecoveryFactor(
    const std::string& icloud_keychain_access_group_prefix,
    SecurityDomainId security_domain_id,
    ICloudKeychainStorage* storage,
    KeyStorage* key_storage,
    TrustedVaultThrottlingConnection* connection,
    CoreAccountInfo primary_account)
    : icloud_keychain_access_group_(
          base::StrCat({icloud_keychain_access_group_prefix,
                        kICloudKeychainRecoveryKeyAccessGroupSuffix})),
      security_domain_id_(security_domain_id),
      storage_(storage),
      key_storage_(key_storage),
      connection_(connection),
      primary_account_(primary_account) {
  CHECK(storage_);
  CHECK(key_storage_);
  CHECK(connection_);
}
ICloudKeychainRecoveryFactor::~ICloudKeychainRecoveryFactor() = default;

LocalRecoveryFactorType ICloudKeychainRecoveryFactor::GetRecoveryFactorType()
    const {
  return LocalRecoveryFactorType::kICloudKeychain;
}

void ICloudKeychainRecoveryFactor::AttemptRecovery(AttemptRecoveryCallback cb) {
  if (key_storage_->HasNonConstantKey(primary_account_.gaia)) {
    // iCloud Keychain is only used to recover keys if there were no
    // non-constant keys available previously.
    FulfillRecoveryWithFailure(
        TrustedVaultDownloadKeysStatusForUMA::kKeyProofVerificationNotSupported,
        std::move(cb));
    return;
  }

  // ICloudRecoveryKey::Retrieve() can't be cancelled, so we use a weak pointer
  // for the callback.
  ICloudRecoveryKey::Retrieve(
      base::BindOnce(
          &ICloudKeychainRecoveryFactor::OnICloudKeysRetrievedForRecovery,
          recovery_weak_ptr_factory_.GetWeakPtr(), std::move(cb)),
      security_domain_id_, icloud_keychain_access_group_);
}

void ICloudKeychainRecoveryFactor::OnICloudKeysRetrievedForRecovery(
    AttemptRecoveryCallback cb,
    std::vector<std::unique_ptr<ICloudRecoveryKey>> local_icloud_keys) {
  if (local_icloud_keys.empty()) {
    MarkAsNotRegistered();
    FulfillRecoveryWithFailure(
        TrustedVaultDownloadKeysStatusForUMA::kDeviceNotRegistered,
        std::move(cb));
    return;
  }

  if (connection_->AreRequestsThrottled(primary_account_)) {
    // Keys download attempt is not possible.
    FulfillRecoveryWithFailure(
        TrustedVaultDownloadKeysStatusForUMA::kThrottledClientSide,
        std::move(cb));
    return;
  }

  ongoing_download_registration_state_request_for_recovery_ =
      connection_->DownloadAuthenticationFactorsRegistrationState(
          primary_account_,
          {trusted_vault_pb::SecurityDomainMember::MEMBER_TYPE_ICLOUD_KEYCHAIN},
          base::BindOnce(&ICloudKeychainRecoveryFactor::
                             OnRecoveryFactorStateDownloadedForRecovery,
                         // `this` outlives `ongoing_request_for_recovery_`.
                         base::Unretained(this), std::move(cb),
                         std::move(local_icloud_keys)),
          base::NullCallback());
  CHECK(ongoing_download_registration_state_request_for_recovery_);
}

void ICloudKeychainRecoveryFactor::OnRecoveryFactorStateDownloadedForRecovery(
    AttemptRecoveryCallback cb,
    std::vector<std::unique_ptr<ICloudRecoveryKey>> local_icloud_keys,
    DownloadAuthenticationFactorsRegistrationStateResult result) {
  // This method should be called only as a result of
  // `ongoing_request_for_recovery_` completion/failure, verify this condition
  // and destroy `ongoing_request_for_recovery_` as it's not  needed anymore.
  CHECK(ongoing_download_registration_state_request_for_recovery_);
  ongoing_download_registration_state_request_for_recovery_ = nullptr;

  if (result.state ==
      DownloadAuthenticationFactorsRegistrationStateResult::State::kError) {
    connection_->RecordFailedRequestForThrottling(primary_account_);
    FulfillRecoveryWithFailure(
        TrustedVaultDownloadKeysStatusForUMA::kNetworkError, std::move(cb));
    return;
  }

  TrustedVaultDownloadKeysStatusForUMA last_status =
      TrustedVaultDownloadKeysStatusForUMA::kDeviceNotRegistered;
  for (const VaultMember& recovery_icloud_key : result.icloud_keys) {
    if (recovery_icloud_key.member_keys.size() == 0) {
      last_status = TrustedVaultDownloadKeysStatusForUMA::kMembershipEmpty;
      continue;
    }

    const std::vector<uint8_t> public_key =
        recovery_icloud_key.public_key->ExportToBytes();

    const auto local_icloud_key_it = std::ranges::find_if(
        local_icloud_keys,
        [&public_key](const auto& key) { return key->id() == public_key; });
    if (local_icloud_key_it != local_icloud_keys.end()) {
      std::optional<std::vector<std::vector<uint8_t>>> new_vault_keys =
          DecryptTrustedVaultWrappedKeys(
              (*local_icloud_key_it)->key()->private_key(),
              recovery_icloud_key.member_keys);

      if (!new_vault_keys) {
        last_status =
            TrustedVaultDownloadKeysStatusForUMA::kMembershipCorrupted;
        continue;
      }

      // Success: all keys were successfully decrypted.
      RecordTrustedVaultDownloadKeysStatus(
          LocalRecoveryFactorType::kICloudKeychain, security_domain_id_,
          TrustedVaultDownloadKeysStatusForUMA::kSuccess);

      int last_vault_key_version =
          std::ranges::max_element(recovery_icloud_key.member_keys, {},
                                   &MemberKeys::version)
              ->version;
      MarkAsRegistered();
      std::move(cb).Run(RecoveryStatus::kSuccess, *new_vault_keys,
                        last_vault_key_version);
      return;
    }
  }

  // None of the retrieved iCloud Keychain keys is in the security domain -
  // fail with the last status. This makes sure to record that status only once
  // per recovery attempt rather than once per vault member, which could skew
  // the metrics.
  MarkAsNotRegistered();
  FulfillRecoveryWithFailure(last_status, std::move(cb));
}

void ICloudKeychainRecoveryFactor::FulfillRecoveryWithFailure(
    TrustedVaultDownloadKeysStatusForUMA status_for_uma,
    AttemptRecoveryCallback cb) {
  RecordTrustedVaultDownloadKeysStatus(LocalRecoveryFactorType::kICloudKeychain,
                                       security_domain_id_, status_for_uma);

  base::BindPostTaskToCurrentDefault(
      base::BindOnce(std::move(cb), RecoveryStatus::kFailure,
                     /*new_vault_keys=*/std::vector<std::vector<uint8_t>>(),
                     /*last_vault_key_version=*/0))
      .Run();
}

bool ICloudKeychainRecoveryFactor::IsRegistered() {
  return storage_->GetICloudKeychainRegistrationInfo(primary_account_.gaia)
      .registered();
}

void ICloudKeychainRecoveryFactor::MarkAsNotRegistered() {
  storage_->MutateICloudKeychainRegistrationInfo(
      primary_account_.gaia,
      [](ICloudKeychainRegistrationInfo& info) { info.set_registered(false); });
}

void ICloudKeychainRecoveryFactor::MarkAsRegistered() {
  storage_->MutateICloudKeychainRegistrationInfo(
      primary_account_.gaia,
      [](ICloudKeychainRegistrationInfo& info) { info.set_registered(true); });
}

TrustedVaultRecoveryFactorRegistrationStateForUMA
ICloudKeychainRecoveryFactor::MaybeRegister(RegisterCallback cb) {
  if (storage_->GetICloudKeychainRegistrationInfo(primary_account_.gaia)
          .registered()) {
    FulfillRegistrationWithFailure(
        TrustedVaultRegistrationStatus::kRegistrationNotAttempted,
        std::move(cb));
    return TrustedVaultRecoveryFactorRegistrationStateForUMA::
        kAlreadyRegisteredV1;
  }

  if (storage_->GetLastRegistrationReturnedLocalDataObsolete(
          primary_account_.gaia)) {
    // Client already knows that existing vault keys (or their absence) isn't
    // sufficient for registration. Fresh keys should be obtained first.
    FulfillRegistrationWithFailure(
        TrustedVaultRegistrationStatus::kRegistrationNotAttempted,
        std::move(cb));
    return TrustedVaultRecoveryFactorRegistrationStateForUMA::
        kLocalKeysAreStale;
  }

  if (connection_->AreRequestsThrottled(primary_account_)) {
    FulfillRegistrationWithFailure(
        TrustedVaultRegistrationStatus::kRegistrationNotAttempted,
        std::move(cb));
    return TrustedVaultRecoveryFactorRegistrationStateForUMA::
        kThrottledClientSide;
  }

  if (!key_storage_->HasNonConstantKey(primary_account_.gaia)) {
    // Registration without non-constant keys isn't supported for iCloud
    // Keychain.
    FulfillRegistrationWithFailure(
        TrustedVaultRegistrationStatus::kRegistrationNotAttempted,
        std::move(cb));
    return TrustedVaultRecoveryFactorRegistrationStateForUMA::
        kRegistrationWithConstantKeyNotSupported;
  }

  if (ongoing_registration_callback_) {
    // Cancel ongoing requests before starting a new one.
    ongoing_download_registration_state_request_for_registration_ = nullptr;
    ongoing_registration_request_ = nullptr;
    registration_weak_ptr_factory_.InvalidateWeakPtrs();
    FulfillRegistrationWithFailure(
        TrustedVaultRegistrationStatus::kRegistrationCancelled,
        std::move(ongoing_registration_callback_));
  }

  ongoing_registration_callback_ = std::move(cb);

  // ICloudRecoveryKey::Retrieve() can't be cancelled, so we use a weak pointer
  // for the callback.
  ICloudRecoveryKey::Retrieve(
      base::BindOnce(
          &ICloudKeychainRecoveryFactor::OnICloudKeysRetrievedForRegistration,
          registration_weak_ptr_factory_.GetWeakPtr()),
      security_domain_id_, icloud_keychain_access_group_);

  // We don't know yet whether there's an existing key pair in iCloud Keychain.
  // However, if there is one that's not yet registered with the security
  // domain, then we have to create a new key pair anyways. Thus, returning
  // `kAttemptingRegistrationWithNewKeyPair` is the most appropriate status
  // here.
  return TrustedVaultRecoveryFactorRegistrationStateForUMA::
      kAttemptingRegistrationWithNewKeyPair;
}

void ICloudKeychainRecoveryFactor::OnICloudKeysRetrievedForRegistration(
    std::vector<std::unique_ptr<ICloudRecoveryKey>> local_icloud_keys) {
  if (local_icloud_keys.empty()) {
    // No local iCloud Keychain key. We need to create a new one and register
    // it.
    ICloudRecoveryKey::Create(
        base::BindOnce(
            &ICloudKeychainRecoveryFactor::OnICloudKeyCreatedForRegistration,
            registration_weak_ptr_factory_.GetWeakPtr()),
        security_domain_id_, icloud_keychain_access_group_);
    return;
  }

  ongoing_download_registration_state_request_for_registration_ =
      connection_->DownloadAuthenticationFactorsRegistrationState(
          primary_account_,
          {trusted_vault_pb::SecurityDomainMember::MEMBER_TYPE_ICLOUD_KEYCHAIN},
          base::BindOnce(&ICloudKeychainRecoveryFactor::
                             OnRecoveryFactorStateDownloadedForRegistration,
                         // `this` outlives `ongoing_request_for_registration_`.
                         base::Unretained(this), std::move(local_icloud_keys)),
          base::NullCallback());
  CHECK(ongoing_download_registration_state_request_for_registration_);
}

void ICloudKeychainRecoveryFactor::
    OnRecoveryFactorStateDownloadedForRegistration(
        std::vector<std::unique_ptr<ICloudRecoveryKey>> local_icloud_keys,
        DownloadAuthenticationFactorsRegistrationStateResult result) {
  // This method should be called only as a result of
  // `ongoing_request_for_registration_` completion/failure, verify this
  // condition and destroy `ongoing_request_for_registration_` as it's not
  // needed anymore.
  CHECK(ongoing_download_registration_state_request_for_registration_);
  ongoing_download_registration_state_request_for_registration_ = nullptr;

  if (result.state ==
      DownloadAuthenticationFactorsRegistrationStateResult::State::kError) {
    connection_->RecordFailedRequestForThrottling(primary_account_);
    FulfillRegistrationWithFailure(
        TrustedVaultRegistrationStatus::kNetworkError,
        std::move(ongoing_registration_callback_));
    return;
  }

  for (const VaultMember& recovery_icloud_key : result.icloud_keys) {
    std::vector<uint8_t> public_key =
        recovery_icloud_key.public_key->ExportToBytes();
    const auto local_icloud_key_it = std::ranges::find_if(
        local_icloud_keys,
        [&public_key](const auto& key) { return key->id() == public_key; });
    if (local_icloud_key_it != local_icloud_keys.end()) {
      MarkAsRegistered();
      int last_vault_key_version =
          std::ranges::max_element(recovery_icloud_key.member_keys, {},
                                   &MemberKeys::version)
              ->version;
      base::BindPostTaskToCurrentDefault(
          base::BindOnce(std::move(ongoing_registration_callback_),
                         TrustedVaultRegistrationStatus::kAlreadyRegistered,
                         last_vault_key_version, /*had_local_keys=*/true))
          .Run();
      return;
    }
  }

  // None of the retrieved iCloud Keychain keys is in the security domain. We
  // need to create a new one and register it.
  ICloudRecoveryKey::Create(
      base::BindOnce(
          &ICloudKeychainRecoveryFactor::OnICloudKeyCreatedForRegistration,
          registration_weak_ptr_factory_.GetWeakPtr()),
      security_domain_id_, icloud_keychain_access_group_);
}

void ICloudKeychainRecoveryFactor::OnICloudKeyCreatedForRegistration(
    std::unique_ptr<ICloudRecoveryKey> local_icloud_key) {
  if (!local_icloud_key) {
    FulfillRegistrationWithFailure(TrustedVaultRegistrationStatus::kOtherError,
                                   std::move(ongoing_registration_callback_));
    return;
  }

  std::vector<std::vector<uint8_t>> vault_keys =
      key_storage_->GetVaultKeys(primary_account_.gaia);
  int last_vault_key_version =
      key_storage_->GetLastKeyVersion(primary_account_.gaia);

  ongoing_registration_request_ = connection_->RegisterAuthenticationFactor(
      primary_account_,
      GetTrustedVaultKeysWithVersions(vault_keys, last_vault_key_version),
      local_icloud_key->key()->public_key(), ICloudKeychain(),
      base::BindOnce(&ICloudKeychainRecoveryFactor::OnRegistered,
                     base::Unretained(this)));
  CHECK(ongoing_registration_request_);
}

void ICloudKeychainRecoveryFactor::OnRegistered(
    TrustedVaultRegistrationStatus status,
    int key_version) {
  // This method should be called only as a result of
  // `ongoing_registration_request_` completion/failure, verify this
  // condition and destroy `ongoing_registration_request_` as it's not
  // needed anymore.
  CHECK(ongoing_registration_request_);
  ongoing_registration_request_ = nullptr;
  CHECK(ongoing_registration_callback_);
  RegisterCallback cb = std::move(ongoing_registration_callback_);

  switch (status) {
    case TrustedVaultRegistrationStatus::kRegistrationNotAttempted:
    case TrustedVaultRegistrationStatus::kRegistrationCancelled:
      NOTREACHED();
    case TrustedVaultRegistrationStatus::kSuccess:
    case TrustedVaultRegistrationStatus::kAlreadyRegistered:
      // kAlreadyRegistered handled as success, because it only means that
      // client doesn't fully handled successful device registration before.
      MarkAsRegistered();
      storage_->SetLastRegistrationReturnedLocalDataObsolete(
          primary_account_.gaia, false);
      break;
    case TrustedVaultRegistrationStatus::kLocalDataObsolete:
      storage_->SetLastRegistrationReturnedLocalDataObsolete(
          primary_account_.gaia, true);
      break;
    case TrustedVaultRegistrationStatus::kTransientAccessTokenFetchError:
    case TrustedVaultRegistrationStatus::kPersistentAccessTokenFetchError:
    case TrustedVaultRegistrationStatus::
        kPrimaryAccountChangeAccessTokenFetchError:
    case TrustedVaultRegistrationStatus::kNetworkError:
    case TrustedVaultRegistrationStatus::kOtherError:
      break;
  }

  std::move(cb).Run(status,
                    /*key_version=*/key_version,
                    /*had_local_keys=*/true);
}

void ICloudKeychainRecoveryFactor::FulfillRegistrationWithFailure(
    TrustedVaultRegistrationStatus status,
    RegisterCallback cb) {
  base::BindPostTaskToCurrentDefault(base::BindOnce(std::move(cb), status,
                                                    /*key_version=*/0,
                                                    /*had_local_keys=*/true))
      .Run();
}

}  // namespace trusted_vault