// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

#pragma once

#include <cassert>
#include <map>
#include <memory>
#include <span>
#include <string>
#include <utility>

#include "arrow/util/secure_string.h"
#include "parquet/exception.h"
#include "parquet/schema.h"
#include "parquet/types.h"

namespace parquet {

static constexpr ParquetCipher::type kDefaultEncryptionAlgorithm =
    ParquetCipher::AES_GCM_V1;
static constexpr int32_t kMaximalAadMetadataLength = 256;
static constexpr bool kDefaultEncryptedFooter = true;
static constexpr bool kDefaultCheckSignature = true;
static constexpr bool kDefaultAllowPlaintextFiles = false;
static constexpr int32_t kAadFileUniqueLength = 8;

class ColumnDecryptionProperties;
using ColumnPathToDecryptionPropertiesMap =
    std::map<std::string, std::shared_ptr<ColumnDecryptionProperties>>;

class ColumnEncryptionProperties;
using ColumnPathToEncryptionPropertiesMap =
    std::map<std::string, std::shared_ptr<ColumnEncryptionProperties>>;

class PARQUET_EXPORT DecryptionKeyRetriever {
 public:
  /// \brief Retrieve a key.
  virtual ::arrow::util::SecureString GetKey(const std::string& key_id) = 0;

  virtual ~DecryptionKeyRetriever() {}
};

/// Simple integer key retriever
class PARQUET_EXPORT IntegerKeyIdRetriever : public DecryptionKeyRetriever {
 public:
  void PutKey(uint32_t key_id, ::arrow::util::SecureString key);

  ::arrow::util::SecureString GetKey(const std::string& key_id_string) override {
    // key_id_string is string but for IntegerKeyIdRetriever it encodes
    // a native-endian 32 bit unsigned integer key_id
    uint32_t key_id;
    assert(key_id_string.size() == sizeof(key_id));
    memcpy(&key_id, key_id_string.data(), sizeof(key_id));

    return key_map_.at(key_id);
  }

 private:
  std::map<uint32_t, ::arrow::util::SecureString> key_map_;
};

// Simple string key retriever
class PARQUET_EXPORT StringKeyIdRetriever : public DecryptionKeyRetriever {
 public:
  void PutKey(std::string key_id, ::arrow::util::SecureString key);
  ::arrow::util::SecureString GetKey(const std::string& key_id) override;

 private:
  std::map<std::string, ::arrow::util::SecureString> key_map_;
};

class PARQUET_EXPORT HiddenColumnException : public ParquetException {
 public:
  explicit HiddenColumnException(const std::string& columnPath)
      : ParquetException(columnPath.c_str()) {}
};

class PARQUET_EXPORT KeyAccessDeniedException : public ParquetException {
 public:
  explicit KeyAccessDeniedException(const std::string& columnPath)
      : ParquetException(columnPath.c_str()) {}
};

inline std::span<const uint8_t> str2span(const std::string& str) {
  if (str.empty()) {
    return {};
  }

  return {reinterpret_cast<const uint8_t*>(str.data()), str.size()};
}

class PARQUET_EXPORT ColumnEncryptionProperties {
 public:
  class PARQUET_EXPORT Builder {
   public:
    PARQUET_DEPRECATED("name argument is ignored, use default constructor instead")
    explicit Builder(const std::string& name) : encrypted_(true) {}

    PARQUET_DEPRECATED("path argument is ignored, use default constructor instead")
    explicit Builder(const schema::ColumnPath& path) : encrypted_(true) {}

    Builder() = default;

    /// Set a column-specific key.
    /// If key is not set on an encrypted column, the column will
    /// be encrypted with the footer key.
    /// keyBytes Key length must be either 16, 24 or 32 bytes.
    /// Caller is responsible for wiping out the input key array.
    Builder* key(::arrow::util::SecureString column_key);

    /// Set a key retrieval metadata.
    /// use either key_metadata() or key_id(), not both
    Builder* key_metadata(std::string key_metadata);

    /// A convenience function to set key metadata using a string id.
    /// Set a key retrieval metadata (converted from String).
    /// use either key_metadata() or key_id(), not both
    /// key_id will be converted to metadata (UTF-8 array).
    Builder* key_id(std::string key_id);

    std::shared_ptr<ColumnEncryptionProperties> build() {
      return std::shared_ptr<ColumnEncryptionProperties>(
          new ColumnEncryptionProperties(encrypted_, key_, key_metadata_));
    }

   private:
    bool encrypted_ = true;
    ::arrow::util::SecureString key_;
    std::string key_metadata_;
  };

  bool is_encrypted() const { return encrypted_; }
  bool is_encrypted_with_footer_key() const { return encrypted_with_footer_key_; }
  const ::arrow::util::SecureString& key() const { return key_; }
  const std::string& key_metadata() const { return key_metadata_; }

  static std::shared_ptr<ColumnEncryptionProperties> Unencrypted();
  static std::shared_ptr<ColumnEncryptionProperties> WithFooterKey();
  static std::shared_ptr<ColumnEncryptionProperties> WithColumnKey(
      ::arrow::util::SecureString key, std::string key_metadata = "");

 private:
  bool encrypted_;
  bool encrypted_with_footer_key_;
  ::arrow::util::SecureString key_;
  std::string key_metadata_;
  explicit ColumnEncryptionProperties(bool encrypted, ::arrow::util::SecureString key,
                                      std::string key_metadata);
};

class PARQUET_EXPORT ColumnDecryptionProperties {
 public:
  class PARQUET_EXPORT Builder {
   public:
    explicit Builder(std::string name) : column_path_(std::move(name)) {}

    explicit Builder(const schema::ColumnPath& path) : Builder(path.ToDotString()) {}

    /// Set an explicit column key. If applied on a file that contains
    /// key metadata for this column the metadata will be ignored,
    /// the column will be decrypted with this key.
    /// key length must be either 16, 24 or 32 bytes.
    Builder* key(::arrow::util::SecureString key);

    std::shared_ptr<ColumnDecryptionProperties> build();

   private:
    std::string column_path_;
    ::arrow::util::SecureString key_;
  };

  const std::string& column_path() const { return column_path_; }
  const ::arrow::util::SecureString& key() const { return key_; }

 private:
  std::string column_path_;
  ::arrow::util::SecureString key_;

  /// This class is only required for setting explicit column decryption keys -
  /// to override key retriever (or to provide keys when key metadata and/or
  /// key retriever are not available)
  explicit ColumnDecryptionProperties(std::string column_path,
                                      ::arrow::util::SecureString key);
};

class PARQUET_EXPORT AADPrefixVerifier {
 public:
  /// Verifies identity (AAD Prefix) of individual file,
  /// or of file collection in a data set.
  /// Throws exception if an AAD prefix is wrong.
  /// In a data set, AAD Prefixes should be collected,
  /// and then checked for missing files.
  virtual void Verify(const std::string& aad_prefix) = 0;
  virtual ~AADPrefixVerifier() {}
};

class PARQUET_EXPORT FileDecryptionProperties {
 public:
  class PARQUET_EXPORT Builder {
   public:
    Builder() {
      check_plaintext_footer_integrity_ = kDefaultCheckSignature;
      plaintext_files_allowed_ = kDefaultAllowPlaintextFiles;
    }

    /// Set an explicit footer key. If applied on a file that contains
    /// footer key metadata the metadata will be ignored, the footer
    /// will be decrypted/verified with this key.
    /// If explicit key is not set, footer key will be fetched from
    /// key retriever.
    /// With explicit keys or AAD prefix, new encryption properties object must be
    /// created for each encrypted file.
    /// Explicit encryption keys (footer and column) are cloned.
    /// Upon completion of file reading, the cloned encryption keys in the properties
    /// will be wiped out (array values set to 0).
    /// Caller is responsible for wiping out the input key array.
    /// param footerKey Key length must be either 16, 24 or 32 bytes.
    Builder* footer_key(::arrow::util::SecureString footer_key);

    /// Set explicit column keys (decryption properties).
    /// Its also possible to set a key retriever on this property object.
    /// Upon file decryption, availability of explicit keys is checked before
    /// invocation of the retriever callback.
    /// If an explicit key is available for a footer or a column,
    /// its key metadata will be ignored.
    Builder* column_keys(
        ColumnPathToDecryptionPropertiesMap column_decryption_properties);

    /// Set a key retriever callback. Its also possible to
    /// set explicit footer or column keys on this file property object.
    /// Upon file decryption, availability of explicit keys is checked before
    /// invocation of the retriever callback.
    /// If an explicit key is available for a footer or a column,
    /// its key metadata will be ignored.
    Builder* key_retriever(std::shared_ptr<DecryptionKeyRetriever> key_retriever);

    /// Skip integrity verification of plaintext footers.
    /// If not called, integrity of plaintext footers will be checked in runtime,
    /// and an exception will be thrown in the following situations:
    /// - footer signing key is not available
    /// (not passed, or not found by key retriever)
    /// - footer content and signature don't match
    Builder* disable_footer_signature_verification() {
      check_plaintext_footer_integrity_ = false;
      return this;
    }

    /// Explicitly supply the file AAD prefix.
    /// A must when a prefix is used for file encryption, but not stored in file.
    /// If AAD prefix is stored in file, it will be compared to the explicitly
    /// supplied value and an exception will be thrown if they differ.
    Builder* aad_prefix(std::string aad_prefix);

    /// Set callback for verification of AAD Prefixes stored in file.
    Builder* aad_prefix_verifier(std::shared_ptr<AADPrefixVerifier> aad_prefix_verifier);

    /// By default, reading plaintext (unencrypted) files is not
    /// allowed when using a decryptor
    /// - in order to detect files that were not encrypted by mistake.
    /// However, the default behavior can be overridden by calling this method.
    /// The caller should use then a different method to ensure encryption
    /// of files with sensitive data.
    Builder* plaintext_files_allowed() {
      plaintext_files_allowed_ = true;
      return this;
    }

    std::shared_ptr<FileDecryptionProperties> build() {
      return std::shared_ptr<FileDecryptionProperties>(new FileDecryptionProperties(
          footer_key_, key_retriever_, check_plaintext_footer_integrity_, aad_prefix_,
          aad_prefix_verifier_, column_decryption_properties_, plaintext_files_allowed_));
    }

   private:
    ::arrow::util::SecureString footer_key_;
    std::string aad_prefix_;
    std::shared_ptr<AADPrefixVerifier> aad_prefix_verifier_;
    ColumnPathToDecryptionPropertiesMap column_decryption_properties_;

    std::shared_ptr<DecryptionKeyRetriever> key_retriever_;
    bool check_plaintext_footer_integrity_;
    bool plaintext_files_allowed_;
  };

  const ::arrow::util::SecureString& column_key(const std::string& column_path) const;

  const ::arrow::util::SecureString& footer_key() const { return footer_key_; }

  const std::string& aad_prefix() const { return aad_prefix_; }

  const std::shared_ptr<DecryptionKeyRetriever>& key_retriever() const {
    return key_retriever_;
  }

  bool check_plaintext_footer_integrity() const {
    return check_plaintext_footer_integrity_;
  }

  bool plaintext_files_allowed() const { return plaintext_files_allowed_; }

  const std::shared_ptr<AADPrefixVerifier>& aad_prefix_verifier() const {
    return aad_prefix_verifier_;
  }

 private:
  ::arrow::util::SecureString footer_key_;
  std::string aad_prefix_;
  std::shared_ptr<AADPrefixVerifier> aad_prefix_verifier_;
  ColumnPathToDecryptionPropertiesMap column_decryption_properties_;
  std::shared_ptr<DecryptionKeyRetriever> key_retriever_;
  bool check_plaintext_footer_integrity_;
  bool plaintext_files_allowed_;

  FileDecryptionProperties(
      ::arrow::util::SecureString footer_key,
      std::shared_ptr<DecryptionKeyRetriever> key_retriever,
      bool check_plaintext_footer_integrity, std::string aad_prefix,
      std::shared_ptr<AADPrefixVerifier> aad_prefix_verifier,
      ColumnPathToDecryptionPropertiesMap column_decryption_properties,
      bool plaintext_files_allowed);
};

class PARQUET_EXPORT FileEncryptionProperties {
 public:
  class PARQUET_EXPORT Builder {
   public:
    explicit Builder(::arrow::util::SecureString footer_key)
        : parquet_cipher_(kDefaultEncryptionAlgorithm),
          encrypted_footer_(kDefaultEncryptedFooter),
          footer_key_(std::move(footer_key)) {
      store_aad_prefix_in_file_ = false;
    }

    /// Create files with plaintext footer.
    /// If not called, the files will be created with encrypted footer (default).
    Builder* set_plaintext_footer() {
      encrypted_footer_ = false;
      return this;
    }

    /// Set encryption algorithm.
    /// If not called, files will be encrypted with AES_GCM_V1 (default).
    Builder* algorithm(ParquetCipher::type parquet_cipher) {
      parquet_cipher_ = parquet_cipher;
      return this;
    }

    /// Set a key retrieval metadata (converted from String).
    /// use either footer_key_metadata or footer_key_id, not both.
    Builder* footer_key_id(std::string key_id);

    /// Set a key retrieval metadata.
    /// use either footer_key_metadata or footer_key_id, not both.
    Builder* footer_key_metadata(std::string footer_key_metadata);

    /// Set the file AAD Prefix.
    Builder* aad_prefix(std::string aad_prefix);

    /// Skip storing AAD Prefix in file.
    /// If not called, and if AAD Prefix is set, it will be stored.
    Builder* disable_aad_prefix_storage();

    /// Set the list of encrypted columns and their properties (keys etc).
    /// If not called, all columns will be encrypted with the footer key.
    /// If called, the file columns not in the list will be left unencrypted.
    Builder* encrypted_columns(ColumnPathToEncryptionPropertiesMap encrypted_columns);

    std::shared_ptr<FileEncryptionProperties> build() {
      return std::shared_ptr<FileEncryptionProperties>(new FileEncryptionProperties(
          parquet_cipher_, footer_key_, footer_key_metadata_, encrypted_footer_,
          aad_prefix_, store_aad_prefix_in_file_, encrypted_columns_));
    }

   private:
    ParquetCipher::type parquet_cipher_;
    bool encrypted_footer_;
    ::arrow::util::SecureString footer_key_;
    std::string footer_key_metadata_;

    std::string aad_prefix_;
    bool store_aad_prefix_in_file_;
    ColumnPathToEncryptionPropertiesMap encrypted_columns_;
  };

  bool encrypted_footer() const { return encrypted_footer_; }

  EncryptionAlgorithm algorithm() const { return algorithm_; }

  const ::arrow::util::SecureString& footer_key() const { return footer_key_; }

  const std::string& footer_key_metadata() const { return footer_key_metadata_; }

  const std::string& file_aad() const { return file_aad_; }

  std::shared_ptr<ColumnEncryptionProperties> column_encryption_properties(
      const std::string& column_path);

  const ColumnPathToEncryptionPropertiesMap& encrypted_columns() const {
    return encrypted_columns_;
  }

 private:
  EncryptionAlgorithm algorithm_;
  ::arrow::util::SecureString footer_key_;
  std::string footer_key_metadata_;
  bool encrypted_footer_;
  std::string file_aad_;
  std::string aad_prefix_;
  bool store_aad_prefix_in_file_;
  ColumnPathToEncryptionPropertiesMap encrypted_columns_;

  FileEncryptionProperties(ParquetCipher::type cipher,
                           ::arrow::util::SecureString footer_key,
                           std::string footer_key_metadata, bool encrypted_footer,
                           std::string aad_prefix, bool store_aad_prefix_in_file,
                           ColumnPathToEncryptionPropertiesMap encrypted_columns);
};

}  // namespace parquet
