#include <crypto/group.h>
#include <crypto/crypto.h>
#include <NTL/GF2EXFactoring.h>

NTL_CLIENT


// ------- mod p groups -------


modp_group :: modp_group(int lmu, int lp)
{
        ZZ lb, ub, gp;

	if (lmu < 10 || lp < 2*lmu)
		throw Key_error();

        do
        {
                mu = GenPrime_ZZ(lmu, 80);
                lb = (power_ZZ(2, lp-1) - 1 + (mu-1)) / mu;
                ub = (power_ZZ(2, lp) - 1) / mu;
                nu = RandomBnd(ub - lb) + lb;
                p = mu * nu + 1;
        }
        while(!(ProbPrime(p, 20)));

	// check:

	if (NumBits(mu) != lmu || NumBits(p) != lp)
		throw General_error();

        do
        {
                gp = RandomBnd(p - 1) + 1;
                g = PowerMod(gp, nu, p);

                if(g != 1)
                        break;

                if(g <= 0 && PowerMod(g, mu, p) != 1)
                        throw General_error();
        }
        while(1);
}


modp_group :: modp_group(istream& ifile)
{
	ifile >> mu >> nu >> p >> g;
}

void modp_group :: PrettyPrint(ostream& tofile)
{
	tofile << "Modp-Group:\n\np = "; prettyprint(tofile, p);
	tofile << "\n\ng = "; prettyprint(tofile, g);
	tofile << "\n\nmu = "; prettyprint(tofile, mu);
	tofile << "\n\nnu = "; prettyprint(tofile, nu);
}

void modp_group :: writetofile(ostream& ofile) const
{
	ofile << mu << "\n" << nu << "\n" << p << "\n" << g << "\n";
}


ostring modp_group :: encode(const ZZ& a) const
{
	return i2osp(a, NumBytes(p));
}

ostring modp_group :: pencode(const ZZ& a) const
{
	return i2osp(a, NumBytes(p));
}

int modp_group :: 
   decode1(ZZ& x, const ostring& s, int pos, unsigned int& fmt) const
{
	if (pos + NumBytes(p) > s.size())
		return 0;
	else {
		x = os2ip(s.substr(pos, NumBytes(p)));
		if (x == 0 || x >= p) return 0;
		fmt = 0;
		return NumBytes(p);
	}
}

ZZ modp_group :: ident() const
{
	return to_ZZ(1);
}

int modp_group :: is_ident(const ZZ& a) const
{
	return a == 1;
}

int modp_group :: equal(const ZZ& a, const ZZ& b) const
{
	return a == b;
}

ZZ modp_group :: add(const ZZ& a, const ZZ& b) const
{
	return (a * b) % p;
}




// ------- elliptic curve mod p groups -------

void ecmodp_group :: writetofile(ostream& ofile) const
{
	ofile << p << "\n";
	ofile << a << "\n";
	ofile << b << "\n";
	ofile << mu << "\n";
	ofile << nu << "\n";
	writetofile(ofile, g);
}

ecmodp_group :: ecmodp_group(istream& ifile) 
{
	ifile >> p >> a >> b >> mu >> nu;
	readfromfile(ifile, g);
	encoding_fmt = uncompressed_fmt;
}

void ecmodp_group :: PrettyPrint(std::ostream& tofile)
{
	tofile << "ECModp-Group:\n\np = "; prettyprint(tofile, p);
	tofile << "\n\na = "; prettyprint(tofile, a % p);
	tofile << "\n\nb = "; prettyprint(tofile, b % p);
	tofile << "\n\nmu = "; prettyprint(tofile, mu);
	tofile << "\n\nnu = "; prettyprint(tofile, nu);
	tofile << "\n\ng(x) = "; prettyprint(tofile, g.x);
	tofile << "\n\ng(y) = "; prettyprint(tofile, g.y);
}

ostring ecmodp_group :: encode(const ecmodp_point& Q) const
{
	octet H;
	octet ytilde, U, C;

	ostring X, Y;

	if (Q.inf) return ostring(octet(0));


	if (encoding_fmt == uncompressed_fmt || encoding_fmt == hybrid_fmt)
		U = 1;
	else
		U = 0;

	if (encoding_fmt == compressed_fmt || encoding_fmt == hybrid_fmt)
		C = 1;
	else
		C = 0;

	if (!C)
		ytilde  = 0;
	else
		ytilde = IsOdd(Q.y);

	H = 4*U + C*(2 + ytilde);

	X = i2osp(Q.x, NumBytes(p));

	if (U) 
		Y = i2osp(Q.y, NumBytes(p));
	else
		Y = ostring();

	return ostring(H) + X + Y;
}

ostring ecmodp_group :: pencode(const ecmodp_point& Q) const
{
	if (Q.inf)
		return ostring(NumBytes(p), octet(0));
	else
		return i2osp(Q.x, NumBytes(p));
}




int ecmodp_group :: decode1(ecmodp_point& Q, const ostring& s, int pos,
			unsigned int& fmt) const
{
	octet H;
	octet U, C, ytilde, ytilde1;
	int curpos;
	ostring X, Y;
	ZZ x, y;
	ZZ alpha;

	if (pos + 1 > s.size())
		return 0;

	H = s[pos];

	if (!H) {
		Q.inf = 1;
		fmt = 0;
		return 1;
	}

	if (H & (~7U)) return 0;

	U = H >> 2;
	C = (H >> 1) & 1U;
	ytilde = H & 1U;

	if (!C && ytilde) return 0;

	curpos = pos + 1;
	if (curpos + NumBytes(p) > s.size())
		return 0;
	X = s.substr(curpos, NumBytes(p));
	x = os2ip(X);
	if (x >= p)
		return 0;
	curpos += NumBytes(p);

	if (U) {
		if (curpos + NumBytes(p) > s.size())
			return 0;
		Y = s.substr(curpos, NumBytes(p));
		y = os2ip(Y);
		if (y >= p)
			return 0;
		curpos += NumBytes(p);
		
		if (((y*y - (x*x*x + a*x + b)) % p) != 0) // check on curve
			return 0;

		// silly consistency check for hybrid fmt

		if (C) {
			ytilde1 = IsOdd(y);
			if (ytilde != ytilde1) 
				return 0;
		}
		
	}
	else { 
		alpha = (x*x*x + a*x + b) % p;
		if (alpha == 0) {
			y = 0;
			if (ytilde) // another silly consistenncy check
				return 0;
		}
		else if (Jacobi(alpha, p) == 1) {
			y = SqrRootMod(alpha, p);
			ytilde1 = IsOdd(y);
			if (ytilde != ytilde1)
				y = p - y;
		}
		else
			return 0;
		
	}

	if (U && !C)
		fmt = 1U;
	if (!U && C)
		fmt = 2U;
	if (U && C)
		fmt = 4U;

	Q.inf = 0;
	Q.x = x;
	Q.y = y;
	return curpos - pos;
}



ecmodp_point ecmodp_group :: ident() const
{
	ecmodp_point P;
	P.inf = 1;
	return P;
}

int ecmodp_group :: is_ident(const ecmodp_point& Q) const
{
	return Q.inf == 1;
}

int ecmodp_group :: equal(const ecmodp_point& Q1, const ecmodp_point& Q2) const
{
	if (Q1.inf && !Q2.inf)
		return 0;
	else if (!Q1.inf && Q2.inf)
		return 0;
	else if (Q1.inf && Q2.inf)
		return 1;
	else
		return Q1.x == Q2.x && Q1.y == Q2.y;
}

ecmodp_point ecmodp_group :: add(const ecmodp_point& Q1, const ecmodp_point& Q2) const
{

	ecmodp_point Q;
	ZZ r, x1, x2, y1, y2;

	if (Q1.inf == 1)
		return Q2;

	if (Q2.inf == 1)
		return Q1;

	x1 = Q1.x;
	y1 = Q1.y;
        x2 = Q2.x;
	y2 = Q2.y;

	if (x1 == x2 && y1 == (-y2 % p))
	{
		Q.inf = 1;
		return Q;
	}
	else
	{
		// if Q1 == Q2
		if (x1 == x2 && y1 == y2)
			r = (((3*x1*x1 + a) % p) *  InvMod((2 * y1) % p, p)) % p;

		// Q1 != Q2
		else
			r = (((y2 - y1) % p) * InvMod((x2 - x1) % p, p)) % p;
	}

	Q.inf = 0;
	Q.x = (r*r - x1 - x2) % p;
	Q.y = (r * (x1 - Q.x) - y1) % p;

	return Q;
}





// ------- elliptic curve over GF(2^n) groups -------

void ecgf2_group :: writetofile(ostream& ofile) const
{
	ofile << p << "\n";
	ofile << a << "\n";
	ofile << b << "\n";
	ofile << mu << "\n";
	ofile << nu << "\n";
	writetofile(ofile, g);
}

ecgf2_group :: ecgf2_group(istream& ifile) 
{
	ifile >> p >> a >> b >> mu >> nu;
	readfromfile(ifile, g);
	encoding_fmt = uncompressed_fmt;
}

void ecgf2_group :: PrettyPrint(std::ostream& tofile)
{
        tofile << "ECGF2-Group:\n\np = "; prettyprint(tofile, p);
        tofile << "\n\na = "; prettyprint(tofile, a % p);
        tofile << "\n\nb = "; prettyprint(tofile, b % p);
        tofile << "\n\nmu = "; prettyprint(tofile, mu);
        tofile << "\n\nnu = "; prettyprint(tofile, nu);
        tofile << "\n\ng(x) = "; prettyprint(tofile, g.x);
        tofile << "\n\ng(y) = "; prettyprint(tofile, g.y);
}

ostring ecgf2_group :: encode(const ecgf2_point& Q) const
{
	octet H;
	octet ytilde, U, C;

	ostring X, Y;

	if (Q.inf) return ostring(octet(0));

	if (encoding_fmt == uncompressed_fmt || encoding_fmt == hybrid_fmt)
		U = 1;
	else
		U = 0;

	if (encoding_fmt == compressed_fmt || encoding_fmt == hybrid_fmt)
		C = 1;
	else
		C = 0;

	if (!C)
		ytilde = 0;
	else if (Q.x == 0)
		ytilde = 0;
	else 
		ytilde = rep(ConstTerm((Q.y * InvMod(Q.x, p)) % p));

	H = 4*U + C*(2 + ytilde);

	X = poly2osp(Q.x, NumBytes1(p));

	if (U) 
		Y = poly2osp(Q.y, NumBytes1(p));
	else
		Y = ostring();

	return ostring(H) + X + Y;
}

ostring ecgf2_group :: pencode(const ecgf2_point& Q) const
{
	if (Q.inf)
		return ostring(NumBytes1(p), octet(0));
	else
		return poly2osp(Q.x, NumBytes1(p));
}


// computes the square root of a mod p, which always
// exists and is unique

static
GF2X SqrRootMod(const GF2X& a, const GF2X& p)
{
	GF2X r;
	GF2EBak bak;
	bak.save();
	GF2E::init(p);

	GF2EX f;
	SetCoeff(f, 2);
	SetCoeff(f, 0, to_GF2E(a));

	vec_pair_GF2EX_long v = CanZass(f);

	// a square root is guaranteed to exist

	r =  rep(ConstTerm(v[0].a));

	if ( (r*r + a) % p != 0 )
		Error("oops----1");

	return r;
}

// if there exists z such that z^2 + z + a = 0 (mod p), then
// r is set to z and 1 is returned; otherwise, 0 is returned

static
int solve1(GF2X& r, const GF2X& a, const GF2X& p)
{
	GF2EBak bak;
	bak.save();
	GF2E::init(p);

	GF2EX f;
	SetCoeff(f, 2);
	SetCoeff(f, 1);
	SetCoeff(f, 0, to_GF2E(a));

	vec_pair_GF2EX_long v = CanZass(f);

	// is either irreducible, or splits into two distinct
	// linear factors

	if (v.length() == 1) // no roots!
		return 0;

	r =  rep(ConstTerm(v[0].a));

	if ( (r*r + r + a) % p != 0 )
		Error("oops----2");

	return 1;
}


int ecgf2_group :: decode1(ecgf2_point& Q, const ostring& s, int pos,
			unsigned int& fmt) const
{
	octet H;
	octet U, C, ytilde, ytilde1;
	int curpos;
	ostring X, Y;
	GF2X x, y;
	GF2X alpha, beta, x_inv;

	if (pos + 1 > s.size())
		return 0;

	H = s[pos];

	if (!H) {
		Q.inf = 1;
		fmt = 0;
		return 1;
	}

	if (H & (~7U)) return 0;

	U = H >> 2;
	C = (H >> 1) & 1U;
	ytilde = H & 1U;

	if (!C && ytilde) return 0;

	curpos = pos + 1;
	if (curpos + NumBytes1(p) > s.size())
		return 0;
	X = s.substr(curpos, NumBytes1(p));
	x = os2polyp(X);
	curpos += NumBytes1(p);

	if (U) {
		if (curpos + NumBytes1(p) > s.size())
			return 0;
		Y = s.substr(curpos, NumBytes1(p));
		y = os2polyp(Y);
		curpos += NumBytes1(p);
		
		// check on curve
		if ((( y*y + x*y + x*x*x + a*x*x + b ) % p) != 0) 
			return 0;

		if (C) {
			// silly consistency check in hybrid fmt

			if (x == 0)
				ytilde1 = 0;
			else 
				ytilde1 = rep(ConstTerm((y * InvMod(x, p)) % p));

			if (ytilde != ytilde1)
				return 0;
		}
		
	}
	else { 
		if (x == 0) {
			y = SqrRootMod(b, p);
			if (ytilde) // another silly consistency check
				return 0;
		}
		else { 

			x_inv = InvMod(x, p);
			alpha = (x + a + b*x_inv*x_inv) % p;

			if (!solve1(beta, alpha, p))
				return 0;

			ytilde1 = rep(ConstTerm(beta));

			if (ytilde != ytilde1)
				beta = beta + 1;

			y = (x * beta) % p;
		}
	}

	if (U && !C)
		fmt = 1U;
	if (!U && C)
		fmt = 2U;
	if (U && C)
		fmt = 4U;

	Q.inf = 0;
	Q.x = x;
	Q.y = y;
	return curpos - pos;
}


ecgf2_point ecgf2_group :: ident() const
{
	ecgf2_point P;
	P.inf = 1;
	return P;
}

int ecgf2_group :: is_ident(const ecgf2_point& Q) const
{
	return Q.inf == 1;
}

int ecgf2_group :: equal(const ecgf2_point& Q1, const ecgf2_point& Q2) const
{
	if (Q1.inf && !Q2.inf)
		return 0;
	else if (!Q1.inf && Q2.inf)
		return 0;
	else if (Q1.inf && Q2.inf)
		return 1;
	else
		return Q1.x == Q2.x && Q1.y == Q2.y;
}

// TODO: modify routine below

ecgf2_point ecgf2_group :: add(const ecgf2_point& Q1, const ecgf2_point& Q2) const
{

	ecgf2_point Q;
	GF2X r, x1, x2, y1, y2, x3, y3;

	if (Q1.inf == 1)
		return Q2;

	if (Q2.inf == 1)
		return Q1;

	x1 = Q1.x;
	y1 = Q1.y;
        x2 = Q2.x;
	y2 = Q2.y;

	if (x1 == x2 && y1 + y2  == x1)
	{
		Q.inf = 1;
		return Q;
	}
	else
	{
		// if Q1 == Q2
		if (x1 == x2 && y1 == y2) {

			r  = (x1 + y1 * InvMod(x1, p)) % p;
			x3 = (r*r + r + a) % p;
			y3 = (x1*x1 + (r+1)*x3) % p;

		}

		// Q1 != Q2
		else {
			r  = ((y2+y1) * InvMod(x2+x1, p)) % p;
			x3 = (r*r + r + x1 + x2 + a) % p;
			y3 = (r*(x1+x3) + x3 + y1) % p;


		}
	}

	Q.inf = 0;
	Q.x = x3;
	Q.y = y3;

	return Q;
}

//-------------------------------------------------------------------

void prettyprint(const char* label, const ZZ& Q, ostream& tofile)
{
        tofile << "\n" << label << " = ";
        prettyprint(tofile, Q);
}


void prettyprint(const char* label, const ecmodp_point& Q, ostream& tofile)
{
        if(Q.inf)
                tofile << "\n" << label << " = Infinity";
        else
        {
                tofile << "\n" << label << "(x) = ";
                prettyprint(tofile, Q.x);
                tofile << "\n" << label << "(y) = ";
                prettyprint(tofile, Q.y);
        }
}

void prettyprint(const char* label, const ecgf2_point& Q, ostream& tofile)
{
        if(Q.inf)
                tofile << "\n" << label << " = Infinity";
        else
        {
                tofile << "\n" << label << "(x) = ";
                prettyprint(tofile, Q.x);
                tofile << "\n" << label << "(y) = ";
                prettyprint(tofile, Q.y);
        }
}

