#include <crypto/kem.h>

NTL_CLIENT

RSAKEM :: RSAKEM(KDF& K, int Keylen, int len) : k(K), KEM(Keylen)
{
	if (len < 20)
		throw General_error();
		
        ZZ p = GenPrime_ZZ(len, 80);
        ZZ q = GenPrime_ZZ(len, 80);
        ZZ phi;

        while(p == q)
                q = GenPrime_ZZ(len, 80);

        e = (1L << 16) + 1;
        n = p*q;
        phi = (p-1)*(q-1);
        d = InvMod(e, phi);
}

RSAKEM :: RSAKEM(KDF& K, int Keylen, istream& ifile) : k(K), KEM(Keylen)
{
	ifile >> n >> e >> d;
}

void RSAKEM :: writetofile(ostream& ofile)
{
	ofile << n << "\n" << e << "\n" << d;
}

void RSAKEM :: PrettyPrint(ostream& tofile)
{
	tofile << "\n-------------";
	tofile << "\nRSA-KEM";
	tofile << "\n-------------";
	tofile << "\n\nKdf=";
	k.PrettyPrint(tofile);
	tofile << "\nkeylen=" << keylen;
	tofile << "\n\n-----\n";
	tofile << "\nPublic Key";
	tofile << "\n\nn = " << n << "\n\ne = " << e;
	tofile << "\n\n-----\n";
	tofile << "\nPrivate Key";
	tofile << "\n\nn = " << n << "\n\nd = " << d;
	tofile << "\n\n";
}

ostring RSAKEM :: rsatransform(const ostring& X, ZZ& a)
{
	long l = NumBytes(n);
	ZZ x, y;
	ostring Y;

	if(X.size() != l)
		throw Decrypt_error();

	// note -- this decryption error (and the one below)
	// can only be triggered during decryption (not encryption)

	x = os2ip(X);

	if(!(x < n))
		throw Decrypt_error();

	y = PowerMod(x,a,n);

	Y = i2osp(y,l);

	return Y;
}

ostring RSAKEM :: encrypt(ostring& cipher)
{
	ZZ r, x;
	ostring R, K;

	r = RandomBnd(n); // ultimately, we want r to be input...or output?

	R = i2osp(r, NumBytes(n));

	cipher = rsatransform(R,e);

	K = k.eval(R,keylen);

	if (logfile != NULL)
	{
		*logfile << "\n-------------------------------";
		*logfile << "\nTrace for RSA-KEM encrypt ";
		*logfile << "\n-------------------------------";
		*logfile << "\n\nr = ";
		prettyprint(*logfile, r);
		*logfile << "\n\nR = ";
		prettyprint(*logfile, R);
		*logfile << "\nC0 = ";
		prettyprint(*logfile, cipher);
		*logfile << "\n\nK = ";
		prettyprint(*logfile, K);
		*logfile << "\n\n";
	}	

	return K;
}

ostring RSAKEM :: decrypt(const ostring& cipher, int& nbytes)
{
	ostring R, K;

	if (cipher.size() < NumBytes(n))
		throw Decrypt_error();

	nbytes = NumBytes(n);

	R = rsatransform(cipher.substr(0, nbytes), d);

	K = k.eval(R,keylen);

	return K;
}

