/*
 * Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License"). You may not use
 * this file except in compliance with the License. A copy of the License is
 * located at
 *
 *     http://aws.amazon.com/apache2.0/
 *
 * or in the "license" file accompanying this file. This file 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.
 */
#include <aws/cryptosdk/private/kms_keyring.h>

#include <aws/core/utils/ARN.h>
#include <aws/core/utils/Outcome.h>
#include <aws/core/utils/logging/LogMacros.h>
#include <aws/core/utils/memory/MemorySystemInterface.h>
#include <aws/core/utils/memory/stl/AWSAllocator.h>
#include <aws/cryptosdk/list_utils.h>
#include <aws/cryptosdk/private/cpputils.h>
#include <aws/cryptosdk/private/user_agent.h>
#include <aws/kms/model/DecryptRequest.h>
#include <aws/kms/model/DecryptResult.h>
#include <aws/kms/model/EncryptRequest.h>
#include <aws/kms/model/EncryptResult.h>
#include <aws/kms/model/GenerateDataKeyRequest.h>
#include <aws/kms/model/GenerateDataKeyResult.h>

namespace Aws {
namespace Cryptosdk {

using Private::append_key_dup_to_edks;
using Private::aws_byte_buf_dup_from_aws_utils;
using Private::aws_map_from_c_aws_hash_table;
using Private::aws_utils_byte_buffer_from_c_aws_byte_buf;

static const char *AWS_CRYPTO_SDK_KMS_CLASS_TAG              = "KmsKeyring";
static const char *AWS_CRYPTO_SDK_DISCOVERY_FILTER_CLASS_TAG = "DiscoveryFilter";
static const char *KEY_PROVIDER_STR                          = "aws-kms";

static void DestroyKeyring(struct aws_cryptosdk_keyring *keyring) {
    auto keyring_data_ptr = static_cast<Aws::Cryptosdk::Private::KmsKeyringImpl *>(keyring);
    Aws::Delete(keyring_data_ptr);
}

static int OnDecrypt(
    struct aws_cryptosdk_keyring *keyring,
    struct aws_allocator *request_alloc,
    struct aws_byte_buf *unencrypted_data_key,
    struct aws_array_list *keyring_trace,
    const struct aws_array_list *edks,
    const struct aws_hash_table *enc_ctx,
    enum aws_cryptosdk_alg_id alg) {
    (void)alg;

    auto self = static_cast<Aws::Cryptosdk::Private::KmsKeyringImpl *>(keyring);
    if (!self || !request_alloc || !unencrypted_data_key || !edks || !enc_ctx) {
        abort();
    }

    Aws::StringStream error_buf;
    const auto enc_ctx_cpp = aws_map_from_c_aws_hash_table(enc_ctx);

    size_t num_elems = aws_array_list_length(edks);
    for (unsigned int idx = 0; idx < num_elems; idx++) {
        struct aws_cryptosdk_edk *edk;
        int rv = aws_array_list_get_at_ptr(edks, (void **)&edk, idx);
        if (rv != AWS_OP_SUCCESS) {
            continue;
        }

        if (!aws_byte_buf_eq(&edk->provider_id, &self->key_provider)) {
            // EDK belongs to a different non KMS keyring. Skip.
            continue;
        }

        const Aws::String key_arn = Private::aws_string_from_c_aws_byte_buf(&edk->provider_info);

        /* If there are no key IDs in the list, keyring is in "discovery" mode and will attempt KMS calls with
         * every key ARN it comes across in the message, so long as the key ARN is authorized by the
         * DiscoveryFilter (matches the partition and an account ID).
         *
         * If there are key IDs in the list, it will cross check the ARN it reads with that list
         * before attempting KMS calls. Note that if caller provided key IDs in anything other than
         * a CMK ARN format, the SDK will not attempt to decrypt those data keys, because the EDK
         * data format always specifies the CMK with the full (non-alias) ARN.
         */
        if (self->key_ids.size() &&
            std::find(self->key_ids.begin(), self->key_ids.end(), key_arn) == self->key_ids.end()) {
            // This keyring does not have access to the CMK used to encrypt this data key. Skip.
            continue;
        }
        // self->discovery_filter is non-null only if self was constructed via BuildDiscovery, which
        // in turn implies discovery mode
        if (self->discovery_filter && !self->discovery_filter->IsAuthorized(key_arn)) {
            // The DiscoveryFilter blocks the CMK used to encrypt this data key. Skip.
            continue;
        }

        Aws::String kms_region = Private::parse_region_from_kms_key_arn(key_arn);
        if (kms_region.empty()) {
            error_buf << "Error: Malformed ciphertext. Provider ID field of KMS EDK is invalid KMS CMK ARN: " << key_arn
                      << " ";
            continue;
        }

        std::function<void()> report_success;
        auto kms_client = self->kms_client_supplier->GetClient(kms_region, report_success);
        if (!kms_client) {
            // Client supplier does not serve this region. Skip.
            continue;
        }

        Aws::KMS::Model::DecryptRequest kms_request;
        kms_request.WithGrantTokens(self->grant_tokens)
            .WithKeyId(key_arn)
            .WithCiphertextBlob(aws_utils_byte_buffer_from_c_aws_byte_buf(&edk->ciphertext))
            .WithEncryptionContext(enc_ctx_cpp);

        Aws::KMS::Model::DecryptOutcome outcome = kms_client->Decrypt(kms_request);
        if (!outcome.IsSuccess()) {
            // Failing on this call is normal behavior in "discovery" mode, but not in standard mode.
            if (self->key_ids.size()) {
                error_buf << "Error: " << outcome.GetError().GetExceptionName()
                          << " Message:" << outcome.GetError().GetMessage() << " ";
            }
            continue;
        }
        report_success();

        const Aws::String &outcome_key_id = outcome.GetResult().GetKeyId();
        if (outcome_key_id == key_arn) {
            int ret = aws_byte_buf_dup_from_aws_utils(
                request_alloc, unencrypted_data_key, outcome.GetResult().GetPlaintext());
            if (ret == AWS_OP_SUCCESS) {
                aws_cryptosdk_keyring_trace_add_record_c_str(
                    request_alloc,
                    keyring_trace,
                    KEY_PROVIDER_STR,
                    key_arn.c_str(),
                    AWS_CRYPTOSDK_WRAPPING_KEY_DECRYPTED_DATA_KEY | AWS_CRYPTOSDK_WRAPPING_KEY_VERIFIED_ENC_CTX);
            }
            return ret;
        } else {
            // Since we specified the key ARN explicitly in the request,
            // KMS had better use that key to decrypt
            return aws_raise_error(AWS_ERROR_INVALID_STATE);
        }
    }

    AWS_LOGSTREAM_ERROR(
        AWS_CRYPTO_SDK_KMS_CLASS_TAG,
        "Could not find any data key that can be decrypted by KMS. Errors:" << error_buf.str());
    // According to materials.h we should return success when no key was found
    return AWS_OP_SUCCESS;
}

static int OnEncrypt(
    struct aws_cryptosdk_keyring *keyring,
    struct aws_allocator *request_alloc,
    struct aws_byte_buf *unencrypted_data_key,
    struct aws_array_list *keyring_trace,
    struct aws_array_list *edk_list,
    const struct aws_hash_table *enc_ctx,
    enum aws_cryptosdk_alg_id alg) {
    if (!keyring || !request_alloc || !unencrypted_data_key || !edk_list || !enc_ctx) {
        abort();
    }
    auto self = static_cast<Aws::Cryptosdk::Private::KmsKeyringImpl *>(keyring);

    /* When no key IDs are configured, i.e. "discovery mode", this function is a no-op. */
    if (!self->key_ids.size()) {
        return AWS_OP_SUCCESS;
    }

    Private::ListRaii my_edks(aws_cryptosdk_edk_list_init, aws_cryptosdk_edk_list_clean_up);
    Private::ListRaii my_keyring_trace(aws_cryptosdk_keyring_trace_init, aws_cryptosdk_keyring_trace_clean_up);
    int rv = my_edks.Create(request_alloc);
    if (rv) return rv;
    rv = my_keyring_trace.Create(request_alloc);
    if (rv) return rv;

    const auto enc_ctx_cpp = aws_map_from_c_aws_hash_table(enc_ctx);

    bool generated_new_data_key = false;
    if (!unencrypted_data_key->buffer) {
        const struct aws_cryptosdk_alg_properties *alg_prop = aws_cryptosdk_alg_props(alg);
        if (alg_prop == NULL) {
            AWS_LOGSTREAM_ERROR(AWS_CRYPTO_SDK_KMS_CLASS_TAG, "Invalid encryption materials algorithm properties");
            return aws_raise_error(AWS_CRYPTOSDK_ERR_CRYPTO_UNKNOWN);
        }
        Aws::String key_id = self->key_ids.front();

        // Already checked on keyring build that this will succeed.
        Aws::String kms_region = Private::parse_region_from_kms_key_arn(key_id);

        std::function<void()> report_success;
        auto kms_client = self->kms_client_supplier->GetClient(kms_region, report_success);
        if (!kms_client) {
            /* Client supplier is allowed to return NULL if, for example, user wants to exclude particular
             * regions. But if we are here it means that user configured keyring with a KMS key that was
             * incompatible with the client supplier in use.
             */
            return aws_raise_error(AWS_CRYPTOSDK_ERR_BAD_STATE);
        }
        Aws::KMS::Model::GenerateDataKeyRequest kms_request;
        kms_request.WithKeyId(key_id)
            .WithGrantTokens(self->grant_tokens)
            .WithNumberOfBytes((int)alg_prop->data_key_len)
            .WithEncryptionContext(enc_ctx_cpp);

        Aws::KMS::Model::GenerateDataKeyOutcome outcome = kms_client->GenerateDataKey(kms_request);
        if (!outcome.IsSuccess()) {
            AWS_LOGSTREAM_ERROR(AWS_CRYPTO_SDK_KMS_CLASS_TAG, "Invalid encryption materials algorithm properties");
            return aws_raise_error(AWS_CRYPTOSDK_ERR_KMS_FAILURE);
        }
        report_success();
        rv = append_key_dup_to_edks(
            request_alloc,
            &my_edks.list,
            &outcome.GetResult().GetCiphertextBlob(),
            &outcome.GetResult().GetKeyId(),
            &self->key_provider);
        if (rv != AWS_OP_SUCCESS) return rv;

        rv = aws_byte_buf_dup_from_aws_utils(request_alloc, unencrypted_data_key, outcome.GetResult().GetPlaintext());
        if (rv != AWS_OP_SUCCESS) return rv;
        generated_new_data_key = true;
        aws_cryptosdk_keyring_trace_add_record_c_str(
            request_alloc,
            &my_keyring_trace.list,
            KEY_PROVIDER_STR,
            key_id.c_str(),
            AWS_CRYPTOSDK_WRAPPING_KEY_GENERATED_DATA_KEY | AWS_CRYPTOSDK_WRAPPING_KEY_ENCRYPTED_DATA_KEY |
                AWS_CRYPTOSDK_WRAPPING_KEY_SIGNED_ENC_CTX);
    }

    const auto unencrypted_data_key_cpp = aws_utils_byte_buffer_from_c_aws_byte_buf(unencrypted_data_key);

    size_t num_key_ids = self->key_ids.size();
    for (size_t key_id_idx = generated_new_data_key ? 1 : 0; key_id_idx < num_key_ids; ++key_id_idx) {
        Aws::String key_id = self->key_ids[key_id_idx];

        // Already checked on keyring build that this will succeed.
        Aws::String kms_region = Private::parse_region_from_kms_key_arn(key_id);

        std::function<void()> report_success;
        auto kms_client = self->kms_client_supplier->GetClient(kms_region, report_success);
        if (!kms_client) {
            /* Client supplier is allowed to return NULL if, for example, user wants to exclude particular
             * regions. But if we are here it means that user configured keyring with a KMS key that was
             * incompatible with the client supplier in use.
             */
            rv = aws_raise_error(AWS_CRYPTOSDK_ERR_BAD_STATE);
            goto out;
        }
        Aws::KMS::Model::EncryptRequest kms_request;
        kms_request.WithKeyId(key_id)
            .WithGrantTokens(self->grant_tokens)
            .WithPlaintext(unencrypted_data_key_cpp)
            .WithEncryptionContext(enc_ctx_cpp);

        Aws::KMS::Model::EncryptOutcome outcome = kms_client->Encrypt(kms_request);
        if (!outcome.IsSuccess()) {
            AWS_LOGSTREAM_ERROR(
                AWS_CRYPTO_SDK_KMS_CLASS_TAG,
                "KMS encryption error : " << outcome.GetError().GetExceptionName()
                                          << " Message: " << outcome.GetError().GetMessage());
            rv = aws_raise_error(AWS_CRYPTOSDK_ERR_KMS_FAILURE);
            goto out;
        }
        report_success();
        rv = append_key_dup_to_edks(
            request_alloc,
            &my_edks.list,
            &outcome.GetResult().GetCiphertextBlob(),
            &outcome.GetResult().GetKeyId(),
            &self->key_provider);
        if (rv != AWS_OP_SUCCESS) {
            goto out;
        }
        aws_cryptosdk_keyring_trace_add_record_c_str(
            request_alloc,
            &my_keyring_trace.list,
            KEY_PROVIDER_STR,
            key_id.c_str(),
            AWS_CRYPTOSDK_WRAPPING_KEY_ENCRYPTED_DATA_KEY | AWS_CRYPTOSDK_WRAPPING_KEY_SIGNED_ENC_CTX);
    }
    rv = aws_cryptosdk_transfer_list(edk_list, &my_edks.list);
    if (rv == AWS_OP_SUCCESS) {
        aws_cryptosdk_transfer_list(keyring_trace, &my_keyring_trace.list);
    }
out:
    if (rv != AWS_OP_SUCCESS && generated_new_data_key) {
        aws_byte_buf_clean_up(unencrypted_data_key);
    }
    return rv;
}

Aws::Cryptosdk::Private::KmsKeyringImpl::~KmsKeyringImpl() {}

Aws::Cryptosdk::Private::KmsKeyringImpl::KmsKeyringImpl(
    const Aws::Vector<Aws::String> &key_ids,
    const Aws::Vector<Aws::String> &grant_tokens,
    std::shared_ptr<Aws::Cryptosdk::KmsKeyring::ClientSupplier> client_supplier)
    : key_provider(aws_byte_buf_from_c_str(KEY_PROVIDER_STR)),
      kms_client_supplier(client_supplier),
      grant_tokens(grant_tokens),
      key_ids(key_ids) {
    static const aws_cryptosdk_keyring_vt kms_keyring_vt = {
        sizeof(struct aws_cryptosdk_keyring_vt), KEY_PROVIDER_STR, &DestroyKeyring, &OnEncrypt, &OnDecrypt
    };

    aws_cryptosdk_keyring_base_init(this, &kms_keyring_vt);
}

static std::shared_ptr<KMS::KMSClient> CreateDefaultKmsClient(const char *allocationTag, const Aws::String &region) {
    Aws::Client::ClientConfiguration client_configuration;
    if (!region.empty()) {
        client_configuration.region = region;
    }
    client_configuration.userAgent += " " AWS_CRYPTOSDK_PRIVATE_VERSION_UA "/kms-keyring-cpp";
#ifdef VALGRIND_TESTS
    // When running under valgrind, the default timeouts are too slow
    client_configuration.requestTimeoutMs = 10000;
    client_configuration.connectTimeoutMs = 10000;
#endif
    return Aws::MakeShared<Aws::KMS::KMSClient>(allocationTag, client_configuration);
}

std::shared_ptr<KmsKeyring::SingleClientSupplier> KmsKeyring::SingleClientSupplier::Create(
    const std::shared_ptr<KMS::KMSClient> &kms_client) {
    return Aws::MakeShared<SingleClientSupplier>(AWS_CRYPTO_SDK_KMS_CLASS_TAG, kms_client);
}

std::shared_ptr<KMS::KMSClient> KmsKeyring::SingleClientSupplier::GetClient(
    const Aws::String &, std::function<void()> &report_success) {
    report_success = [] {};  // no-op lambda
    return this->kms_client;
}

std::shared_ptr<KmsKeyring::CachingClientSupplier> KmsKeyring::CachingClientSupplier::Create() {
    return Aws::MakeShared<KmsKeyring::CachingClientSupplier>(AWS_CRYPTO_SDK_KMS_CLASS_TAG);
}

std::shared_ptr<KMS::KMSClient> KmsKeyring::CachingClientSupplier::GetClient(
    const Aws::String &region, std::function<void()> &report_success) {
    {
        std::unique_lock<std::mutex> lock(cache_mutex);
        if (cache.find(region) != cache.end()) {
            report_success = [] {};  // no-op lambda
            return cache.at(region);
        }
    }
    auto client    = CreateDefaultKmsClient(AWS_CRYPTO_SDK_KMS_CLASS_TAG, region);
    report_success = [this, region, client] {
        std::unique_lock<std::mutex> lock(this->cache_mutex);
        this->cache[region] = client;
    };
    return client;
}

std::shared_ptr<KmsKeyring::ClientSupplier> Aws::Cryptosdk::Private::BuildClientSupplier(
    const Aws::Vector<Aws::String> &key_ids,
    const std::shared_ptr<Aws::KMS::KMSClient> kms_client,
    std::shared_ptr<KmsKeyring::ClientSupplier> client_supplier) {
    // Postcondition: If the caller provides a KMS client, then BuildClientSupplier MUST return a client supplier wrapping the provided client.
    if (kms_client) {
        return KmsKeyring::SingleClientSupplier::Create(kms_client);
    }

    if (key_ids.size() == 1) {
        Aws::String region = Private::parse_region_from_kms_key_arn(key_ids.front());
        std::shared_ptr<Aws::KMS::KMSClient> single_kms_client;
        if (client_supplier) {
            // Postcondition: If the caller provides only one key ID and a client supplier, then BuildClientSupplier MUST call the client supplier to obtain a KMS client in the key's region. BuildClientSupplier MUST return a client supplier that only supplies the obtained KMS client.
            std::function<void()> report_success;
            single_kms_client = client_supplier->GetClient(region, report_success);
            report_success();
        } else {
            // Postcondition: If the caller provides only one key ID and no client supplier, then BuildClientSupplier MUST create a KMS client in the key's region, and it MUST return a client supplier that only supplies the created KMS client.
            single_kms_client = CreateDefaultKmsClient(AWS_CRYPTO_SDK_KMS_CLASS_TAG, region);
        }
        return KmsKeyring::SingleClientSupplier::Create(single_kms_client);
    }

    return client_supplier
               // Postcondition: If the caller provides a client supplier, and provides zero or at least two key IDs, then BuildClientSupplier MUST return the provided client supplier.
               ? client_supplier
               // Postcondition: If the caller does not provide a client supplier, and provides zero or at least two key IDs, then BuildClientSupplier MUST return a default client supplier.
               : KmsKeyring::CachingClientSupplier::Create();
}

aws_cryptosdk_keyring *KmsKeyring::Builder::Build(
    const Aws::String &generator_key_id, const Aws::Vector<Aws::String> &additional_key_ids) const {
    if (Private::parse_region_from_kms_key_arn(generator_key_id).empty()) {
        AWS_LOGSTREAM_ERROR(AWS_CRYPTO_SDK_KMS_CLASS_TAG, "Unable to parse key ARN");
        aws_raise_error(AWS_ERROR_INVALID_ARGUMENT);
        return NULL;
    }

    Aws::Vector<Aws::String> my_key_ids = { generator_key_id };
    for (auto &key : additional_key_ids) {
        if (Private::parse_region_from_kms_key_arn(key).empty()) {
            AWS_LOGSTREAM_ERROR(AWS_CRYPTO_SDK_KMS_CLASS_TAG, "Unable to parse key ARN");
            aws_raise_error(AWS_ERROR_INVALID_ARGUMENT);
            return NULL;
        }
        my_key_ids.push_back(key);
    }

    return Aws::New<Private::KmsKeyringImpl>(
        AWS_CRYPTO_SDK_KMS_CLASS_TAG,
        my_key_ids,
        grant_tokens,
        Private::BuildClientSupplier(my_key_ids, kms_client, client_supplier));
}

aws_cryptosdk_keyring *KmsKeyring::Builder::BuildDiscovery() const {
    Aws::Vector<Aws::String> empty_key_ids_list;
    return Aws::New<Private::KmsKeyringImpl>(
        AWS_CRYPTO_SDK_KMS_CLASS_TAG,
        empty_key_ids_list,
        grant_tokens,
        Private::BuildClientSupplier(empty_key_ids_list, kms_client, client_supplier));
}

aws_cryptosdk_keyring *KmsKeyring::Builder::BuildDiscovery(std::shared_ptr<DiscoveryFilter> discovery_filter) const {
    if (!discovery_filter) {
        return nullptr;
    }

    Aws::Vector<Aws::String> empty_key_ids_list;
    return Aws::New<Private::KmsKeyringImpl>(
        AWS_CRYPTO_SDK_KMS_CLASS_TAG,
        empty_key_ids_list,
        grant_tokens,
        Private::BuildClientSupplier(empty_key_ids_list, kms_client, client_supplier),
        discovery_filter);
}

KmsKeyring::Builder &KmsKeyring::Builder::WithGrantTokens(const Aws::Vector<Aws::String> &grant_tokens) {
    this->grant_tokens.insert(this->grant_tokens.end(), grant_tokens.begin(), grant_tokens.end());
    return *this;
}

KmsKeyring::Builder &KmsKeyring::Builder::WithGrantToken(const Aws::String &grant_token) {
    this->grant_tokens.push_back(grant_token);
    return *this;
}

KmsKeyring::Builder &KmsKeyring::Builder::WithClientSupplier(
    const std::shared_ptr<KmsKeyring::ClientSupplier> &client_supplier) {
    this->client_supplier = client_supplier;
    return *this;
}

KmsKeyring::Builder &KmsKeyring::Builder::WithKmsClient(const std::shared_ptr<KMS::KMSClient> &kms_client) {
    this->kms_client = kms_client;
    return *this;
}

bool KmsKeyring::DiscoveryFilter::IsAuthorized(const Aws::String &key_arn) const {
    Utils::ARN arn(key_arn);
    if (!arn) {
        return false;
    }

    bool matching_partition = arn.GetPartition() == partition;
    bool matching_account   = account_ids.find(arn.GetAccountId()) != account_ids.end();
    return matching_partition && matching_account;
}

KmsKeyring::DiscoveryFilterBuilder KmsKeyring::DiscoveryFilter::Builder(Aws::String partition) {
    KmsKeyring::DiscoveryFilterBuilder builder(partition);
    return builder;
}

KmsKeyring::DiscoveryFilterBuilder &KmsKeyring::DiscoveryFilterBuilder::AddAccount(const Aws::String &account_id) {
    this->account_ids.insert(account_id);
    return *this;
}

KmsKeyring::DiscoveryFilterBuilder &KmsKeyring::DiscoveryFilterBuilder::AddAccounts(
    const Aws::Vector<Aws::String> &account_ids) {
    this->account_ids.insert(account_ids.begin(), account_ids.end());
    return *this;
}

KmsKeyring::DiscoveryFilterBuilder &KmsKeyring::DiscoveryFilterBuilder::WithAccounts(
    const Aws::Vector<Aws::String> &account_ids) {
    this->account_ids.clear();
    return this->AddAccounts(account_ids);
}

std::shared_ptr<KmsKeyring::DiscoveryFilter> KmsKeyring::DiscoveryFilterBuilder::Build() const {
    // Must have at least one account ID, and partition and account IDs cannot be the empty string
    if (account_ids.empty()) {
        AWS_LOGSTREAM_ERROR(
            AWS_CRYPTO_SDK_KMS_CLASS_TAG, "Invalid DiscoveryFilterBuilder: account IDs cannot be empty");
        return nullptr;
    }
    if (partition.empty()) {
        AWS_LOGSTREAM_ERROR(AWS_CRYPTO_SDK_KMS_CLASS_TAG, "Invalid DiscoveryFilterBuilder: partition cannot be blank");
        return nullptr;
    }
    if (account_ids.find("") != account_ids.end()) {
        AWS_LOGSTREAM_ERROR(
            AWS_CRYPTO_SDK_KMS_CLASS_TAG, "Invalid DiscoveryFilterBuilder: account IDs cannot be blank");
        return nullptr;
    }

    return Aws::MakeShared<Private::DiscoveryFilterImpl>(
        AWS_CRYPTO_SDK_DISCOVERY_FILTER_CLASS_TAG, partition, account_ids);
}

}  // namespace Cryptosdk
}  // namespace Aws