#include <stdio.h>
#include <stdarg.h>
#include <ctype.h>
#include <stdlib.h>
#include <string.h>
#include <syslog.h>
#include <signal.h>
#include <errno.h>
#include <netdb.h>
#include <unistd.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <time.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <limits.h>
#include <dialog.h>
#include <tlmplib.h>
#include <tlmpnet.h>
#include "fdpass.h"
#include <vector>
#include <string>

using namespace std;

static DEBUG_KEY D_UDPPROXY ("udpproxy","udpproxy program");

struct DAEMON_INFO{
	bool daemon;
	const char *pidfile;
	const char *user;
	const char *unixsocket;
	DAEMON_INFO (){
		daemon = false;
		pidfile = "/var/run/udpproxy.pid";
		user = "blackhole";
		unixsocket = "/var/run/blackhole/udpproxy.sock";
	}
};

enum CONNECTION_TYPE{
	TYPE_UNKNOWN,
	TYPE_CLIENT,
	TYPE_CONTROL
};

struct HANDLE_INFO: public ARRAY_OBJ{
	CONNECTION_TYPE type;
	HANDLE_INFO(){
		type = TYPE_UNKNOWN;
	}
};

struct UDPREQUEST: public HANDLE_INFO{
	UDPDEST from;
	UDPREQUEST(const UDPDEST &_from){
		from = _from;
		type = TYPE_CLIENT;
	}
};

<mod>
static int udp_server(
	DAEMON_INFO &dinfo,
	const char *udphost,	// Listen on udpport on this address
	const char *udpport,
	const char *bindaddr,	// Connect to tcphost/tcpport bound to this address
	const char *tcphost,
	const char *tcpport)
{
	glocal const char *bindaddr = bindaddr;
	glocal const char *tcphost = tcphost;
	glocal const char *tcpport = tcpport;
	glocal _F_UDPSERVER *userv = NULL;
	glocal TCPSERVER_V1 *tserv = NULL;
	glocal int nbreq = 0;
	<obj UDPSERVER userv>(udphost,udpport);
	<f receive>
		glocal.userv = this;
		glocal.nbreq++;
		debug_printf (D_UDPPROXY,"len = %d from %lu.%lu.%lu.%lu\n",len
			,(from.addr>>24)&0xff
			,(from.addr>>16)&0xff
			,(from.addr>>8)&0xff
			,from.addr&0xff);
		int fd = cmdsock_connect (glocal.bindaddr,glocal.tcphost,glocal.tcpport,1,1);
		if (fd == -1){
			tlmp_error ("udp_server: Can't connect to %s port %s (%s)\n"
				,glocal.tcphost,glocal.tcpport,strerror(errno));
		}else{
			write (fd,buf,len);
			UDPREQUEST *r = new UDPREQUEST(from);
			glocal.tserv->inject(fd,r);
		}
	</f>
	</obj>
	int ret = -1;
	if (userv.is_ok()){
		<obj TCPSERVER tserv>();
		<f newclient>
			HANDLE_INFO *c = new HANDLE_INFO;
			info.data = c;
			if (strncmp(info.port,"unix:",5)==0){
				setrawmode (false);
				c->type = TYPE_CONTROL;
			}else{
				c->type = TYPE_CLIENT;
			}
		</f>
		<f endclient>
		</f>
		<f receive>
			HANDLE_INFO *c = (HANDLE_INFO*)info.data;
			if (c->type == TYPE_CONTROL){
				if (strcmp(line,"status")==0){
					sendf ("nbreq=%d\n",glocal.nbreq);
					send ("Ok\n");
				}else{
					send ("Invalid command\n");
				}
			}else{
				UDPREQUEST *r = (UDPREQUEST*)info.data;
				int ok = glocal.userv->sendto (r->from,line,info.linelen);
				debug_printf (D_UDPPROXY,"sendto ok=%d\n",ok);
			}
			endclient = true;
		</f>
		</obj>
		if (fdpass_setcontrol(tserv,dinfo.unixsocket,dinfo.user)!=-1){
			tserv.setrawmode(true);
			ret = 0;
			if (dinfo.daemon){                                                                                                                                                                                                           
				daemon_init(dinfo.pidfile,dinfo.user);                                                                                                                                                                               
			}  
			<obj NETEVENT_MANAGER ev>();
			<f idle>
			</f>
			</obj>
			ev.add (userv);
			ev.add (tserv);
			glocal.tserv = &tserv;
			ev.loop(5);
		}
	}
	return ret;
}
</mod>

static void fsig (int)
{
        // fprintf (stderr,"signal\n");
}

/*
        Install the signal handler and set the alarm
*/
static void udpproxy_setsig(int time_out)
{
        static struct sigaction siga;
        siga.sa_handler = fsig;
        sigaction (SIGALRM,&siga,NULL);
        alarm (time_out);
}

/*
        Remove the signal handler and reset the alarm
*/
static void udpproxy_resetsig ()
{
        static struct sigaction siga;
        siga.sa_handler = SIG_DFL;
        sigaction (SIGALRM,&siga,NULL);
        alarm (0);
}

<mod>
static int udp_client (
	DAEMON_INFO &dinfo,
	const char *tcphost,	// Listen on TCP tcphost/tcpport
	const char *tcpport,
	const vector<string> &udphost,	// Connects to udphost/udpport	(round robin)
	const char *from,	// When connecting to udphost, bind to address from
	const char *udpport)
{
	glocal const vector<string> *udphost = &udphost;
	glocal const char *udpport = udpport;
	glocal const char *from = from;
	glocal int nbreq = 0;
	glocal int nbtimeout = 0;	// How many time recvfrom failed.
	glocal int nbsendtoerr = 0;	// How many time sendto failed.
	glocal unsigned sel_udphost = 0;


	<obj TCPSERVER o>(tcphost,tcpport,5);
	<f newclient>
		HANDLE_INFO *c = new HANDLE_INFO;
		info.data = c;
		if (strncmp(info.port,"unix:",5)==0){
			setrawmode (false);
			c->type = TYPE_CONTROL;
		}else{
			c->type = TYPE_CLIENT;
		}
	</f>
	<f endclient>
	</f>
	<f receive>
		HANDLE_INFO *c = (HANDLE_INFO*)info.data;
		if (c->type == TYPE_CONTROL){
			if (strcmp(line,"status")==0){
				sendf ("nbreq=%d\n",glocal.nbreq);
				sendf ("nbtimeout=%d\n",glocal.nbtimeout);
				sendf ("nbsendtoerr=%d\n",glocal.nbsendtoerr);
				send ("Ok\n");
			}else{
				send ("Invalid command\n");
			}
			endclient = true;
		}else if (c->type == TYPE_CLIENT){
			glocal.nbreq++;
			struct sockaddr_in sin;
			sin.sin_family = AF_INET;
			struct servent *s = getservbyname(glocal.udpport,"udp");
			if (s != NULL){
				sin.sin_port = s->s_port;
			}else{
				sin.sin_port = htons (atoi(glocal.udpport));
			}
			const char *udphost = (*glocal.udphost)[glocal.sel_udphost].c_str();
			glocal.sel_udphost++;
			if (glocal.sel_udphost == glocal.udphost->size()) glocal.sel_udphost=0;
			struct hostent *h = gethostbyname (udphost);
			memcpy  (&sin.sin_addr,h->h_addr, h->h_length);
	
			int opt = 1;
			int fd = socket (AF_INET, SOCK_DGRAM, 0);
			if (fd == -1){
				tlmp_error ("socket error (%s)\n",strerror(errno));
				endclient = true;
			}else if (setsockopt(fd,SOL_SOCKET,SO_BROADCAST,&opt,sizeof(opt))==-1){
				tlmp_error ("setsockopt error (%s)\n",strerror(errno));
				endclient = true;
			}else{
				if (glocal.from != NULL){
					struct hostent *f = gethostbyname (glocal.from);
					struct sockaddr_in s;
					s.sin_family = AF_INET;
					memcpy  (&s.sin_addr,f->h_addr, f->h_length);
					if (bind (fd,(struct sockaddr *) &s, sizeof (s)) == -1){
						tlmp_error ("Can't bind %d(%s)\n",errno,strerror(errno));
						endclient=true;
					}
				}
				int ok = ::sendto (fd,line,info.linelen,0,(struct sockaddr *)&sin,sizeof(sin));
				if (ok <= 0){
					tlmp_error ("sendto %d (%s)\n",ok,strerror(errno));
					glocal.nbsendtoerr++;
				}else{
					struct sockaddr_in cli;
					socklen_t len = sizeof(cli);
					char buf[1000];
					udpproxy_setsig(10);
					ok = recvfrom (fd,buf,sizeof(buf),0,(struct sockaddr*)&cli,&len);
					udpproxy_resetsig();
					debug_printf (D_UDPPROXY,"recvfrom ok=%d\n",ok);
					if (ok >= 0){
						send (buf,ok);
					}else if (ok == -1){
						//tlmp_error ("recvfrom ends with signal\n");
						glocal.nbtimeout++;
					}
				}
			}
			if (fd != -1) close (fd);
			endclient = true;
		}else{
			tlmp_error ("udp_client, unknown connection type ???\n");
			endclient = true;
		}
	</f>
	</obj>
	int ret = -1;
	if (o.is_ok()){
		if (fdpass_setcontrol(o,dinfo.unixsocket,dinfo.user)!=-1){
			ret = 0;
			o.setrawmode (true);
			if (dinfo.daemon){
				daemon_init (dinfo.pidfile,dinfo.user);
			}
			o.loop();
		}
	}
	return ret;
}
</mod>


<mod>
int main (int argc, char *argv[])
{
	glocal int ret = -1;
	glocal bool udp2tcp = false;
	glocal bool tcp2udp = false;
	glocal bool status = false;
	glocal const char *udpport = NULL;
	glocal vector<string> udphost;
	glocal const char *bindaddr = NULL;
	glocal const char *tcpport = NULL;
	glocal const char *tcphost = NULL;
	glocal DAEMON_INFO dinfo;

	glocal.ret = <call tlmpprogram>(argc,argv);
	<f init>
		setproginfo ("udpproxy",VERSION,"UDP proxy through blackhole, back and forth");
		setgrouparg ("Operation mode");
		setarg (' ',"udp2tcp","Proxy UDP request through TCP",glocal.udp2tcp,false);
		setarg (' ',"tcp2udp","Proxy TCP request through UDP",glocal.tcp2udp,false);
		setarg (' ',"status","get daemon status",glocal.status,false);
		setgrouparg ("Networking");
		setarg (' ',"udpport","UDP port to listen or connect",glocal.udpport,false);
		setarg (' ',"tcpport","TCP port to listen or connect",glocal.tcpport,false);
		setarg (' ',"udphost","UDP host to listen or connect",glocal.udphost,false);
		setarg (' ',"tcphost","TCP host to connect",glocal.tcphost,false);
		setarg (' ',"bindaddr","Client address for udp2tcp",glocal.bindaddr,false);

		setgrouparg ("Misc.");
		setarg (' ',"daemon","Operate in background",glocal.dinfo.daemon,false);
		setarg (' ',"pidfile","In daemon mode, will hold the PID",glocal.dinfo.pidfile,false);
		setarg (' ',"control","Unix socket to get status",glocal.dinfo.unixsocket,false);
		setarg (' ',"user","Run as this user",glocal.dinfo.user,false);
	</f>
	<f showerror>
		if (glocal.dinfo.daemon){
			syslog (LOG_ERR,"%s",msg);
		}else{
			fprintf (stderr,"%s",msg);
		}
	</f>
	<f main_noarg>
		int ret = -1;
		if (glocal.status){
			ret = <call tcpconnect>("unix:",glocal.dinfo.unixsocket,5);
			<f init>
				sendf ("status\n");
			</f>
			<f oneline>
				if (strcmp(line,"Ok")==0){
					end = true;
				}else{
					printf ("%s\n",line);
				}
			</f>
			<f end>
			</f>
			</call>
		}else if (glocal.udp2tcp && glocal.tcp2udp){
			tlmp_error ("udp2tcp may not be used with tcp2udp\n");
		}else if (!glocal.udp2tcp && !glocal.tcp2udp){
			tlmp_error ("One of the following option is required: --udp2tcp or --tcp2udp\n");
			usage();
		}else if (glocal.tcpport == NULL){
			tlmp_error ("Option --tcpport is mandatory\n");
		}else if (glocal.udpport == NULL){
			tlmp_error ("Option --udpport is mandatory\n");
		}else{
			if (glocal.udp2tcp){
				const char *udphost = "0.0.0.0";
				if (glocal.udphost.size()==1) udphost = glocal.udphost[0].c_str();
				ret = udp_server (glocal.dinfo,udphost,glocal.udpport,glocal.bindaddr,glocal.tcphost,glocal.tcpport);
			}else if (glocal.tcp2udp){
				if (glocal.tcphost == NULL) glocal.tcphost = "0.0.0.0";
				if (glocal.udphost.size()==0){
					tlmp_error ("udphost must be defined\n");	
				}else{
					ret = udp_client (glocal.dinfo,glocal.tcphost,glocal.tcpport,glocal.udphost,NULL,glocal.udpport);
				}
			}
		}
		return ret;
	</f>
	</call>
	return glocal.ret;
}
</mod>