/* file: rsa.c
 * author: Marcus Fritzsch <m ATfritschydotde>
 * date: 20050307
 * desc: simple, non serious implementation of the RSA algorithm
 *
 * compile:
 * gcc-4.0 -O3 -pedantic -Wall -std=c99 -lgmp -o rsa rsa.c
 */

#include <gmp.h>
#include <stdio.h>
#include <string.h>
#include <time.h>
#include <unistd.h>

/* my rsa context */
typedef struct RsaCtx_S {
	mpz_t p;
	mpz_t q;
	mpz_t n;	/* p*q, the modulus */
	mpz_t m;	/* (p-1) * (q-1) */
	mpz_t d;	/* private exponent */
	mpz_t e;	/* public exponent */

	gmp_randstate_t r;	/* gmp random state */
} RsaCtx;

void
initRsaCtx (RsaCtx * ctx)
{
	mpz_init (ctx->p);
	mpz_init (ctx->q);
	mpz_init (ctx->n);
	mpz_init (ctx->m);
	mpz_init (ctx->e);
	mpz_init (ctx->d);
	gmp_randinit_default (ctx->r);
}

void
clearRsaCtx (RsaCtx* ctx)
{
	mpz_clear (ctx->p);
	mpz_clear (ctx->q);
	mpz_clear (ctx->n);
	mpz_clear (ctx->m);
	mpz_clear (ctx->e);
	mpz_clear (ctx->d);
}

void
getPrime (mpz_t p, RsaCtx ctx, unsigned int bits)
{
	/* set p to a randomly chosen prime of approximately 'bits' bits */
	for (;;)
		{
			mpz_urandomb (p, ctx.r, bits);
			if (mpz_probab_prime_p (p, 10))
				return;
		}
}

RsaCtx
generateRsaCtx (unsigned int bits) {
	mpz_t t;	/* temp */
	mpz_t t2; /* temp2 */
	mpz_init (t);
	mpz_init (t2);

	RsaCtx ctx;
	initRsaCtx (&ctx);

	/* init mp random */
	gmp_randseed_ui (ctx.r, getpid () * time (NULL));

	/* generate P and Q
	 * P and Q should be of same size */
	mpz_set_ui (ctx.p, 0u);
	mpz_set_ui (ctx.q, 0u);
	for (;;)
		{
			mpz_mul (t, ctx.p, ctx.q);
			if (mpz_sizeinbase (t, 2) >= bits)
				break;
			getPrime (ctx.p, ctx, bits >> 1);
			getPrime (ctx.q, ctx, bits >> 1);
		}

	/* initialise N = PQ */
	mpz_mul (ctx.n, ctx.p, ctx.q);

	/* initialise M = (P-1)(Q-1) */
	mpz_sub_ui (ctx.p, ctx.p, 1u);
	mpz_sub_ui (ctx.q, ctx.q, 1u);
	mpz_mul (ctx.m, ctx.p, ctx.q);
	mpz_add_ui (ctx.p, ctx.p, 1u);
	mpz_add_ui (ctx.q, ctx.q, 1u);

	/* calculate E , E is odd and gcd (M, E) == 1 */
	mpz_set_ui (ctx.e, 0xffffu + 2);
	for (;;)
		{
			mpz_gcd (t, ctx.m, ctx.e);
			if (mpz_cmp_ui (t, 1u) == 0)
				break;
			mpz_add_ui (ctx.e, ctx.e, 2u);
		}

	/* calculate D , D = (X*M + 1) / E
	 * D is integer and D > 1
	 * note: this is called the mod inverse of (E,M) */
	mpz_invert (ctx.d, ctx.e, ctx.m);

	mpz_clear (t);
	mpz_clear (t2);

	return ctx;
}

void
encrypt (mpz_t r, RsaCtx c, mpz_t t)
{
	mpz_powm (r, t, c.e, c.n);
}

void
decrypt (mpz_t r, RsaCtx c, mpz_t t)
{
	mpz_powm (r, t, c.d, c.n);
}

void
printZ (char * name, mpz_t n)
{
	printf ("%s = ", name);
	mpz_out_str (stdout, 10, n);
	printf ("\n");
}

void
printCipherText (mpz_t * a, int l)
{
	int i = 0;
	while (i < l)
		{
			mpz_out_str (stdout, 10, a [i]);
			printf (" ");
			i++;
		}
	printf ("\n");
}

int
main ()
{
	/* generate an RSA contect with n bit keys */
	RsaCtx ctx = generateRsaCtx (1024);
	mpz_t t;
	mpz_t t2;
	mpz_t c;

	char plaintext   [] =
		"Das ist ein Testtext! Mit vielen Zeichen ;o) -- by Fritschy";
	mpz_t ciphertext [strlen (plaintext)];
	char plaintext2  [strlen (plaintext)+1];
	int i;

	printZ ("P", ctx.p);
	printZ ("Q", ctx.q);
	printZ ("N", ctx.n);
	printZ ("M", ctx.m);
	printZ ("E", ctx.e);
	printZ ("D", ctx.d);

	printf ("\ntesting:\n");

	mpz_init_set_ui (t, 100010001);
	mpz_init (c);
	mpz_init (t2);

	encrypt (c, ctx, t);
	decrypt (t2, ctx, c);

	printZ ("plaintext ", t);
	printZ ("ciphertext", c);
	printZ ("decrypted ", t2);

	printf ("\n");

	/* encrypt string */
	for (i = 0; i < strlen (plaintext); i++) {
		/* encrypt */
		mpz_init (ciphertext [i]);
		mpz_set_ui (t, plaintext [i]);
		encrypt (ciphertext [i], ctx, t);
		/* decrypt */
		decrypt (t, ctx, ciphertext [i]);
		plaintext2 [i] = mpz_get_ui (t);
	}
	/* terminate */
	plaintext2 [i] = 0;

	printf ("plaintext  = %s\nciphertext = ", plaintext);
	printCipherText (ciphertext, strlen (plaintext));
	printf ("decrypted  = %s\n", plaintext2);

	/* do cleanup */
	mpz_clear (t);
	mpz_clear (t2);
	mpz_clear (c);

	for (i = 0; i < strlen (plaintext); i++)
		mpz_clear (ciphertext [i]);

	clearRsaCtx (&ctx);

	return 0;
}