#include <config.h>

#include <stdint.h>

#include <stdio.h>

#include <string.h>

#include <utils.h>

#include <stdlib.h>

#include <gnutls/gnutls.h>

#include <gnutls/crypto.h>

#include <gnutls/self-test.h>

#ifndef HAVE_LIBNETTLE

int main(int argc, char **argv)

{

	exit(77);

}

#else

# include <nettle/aes.h>

# include <nettle/cbc.h>

# include <nettle/gcm.h>

# include <assert.h>

/* this tests whether the API to override ciphers works sanely.

 */

static int used = 0;

static int aead_used = 0;

static void tls_log_func(int level, const char *str)

{

	fprintf(stderr, "<%d>| %s", level, str);

}

struct myaes_ctx {

	struct aes128_ctx aes;

	unsigned char iv[16];

	int enc;

};

static int

myaes_init(gnutls_cipher_algorithm_t algorithm, void **_ctx, int enc)

{

	/* we use key size to distinguish */

	if (algorithm != GNUTLS_CIPHER_AES_128_CBC)

		return GNUTLS_E_INVALID_REQUEST;

	*_ctx = calloc(1, sizeof(struct myaes_ctx));

	if (*_ctx == NULL) {

		return GNUTLS_E_MEMORY_ERROR;

	}

	((struct myaes_ctx *) (*_ctx))->enc = enc;

	return 0;

}

static int

myaes_setkey(void *_ctx, const void *userkey, size_t keysize)

{

	struct myaes_ctx *ctx = _ctx;

	assert(keysize == 16);

	if (ctx->enc)

		aes128_set_encrypt_key(&ctx->aes, userkey);

	else

		aes128_set_decrypt_key(&ctx->aes, userkey);

	return 0;

}

static int myaes_setiv(void *_ctx, const void *iv, size_t iv_size)

{

	struct myaes_ctx *ctx = _ctx;

	memcpy(ctx->iv, iv, 16);

	return 0;

}

static int

myaes_encrypt(void *_ctx, const void *src, size_t src_size,

	    void *dst, size_t dst_size)

{

	struct myaes_ctx *ctx = _ctx;

	used++;

	cbc_encrypt(&ctx->aes, (nettle_cipher_func*)aes128_encrypt, 16, ctx->iv, src_size, dst, src);

	return 0;

}

static int

myaes_decrypt(void *_ctx, const void *src, size_t src_size,

	    void *dst, size_t dst_size)

{

	struct myaes_ctx *ctx = _ctx;

	used++;

	cbc_decrypt(&ctx->aes, (nettle_cipher_func*)aes128_decrypt, 16, ctx->iv, src_size, dst, src);

	return 0;

}

static void myaes_deinit(void *_ctx)

{

	free(_ctx);

}

/* AES-GCM */

struct myaes_gcm_ctx {

	struct gcm_aes128_ctx aes;

};

static int

myaes_gcm_init(gnutls_cipher_algorithm_t algorithm, void **_ctx, int enc)

{

	/* we use key size to distinguish */

	if (algorithm != GNUTLS_CIPHER_AES_128_GCM

	    && algorithm != GNUTLS_CIPHER_AES_256_GCM)

		return GNUTLS_E_INVALID_REQUEST;

	*_ctx = calloc(1, sizeof(struct myaes_gcm_ctx));

	if (*_ctx == NULL) {

		return GNUTLS_E_MEMORY_ERROR;

	}

	return 0;

}

static int

myaes_gcm_setkey(void *_ctx, const void *userkey, size_t keysize)

{

	struct myaes_gcm_ctx *ctx = _ctx;

	gcm_aes128_set_key(&ctx->aes, userkey);

	return 0;

}

static void myaes_gcm_deinit(void *_ctx)

{

	free(_ctx);

}

static int

myaes_gcm_encrypt(void *_ctx,

		  const void *nonce, size_t nonce_size,

		  const void *auth, size_t auth_size,

		  size_t tag_size,

		  const void *plain, size_t plain_size,

		  void *encr, size_t encr_size)

{

	/* proper AEAD cipher */

	struct myaes_gcm_ctx *ctx = _ctx;

	if (encr_size < plain_size + tag_size)

		return GNUTLS_E_SHORT_MEMORY_BUFFER;

	aead_used++;

	gcm_aes128_set_iv(&ctx->aes, nonce_size, nonce);

	gcm_aes128_update(&ctx->aes, auth_size, auth);

	gcm_aes128_encrypt(&ctx->aes, plain_size, encr, plain);

	gcm_aes128_digest(&ctx->aes, tag_size, ((uint8_t*)encr) + plain_size);

	return 0;

}

static int

myaes_gcm_decrypt(void *_ctx,

		  const void *nonce, size_t nonce_size,

		  const void *auth, size_t auth_size,

		  size_t tag_size,

		  const void *encr, size_t encr_size,

		  void *plain, size_t plain_size)

{

	uint8_t tag[16];

	struct myaes_gcm_ctx *ctx = _ctx;

	if (encr_size < tag_size)

		return GNUTLS_E_DECRYPTION_FAILED;

	aead_used++;

	gcm_aes128_set_iv(&ctx->aes, nonce_size, nonce);

	gcm_aes128_update(&ctx->aes, auth_size, auth);

	encr_size -= tag_size;

	gcm_aes128_decrypt(&ctx->aes, encr_size, plain, encr);

	gcm_aes128_digest(&ctx->aes, tag_size, tag);

	if (gnutls_memcmp(((uint8_t*)encr)+encr_size, tag, tag_size) != 0)

		return GNUTLS_E_DECRYPTION_FAILED;

	return 0;

}

int main(int argc, char **argv)

{

	int ret;

	gnutls_global_set_log_function(tls_log_func);

	if (argc > 1)

		gnutls_global_set_log_level(4711);

	ret = gnutls_crypto_register_cipher(GNUTLS_CIPHER_AES_128_CBC, 1,

		myaes_init,

		myaes_setkey,

		myaes_setiv,

		myaes_encrypt,

		myaes_decrypt,

		myaes_deinit);

	if (ret < 0) {

		fprintf(stderr, "%d: cannot register cipher\n", __LINE__);

		exit(1);

	}

	ret = gnutls_crypto_register_aead_cipher(GNUTLS_CIPHER_AES_128_GCM, 1,

		myaes_gcm_init,

		myaes_gcm_setkey,

		myaes_gcm_encrypt,

		myaes_gcm_decrypt,

		myaes_gcm_deinit);

	if (ret < 0) {

		fprintf(stderr, "%d: cannot register cipher\n", __LINE__);

		exit(1);

	}

	global_init();

	if (gnutls_cipher_self_test(GNUTLS_SELF_TEST_FLAG_ALL|GNUTLS_SELF_TEST_FLAG_NO_COMPAT, 0) < 0)

		return 1;

	if (used == 0) {

		fprintf(stderr, "The CBC cipher was not used\n");

		exit(1);

	}

	if (aead_used == 0) {

		fprintf(stderr, "The AEAD cipher was not used\n");

		exit(1);

	}

	gnutls_global_deinit();

	return 0;

}

#endif