/*
 * %CopyrightBegin%
 *
 * Copyright Ericsson AB 2010-2018. 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.
 * 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.
 *
 * %CopyrightEnd%
 */

#include "aead.h"
#include "aes.h"
#include "cipher.h"

ERL_NIF_TERM aead_encrypt(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
{/* (Type,Key,Iv,AAD,In) */
#if defined(HAVE_AEAD)
    const struct cipher_type_t *cipherp;
    EVP_CIPHER_CTX *ctx = NULL;
    const EVP_CIPHER *cipher = NULL;
    ErlNifBinary key, iv, aad, in;
    unsigned int tag_len;
    unsigned char *outp, *tagp;
    ERL_NIF_TERM type, out, out_tag, ret;
    int len, ctx_ctrl_set_ivlen, ctx_ctrl_get_tag, ctx_ctrl_set_tag;

    type = argv[0];

    ASSERT(argc == 6);

    if (!enif_is_atom(env, type))
        {ret = EXCP_BADARG(env, "non-atom cipher type"); goto done;}
    if (!enif_inspect_iolist_as_binary(env, argv[1], &key))
        {ret = EXCP_BADARG(env, "non-binary key"); goto done;}
    if (!enif_inspect_binary(env, argv[2], &iv))
        {ret = EXCP_BADARG(env, "non-binary iv"); goto done;}
    if (!enif_inspect_iolist_as_binary(env, argv[3], &aad))
        {ret = EXCP_BADARG(env, "non-binary AAD"); goto done;}
    if (!enif_inspect_iolist_as_binary(env, argv[4], &in))
        {ret = EXCP_BADARG(env, "non-binary text"); goto done;}
    if (!enif_get_uint(env, argv[5], &tag_len))
        {ret = EXCP_BADARG(env, ""); goto done;}

    if (tag_len > INT_MAX
        || iv.size > INT_MAX
        || in.size > INT_MAX
        || aad.size > INT_MAX)
        {ret = EXCP_BADARG(env, "binary too long"); goto done;}

    if ((cipherp = get_cipher_type(type, key.size)) == NULL)
        {ret = EXCP_BADARG(env, "Unknown cipher"); goto done;}
    if (cipherp->flags & NON_EVP_CIPHER)
        {ret = EXCP_BADARG(env, "Bad cipher"); goto done;}
    if (! (cipherp->flags & AEAD_CIPHER) )
        {ret = EXCP_BADARG(env, "Not aead cipher"); goto done;}
    if ((cipher = cipherp->cipher.p) == NULL)
        {ret = EXCP_NOTSUP(env, "Cipher not supported in this libcrypto version"); goto done;}

    ctx_ctrl_set_ivlen = cipherp->extra.aead.ctx_ctrl_set_ivlen;
    ctx_ctrl_get_tag =  cipherp->extra.aead.ctx_ctrl_get_tag;
    ctx_ctrl_set_tag =  cipherp->extra.aead.ctx_ctrl_set_tag;
    
    if ((ctx = EVP_CIPHER_CTX_new()) == NULL)
        {ret = EXCP_ERROR(env, ""); goto done;}

    if (EVP_EncryptInit_ex(ctx, cipher, NULL, NULL, NULL) != 1)
        {ret = EXCP_ERROR(env, ""); goto done;}
    if (EVP_CIPHER_CTX_ctrl(ctx, ctx_ctrl_set_ivlen, (int)iv.size, NULL) != 1)
        {ret = EXCP_ERROR(env, ""); goto done;}

#if defined(HAVE_CCM)
    if (type == atom_aes_ccm) {
        if (EVP_CIPHER_CTX_ctrl(ctx, ctx_ctrl_set_tag, (int)tag_len, NULL) != 1)
            {ret = EXCP_ERROR(env, ""); goto done;}
        if (EVP_EncryptInit_ex(ctx, NULL, NULL, key.data, iv.data) != 1)
            {ret = EXCP_ERROR(env, ""); goto done;}
        if (EVP_EncryptUpdate(ctx, NULL, &len, NULL, (int)in.size) != 1)
            {ret = EXCP_ERROR(env, ""); goto done;}
    } else
#endif
        {
            if (EVP_EncryptInit_ex(ctx, NULL, NULL, key.data, iv.data) != 1)
                {ret = EXCP_ERROR(env, ""); goto done;}
        }

    if (EVP_EncryptUpdate(ctx, NULL, &len, aad.data, (int)aad.size) != 1)
        {ret = EXCP_ERROR(env, ""); goto done;}

    if ((outp = enif_make_new_binary(env, in.size, &out)) == NULL)
        {ret = EXCP_ERROR(env, ""); goto done;}

    if (EVP_EncryptUpdate(ctx, outp, &len, in.data, (int)in.size) != 1)
        {ret = EXCP_ERROR(env, ""); goto done;}
    if (EVP_EncryptFinal_ex(ctx, outp/*+len*/, &len) != 1)
        {ret = EXCP_ERROR(env, ""); goto done;}

    if ((tagp = enif_make_new_binary(env, tag_len, &out_tag)) == NULL)
        {ret = EXCP_ERROR(env, ""); goto done;}

    if (EVP_CIPHER_CTX_ctrl(ctx, ctx_ctrl_get_tag, (int)tag_len, tagp) != 1)
        {ret = EXCP_ERROR(env, ""); goto done;}

    CONSUME_REDS(env, in);
    ret = enif_make_tuple2(env, out, out_tag);

 done:
    if (ctx)
        EVP_CIPHER_CTX_free(ctx);
    return ret;

#else
    return EXCP_NOTSUP(env, "");
#endif
}

ERL_NIF_TERM aead_decrypt(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
{/* (Type,Key,Iv,AAD,In,Tag) */
#if defined(HAVE_AEAD)
    const struct cipher_type_t *cipherp;
    EVP_CIPHER_CTX *ctx = NULL;
    const EVP_CIPHER *cipher = NULL;
    ErlNifBinary key, iv, aad, in, tag;
    unsigned char *outp;
    ERL_NIF_TERM type, out, ret;
    int len, ctx_ctrl_set_ivlen, ctx_ctrl_set_tag;

    ASSERT(argc == 6);

    type = argv[0];
#if defined(HAVE_GCM_EVP_DECRYPT_BUG)
    if (type == atom_aes_gcm)
        return aes_gcm_decrypt_NO_EVP(env, argc, argv);
#endif

    if (!enif_is_atom(env, type))
        {ret = EXCP_BADARG(env, "non-atom cipher type"); goto done;}
    if (!enif_inspect_iolist_as_binary(env, argv[1], &key))
        {ret = EXCP_BADARG(env, "non-binary key"); goto done;}
    if (!enif_inspect_binary(env, argv[2], &iv))
        {ret = EXCP_BADARG(env, "non-binary iv"); goto done;}
    if (!enif_inspect_iolist_as_binary(env, argv[3], &aad))
        {ret = EXCP_BADARG(env, "non-binary AAD"); goto done;}
    if (!enif_inspect_iolist_as_binary(env, argv[4], &in))
        {ret = EXCP_BADARG(env, ""); goto done;}
    if (!enif_inspect_iolist_as_binary(env, argv[5], &tag))
        {ret = EXCP_BADARG(env, "non-binary text"); goto done;}

    if (tag.size > INT_MAX
        || key.size > INT_MAX
        || iv.size > INT_MAX
        || in.size > INT_MAX
        || aad.size > INT_MAX)
        {ret = EXCP_BADARG(env, "binary too long"); goto done;}

    if ((cipherp = get_cipher_type(type, key.size)) == NULL)
        {ret = EXCP_BADARG(env, "Unknown cipher"); goto done;}
    if (cipherp->flags & NON_EVP_CIPHER)
        {ret = EXCP_BADARG(env, "Bad cipher"); goto done;}
    if ( !(cipherp->flags & AEAD_CIPHER) )
        {ret = EXCP_BADARG(env, "Not aead cipher"); goto done;}
    if ((cipher = cipherp->cipher.p) == NULL)
        {ret = EXCP_NOTSUP(env, "Cipher not supported in this libcrypto version"); goto done;}

    ctx_ctrl_set_ivlen = cipherp->extra.aead.ctx_ctrl_set_ivlen;
    ctx_ctrl_set_tag =  cipherp->extra.aead.ctx_ctrl_set_tag;

    if ((outp = enif_make_new_binary(env, in.size, &out)) == NULL)
        {ret = EXCP_ERROR(env, ""); goto done;}

    if ((ctx = EVP_CIPHER_CTX_new()) == NULL)
        {ret = EXCP_ERROR(env, ""); goto done;}
    if (EVP_DecryptInit_ex(ctx, cipher, NULL, NULL, NULL) != 1)
        {ret = EXCP_ERROR(env, ""); goto done;}
    if (EVP_CIPHER_CTX_ctrl(ctx,  ctx_ctrl_set_ivlen, (int)iv.size, NULL) != 1)
        {ret = EXCP_ERROR(env, ""); goto done;}

#if defined(HAVE_CCM)
    if (type == atom_aes_ccm) {
        if (EVP_CIPHER_CTX_ctrl(ctx, ctx_ctrl_set_tag, (int)tag.size, tag.data) != 1)
            {ret = EXCP_ERROR(env, ""); goto done;}
        if (EVP_DecryptInit_ex(ctx, NULL, NULL, key.data, iv.data) != 1)
            {ret = EXCP_ERROR(env, ""); goto done;}
        if (EVP_DecryptUpdate(ctx, NULL, &len, NULL, (int)in.size) != 1)
            {ret = EXCP_ERROR(env, ""); goto done;}
    }
    else
#endif
        {
            if (EVP_DecryptInit_ex(ctx, NULL, NULL, key.data, iv.data) != 1)
                {ret = EXCP_ERROR(env, ""); goto done;}
        }

    if (EVP_DecryptUpdate(ctx, NULL, &len, aad.data, (int)aad.size) != 1)
        {ret = EXCP_ERROR(env, ""); goto done;}
    if (EVP_DecryptUpdate(ctx, outp, &len, in.data, (int)in.size) != 1)
        {ret = EXCP_ERROR(env, ""); goto done;}

#if defined(HAVE_GCM)
    if (type == atom_aes_gcm) {
        if (EVP_CIPHER_CTX_ctrl(ctx, ctx_ctrl_set_tag, (int)tag.size, tag.data) != 1)
            goto err;
        if (EVP_DecryptFinal_ex(ctx, outp+len, &len) != 1)
            goto err;
    }
#endif
    CONSUME_REDS(env, in);
    ret = out;
    goto done;

 err:
    /* Decrypt failed, that is, wrong tag */
    ret = atom_error;

 done:
    if (ctx)
        EVP_CIPHER_CTX_free(ctx);
    return ret;

#else
    return EXCP_NOTSUP(env, "");
#endif
}