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

#ifndef CRYPTOHOME_USER_SECRET_STASH_H_
#define CRYPTOHOME_USER_SECRET_STASH_H_

#include <brillo/secure_blob.h>
#include <stdint.h>

#include <map>
#include <memory>
#include <optional>
#include <string>

#include "cryptohome/error/cryptohome_error.h"
#include "cryptohome/flatbuffer_schemas/user_secret_stash_container.h"
#include "cryptohome/storage/file_system_keyset.h"

namespace cryptohome {

// Returns the UserSecretStash experiment version. This will compared with the
// `last_invalid` field in the fetched experiment config to determine whether
// this version should enable the experiment. Will be incremented whenever a
// known issue that blocks the experiment is fixed so that it can be enabled
// again.
int UserSecretStashExperimentVersion();

// This is used by the UssExperimentConfigFetcher to set the experiment flag to
// enabled or disabled based on whether this USS version is valid and how much
// of the population should have the experiment enabled.
void SetUserSecretStashExperimentFlag(bool enabled);

// Returns whether the UserSecretStash experiment (using the USS instead of
// vault keysets) is enabled.
// The experiment is controlled by fetching a config file from gstatic. It
// matches the local USS version returned by
// `UserSecretStashExperimentVersion()` and the `last_invalid` version specified
// in the config file. If our version is greater, the experiment is enabled with
// `population` probability, and disabled otherwise. Whether the experiment is
// enabled can be overridden by creating the /var/lib/cryptohome/uss_enabled (to
// enable) or the /var/lib/cryptohome/uss_disabled (to disable) file. Unit tests
// can furthermore override this behavior using
// `SetUserSecretStashExperimentForTesting()`.
bool IsUserSecretStashExperimentEnabled();
// Allows to toggle the experiment state in tests. Passing nullopt reverts to
// the default behavior.
void SetUserSecretStashExperimentForTesting(std::optional<bool> enabled);

// This wraps the UserSecretStash flatbuffer message, and is the only way that
// the UserSecretStash is accessed. Don't pass the raw flatbuffer around.
class UserSecretStash {
 public:
  // Container for a wrapped (encrypted) USS main key.
  struct WrappedKeyBlock {
    // The algorithm used for wrapping the USS main key.
    UserSecretStashEncryptionAlgorithm encryption_algorithm;
    // This is the encrypted USS main key.
    brillo::Blob encrypted_key;
    // The random IV used in the USS main key encryption.
    brillo::Blob iv;
    // The GCM tag generated by the block cipher.
    brillo::Blob gcm_tag;
  };

  // Sets up a UserSecretStash with random contents (reset secret, etc.) and the
  // values from the specified file system keyset.
  static CryptohomeStatusOr<std::unique_ptr<UserSecretStash>> CreateRandom(
      const FileSystemKeyset& file_system_keyset);
  // This deserializes the |flatbuffer| into a UserSecretStashContainer table.
  // Besides unencrypted data, that table contains a ciphertext, which is
  // decrypted with the |main_key| using AES-GCM-256. It doesn't return the
  // plaintext, it populates the fields of the class with the encrypted message.
  static CryptohomeStatusOr<std::unique_ptr<UserSecretStash>>
  FromEncryptedContainer(const brillo::Blob& flatbuffer,
                         const brillo::SecureBlob& main_key);
  // Same as |FromEncryptedContainer()|, but the main key is unwrapped from the
  // USS container using the given wrapping key. The |main_key| output argument
  // is populated with the unwrapped main key on success.
  static CryptohomeStatusOr<std::unique_ptr<UserSecretStash>>
  FromEncryptedContainerWithWrappingKey(const brillo::Blob& flatbuffer,
                                        const std::string& wrapping_id,
                                        const brillo::SecureBlob& wrapping_key,
                                        brillo::SecureBlob* main_key);

  // Randomly generates a USS Main Key. This is intended to be used when
  // creating a fresh USS via |CreateRandom()|.
  static brillo::SecureBlob CreateRandomMainKey();

  virtual ~UserSecretStash() = default;

  // Because this class contains raw secrets, it should never be copy-able.
  UserSecretStash(const UserSecretStash&) = delete;
  UserSecretStash& operator=(const UserSecretStash&) = delete;

  const FileSystemKeyset& GetFileSystemKeyset() const;

  // This gets the reset secret for the auth factor with the associated
  // |label|.
  std::optional<brillo::SecureBlob> GetResetSecretForLabel(
      const std::string& label) const;

  // This sets the reset secret for an auth factor with the associated |label|.
  // This does not overwrite an existing reset secret. It returns if the
  // insertion succeeded.
  [[nodiscard]] bool SetResetSecretForLabel(const std::string& label,
                                            const brillo::SecureBlob& secret);

  // This removes the reset secret for an auth factor with the associated
  // |label|. Returns false if reset secret wasn't present for provided |label|,
  // true otherwise.
  // TODO(b/238897234): Move this to RemoveWrappedMainKey.
  bool RemoveResetSecretForLabel(const std::string& label);

  // The OS version on which this particular user secret stash was originally
  // created. The format is the one of the CHROMEOS_RELEASE_VERSION field in
  // /etc/lsb-release, e.g.: "11012.0.2018_08_28_1422". Empty if the version
  // fetch failed at the creation time.
  // !!!WARNING!!!: This value is not authenticated nor validated. It must not
  // be used for security-critical features.
  const std::string& GetCreatedOnOsVersion() const;

  // Returns whether there's a wrapped key block with the given wrapping ID.
  bool HasWrappedMainKey(const std::string& wrapping_id) const;
  // Unwraps (decrypts) the USS main key from the wrapped key block with the
  // given wrapping ID. Returns a status if it doesn't exist or the unwrapping
  // fails.
  CryptohomeStatusOr<brillo::SecureBlob> UnwrapMainKey(
      const std::string& wrapping_id,
      const brillo::SecureBlob& wrapping_key) const;
  // Wraps (encrypts) the USS main key using the given wrapped key. The wrapped
  // data is added into the USS as a wrapped key block with the given wrapping
  // ID. |main_key| must be non-empty, and |wrapping_key| - of
  // |kAesGcm256KeySize| length. Returns a status if the wrapping ID is already
  // used or the wrapping fails.
  CryptohomeStatus AddWrappedMainKey(const brillo::SecureBlob& main_key,
                                     const std::string& wrapping_id,
                                     const brillo::SecureBlob& wrapping_key);
  // Removes the wrapped key with the given ID. If it doesn't exist, returns
  // false.
  bool RemoveWrappedMainKey(const std::string& wrapping_id);

  // This uses the |main_key|, which should be 256-bit as of right now, to
  // encrypt this UserSecretStash class. The object is converted to a
  // UserSecretStashPayload table, serialized, encrypted with AES-GCM-256, and
  // serialized as a UserSecretStashContainer table.
  CryptohomeStatusOr<brillo::Blob> GetEncryptedContainer(
      const brillo::SecureBlob& main_key);

 private:
  // Decrypts the USS payload flatbuffer using the passed main key and
  // constructs the USS instance from it. Returns null on decryption or
  // validation failure.
  static CryptohomeStatusOr<std::unique_ptr<UserSecretStash>>
  FromEncryptedPayload(
      const brillo::Blob& ciphertext,
      const brillo::Blob& iv,
      const brillo::Blob& gcm_tag,
      const std::map<std::string, WrappedKeyBlock>& wrapped_key_blocks,
      const std::string& created_on_os_version,
      const brillo::SecureBlob& main_key);

  UserSecretStash(
      const FileSystemKeyset& file_system_keyset,
      const std::map<std::string, brillo::SecureBlob>& reset_secrets);

  explicit UserSecretStash(const FileSystemKeyset& file_system_keyset);

  // Keys registered with the kernel to decrypt files and file names, together
  // with corresponding salts and signatures.
  const FileSystemKeyset file_system_keyset_;
  // Stores multiple wrapped (encrypted) representations of the main key, each
  // wrapped using a different intermediate key. The map's index is the wrapping
  // ID, which is an opaque string (although upper programmatic layers can add
  // semantics to it, in order to map it to the authentication method).
  std::map<std::string, WrappedKeyBlock> wrapped_key_blocks_;
  // The OS version on which this particular user secret stash was originally
  // created.
  std::string created_on_os_version_;
  // The reset secrets corresponding to each auth factor, by label.
  std::map<std::string, brillo::SecureBlob> reset_secrets_;
};

}  // namespace cryptohome

#endif  // CRYPTOHOME_USER_SECRET_STASH_H_