#include "testutils.h"

#include <assert.h>

#include "rsa.h"

#define KEY_COUNT 20
#define COUNT 100

static void
random_fn (void *ctx, size_t n, uint8_t *dst)
{
  gmp_randstate_t *rands = (gmp_randstate_t *)ctx;
  mpz_t r;

  mpz_init (r);
  mpz_urandomb (r, *rands, n*8);
  nettle_mpz_get_str_256 (n, dst, r);
  mpz_clear (r);
}

static void
test_one (gmp_randstate_t *rands, struct rsa_public_key *pub,
          struct rsa_private_key *key, mpz_t plaintext)
{
  mpz_t ciphertext;
  mpz_t decrypted;

  mpz_init (ciphertext);
  mpz_init (decrypted);

  mpz_powm (ciphertext, plaintext, pub->e, pub->n);
  rsa_compute_root_tr (pub, key, rands, random_fn, decrypted, ciphertext);
  if (mpz_cmp (plaintext, decrypted)) {
    fprintf (stderr, "rsa_compute_root_tr failed\n");

    fprintf(stderr, "Public key: size=%u\n n:", (unsigned) pub->size);
    mpz_out_str (stderr, 10, pub->n);
    fprintf(stderr, "\n e:");
    mpz_out_str (stderr, 10, pub->e);
    fprintf(stderr, "\nPrivate key: size=%u\n p:", (unsigned) key->size);
    mpz_out_str (stderr, 10, key->p);
    fprintf(stderr, "\n q:");
    mpz_out_str (stderr, 10, key->q);
    fprintf(stderr, "\n a:");
    mpz_out_str (stderr, 10, key->a);
    fprintf(stderr, "\n b:");
    mpz_out_str (stderr, 10, key->b);
    fprintf(stderr, "\n c:");
    mpz_out_str (stderr, 10, key->c);
    fprintf(stderr, "\n d:");
    mpz_out_str (stderr, 10, key->d);
    fprintf(stderr, "\n");

    fprintf (stderr, "plaintext(%u) = ", (unsigned) mpz_sizeinbase (plaintext, 2));
    mpz_out_str (stderr, 10, plaintext);
    fprintf (stderr, "\n");
    fprintf (stderr, "ciphertext(%u) = ", (unsigned) mpz_sizeinbase (ciphertext, 2));
    mpz_out_str (stderr, 10, ciphertext);
    fprintf (stderr, "\n");
    fprintf (stderr, "decrypted(%u) = ", (unsigned) mpz_sizeinbase (decrypted, 2));
    mpz_out_str (stderr, 10, decrypted);
    fprintf (stderr, "\n");
    abort();
  }

  mpz_clear (ciphertext);
  mpz_clear (decrypted);
}

#if !NETTLE_USE_MINI_GMP
/* We want to generate keypairs that are not "standard" but have more size
 * variance between q and p.
 */
static void
generate_keypair (gmp_randstate_t rands,
                  struct rsa_public_key *pub, struct rsa_private_key *key)
{
  unsigned long int psize;
  unsigned long int qsize;
  mpz_t p1;
  mpz_t q1;
  mpz_t phi;
  mpz_t tmp;
  int res;

  mpz_init (p1);
  mpz_init (q1);
  mpz_init (phi);
  mpz_init (tmp);

  psize = 100 + gmp_urandomm_ui (rands, 400);
  qsize = 100 + gmp_urandomm_ui (rands, 400);

  mpz_set_ui (pub->e, 65537);

  for (;;)
    {
      for (;;)
        {
          mpz_rrandomb (key->p, rands, psize);
          mpz_nextprime (key->p, key->p);
          mpz_sub_ui (p1, key->p, 1);
          mpz_gcd (tmp, pub->e, p1);
          if (mpz_cmp_ui (tmp, 1) == 0)
            break;
        }

      for (;;)
        {
          mpz_rrandomb (key->q, rands, qsize);
          mpz_nextprime (key->q, key->q);
          mpz_sub_ui (q1, key->q, 1);
          mpz_gcd (tmp, pub->e, q1);
          if (mpz_cmp_ui (tmp, 1) == 0)
            break;
        }

      if (mpz_invert (key->c, key->q, key->p))
        break;
    }

  mpz_mul(phi, p1, q1);
  res = mpz_invert(key->d, pub->e, phi);
  assert (res);

  mpz_fdiv_r (key->a, key->d, p1);
  mpz_fdiv_r (key->b, key->d, q1);

  mpz_mul (pub->n, key->p, key->q);

  pub->size = key->size = mpz_size(pub->n) * sizeof(mp_limb_t);

  mpz_clear (tmp);
  mpz_clear (phi);
  mpz_clear (q1);
  mpz_clear (p1);
}
#endif

void
test_main (void)
{
  gmp_randstate_t rands;
  struct rsa_public_key pub;
  struct rsa_private_key key;
  mpz_t plaintext;
  unsigned i, j;

  rsa_private_key_init(&key);
  rsa_public_key_init(&pub);
  mpz_init (plaintext);

  gmp_randinit_default (rands);
  test_randomize(rands);

  for (j = 0; j < KEY_COUNT; j++)
    {
#if !NETTLE_USE_MINI_GMP
      generate_keypair(rands, &pub, &key);
#else
      rsa_generate_keypair(&pub, &key, &rands, random_fn, NULL, NULL, 512, 16);
#endif /* !NETTLE_USE_MINI_GMP */

      for (i = 0; i < COUNT; i++)
	{
	  mpz_urandomb(plaintext, rands, mpz_sizeinbase(pub.n, 2) - 1);
	  test_one(&rands, &pub, &key, plaintext);
	}
      for (i = 0; i < COUNT; i++)
	{
	  mpz_rrandomb(plaintext, rands, mpz_sizeinbase(pub.n, 2) - 1);
	  test_one(&rands, &pub, &key, plaintext);
	}
    }
  mpz_clear (plaintext);
  rsa_public_key_clear (&pub);
  rsa_private_key_clear (&key);

  gmp_randclear (rands);
}