/* -*- 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_encrypt_impl.h"

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

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

namespace gr {
  namespace satlab {

    srs3_encrypt::sptr
    srs3_encrypt::make(const char *key, bool encrypt, bool auth)
    {
      return gnuradio::get_initial_sptr
        (new srs3_encrypt_impl(key, encrypt, auth));
    }

    int
    srs3_encrypt_impl::aes256gcm_encrypt(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, uint8_t *tag)
    {
        int outlen;
        EVP_CIPHER_CTX *ctx;

        ctx = EVP_CIPHER_CTX_new();

        /* Set cipher type and mode */
        EVP_EncryptInit_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_EncryptInit_ex(ctx, NULL, NULL, key, iv);

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

        /* Encrypt plaintext */
        EVP_EncryptUpdate(ctx, out, &outlen, in, length);

        /* Finalise */
        EVP_EncryptFinal_ex(ctx, out, &outlen);

        /* Get tag */
        if (tag)
            EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_AEAD_GET_TAG, SPACEFRAME_CRYPTO_TAG_LENGTH, tag);

        /* Free context */
        EVP_CIPHER_CTX_free(ctx);

        return 0;
    }

    void
    srs3_encrypt_impl::encrypt(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_encrypt && !d_auth) {
            message_port_pub(pmt::mp("output"), pdu);
            return;
        }

        /* Calculate output buffer length */
        size_t out_length = SPACEFRAME_CRYPTO_IV_LENGTH + length;
        if (d_auth)
            out_length += SPACEFRAME_CRYPTO_TAG_LENGTH;

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

        /* Create random IV */
        uint8_t iv[SPACEFRAME_CRYPTO_IV_LENGTH];
        int rc = RAND_bytes(iv, sizeof(iv));
        if (rc != 1) {
            free(out_data);
            return;
        }

        /* Add IV to beginning of frame */
        memcpy(out_data, iv, sizeof(iv));

        /* Encrypt */
        uint8_t *tag = d_auth ? &out_data[SPACEFRAME_CRYPTO_IV_LENGTH + length] : NULL;
        uint8_t *out = d_encrypt ? &out_data[SPACEFRAME_CRYPTO_IV_LENGTH] : NULL;
        int ret = aes256gcm_encrypt(d_key, iv, NULL, 0, data, out, length, tag);
        if (ret < 0) {
            free(out_data);
            return;
        }

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

        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_encrypt_impl::set_key(const char *key)
    {
        memcpy(d_key, key, sizeof(d_key));
    }

    void
    srs3_encrypt_impl::set_encrypt(bool encrypt)
    {
        d_encrypt = encrypt;
    }

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

    /*
     * The private constructor
     */
    srs3_encrypt_impl::srs3_encrypt_impl(const char *key, bool encrypt, bool auth)
      : gr::block("srs3_encrypt",
              gr::io_signature::make(0, 0, 0),
              gr::io_signature::make(0, 0, 0)),
        d_encrypt(encrypt),
        d_auth(auth)
    {
        memcpy(d_key, key, sizeof(d_key));
        message_port_register_out(pmt::mp("output"));

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

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

