/* -*- c++ -*- */
/* 
 * Copyright (c) 2021 Satlab A/S <satlab@satlab.com>
 * 
 * This is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 3, or (at your option)
 * any later version.
 * 
 * This software is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this software; see the file COPYING.  If not, write to
 * the Free Software Foundation, Inc., 51 Franklin Street,
 * Boston, MA 02110-1301, USA.
 */

#ifdef HAVE_CONFIG_H
#include "config.h"
#endif

#include <endian.h>
#include <gnuradio/io_signature.h>
#include "srs3_decrypt_impl.h"

#include <openssl/rand.h>
#include <openssl/bio.h>
#include <openssl/evp.h>

extern "C" {
  #include "spaceframe.h"
}

namespace gr {
  namespace satlab {

    srs3_decrypt::sptr
    srs3_decrypt::make(const char *key, bool decrypt, bool auth)
    {
      return gnuradio::get_initial_sptr
        (new srs3_decrypt_impl(key, decrypt, auth));
    }

    int
    srs3_decrypt_impl::aes256gcm_decrypt(const uint8_t *key, const uint8_t *iv, const uint8_t *aad, size_t aad_length, const uint8_t *in, uint8_t *out, size_t length, const uint8_t *tag)
    {
        int ret, outlen;
        EVP_CIPHER_CTX *ctx;

        ctx = EVP_CIPHER_CTX_new();

        /* Set cipher type and mode */
        EVP_DecryptInit_ex(ctx, EVP_aes_256_gcm(), NULL, NULL, NULL);

        /* Set IV length */
        EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_AEAD_SET_IVLEN, SPACEFRAME_CRYPTO_IV_LENGTH, NULL);

        /* Specify key and IV */
        EVP_DecryptInit_ex(ctx, NULL, NULL, key, iv);

        /* Add AAD */
        if (aad && aad_length)
            EVP_DecryptUpdate(ctx, NULL, &outlen, aad, aad_length);

        /* Decrypt ciphertext */
        EVP_DecryptUpdate(ctx, out, &outlen, in, length);

        /* Set expected tag value */
        if (tag)
            EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_AEAD_SET_TAG, SPACEFRAME_CRYPTO_TAG_LENGTH, (void *)tag);

        /* Finalise and verify tag */
        ret = EVP_DecryptFinal_ex(ctx, out, &outlen);

        /* Free context */
        EVP_CIPHER_CTX_free(ctx);

        return (!tag || ret > 0) ? 0 : -EINVAL;
    }

    void
    srs3_decrypt_impl::decrypt(pmt::pmt_t pdu)
    {
        pmt::pmt_t vector = pmt::cdr(pdu);
        size_t length = pmt::blob_length(vector);
        const uint8_t *data = (const uint8_t *)pmt::blob_data(vector);

        /* Fast path if unencrypted */
        if (!d_decrypt && !d_auth) {
            message_port_pub(pmt::mp("output"), pdu);
            return;
        }

        /* Encrypted data must contain at least the IV */
        if (length < SPACEFRAME_CRYPTO_IV_LENGTH)
            return;

        /* Calculate output buffer length */
        size_t out_length = length - SPACEFRAME_CRYPTO_IV_LENGTH;

        /* Check that we have a tag if authenticating */
        if (d_auth) {
            if (out_length < SPACEFRAME_CRYPTO_TAG_LENGTH)
                return;
            out_length -= SPACEFRAME_CRYPTO_TAG_LENGTH;
        }

        /* Allocate output buffer */
        uint8_t *out_data = (uint8_t *)malloc(out_length);
        if (!out_data)
            return;

        const uint8_t *tag = d_auth ? &data[length - SPACEFRAME_CRYPTO_TAG_LENGTH] : NULL;
        uint8_t *out = d_decrypt ? out_data : NULL;

        /* Decrypt */
        int ret = aes256gcm_decrypt(d_key, data, NULL, 0, &data[SPACEFRAME_CRYPTO_IV_LENGTH], out, out_length, tag);
        if (ret < 0) {
            free(out_data);
            return;
        }

        /* Copy input to output in MAC-only mode */
        if (!d_decrypt)
            memcpy(out_data, &data[SPACEFRAME_CRYPTO_IV_LENGTH], out_length);

        /* Create output PDU from buffer */
        pmt::pmt_t vector_out = pmt::make_blob(out_data, out_length);
        pmt::pmt_t pdu_enc = pmt::cons(pmt::car(pdu), vector_out);
        
        free(out_data);

        message_port_pub(pmt::mp("output"), pdu_enc);
    }

    void
    srs3_decrypt_impl::set_key(const char *key)
    {
        memcpy(d_key, key, sizeof(d_key));
    }

    void
    srs3_decrypt_impl::set_decrypt(bool decrypt)
    {
        d_decrypt = decrypt;
    }

    void
    srs3_decrypt_impl::set_auth(bool auth)
    {
        d_auth = auth;
    }

    /*
     * The private constructor
     */
    srs3_decrypt_impl::srs3_decrypt_impl(const char *key, bool decrypt, bool auth)
      : gr::block("srs3_decrypt",
              gr::io_signature::make(0, 0, 0),
              gr::io_signature::make(0, 0, 0)),
        d_decrypt(decrypt),
        d_auth(auth)
    {
        memcpy(d_key, key, sizeof(d_key));
        message_port_register_out(pmt::mp("output"));

        message_port_register_in(pmt::mp("decrypt"));
        set_msg_handler(pmt::mp("decrypt"), [this](pmt::pmt_t msg) {
            this->decrypt(msg);
        });
    }

    /*
     * Our virtual destructor.
     */
    srs3_decrypt_impl::~srs3_decrypt_impl()
    {
    }
  } /* namespace satlab */
} /* namespace gr */

