/*	 
 *	 Miracl Library by Shamus Software Ltd.
 *
 *   Requires: BIG.cpp, CRT.cpp from miracl library
 *
 *   James O' Connor 50703338 and Feargal Gallaghar 50729035
 */
#include <winsock.h>	//sockets
#include <iostream>		//IO
#include <windows.h>	//system time functionsfor rnd
#include <fstream>		//fileIO	
#include <big.h>		//include MIRACL system 
#include <crt.h>		//chinese remainder thereom    

using namespace std;
int genRSAKeys(Big privateKey[2], Big * publicKey);
void initRandom();
void encode(Big pubKey, int mesLen, char message[], char cryptMessage[]);
void decode(Big privKey[], Big pubKey, char message[], char cryptMessage[], int totalBlocks);
long randLong();
char randByte();
void genIDEA(char idea[16]);
void sendKey(Big pubKey, SOCKET conn);
void recvKey(Big *foreignPublicKey, SOCKET conn);
void sendBlock(char block[], SOCKET conn);
void recvBlock(char block[], SOCKET conn);
void StreamServer(short nPortBig, Big privKey[], Big pubKey);
void StreamClient(char *szServer, short nPort, Big privKey[], Big pubKey);
Big strongPrime(int n,long seed1,long seed2);

Miracl precision(2048, 256);


csprng *rng;
ifstream fin("seeds.txt");

#define PRINTERROR(s)	\
		fprintf(stderr,"\n%: %d\n", s, WSAGetLastError())


BOOL server = false;
void main(int argc, char **argv)
{
	WORD wVersionRequested = MAKEWORD(1,1);
	WSADATA wsaData;
	int nRet;
	short nPort;
	char hostname[256];
	int set = 0;

	if(argc == 1)
	while(true)
	{
		
		cout << "Select Server or Client:" << endl;
		cout << "Type 1 for server. 2 for client:" << endl;
		cin >> set;

		if(set==1)
		{
			cout << "You have chosen server !"<<endl;
			break;
		}

		else if(set==2)
		{
			cout << "You have chosen client !"<<endl;
			break;
		}
	
		else
			cout << "# INVALID CHOICE #"<<endl;
	}

	else
		set = 2;


	nRet = WSAStartup(wVersionRequested, &wsaData);
	if (wsaData.wVersion != wVersionRequested)
	{	
		fprintf(stderr,"\n Wrong version\n");
		return;
	}

	int i =0;
	initRandom();
	Big privateKey[2];
	privateKey[0] =0;
	privateKey[1] =0;
	Big * publicKey = new Big();
	*publicKey =0;
	// generate new RSA keys
	genRSAKeys(privateKey, publicKey);

	Big pubKey = *publicKey;
	Big privKey[2];
	privKey[0] = privateKey[0];
	privKey[1] = privateKey[1];

	if (set == 1)
	{
		cout << "Enter port number to listen on:" << endl;
		cin >> nPort;
		server = true;
		StreamServer(nPort, privKey, pubKey);

	}
	else if (set == 2 && argc ==1)
	{
		cout << "Enter host name:" << endl;
		cin >> hostname;
		cout << "Enter port number to connect to:" << endl;
		cin >> nPort;
		StreamClient(hostname, nPort, privKey, pubKey);
	}

	else
	{
		char *temp = argv[1];
		strcpy(&hostname[0], temp);
		nPort = atoi(argv[2]);
		StreamClient(hostname, nPort, privKey, pubKey);
	}

	WSACleanup();
}


void StreamServer(short nPort, Big privKey[], Big pubKey)
{
	// Create a TCP/IP stream socket to "listen" with
	SOCKET	listenSocket;

	listenSocket = socket(AF_INET,			// Address family
						  SOCK_STREAM,		// Socket type
						  IPPROTO_TCP);		// Protocol
	if (listenSocket == INVALID_SOCKET)
	{
		cout << "Socket failed: ";
		PRINTERROR("socket()");
		return;
	}

	// Fill in the address structure
	SOCKADDR_IN saServer;		

	saServer.sin_family = AF_INET;
	saServer.sin_addr.s_addr = INADDR_ANY;	// Let WinSock supply address
	saServer.sin_port = htons(nPort);		// Use port from command line

	// bind the name to the socket
	int nRet;

	nRet = bind(listenSocket,				// Socket 
				(LPSOCKADDR)&saServer,		// Our address
				sizeof(struct sockaddr));	// Size of address structure
	if (nRet == SOCKET_ERROR)
	{
		cout << "Server Socket failed: ";
		PRINTERROR("bind()");
		closesocket(listenSocket);
		return;
	}


	int nLen;
	nLen = sizeof(SOCKADDR);
	char szBuf[256];

	nRet = gethostname(szBuf, sizeof(szBuf));
	if (nRet == SOCKET_ERROR)
	{
		PRINTERROR("gethostname()");
		closesocket(listenSocket);
		return;
	}

	// Show the server name and port number
	printf("\nServer named %s waiting on port %d\n",
			szBuf, nPort);

	// Set the socket to listen
	nRet = listen(listenSocket,		
				  SOMAXCONN);					
	if (nRet == SOCKET_ERROR)
	{
		PRINTERROR("listen()");
		closesocket(listenSocket);
		return;
	}

	


while(true)
{
	// Wait for an incoming request
	SOCKET	remoteSocket;

	printf("\nWaiting for connection...");
	remoteSocket = accept(listenSocket,		
						  NULL,				
						  NULL);
	if (remoteSocket == INVALID_SOCKET)
	{
		PRINTERROR("accept()");
		closesocket(listenSocket);
		return;
	}

	Big *clientKey = new Big();
	*clientKey =0;

	//wait for client public key to be transmitted
	recvKey(clientKey, remoteSocket);

	//send own public key
	sendKey(pubKey, remoteSocket);

	char random[256];
	char clear[256];
	char compare[256];
	int i;
	
	//generate random number (challenge)
	for(i=0; i<256; i++)
	{
		clear[i] = randByte();
		compare[i] = clear[i];
	}


	//encrypt clear number
	encode(*clientKey, 256, clear, random);
	
	//send challenge
	sendBlock(random, remoteSocket);

	//receive response
	recvBlock(random, remoteSocket);

	//extra security, in case decode fails
	for(i=0; i<256; i++)
		clear[i] = '\0';

	//decode response
	decode(privKey, pubKey, clear, random, 1);

	BOOL same = TRUE;
	// compare random numbers for auth
	for(i=0; i<256; i++)
	{
		//cout <<" compare: " << clear[i] << " " << compare[i]<<endl;
		if(! (clear[i] == compare[i]))
		{
			same = FALSE;
			cout << "Authentication Failed." << endl;
			break;
		}
	}

	if(same)
	{
		//receive client's challenge
		recvBlock(random, remoteSocket);

		//decode into clear
		decode(privKey, pubKey, clear, random, 1);

		//recrypt clear number
		encode(*clientKey, 256, clear, random);

		//send response to clients challenge
		sendBlock(random, remoteSocket);

		//receive the idea key from the client
		recvBlock(random, remoteSocket);

		char idea[256] = {'\0'};
		//decrypt idea key
		decode(privKey, pubKey, idea, random, 1);
	
		miracl *mip = &precision;
		//print out key in hex
		mip->IOBASE = 16;
		//key is first 16 bytes of data
		cout << "IDEA key:\n"<< from_binary(16, idea) <<endl;
		mip->IOBASE = 256;

		//close socket
		closesocket(remoteSocket);
		// generate new RSA keys
		genRSAKeys(privKey, &pubKey);
	}

	else
	{	//authentication failed
		//close socket
		closesocket(remoteSocket);
		// generate new RSA keys
		genRSAKeys(privKey, &pubKey);
	}
}
	// Close server socket before exiting	
	closesocket(listenSocket);
	return;
}


void StreamClient(char *szServer, short nPort, Big privKey[], Big pubKey)
{
	int i=0;
	printf("\nStream Client connecting to server: %s on port: %d",
				szServer, nPort);

	// Find the server
    LPHOSTENT lpHostEntry;

	lpHostEntry = gethostbyname(szServer);
    if (lpHostEntry == NULL)
    {
        PRINTERROR("gethostbyname()");
        return;
    }

	// Create a TCP/IP stream socket
	SOCKET	theSocket;

	theSocket = socket(AF_INET,				// Address family
					   SOCK_STREAM,			// Socket type
					   IPPROTO_TCP);		// Protocol
	if (theSocket == INVALID_SOCKET)
	{	cout << "Socket failed: ";
		PRINTERROR("socket()");
		return;
	}

	// Fill in the address structure
	SOCKADDR_IN saServer;
	saServer.sin_family = AF_INET;
	saServer.sin_addr = *((LPIN_ADDR)*lpHostEntry->h_addr_list);
	saServer.sin_port = htons(nPort);	// Port number

	// connect to the server
	int nRet;

	nRet = connect(theSocket,				// Socket
				   (LPSOCKADDR)&saServer,	// Server address
				   sizeof(struct sockaddr));// Length of server address structure
	if (nRet == SOCKET_ERROR)
	{
		cout << "Socket failed: ";
		PRINTERROR("socket()");
		closesocket(theSocket);
		return;
	}

	//IDEA key is 16 bytes, public key is 256 bytes == 1 block
	int numBlocks = 1; 

	//send public key to server
	sendKey(pubKey, theSocket);


	Big *serverPubKey = new Big();
	*serverPubKey =0;
	//receive servers public key
	recvKey(serverPubKey, theSocket);

	char random[256];
	//receive servers challenge
	recvBlock(random, theSocket);

	char clear[256];

	//decrypt random number
	decode(privKey, pubKey, clear, random, 1);
	
	//encrypt same random number with servers key
	encode(*serverPubKey, 256, clear, random);

	//send back random number (answer challenge)
	sendBlock(random, theSocket);	

	//create new random number
	char compare[256];
	
	for(i=0; i<256; i++)
	{	clear[i] = randByte();
		compare[i] = clear[i];
	}
	
	//encrypt random number
	encode(*serverPubKey, 256, clear, random);

	//send new challenge
	sendBlock(random, theSocket);

	//receive number back
	recvBlock(random, theSocket);

	//extra security, in case decode fails
	for(i=0; i<256; i++)
		clear[i] = '\0';

	//decrypt random number
	decode(privKey, pubKey, clear, random, 1);
	
	BOOL same = TRUE;
	// compare random numbers for auth
	for(i=0; i<256; i++)
	{
		//cout <<" compare: " << clear[i] << " " << compare[i]<<endl;
		if(! (clear[i] == compare[i]))
		{
			same = FALSE;
			cout << "Authentication Failed." << endl;
			break;
		}
	}

	if(same)
	{
		//idea key is 16 bytes, rest is filler for blcok
		char idea[256] = {'0'};
		genIDEA(idea);	//fill IDEA key with random bytes

		//encrypt idea key
		encode(*serverPubKey, 256, idea, random);

		//send encrypted IDEA key to server
		sendBlock(random, theSocket);
		miracl *mip= &precision;

		mip->IOBASE = 16;
		cout <<"IDEA key is:\n" << from_binary(16, idea) <<endl;
		mip->IOBASE = 256;

	}

	closesocket(theSocket);
	return;
}

//return random long number
long randLong()
{
	char temp[sizeof(long)];
	for(int i=0; i< sizeof(long); i++)
		temp[i] = strong_rng(rng);

	long l = (long) *temp;
    return l;
}

//return random byte
char randByte()
{ 
    return strong_rng(rng);
}

//generate strong prime number 
Big strongPrime(int n,long seed1,long seed2)
{	Big p =0;
	Big pd =0;
	Big pl =0;
	Big ph =0;
    int r=0;
	int r1 =0;
	int r2 =0;
    irand(seed1);
    pd=rand(2*n/3,2);
    pd=nextprime(pd);
    ph=pow((Big)2,n-1)/pd;  
    pl=pow((Big)2,n-2)/pd;
    ph-=pl;
    irand(seed2);
    ph=rand(ph);
    ph+=pl;
    r1=pd%12;
    r2=ph%12;
    r=0;
    while ((r1*(r2+r))%12!=5) 
		r++;

    ph+=r;

    do 
    { 
        p=2*ph*pd+1;
        ph+=12;

    } while (!prime(p));

    return p;
}


void initRandom()
{
	char random[128];
	for(int i=0; i<128; i++)
		fin.get(random[i]);

	SYSTEMTIME t;
	GetSystemTime(&t);	
	unsigned long tmp1 = (unsigned long) (t.wMilliseconds 
		+ (1000 * t.wSecond) 
		+ (60000 * t.wMinute)
		+ (360000 * t.wHour));
	mr_unsign32  tod = (mr_unsign32) tmp1;
	rng = new csprng();
	strong_init(rng, 128, random, tod);
}

//  generate RSA keys 
int genRSAKeys(Big p[], Big *publicKey)
{  
	cout << "Please wait, generating keys...." <<endl;
    int k,i;
    long seed[4];
    ofstream public_key("public.key");
    ofstream private_key("private.key");
    miracl *mip=&precision;
	Big ke =0;
    gprime(15000);  
	k = 1024;	//size of primes
    for (i=0;i<4;i++)
        seed[i]=randLong();
    ke=1;

    for (i=0;i<2;i++) 
    {
        p[i]=strongPrime(k,seed[2*i],seed[2*i+1]);
        ke*=p[i];
    }

    mip->IOBASE=16;
    public_key << ke << endl;
    for (i=0;i<2;i++)
        private_key << p[i] << endl;

	fin.close();
	
	//set new seeds ;)
	ofstream fout("seeds.txt");
	int v=0;
	for(v=0; v<128; v++)
		fout << randByte();

	fout.close();
	*publicKey = ke;
    return 0;
}

//  encode using public key  
void encode(Big pubKey, int mesLen, char message[], char cryptMessage[])
{  
	cout << "Encoding. Keysize: " << bits(pubKey) <<" bits" << endl;
    Big e =0;
	Big m =0;
	Big mn =0;
	Big mx =0;
	Big ep[2];
	ep[0] =0;
	ep[1] =0;
	Big kp[2];
	kp[0] =0;
	kp[1] =0;
	int i =0;
	int j =0;
	int count =0;
	int klen =0;
	int numBlocks=0;
	char buff[256] = {'\0'};
	char block[256] = {'\0'};
	miracl *mip=&precision;
	mip->IOBASE=256;
	mn=root(pubKey, 3);
    mx=mn*mn*mn-mn*mn;

    klen=256;
	BOOL last = FALSE;
	if(mesLen < klen) 
		last = TRUE;

	count = 0;
	numBlocks = 0;

	// chop up into klen-sized chunks and encode 	
    while (!last)
    {	
		for (i=0;i<256;i++)
			buff[i]=message[count + i];

        m= from_binary(256, buff);
        e=pow(m,3, pubKey);

		to_binary(e, 256, block, FALSE); 
		
		for(j=0; j<256; j++)
			cryptMessage[count + j] = block[j];

		
		count = count + 256;
		numBlocks++;
		
		if((mesLen - count) < klen)
		{	last = TRUE;
			
		}
	
    }

	// now deal with left overs 
	if (last && (mesLen%klen)!=0)
    {	
		cout << "in if extra"<<endl;
		for (i=0;i<(mesLen - count);i++)
			buff[i]=message[i+count];

		//add random characters to fill block
		for (j=i; j<256; j++)
			buff[j]=randByte();	

		//cast buffer into Big
        m= from_binary(256, buff); 
        
		e=pow(m,3,pubKey);

		to_binary(e, 256, block, FALSE); 
		for(int j=0; j<256; j++)
			cryptMessage[count + j] = block[j];
		
		count= count + 256;
		numBlocks++;
		
    }

}   

void decode(Big privKey[], Big pubKey, char message[], char cryptMessage[], int totalBlocks)
{  //  decode using private key  
    Big e =0;
	Big m =0;
	Big mn =0;
	Big mx =0;
	Big ep[2];
	ep[0] =0;
	ep[1] =0;
	Big kp[2];
	kp[0] =0;
	kp[1] =0;
	int i =0;
	int j =0;
	int count =0;
	int klen =0;
	int numBlocks=0;
	
	miracl *mip=&precision;
	mip->IOBASE=256;

    klen=256;
	count = 0;
	numBlocks = 0;
	char block[256] = {'\0'};
	char buff[256]= {'\0'};

	// use chinese remainder thereom 
	Crt chinese(2, privKey);   

    // prepare for Chinese remainder thereom 
    // Find 1/3 mod p[i]-1  Euler TF       
    kp[0]=(2*(privKey[0]-1)+1)/3;
	kp[1]=(2*(privKey[1]-1)+1)/3;

    mn=root(pubKey,3);
    mx=mn*mn*mn-mn*mn;

	// chop up into klen-sized chunks and decode 
    for(int jk = 0; jk < totalBlocks; jk++)
    {	
		for (i=0;i<256;i++)
			buff[i]=cryptMessage[count + i];

        m= from_binary(256, buff);

		ep[0]=pow(m,kp[0],privKey[0]);
		ep[1]=pow(m,kp[1],privKey[1]);
        e=chinese.eval(ep);
        if (e>=mx) 
			e%=mn;

		to_binary(e, 256, block, FALSE); 	
		for(j=0; j<256; j++)
			message[count + j] = block[j];

        count = count + 256;
		numBlocks++;
	
    }

}   

void genIDEA(char idea[16])
{
	for(int i=0; i<16; i++)
		idea[i] = randByte();
}

void sendKey(Big pubKey, SOCKET conn)
{	
	// Send data to the server
	cout << "sending key.." << endl;
	miracl *mip =  & precision;
	mip->IOBASE=16;
	cout << pubKey<< "\n" << endl;
	mip->IOBASE = 256;
	int nRet =0;
	char key[256];
	to_binary(pubKey, 256, key, FALSE);

	nRet = send(conn,					// Connected socket
				key,					// Data buffer
				256,					// Length of data
				0);						// Flags
	if (nRet == SOCKET_ERROR)
	{
		PRINTERROR("send()");
		closesocket(conn);
		return;
	}

}

void recvKey(Big *foreignPublicKey, SOCKET conn)
{
	// Wait for a reply
	int nRet =0;
	char key[256];
	nRet = recv(conn,		// Connected socket
				key,		// Receive buffer
				256,		// Size of receive buffer
				0);			// Flags
	if (nRet == SOCKET_ERROR)
	{
		PRINTERROR("recv()");
		closesocket(conn);
		return;
	}

	*foreignPublicKey = from_binary(256, key);

	cout << "received key.." << endl;
	miracl *mip =  & precision;
	mip->IOBASE=16;
	cout << *foreignPublicKey << "\n" << endl;
	mip->IOBASE = 256;
}

void sendBlock(char block[], SOCKET conn)
{
	// Send data to the server
	cout << "sending block..." << endl;
	miracl *mip =  & precision;
	mip->IOBASE=16;
	cout << from_binary(256, block) << "\n" << endl;
	mip->IOBASE = 256;
	int nRet =0;

	nRet = send(conn,		// Connected socket
				block,		// Data buffer
				256,		// Length of data
				0);			// Flags
	if (nRet == SOCKET_ERROR)
	{
		PRINTERROR("send()");
		closesocket(conn);
		return;
	}
}

void recvBlock(char block[], SOCKET conn)
{
	// Wait for a reply
	int nRet =0;
	nRet = recv(conn,		// Connected socket
				block,		// Receive buffer
				256,		// Size of receive buffer
				0);			// Flags
	if (nRet == SOCKET_ERROR)
	{
		PRINTERROR("recv()");
		closesocket(conn);
		return;
	}

	cout << "received block..." << endl;

	miracl *mip =  & precision;
	mip->IOBASE=16;
	cout << from_binary(256, block) << "\n" << endl;
	mip->IOBASE = 256;
}