/*
 * Copyright (C) 2010-2012 Free Software Foundation, Inc.
 *
 * Author: Nikos Mavrogiannopoulos
 *
 * This file is part of GNUTLS.
 *
 * The GNUTLS library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public License
 * as published by the Free Software Foundation; either version 2.1 of
 * the License, or (at your option) any later version.
 *
 * This library 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
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>
 *
 */

/* Here lie everything that has to do with large numbers, gmp.
 */

#include <gnutls_int.h>
#include <gnutls_errors.h>
#include <algorithms.h>
#include <gnutls_num.h>
#include <gnutls_mpi.h>
#include <nettle/bignum.h> /* includes gmp.h */
#include <gnettle.h>
#include <random.h>

static int
wrap_nettle_mpi_print(const bigint_t a, void *buffer, size_t * nbytes,
		      gnutls_bigint_format_t format)
{
	unsigned int size;
	mpz_t *p = (void *) a;

	if (format == GNUTLS_MPI_FORMAT_USG) {
		size = nettle_mpz_sizeinbase_256_u(*p);
	} else if (format == GNUTLS_MPI_FORMAT_STD) {
		size = nettle_mpz_sizeinbase_256_s(*p);
	} else if (format == GNUTLS_MPI_FORMAT_PGP) {
		size = nettle_mpz_sizeinbase_256_u(*p) + 2;
	} else {
		gnutls_assert();
		return GNUTLS_E_INVALID_REQUEST;
	}

	if (buffer == NULL || size > *nbytes) {
		*nbytes = size;
		return GNUTLS_E_SHORT_MEMORY_BUFFER;
	}

	if (format == GNUTLS_MPI_FORMAT_PGP) {
		uint8_t *buf = buffer;
		unsigned int nbits = _gnutls_mpi_get_nbits(a);
		buf[0] = (nbits >> 8) & 0xff;
		buf[1] = (nbits) & 0xff;
		nettle_mpz_get_str_256(size - 2, buf + 2, *p);
	} else {
		nettle_mpz_get_str_256(size, buffer, *p);
	}
	*nbytes = size;

	return 0;
}

static bigint_t wrap_nettle_mpi_new(int nbits)
{
	bigint_t p;

	p = gnutls_malloc(SIZEOF_MPZT);
	if (p == NULL) {
		gnutls_assert();
		return NULL;
	}
	if (nbits == 0)
		mpz_init(TOMPZ(p));
	else
		mpz_init2(TOMPZ(p), nbits);

	return p;
}

static bigint_t
wrap_nettle_mpi_scan(const void *buffer, size_t nbytes,
		     gnutls_bigint_format_t format)
{
	bigint_t r = wrap_nettle_mpi_new(nbytes * 8);

	if (r == NULL) {
		gnutls_assert();
		return r;
	}

	if (format == GNUTLS_MPI_FORMAT_USG) {
		nettle_mpz_set_str_256_u(TOMPZ(r), nbytes, buffer);
	} else if (format == GNUTLS_MPI_FORMAT_STD) {
		nettle_mpz_set_str_256_s(TOMPZ(r), nbytes, buffer);
	} else if (format == GNUTLS_MPI_FORMAT_PGP) {
		const uint8_t *buf = buffer;
		size_t size;

		if (nbytes < 3) {
			gnutls_assert();
			goto fail;
		}

		size = (buf[0] << 8) | buf[1];
		size = (size + 7) / 8;

		if (size > nbytes - 2) {
			gnutls_assert();
			goto fail;
		}
		nettle_mpz_set_str_256_u(TOMPZ(r), size, buf + 2);
	} else {
		gnutls_assert();
		goto fail;
	}

	return r;
      fail:
	_gnutls_mpi_release(&r);
	return NULL;

}

static int wrap_nettle_mpi_cmp(const bigint_t u, const bigint_t v)
{
	mpz_t *i1 = u, *i2 = v;

	return mpz_cmp(*i1, *i2);
}

static int wrap_nettle_mpi_cmp_ui(const bigint_t u, unsigned long v)
{
	mpz_t *i1 = u;

	return mpz_cmp_ui(*i1, v);
}

static bigint_t wrap_nettle_mpi_set(bigint_t w, const bigint_t u)
{
	mpz_t *i1, *i2 = u;

	if (w == NULL)
		w = _gnutls_mpi_alloc_like(u);
	i1 = w;

	mpz_set(*i1, *i2);

	return i1;
}

static bigint_t wrap_nettle_mpi_set_ui(bigint_t w, unsigned long u)
{
	mpz_t *i1;

	if (w == NULL)
		w = wrap_nettle_mpi_new(32);

	i1 = w;

	mpz_set_ui(*i1, u);

	return i1;
}

static unsigned int wrap_nettle_mpi_get_nbits(bigint_t a)
{
	return mpz_sizeinbase(TOMPZ(a), 2);
}

static void wrap_nettle_mpi_release(bigint_t a)
{
	mpz_clear(TOMPZ(a));
	gnutls_free(a);
}

static void wrap_nettle_mpi_clear(bigint_t a)
{
	memset(TOMPZ(a)[0]._mp_d, 0,
	       TOMPZ(a)[0]._mp_alloc * sizeof(mp_limb_t));
}

static bigint_t wrap_nettle_mpi_mod(const bigint_t a, const bigint_t b)
{
	bigint_t r = wrap_nettle_mpi_new(wrap_nettle_mpi_get_nbits(b));

	if (r == NULL)
		return NULL;

	mpz_mod(TOMPZ(r), TOMPZ(a), TOMPZ(b));

	return r;
}

static bigint_t
wrap_nettle_mpi_powm(bigint_t w, const bigint_t b, const bigint_t e,
		     const bigint_t m)
{
	if (w == NULL)
		w = wrap_nettle_mpi_new(wrap_nettle_mpi_get_nbits(m));

	if (w == NULL)
		return NULL;

	mpz_powm(TOMPZ(w), TOMPZ(b), TOMPZ(e), TOMPZ(m));

	return w;
}

static bigint_t
wrap_nettle_mpi_addm(bigint_t w, const bigint_t a, const bigint_t b,
		     const bigint_t m)
{
	if (w == NULL)
		w = wrap_nettle_mpi_new(wrap_nettle_mpi_get_nbits(a));

	if (w == NULL)
		return NULL;

	mpz_add(TOMPZ(w), TOMPZ(b), TOMPZ(a));
	mpz_fdiv_r(TOMPZ(w), TOMPZ(w), TOMPZ(m));

	return w;
}

static bigint_t
wrap_nettle_mpi_subm(bigint_t w, const bigint_t a, const bigint_t b,
		     const bigint_t m)
{
	if (w == NULL)
		w = wrap_nettle_mpi_new(wrap_nettle_mpi_get_nbits(a));

	if (w == NULL)
		return NULL;

	mpz_sub(TOMPZ(w), TOMPZ(a), TOMPZ(b));
	mpz_fdiv_r(TOMPZ(w), TOMPZ(w), TOMPZ(m));

	return w;
}

static bigint_t
wrap_nettle_mpi_mulm(bigint_t w, const bigint_t a, const bigint_t b,
		     const bigint_t m)
{
	if (w == NULL)
		w = wrap_nettle_mpi_new(wrap_nettle_mpi_get_nbits(m));

	if (w == NULL)
		return NULL;

	mpz_mul(TOMPZ(w), TOMPZ(a), TOMPZ(b));
	mpz_fdiv_r(TOMPZ(w), TOMPZ(w), TOMPZ(m));

	return w;
}

static bigint_t
wrap_nettle_mpi_add(bigint_t w, const bigint_t a, const bigint_t b)
{
	if (w == NULL)
		w = wrap_nettle_mpi_new(wrap_nettle_mpi_get_nbits(b));

	if (w == NULL)
		return NULL;

	mpz_add(TOMPZ(w), TOMPZ(a), TOMPZ(b));

	return w;
}

static bigint_t
wrap_nettle_mpi_sub(bigint_t w, const bigint_t a, const bigint_t b)
{
	if (w == NULL)
		w = wrap_nettle_mpi_new(wrap_nettle_mpi_get_nbits(a));

	if (w == NULL)
		return NULL;

	mpz_sub(TOMPZ(w), TOMPZ(a), TOMPZ(b));

	return w;
}

static bigint_t
wrap_nettle_mpi_mul(bigint_t w, const bigint_t a, const bigint_t b)
{
	if (w == NULL)
		w = wrap_nettle_mpi_new(wrap_nettle_mpi_get_nbits(a));

	if (w == NULL)
		return NULL;

	mpz_mul(TOMPZ(w), TOMPZ(a), TOMPZ(b));

	return w;
}

/* q = a / b */
static bigint_t
wrap_nettle_mpi_div(bigint_t q, const bigint_t a, const bigint_t b)
{
	if (q == NULL)
		q = wrap_nettle_mpi_new(wrap_nettle_mpi_get_nbits(a));

	if (q == NULL)
		return NULL;

	mpz_cdiv_q(TOMPZ(q), TOMPZ(a), TOMPZ(b));

	return q;
}

static bigint_t
wrap_nettle_mpi_add_ui(bigint_t w, const bigint_t a, unsigned long b)
{
	if (w == NULL)
		w = wrap_nettle_mpi_new(wrap_nettle_mpi_get_nbits(a));

	if (w == NULL)
		return NULL;

	mpz_add_ui(TOMPZ(w), TOMPZ(a), b);

	return w;
}

static bigint_t
wrap_nettle_mpi_sub_ui(bigint_t w, const bigint_t a, unsigned long b)
{
	if (w == NULL)
		w = wrap_nettle_mpi_new(wrap_nettle_mpi_get_nbits(a));

	if (w == NULL)
		return NULL;

	mpz_sub_ui(TOMPZ(w), TOMPZ(a), b);

	return w;

}

static bigint_t
wrap_nettle_mpi_mul_ui(bigint_t w, const bigint_t a, unsigned long b)
{
	if (w == NULL)
		w = wrap_nettle_mpi_new(wrap_nettle_mpi_get_nbits(a));

	if (w == NULL)
		return NULL;

	mpz_mul_ui(TOMPZ(w), TOMPZ(a), b);

	return w;

}

static int wrap_nettle_prime_check(bigint_t pp)
{
	int ret;
	ret = mpz_probab_prime_p(TOMPZ(pp), PRIME_CHECK_PARAM);

	if (ret > 0) {
		return 0;
	}

	return GNUTLS_E_INTERNAL_ERROR;	/* ignored */
}


/* generate a prime of the form p=2qw+1
 * The algorithm is simple but probably it has to be modified to gcrypt's
 * since it is slow. Nature did not want 2qw+1 to be prime.
 * The generator will be the generator of a subgroup of order q-1.
 *
 * Algorithm based on the algorithm in "A Computational Introduction to Number 
 * Theory and Algebra" by V. Shoup, sec 11.1 Finding a generator for Z^{*}_p
 *
 */
inline static int
gen_group(mpz_t * prime, mpz_t * generator, unsigned int nbits,
	  unsigned int *q_bits)
{
	mpz_t q, w, r;
	unsigned int p_bytes = nbits / 8;
	uint8_t *buffer = NULL;
	unsigned int q_bytes, w_bytes, r_bytes, w_bits;
	int ret;

	/* security level enforcement. 
	 * Values for q are selected according to ECRYPT II recommendations.
	 */
	q_bytes = _gnutls_pk_bits_to_subgroup_bits(nbits);
	q_bytes /= 8;

	if (q_bytes == 0 || q_bytes >= p_bytes) {
		gnutls_assert();
		return GNUTLS_E_INVALID_REQUEST;
	}

	if (nbits % 8 != 0)
		p_bytes++;

	w_bits = nbits - q_bytes * 8;
	w_bytes = w_bits / 8;
	if (w_bits % 8 != 0)
		w_bytes++;

	_gnutls_debug_log
	    ("Generating group of prime of %u bits and format of 2wq+1. q_size=%u bits\n",
	     nbits, q_bytes * 8);
	buffer = gnutls_malloc(p_bytes);	/* p_bytes > q_bytes */
	if (buffer == NULL) {
		gnutls_assert();
		return GNUTLS_E_MEMORY_ERROR;
	}

	mpz_init(q);
	mpz_init(w);
	mpz_init(r);

	/* search for a prime. We are not that unlucky so search
	 * forever.
	 */
	for (;;) {
		ret = _gnutls_rnd(GNUTLS_RND_RANDOM, buffer, w_bytes);
		if (ret < 0) {
			gnutls_assert();
			goto fail;
		}

		nettle_mpz_set_str_256_u(w, w_bytes, buffer);
		/* always odd */
		mpz_setbit(w, 0);

		ret = mpz_probab_prime_p(w, PRIME_CHECK_PARAM);
		if (ret > 0) {
			break;
		}
	}

	/* now generate q of size p_bytes - w_bytes */

	_gnutls_debug_log
	    ("Found prime w of %u bits. Will look for q of %u bits...\n",
	     wrap_nettle_mpi_get_nbits(&w), q_bytes * 8);

	for (;;) {
		ret = _gnutls_rnd(GNUTLS_RND_RANDOM, buffer, q_bytes);
		if (ret < 0) {
			gnutls_assert();
			return ret;
		}

		nettle_mpz_set_str_256_u(q, q_bytes, buffer);
		/* always odd */
		mpz_setbit(q, 0);

		ret = mpz_probab_prime_p(q, PRIME_CHECK_PARAM);
		if (ret == 0) {
			continue;
		}

		/* check if 2wq+1 is prime */
		mpz_mul_ui(*prime, w, 2);
		mpz_mul(*prime, *prime, q);
		mpz_add_ui(*prime, *prime, 1);

		ret = mpz_probab_prime_p(*prime, PRIME_CHECK_PARAM);
		if (ret > 0) {
			break;
		}
	}

	*q_bits = wrap_nettle_mpi_get_nbits(&q);
	_gnutls_debug_log
	    ("Found prime q of %u bits. Looking for generator...\n",
	     *q_bits);

	/* finally a prime! Let's calculate generator
	 */

	/* c = r^((p-1)/q), r == random
	 * c = r^(2w)
	 * if c!=1 c is the generator for the subgroup of order q-1
	 * 
	 */
	r_bytes = p_bytes;

	mpz_mul_ui(w, w, 2);	/* w = w*2 */
	mpz_fdiv_r(w, w, *prime);

	for (;;) {
		ret = _gnutls_rnd(GNUTLS_RND_NONCE, buffer, r_bytes);
		if (ret < 0) {
			gnutls_assert();
			return ret;
		}

		nettle_mpz_set_str_256_u(r, r_bytes, buffer);
		mpz_fdiv_r(r, r, *prime);

		/* check if r^w mod n != 1 mod n */
		mpz_powm(*generator, r, w, *prime);

		if (mpz_cmp_ui(*generator, 1) == 0)
			continue;
		else
			break;
	}

	_gnutls_debug_log("Found generator g of %u bits\n",
			  wrap_nettle_mpi_get_nbits(generator));
	_gnutls_debug_log("Prime n is %u bits\n",
			  wrap_nettle_mpi_get_nbits(prime));

	ret = 0;
	goto exit;

      fail:
	mpz_clear(*prime);
	mpz_clear(*generator);

      exit:
	mpz_clear(q);
	mpz_clear(w);
	mpz_clear(r);
	gnutls_free(buffer);

	return ret;
}

static int
wrap_nettle_generate_group(gnutls_group_st * group, unsigned int bits)
{
	int ret;
	bigint_t p = wrap_nettle_mpi_new(bits);
	bigint_t g;
	unsigned int q_bits;

	if (p == NULL) {
		gnutls_assert();
		return GNUTLS_E_MEMORY_ERROR;
	}

	g = wrap_nettle_mpi_new(bits);
	if (g == NULL) {
		gnutls_assert();
		_gnutls_mpi_release(&p);
		return GNUTLS_E_MEMORY_ERROR;
	}

	ret = gen_group(p, g, bits, &q_bits);
	if (ret < 0) {
		_gnutls_mpi_release(&g);
		_gnutls_mpi_release(&p);
		gnutls_assert();
		return ret;
	}

	group->p = p;
	group->g = g;
	group->q_bits = q_bits;

	return 0;
}


int crypto_bigint_prio = INT_MAX;

gnutls_crypto_bigint_st _gnutls_mpi_ops = {
	.bigint_new = wrap_nettle_mpi_new,
	.bigint_cmp = wrap_nettle_mpi_cmp,
	.bigint_cmp_ui = wrap_nettle_mpi_cmp_ui,
	.bigint_mod = wrap_nettle_mpi_mod,
	.bigint_set = wrap_nettle_mpi_set,
	.bigint_set_ui = wrap_nettle_mpi_set_ui,
	.bigint_get_nbits = wrap_nettle_mpi_get_nbits,
	.bigint_powm = wrap_nettle_mpi_powm,
	.bigint_addm = wrap_nettle_mpi_addm,
	.bigint_subm = wrap_nettle_mpi_subm,
	.bigint_add = wrap_nettle_mpi_add,
	.bigint_sub = wrap_nettle_mpi_sub,
	.bigint_add_ui = wrap_nettle_mpi_add_ui,
	.bigint_sub_ui = wrap_nettle_mpi_sub_ui,
	.bigint_mul = wrap_nettle_mpi_mul,
	.bigint_mulm = wrap_nettle_mpi_mulm,
	.bigint_mul_ui = wrap_nettle_mpi_mul_ui,
	.bigint_div = wrap_nettle_mpi_div,
	.bigint_prime_check = wrap_nettle_prime_check,
	.bigint_release = wrap_nettle_mpi_release,
	.bigint_clear = wrap_nettle_mpi_clear,
	.bigint_print = wrap_nettle_mpi_print,
	.bigint_scan = wrap_nettle_mpi_scan,
	.bigint_generate_group = wrap_nettle_generate_group
};
