/*
 * Copyright 2025 The OpenSSL Project Authors. All Rights Reserved.
 *
 * Licensed under the Apache License 2.0 (the "License").  You may not use
 * this file except in compliance with the License.  You can obtain a copy
 * in the file LICENSE in the source distribution or at
 * https://www.openssl.org/source/license.html
 */

#include <openssl/cms.h>
#include <openssl/core_names.h>
#include <openssl/crypto.h>
#include <openssl/err.h>
#include <openssl/evp.h>
#include <openssl/kdf.h>
#include <openssl/x509.h>
#include "cms_local.h"
#include "crypto/evp.h"
#include "internal/sizes.h"

/* KEM Recipient Info (KEMRI) routines */

int ossl_cms_RecipientInfo_kemri_get0_alg(CMS_RecipientInfo *ri,
                                          uint32_t **pkekLength,
                                          X509_ALGOR **pwrap)
{
    if (ri->type != CMS_RECIPINFO_KEM) {
        ERR_raise(ERR_LIB_CMS, CMS_R_NOT_KEM);
        return 0;
    }
    if (pkekLength)
        *pkekLength = &ri->d.ori->d.kemri->kekLength;
    if (pwrap)
        *pwrap = ri->d.ori->d.kemri->wrap;
    return 1;
}

int CMS_RecipientInfo_kemri_cert_cmp(CMS_RecipientInfo *ri, X509 *cert)
{
    if (ri->type != CMS_RECIPINFO_KEM) {
        ERR_raise(ERR_LIB_CMS, CMS_R_NOT_KEM);
        return -2;
    }
    return ossl_cms_SignerIdentifier_cert_cmp(ri->d.ori->d.kemri->rid, cert);
}

int CMS_RecipientInfo_kemri_set0_pkey(CMS_RecipientInfo *ri, EVP_PKEY *pk)
{
    EVP_PKEY_CTX *pctx = NULL;
    CMS_KEMRecipientInfo *kemri;

    if (ri->type != CMS_RECIPINFO_KEM) {
        ERR_raise(ERR_LIB_CMS, CMS_R_NOT_KEM);
        return 0;
    }

    kemri = ri->d.ori->d.kemri;

    EVP_PKEY_CTX_free(kemri->pctx);
    kemri->pctx = NULL;

    if (pk != NULL) {
        pctx = EVP_PKEY_CTX_new_from_pkey(ossl_cms_ctx_get0_libctx(kemri->cms_ctx), pk,
                                          ossl_cms_ctx_get0_propq(kemri->cms_ctx));
        if (pctx == NULL || EVP_PKEY_decapsulate_init(pctx, NULL) <= 0)
            goto err;

        kemri->pctx = pctx;
    }

    return 1;
err:
    EVP_PKEY_CTX_free(pctx);
    return 0;
}

/* Initialise a kemri based on passed certificate and key */

int ossl_cms_RecipientInfo_kemri_init(CMS_RecipientInfo *ri, X509 *recip,
                                      EVP_PKEY *recipPubKey, unsigned int flags,
                                      const CMS_CTX *ctx)
{
    CMS_OtherRecipientInfo *ori;
    CMS_KEMRecipientInfo *kemri;
    int idtype;
    X509_PUBKEY *x_pubkey;
    X509_ALGOR *x_alg;

    ri->d.ori = M_ASN1_new_of(CMS_OtherRecipientInfo);
    if (ri->d.ori == NULL)
        return 0;
    ri->encoded_type = CMS_RECIPINFO_OTHER;
    ri->type = CMS_RECIPINFO_KEM;

    ori = ri->d.ori;
    ori->oriType = OBJ_nid2obj(NID_id_smime_ori_kem);
    if (ori->oriType == NULL)
        return 0;
    ori->d.kemri = M_ASN1_new_of(CMS_KEMRecipientInfo);
    if (ori->d.kemri == NULL)
        return 0;

    kemri = ori->d.kemri;
    kemri->version = 0;
    kemri->cms_ctx = ctx;

    /*
     * Not a typo: RecipientIdentifier and SignerIdentifier are the same
     * structure.
     */

    idtype = (flags & CMS_USE_KEYID) ? CMS_RECIPINFO_KEYIDENTIFIER : CMS_RECIPINFO_ISSUER_SERIAL;
    if (!ossl_cms_set1_SignerIdentifier(kemri->rid, recip, idtype, ctx))
        return 0;

    x_pubkey = X509_get_X509_PUBKEY(recip);
    if (x_pubkey == NULL)
        return 0;
    if (!X509_PUBKEY_get0_param(NULL, NULL, NULL, &x_alg, x_pubkey))
        return 0;
    if (!X509_ALGOR_copy(kemri->kem, x_alg))
        return 0;

    kemri->pctx = EVP_PKEY_CTX_new_from_pkey(ossl_cms_ctx_get0_libctx(ctx),
                                             recipPubKey,
                                             ossl_cms_ctx_get0_propq(ctx));
    if (kemri->pctx == NULL)
        return 0;
    if (EVP_PKEY_encapsulate_init(kemri->pctx, NULL) <= 0)
        return 0;

    return 1;
}

EVP_CIPHER_CTX *CMS_RecipientInfo_kemri_get0_ctx(CMS_RecipientInfo *ri)
{
    if (ri->type == CMS_RECIPINFO_KEM)
        return ri->d.ori->d.kemri->ctx;
    return NULL;
}

X509_ALGOR *CMS_RecipientInfo_kemri_get0_kdf_alg(CMS_RecipientInfo *ri)
{
    if (ri->type == CMS_RECIPINFO_KEM)
        return ri->d.ori->d.kemri->kdf;
    return NULL;
}

int CMS_RecipientInfo_kemri_set_ukm(CMS_RecipientInfo *ri,
                                    const unsigned char *ukm,
                                    int ukmLength)
{
    CMS_KEMRecipientInfo *kemri;
    ASN1_OCTET_STRING *ukm_str;

    if (ri->type != CMS_RECIPINFO_KEM) {
        ERR_raise(ERR_LIB_CMS, CMS_R_NOT_KEM);
        return 0;
    }

    if (ukm == NULL && ukmLength != 0) {
        ERR_raise(ERR_LIB_CMS, ERR_R_PASSED_INVALID_ARGUMENT);
        return 0;
    }

    kemri = ri->d.ori->d.kemri;

    ukm_str = ASN1_OCTET_STRING_new();
    if (ukm_str == NULL)
        return 0;
    if (!ASN1_OCTET_STRING_set(ukm_str, ukm, ukmLength)) {
        ASN1_OCTET_STRING_free(ukm_str);
        return 0;
    }
    ASN1_OCTET_STRING_free(kemri->ukm);
    kemri->ukm = ukm_str;
    return 1;
}

static EVP_KDF_CTX *create_kdf_ctx(CMS_KEMRecipientInfo *kemri)
{
    const ASN1_OBJECT *kdf_oid;
    int ptype;
    char kdf_alg[OSSL_MAX_NAME_SIZE];
    EVP_KDF *kdf = NULL;
    EVP_KDF_CTX *kctx = NULL;

    /*
     * KDFs with algorithm identifier parameters are not supported yet. To
     * support this, EVP_KDF_CTX_set_algor_params from
     * `doc/designs/passing-algorithmidentifier-parameters.md` needs to be
     * implemented.
     */
    X509_ALGOR_get0(&kdf_oid, &ptype, NULL, kemri->kdf);
    if (ptype != V_ASN1_UNDEF && ptype != V_ASN1_NULL) {
        ERR_raise(ERR_LIB_CMS, CMS_R_UNSUPPORTED_KDF_ALGORITHM);
        goto err;
    }
    if (OBJ_obj2txt(kdf_alg, sizeof(kdf_alg), kdf_oid, 1) < 0)
        goto err;

    kdf = EVP_KDF_fetch(ossl_cms_ctx_get0_libctx(kemri->cms_ctx), kdf_alg,
                        ossl_cms_ctx_get0_propq(kemri->cms_ctx));
    if (kdf == NULL)
        goto err;

    kctx = EVP_KDF_CTX_new(kdf);
err:
    EVP_KDF_free(kdf);
    return kctx;
}

static int kdf_derive(unsigned char *kek, size_t keklen,
                      const unsigned char *ss, size_t sslen,
                      CMS_KEMRecipientInfo *kemri)
{
    EVP_KDF_CTX *kctx = NULL;
    OSSL_PARAM params[3];
    unsigned char *infoder = NULL;
    int infolen = 0;
    int rv = 0;

    infolen = CMS_CMSORIforKEMOtherInfo_encode(&infoder, kemri->wrap, kemri->ukm,
                                               kemri->kekLength);
    if (infolen <= 0)
        goto err;

    kctx = create_kdf_ctx(kemri);
    if (kctx == NULL)
        goto err;

    params[0] = OSSL_PARAM_construct_octet_string(OSSL_KDF_PARAM_KEY,
                                                  (unsigned char *)ss, sslen);
    params[1] = OSSL_PARAM_construct_octet_string(OSSL_KDF_PARAM_INFO,
                                                  (char *)infoder, infolen);
    params[2] = OSSL_PARAM_construct_end();

    if (EVP_KDF_derive(kctx, kek, keklen, params) <= 0)
        goto err;

    rv = 1;
err:
    OPENSSL_free(infoder);
    EVP_KDF_CTX_free(kctx);

    return rv;
}

/*
 * Derive KEK and decrypt/encrypt with it to produce either the original CEK
 * or the encrypted CEK.
 */

static int cms_kek_cipher(unsigned char **pout, size_t *poutlen,
                          const unsigned char *ss, size_t sslen,
                          const unsigned char *in, size_t inlen,
                          CMS_KEMRecipientInfo *kemri, int enc)
{
    /* Key encryption key */
    unsigned char kek[EVP_MAX_KEY_LENGTH];
    size_t keklen = kemri->kekLength;
    unsigned char *out = NULL;
    int outlen = 0;
    int rv = 0;

    if (keklen > sizeof(kek)) {
        ERR_raise(ERR_LIB_CMS, CMS_R_INVALID_KEY_LENGTH);
        return 0;
    }

    if (!kdf_derive(kek, keklen, ss, sslen, kemri))
        goto err;

    /* Set KEK in context */
    if (!EVP_CipherInit_ex(kemri->ctx, NULL, NULL, kek, NULL, enc))
        goto err;
    /* obtain output length of ciphered key */
    if (!EVP_CipherUpdate(kemri->ctx, NULL, &outlen, in, (int)inlen))
        goto err;
    out = OPENSSL_malloc(outlen);
    if (out == NULL)
        goto err;
    if (!EVP_CipherUpdate(kemri->ctx, out, &outlen, in, (int)inlen))
        goto err;
    *pout = out;
    out = NULL;
    *poutlen = (size_t)outlen;

    rv = 1;
err:
    OPENSSL_free(out);
    OPENSSL_cleanse(kek, sizeof(kek));
    EVP_CIPHER_CTX_reset(kemri->ctx);
    EVP_PKEY_CTX_free(kemri->pctx);
    kemri->pctx = NULL;
    return rv;
}

/* Encrypt content key in KEM recipient info */

int ossl_cms_RecipientInfo_kemri_encrypt(const CMS_ContentInfo *cms,
                                         CMS_RecipientInfo *ri)
{
    CMS_KEMRecipientInfo *kemri;
    CMS_EncryptedContentInfo *ec;
    unsigned char *kem_ct = NULL;
    size_t kem_ct_len;
    unsigned char *kem_secret = NULL;
    size_t kem_secret_len = 0;
    unsigned char *enckey;
    size_t enckeylen;
    int rv = 0;

    if (ri->type != CMS_RECIPINFO_KEM) {
        ERR_raise(ERR_LIB_CMS, CMS_R_NOT_KEM);
        return 0;
    }

    kemri = ri->d.ori->d.kemri;

    ec = ossl_cms_get0_env_enc_content(cms);
    /* Initialise wrap algorithm parameters */
    if (!ossl_cms_RecipientInfo_wrap_init(ri, ec->cipher))
        return 0;

    /* Initialise KDF algorithm */
    if (!ossl_cms_env_asn1_ctrl(ri, 0))
        return 0;

    if (EVP_PKEY_encapsulate(kemri->pctx, NULL, &kem_ct_len, NULL, &kem_secret_len) <= 0)
        return 0;
    kem_ct = OPENSSL_malloc(kem_ct_len);
    kem_secret = OPENSSL_malloc(kem_secret_len);
    if (kem_ct == NULL || kem_secret == NULL)
        goto err;

    if (EVP_PKEY_encapsulate(kemri->pctx, kem_ct, &kem_ct_len, kem_secret, &kem_secret_len) <= 0)
        goto err;

    ASN1_STRING_set0(kemri->kemct, kem_ct, (int)kem_ct_len);
    kem_ct = NULL;

    if (!cms_kek_cipher(&enckey, &enckeylen, kem_secret, kem_secret_len, ec->key, ec->keylen,
                        kemri, 1))
        goto err;
    ASN1_STRING_set0(kemri->encryptedKey, enckey, (int)enckeylen);

    rv = 1;
err:
    OPENSSL_free(kem_ct);
    OPENSSL_clear_free(kem_secret, kem_secret_len);
    return rv;
}

int ossl_cms_RecipientInfo_kemri_decrypt(const CMS_ContentInfo *cms,
                                         CMS_RecipientInfo *ri)
{
    CMS_KEMRecipientInfo *kemri;
    CMS_EncryptedContentInfo *ec;
    const unsigned char *kem_ct = NULL;
    size_t kem_ct_len;
    unsigned char *kem_secret = NULL;
    size_t kem_secret_len = 0;
    unsigned char *enckey = NULL;
    size_t enckeylen;
    unsigned char *cek = NULL;
    size_t ceklen;
    int ret = 0;

    if (ri->type != CMS_RECIPINFO_KEM) {
        ERR_raise(ERR_LIB_CMS, CMS_R_NOT_KEM);
        return 0;
    }

    kemri = ri->d.ori->d.kemri;

    ec = ossl_cms_get0_env_enc_content(cms);

    if (kemri->pctx == NULL) {
        ERR_raise(ERR_LIB_CMS, CMS_R_NO_PRIVATE_KEY);
        return 0;
    }

    /* Setup all parameters to derive KEK */
    if (!ossl_cms_env_asn1_ctrl(ri, 1))
        goto err;

    kem_ct = ASN1_STRING_get0_data(kemri->kemct);
    kem_ct_len = ASN1_STRING_length(kemri->kemct);

    if (EVP_PKEY_decapsulate(kemri->pctx, NULL, &kem_secret_len, kem_ct, kem_ct_len) <= 0)
        return 0;
    kem_secret = OPENSSL_malloc(kem_secret_len);
    if (kem_secret == NULL)
        goto err;

    if (EVP_PKEY_decapsulate(kemri->pctx, kem_secret, &kem_secret_len, kem_ct, kem_ct_len) <= 0)
        goto err;

    /* Attempt to decrypt CEK */
    enckeylen = kemri->encryptedKey->length;
    enckey = kemri->encryptedKey->data;
    if (!cms_kek_cipher(&cek, &ceklen, kem_secret, kem_secret_len, enckey, enckeylen, kemri, 0))
        goto err;
    ec = ossl_cms_get0_env_enc_content(cms);
    OPENSSL_clear_free(ec->key, ec->keylen);
    ec->key = cek;
    ec->keylen = ceklen;

    ret = 1;
err:
    OPENSSL_clear_free(kem_secret, kem_secret_len);
    return ret;
}