#include <openssl/evp.h>

#include <stdlib.h>
#include <string.h>

#include "deucessh-enc.h"
#include "deucessh.h"
#ifdef DSSH_TESTING
 #include "ssh-internal.h"
#endif

#define AES256_CTR_BLOCK_SIZE 16
#define AES256_CTR_KEY_SIZE   32
#define AES256_CTR_NAME       "aes256-ctr"
#define AES256_CTR_NAME_LEN   10

struct dssh_enc_ctx {
	EVP_CIPHER_CTX *ctx;
};

static int
init_ctx(const uint8_t *key, const uint8_t *iv, bool encrypt_dir, dssh_enc_ctx **ctx)
{
	struct dssh_enc_ctx *cbd = malloc(sizeof(*cbd));

	if (cbd == NULL)
		return DSSH_ERROR_ALLOC;

	cbd->ctx = EVP_CIPHER_CTX_new();
	if (cbd->ctx == NULL) {
		free(cbd);
		return DSSH_ERROR_ALLOC;
	}

	(void)encrypt_dir;
	EVP_CIPHER *cipher = EVP_CIPHER_fetch(NULL, "AES-256-CTR", NULL);

	if (cipher == NULL) {
		EVP_CIPHER_CTX_free(cbd->ctx);
		free(cbd);
		return DSSH_ERROR_INIT;
	}
	int ok = (EVP_EncryptInit_ex(cbd->ctx, cipher, NULL, key, iv) == 1)
	         && (EVP_CIPHER_CTX_set_padding(cbd->ctx, 0) == 1);

	EVP_CIPHER_free(cipher);
	if (!ok) {
		EVP_CIPHER_CTX_free(cbd->ctx);
		free(cbd);
		return DSSH_ERROR_INIT;
	}

	*ctx = cbd;
	return 0;
}

static int
do_crypt(uint8_t *buf, size_t bufsz, dssh_enc_ctx *ctx)
{
	struct dssh_enc_ctx *cbd = ctx;

	if ((cbd == NULL) || (cbd->ctx == NULL))
		return DSSH_ERROR_INIT;
	if (bufsz > INT_MAX)
		return DSSH_ERROR_INVALID;

	int outlen = 0;

	if (EVP_EncryptUpdate(cbd->ctx, buf, &outlen, buf, (int)bufsz) != 1)
		return DSSH_ERROR_INIT;

	return 0;
}

static void
cleanup(dssh_enc_ctx *ctx)
{
	struct dssh_enc_ctx *cbd = ctx;

	if (cbd != NULL) {
		EVP_CIPHER_CTX_free(cbd->ctx);
		free(cbd);
	}
}

DSSH_PUBLIC int
dssh_register_aes256_ctr(void)
{
	static const char  name[] = AES256_CTR_NAME;
	struct dssh_enc_s *enc    = malloc(sizeof(*enc) + sizeof(name));

	if (enc == NULL)
		return DSSH_ERROR_ALLOC;
	enc->next          = NULL;
	enc->init          = init_ctx;
	enc->encrypt       = do_crypt;
	enc->decrypt       = do_crypt;
	enc->cleanup       = cleanup;
	enc->bytes_per_key = UINT64_C(1) << 36; /* RFC 4344 s3.2: 2^32 blocks * 16 = 64 GiB */
	enc->flags         = 0;
	enc->blocksize     = AES256_CTR_BLOCK_SIZE;
	enc->key_size      = AES256_CTR_KEY_SIZE;
	memcpy(enc->name, name, sizeof(name));
	return dssh_transport_register_enc(enc);
}
