/**
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0.
 */
#include <aws/cal/private/rsa.h>

#include <aws/cal/cal.h>
#include <aws/cal/hash.h>
#include <aws/cal/private/der.h>

typedef struct aws_rsa_key_pair *(aws_rsa_key_pair_new_from_public_pkcs1_fn)(struct aws_allocator *allocator,
                                                                             struct aws_byte_cursor public_key);

typedef struct aws_rsa_key_pair *(aws_rsa_key_pair_new_from_private_pkcs1_fn)(struct aws_allocator *allocator,
                                                                              struct aws_byte_cursor private_key);

typedef struct aws_rsa_key_pair *(aws_rsa_key_pair_new_from_private_pkcs8_fn)(struct aws_allocator *allocator,
                                                                              struct aws_byte_cursor private_key);

#ifndef BYO_CRYPTO

extern struct aws_rsa_key_pair *aws_rsa_key_pair_new_from_public_key_pkcs1_impl(
    struct aws_allocator *allocator,
    struct aws_byte_cursor public_key);

extern struct aws_rsa_key_pair *aws_rsa_key_pair_new_from_private_key_pkcs1_impl(
    struct aws_allocator *allocator,
    struct aws_byte_cursor private_key);

struct aws_rsa_key_pair *aws_rsa_key_pair_new_from_private_key_pkcs8_impl(
    struct aws_allocator *allocator,
    struct aws_byte_cursor private_key);

#else  /* BYO_CRYPTO */

struct aws_rsa_key_pair *aws_rsa_key_pair_new_from_public_key_pkcs1_impl(
    struct aws_allocator *allocator,
    struct aws_byte_cursor public_key) {
    (void)allocator;
    (void)public_key;
    abort();
}

struct aws_rsa_key_pair *aws_rsa_key_pair_new_from_private_key_pkcs1_impl(
    struct aws_allocator *allocator,
    struct aws_byte_cursor private_key) {
    (void)allocator;
    (void)private_key;
    abort();
}

struct aws_rsa_key_pair *aws_rsa_key_pair_new_from_private_key_pkcs8_impl(
    struct aws_allocator *allocator,
    struct aws_byte_cursor private_key) {
    (void)allocator;
    (void)private_key;
    abort();
}
#endif /* BYO_CRYPTO */

static aws_rsa_key_pair_new_from_public_pkcs1_fn *s_rsa_key_pair_new_from_public_key_pkcs1_fn =
    aws_rsa_key_pair_new_from_public_key_pkcs1_impl;

static aws_rsa_key_pair_new_from_private_pkcs1_fn *s_rsa_key_pair_new_from_private_key_pkcs1_fn =
    aws_rsa_key_pair_new_from_private_key_pkcs1_impl;

static aws_rsa_key_pair_new_from_private_pkcs8_fn *s_rsa_key_pair_new_from_private_key_pkcs8_fn =
    aws_rsa_key_pair_new_from_private_key_pkcs8_impl;

struct aws_rsa_key_pair *aws_rsa_key_pair_new_from_public_key_pkcs1(
    struct aws_allocator *allocator,
    struct aws_byte_cursor public_key) {
    return s_rsa_key_pair_new_from_public_key_pkcs1_fn(allocator, public_key);
}

struct aws_rsa_key_pair *aws_rsa_key_pair_new_from_private_key_pkcs1(
    struct aws_allocator *allocator,
    struct aws_byte_cursor private_key) {
    return s_rsa_key_pair_new_from_private_key_pkcs1_fn(allocator, private_key);
}

AWS_CAL_API struct aws_rsa_key_pair *aws_rsa_key_pair_new_from_private_key_pkcs8(
    struct aws_allocator *allocator,
    struct aws_byte_cursor key) {
    return s_rsa_key_pair_new_from_private_key_pkcs8_fn(allocator, key);
}

void aws_rsa_key_pair_base_clean_up(struct aws_rsa_key_pair *key_pair) {
    aws_byte_buf_clean_up_secure(&key_pair->priv);
    aws_byte_buf_clean_up_secure(&key_pair->pub);
}

struct aws_rsa_key_pair *aws_rsa_key_pair_acquire(struct aws_rsa_key_pair *key_pair) {
    return aws_ref_count_acquire(&key_pair->ref_count);
}

struct aws_rsa_key_pair *aws_rsa_key_pair_release(struct aws_rsa_key_pair *key_pair) {
    if (key_pair != NULL) {
        aws_ref_count_release(&key_pair->ref_count);
    }
    return NULL;
}

size_t aws_rsa_key_pair_max_encrypt_plaintext_size(
    const struct aws_rsa_key_pair *key_pair,
    enum aws_rsa_encryption_algorithm algorithm) {
    /*
     * Per rfc8017, max size of plaintext for encrypt operation is as follows:
     * PKCS1-v1_5: (key size in bytes) - 11
     * OAEP: (key size in bytes) - 2 * (hash bytes) - 2
     */

    size_t key_size_in_bytes = key_pair->key_size_in_bits / 8;
    switch (algorithm) {
        case AWS_CAL_RSA_ENCRYPTION_PKCS1_5:
            return key_size_in_bytes - 11;
        case AWS_CAL_RSA_ENCRYPTION_OAEP_SHA256:
            return key_size_in_bytes - 2 * (256 / 8) - 2;
        case AWS_CAL_RSA_ENCRYPTION_OAEP_SHA512:
            return key_size_in_bytes - 2 * (512 / 8) - 2;
        default:
            AWS_FATAL_ASSERT("Unsupported RSA Encryption Algorithm");
    }

    return 0;
}

int aws_rsa_key_pair_encrypt(
    const struct aws_rsa_key_pair *key_pair,
    enum aws_rsa_encryption_algorithm algorithm,
    struct aws_byte_cursor plaintext,
    struct aws_byte_buf *out) {
    AWS_PRECONDITION(key_pair);
    AWS_PRECONDITION(out);
    AWS_PRECONDITION(aws_byte_cursor_is_valid(&plaintext));

    if (AWS_UNLIKELY(aws_rsa_key_pair_max_encrypt_plaintext_size(key_pair, algorithm) < plaintext.len)) {
        AWS_LOGF_ERROR(AWS_LS_CAL_RSA, "Unexpected buffer size. For RSA, ciphertext must not exceed block size");
        return aws_raise_error(AWS_ERROR_CAL_BUFFER_TOO_LARGE_FOR_ALGORITHM);
    }

    return key_pair->vtable->encrypt(key_pair, algorithm, plaintext, out);
}

AWS_CAL_API int aws_rsa_key_pair_decrypt(
    const struct aws_rsa_key_pair *key_pair,
    enum aws_rsa_encryption_algorithm algorithm,
    struct aws_byte_cursor ciphertext,
    struct aws_byte_buf *out) {
    AWS_PRECONDITION(key_pair);
    AWS_PRECONDITION(out);
    AWS_PRECONDITION(aws_byte_cursor_is_valid(&ciphertext));

    if (AWS_UNLIKELY(ciphertext.len != (key_pair->key_size_in_bits / 8))) {
        AWS_LOGF_ERROR(AWS_LS_CAL_RSA, "Unexpected buffer size. For RSA, ciphertext is expected to match block size.");
        return aws_raise_error(AWS_ERROR_INVALID_ARGUMENT);
    }

    return key_pair->vtable->decrypt(key_pair, algorithm, ciphertext, out);
}

int aws_rsa_key_pair_sign_message(
    const struct aws_rsa_key_pair *key_pair,
    enum aws_rsa_signature_algorithm algorithm,
    struct aws_byte_cursor digest,
    struct aws_byte_buf *out) {
    AWS_PRECONDITION(key_pair);
    AWS_PRECONDITION(out);
    AWS_PRECONDITION(aws_byte_cursor_is_valid(&digest));

    AWS_FATAL_ASSERT(
        algorithm == AWS_CAL_RSA_SIGNATURE_PKCS1_5_SHA256 || algorithm == AWS_CAL_RSA_SIGNATURE_PKCS1_5_SHA1 ||
        algorithm == AWS_CAL_RSA_SIGNATURE_PSS_SHA256);

    if (digest.len > AWS_SHA256_LEN) {
        AWS_LOGF_ERROR(
            AWS_LS_CAL_RSA, "Unexpected digest size. For RSA, digest length is bound by max size of hash function");
        return aws_raise_error(AWS_ERROR_INVALID_ARGUMENT);
    }

    return key_pair->vtable->sign(key_pair, algorithm, digest, out);
}

int aws_rsa_key_pair_verify_signature(
    const struct aws_rsa_key_pair *key_pair,
    enum aws_rsa_signature_algorithm algorithm,
    struct aws_byte_cursor digest,
    struct aws_byte_cursor signature) {
    AWS_PRECONDITION(key_pair);
    AWS_PRECONDITION(aws_byte_cursor_is_valid(&digest));
    AWS_PRECONDITION(aws_byte_cursor_is_valid(&signature));

    return key_pair->vtable->verify(key_pair, algorithm, digest, signature);
}

size_t aws_rsa_key_pair_block_length(const struct aws_rsa_key_pair *key_pair) {
    AWS_PRECONDITION(key_pair);
    return key_pair->key_size_in_bits / 8;
}

size_t aws_rsa_key_pair_signature_length(const struct aws_rsa_key_pair *key_pair) {
    AWS_PRECONDITION(key_pair);
    return key_pair->key_size_in_bits / 8;
}

int aws_rsa_key_pair_get_public_key(
    const struct aws_rsa_key_pair *key_pair,
    enum aws_rsa_key_export_format format,
    struct aws_byte_buf *out) {
    (void)format; /* ignore format for now, since only pkcs1 is supported. */
    AWS_PRECONDITION(key_pair);
    AWS_PRECONDITION(out);

    if (key_pair->pub.len == 0) {
        return aws_raise_error(AWS_ERROR_PLATFORM_NOT_SUPPORTED);
    }

    aws_byte_buf_init_copy(out, key_pair->allocator, &key_pair->pub);
    return AWS_OP_SUCCESS;
}

int aws_rsa_key_pair_get_private_key(
    const struct aws_rsa_key_pair *key_pair,
    enum aws_rsa_key_export_format format,
    struct aws_byte_buf *out) {
    (void)format; /* ignore format for now, since only pkcs1 is supported. */
    AWS_PRECONDITION(key_pair);
    AWS_PRECONDITION(out);

    if (key_pair->priv.len == 0) {
        return aws_raise_error(AWS_ERROR_PLATFORM_NOT_SUPPORTED);
    }

    aws_byte_buf_init_copy(out, key_pair->allocator, &key_pair->priv);
    return AWS_OP_SUCCESS;
}

int aws_der_decoder_load_private_rsa_pkcs1(struct aws_der_decoder *decoder, struct aws_rsa_private_key_pkcs1 *out) {

    if (!aws_der_decoder_next(decoder) || aws_der_decoder_tlv_type(decoder) != AWS_DER_SEQUENCE) {
        return aws_raise_error(AWS_ERROR_CAL_MALFORMED_ASN1_ENCOUNTERED);
    }

    struct aws_byte_cursor version_cur;
    if (!aws_der_decoder_next(decoder) || aws_der_decoder_tlv_unsigned_integer(decoder, &version_cur)) {
        return aws_raise_error(AWS_ERROR_CAL_MALFORMED_ASN1_ENCOUNTERED);
    }

    if (version_cur.len != 1 || version_cur.ptr[0] != 0) {
        return aws_raise_error(AWS_ERROR_CAL_UNSUPPORTED_KEY_FORMAT);
    }
    out->version = 0;

    if (!aws_der_decoder_next(decoder) || aws_der_decoder_tlv_unsigned_integer(decoder, &(out->modulus))) {
        return aws_raise_error(AWS_ERROR_CAL_MALFORMED_ASN1_ENCOUNTERED);
    }

    if (!aws_der_decoder_next(decoder) || aws_der_decoder_tlv_unsigned_integer(decoder, &out->publicExponent)) {
        return aws_raise_error(AWS_ERROR_CAL_MALFORMED_ASN1_ENCOUNTERED);
    }

    if (!aws_der_decoder_next(decoder) || aws_der_decoder_tlv_unsigned_integer(decoder, &out->privateExponent)) {
        return aws_raise_error(AWS_ERROR_CAL_MALFORMED_ASN1_ENCOUNTERED);
    }

    if (!aws_der_decoder_next(decoder) || aws_der_decoder_tlv_unsigned_integer(decoder, &out->prime1)) {
        return aws_raise_error(AWS_ERROR_CAL_MALFORMED_ASN1_ENCOUNTERED);
    }

    if (!aws_der_decoder_next(decoder) || aws_der_decoder_tlv_unsigned_integer(decoder, &out->prime2)) {
        return aws_raise_error(AWS_ERROR_CAL_MALFORMED_ASN1_ENCOUNTERED);
    }

    if (!aws_der_decoder_next(decoder) || aws_der_decoder_tlv_unsigned_integer(decoder, &out->exponent1)) {
        return aws_raise_error(AWS_ERROR_CAL_MALFORMED_ASN1_ENCOUNTERED);
    }

    if (!aws_der_decoder_next(decoder) || aws_der_decoder_tlv_unsigned_integer(decoder, &out->exponent2)) {
        return aws_raise_error(AWS_ERROR_CAL_MALFORMED_ASN1_ENCOUNTERED);
    }

    if (!aws_der_decoder_next(decoder) || aws_der_decoder_tlv_unsigned_integer(decoder, &out->coefficient)) {
        return aws_raise_error(AWS_ERROR_CAL_MALFORMED_ASN1_ENCOUNTERED);
    }

    return AWS_OP_SUCCESS;
}

int aws_der_decoder_load_public_rsa_pkcs1(struct aws_der_decoder *decoder, struct aws_rsa_public_key_pkcs1 *out) {
    if (!aws_der_decoder_next(decoder) || aws_der_decoder_tlv_type(decoder) != AWS_DER_SEQUENCE) {
        return aws_raise_error(AWS_ERROR_CAL_MALFORMED_ASN1_ENCOUNTERED);
    }

    if (!aws_der_decoder_next(decoder) || aws_der_decoder_tlv_unsigned_integer(decoder, &(out->modulus))) {
        return aws_raise_error(AWS_ERROR_CAL_MALFORMED_ASN1_ENCOUNTERED);
    }

    if (!aws_der_decoder_next(decoder) || aws_der_decoder_tlv_unsigned_integer(decoder, &out->publicExponent)) {
        return aws_raise_error(AWS_ERROR_CAL_MALFORMED_ASN1_ENCOUNTERED);
    }

    return AWS_OP_SUCCESS;
}

int is_valid_rsa_key_size(size_t key_size_in_bits) {
    if (key_size_in_bits < AWS_CAL_RSA_MIN_SUPPORTED_KEY_SIZE_IN_BITS ||
        key_size_in_bits > AWS_CAL_RSA_MAX_SUPPORTED_KEY_SIZE_IN_BITS || key_size_in_bits % 8 != 0) {
        return aws_raise_error(AWS_ERROR_INVALID_ARGUMENT);
    }

    return AWS_OP_SUCCESS;
}

#ifndef BYO_CRYPTO

static uint8_t s_rsa_encryption_oid[] = {0x2a, 0x86, 0x48, 0x86, 0xf7, 0x0d, 0x01, 0x01, 0x01};

/**
 * There are likely tens of algo specifiers for rsa out there.
 * But in practice mostly everyone uses rsaEncryption oid for key at rest.
 * So for now just support that and we can add other ones later if needed.
 * Even though its technically encryption oid, most crypto libs are lax
 * and allow the key to be used for all operations.
 */
static struct aws_byte_cursor s_rsa_encryption_oid_cur = {
    .ptr = (s_rsa_encryption_oid),
    .len = sizeof(s_rsa_encryption_oid),
};

/**
 * Win and Mac dont provide native support for loading pkcs8 keys, so
 * for simplicity and consistency just parse pkcs1 key from pkcs8 and
 * use it with existing pkcs1 loader.
 */
struct aws_rsa_key_pair *aws_rsa_key_pair_new_from_private_key_pkcs8_impl(
    struct aws_allocator *allocator,
    struct aws_byte_cursor private_key) {

    struct aws_der_decoder *decoder = aws_der_decoder_new(allocator, private_key);

    if (decoder == NULL) {
        return NULL;
    }

    /**
     * Format of pkcs8 is as follows.
     * PrivateKeyInfo ::= SEQUENCE {
     *  version         Version,          -- INTEGER (0)
     *  algorithm       AlgorithmIdentifier,
     *  privateKey      OCTET STRING,     -- contains DER encoded key
     *  attributes  [0] IMPLICIT SET OF Attribute OPTIONAL
     * }
     * AlgorithmIdentifier ::= SEQUENCE {
     *  algorithm       OBJECT IDENTIFIER,
     *  parameters      ANY DEFINED BY algorithm OPTIONAL
     * }
     */

    struct aws_rsa_key_pair *key_pair = NULL;

    if (!aws_der_decoder_next(decoder) || aws_der_decoder_tlv_type(decoder) != AWS_DER_SEQUENCE) {
        aws_raise_error(AWS_ERROR_CAL_MALFORMED_ASN1_ENCOUNTERED);
        goto on_done;
    }

    /* version */
    struct aws_byte_cursor version_cur;
    if (!aws_der_decoder_next(decoder) || aws_der_decoder_tlv_unsigned_integer(decoder, &version_cur)) {
        aws_raise_error(AWS_ERROR_CAL_MALFORMED_ASN1_ENCOUNTERED);
        goto on_done;
    }

    if (version_cur.len != 1 || version_cur.ptr[0] != 0) {
        aws_raise_error(AWS_ERROR_CAL_UNSUPPORTED_KEY_FORMAT);
        goto on_done;
    }

    /* oid */
    struct aws_byte_cursor oid_cur;
    if (!aws_der_decoder_next(decoder) || aws_der_decoder_tlv_type(decoder) != AWS_DER_SEQUENCE) {
        aws_raise_error(AWS_ERROR_CAL_MALFORMED_ASN1_ENCOUNTERED);
        goto on_done;
    }

    if (!aws_der_decoder_next(decoder) || aws_der_decoder_tlv_type(decoder) != AWS_DER_OBJECT_IDENTIFIER ||
        aws_der_decoder_tlv_blob(decoder, &oid_cur)) {
        aws_raise_error(AWS_ERROR_CAL_MALFORMED_ASN1_ENCOUNTERED);
        goto on_done;
    }

    if (!aws_byte_cursor_eq(&s_rsa_encryption_oid_cur, &oid_cur)) {
        aws_raise_error(AWS_ERROR_CAL_UNSUPPORTED_KEY_FORMAT);
        goto on_done;
    }

    /* skip additional params */
    if (!aws_der_decoder_next(decoder)) {
        aws_raise_error(AWS_ERROR_CAL_MALFORMED_ASN1_ENCOUNTERED);
        goto on_done;
    }

    /* key */
    struct aws_byte_cursor key;
    if (!aws_der_decoder_next(decoder) || aws_der_decoder_tlv_string(decoder, &key)) {
        aws_raise_error(AWS_ERROR_CAL_MALFORMED_ASN1_ENCOUNTERED);
        goto on_done;
    }

    key_pair = aws_rsa_key_pair_new_from_private_key_pkcs1(allocator, key);

on_done:
    aws_der_decoder_destroy(decoder);
    return key_pair;
}

#endif