#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include <signal.h>
#include <sys/time.h>
#include <misc.h>
#include <tlmplib.h>
#include <tlmpnet.h>
#include <string>
#include <set>

using namespace std;

#include "ddosproof.h"

static DEBUG_KEY D_DDOSPROOF ("ddosproof","Main program");
#if 0
// Replace with the flipset class
/*
	Remove all entries in works which are older than "back" seconds.
	See function above to see how a work is formatted using the time of day at the start

	Return the number of entries removed.
*/
static int ddosproof_purge (
	set<string> &works,
	int back)
{
	time_t now = time(NULL);
	char tmp[100];
	snprintf (tmp,sizeof(tmp),"%ld",now-back);
	set<string> newworks;
	for (set<string>::iterator it=works.begin(); it != works.end(); it++){
		if (strcmp(tmp,it->c_str())<0) newworks.insert (*it);
	}
	int ret = works.size() - newworks.size();
	works = newworks;
	return ret;
}
#endif

/*
	Read the content of the secret file. If missing, generates one.
*/
<mod>
static void ddosproof_readsecret (int fd, const char *fname, string &secret)
{
	if (file_exist(fname)){
		glocal string line;
		<call loadfile>(fname,true);
		<f oneline>
			glocal.line = line;
			return -1;
		</f>
		</call>
		if (glocal.line.size()==0){
			tlmp_error ("Invalid content in file %s, ending\n",fname);
			exit (-1);
		}
		secret = glocal.line;
	}else{
		char tmp[100];
		if (util_getwork(fd,tmp)!=0){
			tlmp_error ("Can't build a secret for file %s, exiting\n",fname);
		}else{
			glocal const char *line = tmp;
			<call savefile>(fname,false);
			<f dowrite>
				fprintf (fout,"%s\n",glocal.line);
				return 0;
			</f>
			<f openfail>
				tlmp_error ("Can't write secretfile %s (%s)\n",fname,strerror(errno));
				exit (-1);
			</f>
			</call>
			secret = tmp;
		}
	}
}
</mod>

// This is a set<string> with two bucket.
// This is used to do an efficient cleanup of a set. The idea is that we insert in one set for a while, then switch to the other.
// Once the other is filled, we clear the first and start to insert in this one. So old entry are deleted.
class flipset {
	int selset;		// Select which set is used to insert
	set<string> set1,set2;
	size_t maxset;		// Maximum size allow for a set
public:
	flipset(){
		selset = 0;
		maxset = 1000000;
	}
	void setmax(int _maxset){
		maxset = _maxset;
	}
	void flip(){
		if (selset == 0){
			selset = 1;
			set2.clear();
		}else{
			selset = 0;
			set1.clear();
		}
	}
	void insert (const string &s){
		if (selset == 0){
			set1.insert (s);
			if (set1.size() > maxset) flip();
		}else{
			set2.insert (s);
			if (set2.size() > maxset) flip();
		}
	}
	int count (const string &s){
		int ret = 0;
		if (selset == 0){
			ret = set1.count(s);
			if (ret == 0) ret = set2.count(s);
		}else{
			ret = set2.count(s);
			if (ret == 0) ret = set1.count(s);
		}
		return ret;
	}
	size_t size() const {
		return set1.size() + set2.size();
	}
	void erase (const string &s){
		set1.erase(s);
		set2.erase(s);
	}
	class iterator{
		set<string>::iterator sit;
		int sel;
		const flipset *f;
	public:
		iterator (int _sel, const flipset *_f){
			sel = _sel;
			f = _f;
			if (sel == 2){
				sit = f->set2.end();
			}else if (f->set1.size() > 0){
				sit = f->set1.begin();
			}else if (f->set2.size() > 0){
				sit = f->set2.begin();
				sel = 1;
			}else{
				sit = f->set2.end();
				sel = 2;
			}
		}
		void operator ++(int v ){
			if (sel == 0){
				sit++;
				if (sit == f->set1.end()){
					sel = 1;
					sit = f->set2.begin();
					if (sit == f->set2.end()){
						sel = 2;
					}
				}
			}else if (sel == 1){
				if (sit != f->set2.end()){
					sit++;
					if (sit == f->set2.end()){
						sel = 2;
					}
				}
			}
		}
		bool operator != (const iterator &s) const {
			return sel != s.sel || sit != s.sit;
		}
		set<string>::iterator operator -> () const{
			return sit;
		}
					
	};
	iterator begin() const {
		return iterator(0,this);
	}
	iterator end() const {
		return iterator(2,this);
	}
		
};

<mod>
int main (int argc, char *argv[])
{
	glocal int ret = -1;
	glocal const char *tcpbind = "0.0.0.0";
	glocal const char *tcpport = "none";
	glocal const char *unixport = "/var/run/ddosproof.sock";
	glocal const char *secretfile = "/var/run/ddosproof.secret";
	glocal int difficulty=10;
	glocal int maxworks = 1000000;
	glocal int maxbadips = 1000000;
	glocal int maxbadcookies = 1000000;
	glocal bool rejectcookies = false;
	glocal.ret = <call tlmpprogram>(argc,argv);
	<f init>
		setarg (' ',"bind","TCP bind on IP",glocal.tcpbind,false);
		setarg (' ',"tcpport","TCP port to manage proofs",glocal.tcpport,false);
		setarg (' ',"unixport","Unix socket to manage proofs",glocal.unixport,false);
		setarg (' ',"secretfile","File holding the secret to assemble cookies",glocal.secretfile,false);

		setgrouparg ("Tuning");
		setarg (' ',"maxworks","Maximum work in progress entries",glocal.maxworks,false);
		setarg (' ',"maxbadips","Maximum bad IP numbers",glocal.maxbadips,false);
		setarg (' ',"maxbadcookies","Maximum bad cookies",glocal.maxbadcookies,false);
		setarg (' ',"difficulty","Initial difficulty for proof of work",glocal.difficulty,false);
		setarg (' ',"rejectcookies","Reject all cookies (test mode)",glocal.rejectcookies,false);
		
	</f>
	<f main_noarg>
		signal (SIGPIPE,SIG_IGN);
		util_filltb0();
		int ret = -1;
		glocal int fdurandom;
		glocal flipset works;
		glocal flipset badips;
		glocal flipset badcookies;
		glocal string secret;
		glocal unsigned hitbadwork=0;
		glocal unsigned hitbadip=0;
		glocal unsigned hitbadcookie=0;

		glocal.works.setmax (glocal.maxworks);
		glocal.badips.setmax (glocal.maxbadips);
		glocal.badcookies.setmax (glocal.maxbadcookies);

		glocal.fdurandom = open ("/dev/urandom",O_RDONLY,0);
		if (glocal.fdurandom == -1){
			tlmp_error ("Can't open /dev/urandom (%s),exiting\n",strerror(errno));
			exit (-1);
		}
		ddosproof_readsecret (glocal.fdurandom,glocal.secretfile,glocal.secret);
		string tmp = string("unix:")+glocal.unixport;
		<obj TCPSERVER o>(tmp.c_str(),1);
		<f newclient>
			settcpnodelay(true);
		</f>
		<f endclient>
		</f>
		<f receive>
			debug_printf (D_DDOSPROOF,"Receive: %s\n",line);
			SSTRINGS tb;
			int n = str_splitline (line,' ',tb);
			if (n > 0){
				const char *cmd = tb.getitem(0)->c_str();
				if (n==1 && strcmp(cmd,"status")==0){
					sendf ("secret %s\n",glocal.secret.c_str());
					sendf ("difficulty %d\n",glocal.difficulty);
					sendf ("badips %u\n",glocal.badips.size());
					sendf ("works %u\n",glocal.works.size());
					sendf ("badcookies %u\n",glocal.badcookies.size());
					sendf ("hitbadip %u\n",glocal.hitbadip);
					sendf ("hitbadcookie %u\n",glocal.hitbadcookie);
					sendf ("hitbadwork %u\n",glocal.hitbadwork);
					sendf ("rejectcookies %s\n",glocal.rejectcookies ? "on" : "off");
					sendf ("%d ok\n",DDOSCODE_OK);
				}else if (n==3 && strcmp(cmd,"getwork")==0){
					// getwork ip cookie
					// Check if the cookie is valid, return "ok go"
					// Check if the IP is banned, if so return "ok ignore". We just end the connection
					// Check if the IP is working already. If so, send "ok ignore"
					// Check the result (proof of work)
					const char *ip = tb.getitem(1)->c_str();
					const char *cookie = tb.getitem(2)->c_str();
					if (glocal.badips.count(ip)>0){
						sendf ("%d reject\n",DDOSCODE_BADIP);
						glocal.hitbadip++;
					}else{
						int countcookie = glocal.badcookies.count(cookie);
						 if (!glocal.rejectcookies
							&& countcookie==0
							&& util_validcookie(cookie,ip,glocal.secret)){
							sendf ("%d go\n",DDOSCODE_COOKIEOK);
						}else{
							if (countcookie > 0) glocal.hitbadcookie++;
							char tmp[100];
							if (util_getwork(glocal.fdurandom,tmp)!=-1){
								glocal.works.insert (tmp);
								sendf ("%d %d %s\n",DDOSCODE_GOTOWORK,glocal.difficulty,tmp);
							}else{
								sendf ("%d fail\n",DDOSCODE_INTERNAL);
							}
						}
					}
				}else if (n==3 && strcmp(cmd,"validate")==0){
					// validade ip work_to_do answer
					const char *ip = tb.getitem(1)->c_str();
					if (glocal.badips.count(ip)>0){
						sendf ("%d reject\n",DDOSCODE_BADIP);
					}else{
						const char *answer = tb.getitem(2)->c_str();
						const char *pt = strchr(answer,',');
						// printf ("answer=%s pt=%p\n",answer,pt);
						if (pt == NULL){
							sendf ("%d ignore no comma\n",DDOSCODE_BADPROOF);
							glocal.badips.insert (ip);
							glocal.hitbadwork++;
						}else{
							string base (answer,pt-answer);
							if (glocal.works.count(base)==0){
								// We never ask for this work
								glocal.badips.insert (ip);
								sendf ("%d reject\n",DDOSCODE_BADPROOF);
								glocal.hitbadwork++;
							}else  if (util_validwork (answer,glocal.difficulty)){
								glocal.works.erase (base);	// Can't be reused
								char tmp[100];
								if (util_makecookie (glocal.fdurandom,glocal.secret,ip,tmp)==-1){
									sendf ("%d fail\n",DDOSCODE_INTERNAL);
								}else{
									sendf ("%d %s\n",DDOSCODE_GOODWORK,tmp);
								}
							}else{
								sendf ("%d invalid\n",DDOSCODE_BADPROOF);
								glocal.badips.insert (ip);
								glocal.hitbadwork++;
							}
						}
					}
				}else if (n==1 && strcmp(cmd,"purge")==0){
					size_t oldsize = glocal.works.size();
					glocal.works.flip();
					size_t newsize = glocal.works.size();
					int removed = oldsize - newsize;
					sendf ("ok %d\n",removed);
				}else if (n==1 && strcmp(cmd,"getcookie")==0){
					// For testing purpose
					char tmp[100];
					if (util_makecookie (glocal.fdurandom,glocal.secret,"127.0.0.1",tmp)==-1){
						sendf ("%d fail\n",DDOSCODE_INTERNAL);
					}else{
						sendf ("%d %s\n",DDOSCODE_OK,tmp);
					}
				}else if (n==2 && strcmp(cmd,"badcookie")==0){
					glocal.badcookies.insert (tb.getitem(1)->c_str());
					sendf ("%d ok\n",DDOSCODE_OK);
				}else if (n==2 && strcmp(cmd,"badip")==0){
					glocal.badips.insert (tb.getitem(1)->c_str());
					sendf ("%d ok\n",DDOSCODE_OK);
				}else if (n==2 && strcmp(cmd,"difficulty")==0){
					glocal.difficulty = atoi(tb.getitem(1)->c_str());
					sendf ("%d ok\n",DDOSCODE_OK);
				}else if (n==1 && strcmp(cmd,"badcookies")==0){
					for (flipset::iterator it=glocal.badcookies.begin(); it != glocal.badcookies.end(); it++){
						sendf ("cookie %s\n",it->c_str());
					}
					sendf ("%d ok\n",DDOSCODE_OK);
				}else if (n==1 && strcmp(cmd,"works")==0){
					for (flipset::iterator it=glocal.works.begin(); it != glocal.works.end(); it++){
						sendf ("work %s\n",it->c_str());
					}
					sendf ("%d ok\n",DDOSCODE_OK);
				}else if (n==1 && strcmp(cmd,"badips")==0){
					for (flipset::iterator it=glocal.badips.begin(); it != glocal.badips.end(); it++){
						sendf ("ip %s\n",it->c_str());
					}
					sendf ("%d ok\n",DDOSCODE_OK);
				}else if (n==1 && strcmp(cmd,"flipworks")==0){
					glocal.works.flip();
					sendf ("%d ok\n",DDOSCODE_OK);
				}else if (n==1 && strcmp(cmd,"flipbadips")==0){
					glocal.badips.flip();
					sendf ("%d ok\n",DDOSCODE_OK);
				}else if (n==1 && strcmp(cmd,"flipbadcookies")==0){
					glocal.badcookies.flip();
					sendf ("%d ok\n",DDOSCODE_OK);
				}else if (n==1 && strcmp(cmd,"resethit")==0){
					glocal.hitbadip = 0;
					glocal.hitbadcookie = 0;
					glocal.hitbadwork = 0;
					sendf ("%d ok\n",DDOSCODE_OK);
				}else if (n==2 && strcmp(cmd,"rejectcookies")==0){
					const char *val = tb.getitem(1)->c_str();
					if (strcmp(val,"on")==0){
						glocal.rejectcookies = true;
						sendf ("%d ok\n",DDOSCODE_OK);
					}else if (strcmp(val,"off")==0){
						glocal.rejectcookies = false;
						sendf ("%d ok\n",DDOSCODE_OK);
					}else{
						sendf ("%d bad value, on/off expected\n",DDOSCODE_INVALID);
					}
				}else if (n==1 && strcmp(cmd,"quit")==0){
					sendf ("%d bye\n",DDOSCODE_OK);
					endclient = true;
				}else if (n==1 && strcmp(cmd,"help")==0){
					send (
						"badip some_ip\n"
						"badips\n"
						"badcookie a_cookie\n"
						"badcookies\n"
						"difficulty value\n"
						"flipbadips\n"
						"flipbadcookies\n"
						"flipworks\n"
						"getwork ip cookie\n"
						"help\n"
						"quit\n"
						"nocookies on/off\n"
						"resethit\n"
						"status\n"
						"validate ip work\n"
						"works\n"
						);
			
					sendf ("%d ok\n",DDOSCODE_OK);
				}else{
					tlmp_error ("Invalid command: %s\n",line);
					endclient = true;
				}
			}
		</f>
		</obj>
		if (strcmp(glocal.tcpport,"none")!=0){
			o.listen (glocal.tcpbind,glocal.tcpport);
		}
		if (o.is_ok()){
			o.loop();
			ret = 0;
		}else{
			tlmp_error ("Can't setup network service, ending\n");
		}
		return ret;
	</f>
	</call>
	return glocal.ret;
}
</mod>