/*
*	Plisten1 is a listen1 that protect itself from brute force attacks.
*	Most of the code came from official distribution from Bell-labs
*	/sys/src/cmd/aux/listen1.c
*	ver.1.2
*	-Kenar-
*/

#include <u.h>
#include <libc.h>
#include <auth.h>

typedef struct laccess laccess;
struct laccess{
	char ip[16];
	long time;
};

int maxconnect = 0;
#define OTIME 10

int verbose;
int trusted;

char *acceptdir = "/sys/log/accept";
char *rejectdir = "/sys/log/reject";
char *listenlog = "plisten1";
char *script;

void
usage(void)
{
	fprint(2, "usage1: plisten1 [-tv] [-m maxconnect] address cmd args...\n"
		"usage2: plisten1 [-tv] [-m maxconnect] -s script address\n"
	);
	exits("usage");
}

void
becomenone(void)
{
	int fd;

	fd = open("#c/user", OWRITE);
	if(fd < 0 || write(fd, "none", strlen("none")) < 0)
		sysfatal("can't become none: %r");
	close(fd);
	if(newns("none", nil) < 0)
		sysfatal("can't build namespace: %r");
}

char*
remoteaddr(char *dir)
{
	static char buf[128];
	char *p;
	int n, fd;

	snprint(buf, sizeof buf, "%s/remote", dir);
	fd = open(buf, OREAD);
	if(fd < 0)
		return "";
	n = read(fd, buf, sizeof(buf));
	close(fd);
	if(n > 0){
		buf[n] = 0;
		p = strchr(buf, '!');
		if(p)
			*p = 0;
		return buf;
	}
	return "";
}

void
main(int argc, char **argv)
{
	char data[60], dir[40], ndir[40], buf[256];
	char *args[6]; 
	int ctl, nctl, fd;
	laccess *lbuf = nil;
	int lnext = 0;
	int lbufsize;
	int accepted;
	long now;
	char *raddr;
	int i,n;

	ARGBEGIN{
	case 't':
		trusted = 1;
		break;
	case 'm':
		maxconnect = atoi(EARGF(usage()));
		break;
	case 'v':
		verbose = 1;
		break;
	case 's':
		script = EARGF(usage());
		break;
	default:
		usage();
	}ARGEND


	if(script){
		char *p,*t;
		if(argc != 1)
			usage();
		args[0] = argv[0]; /* example: "tcp!*!110" */
		args[1] = script;
		args[2] = nil; /* to be serv, "tcp110" */
		t = strdup(argv[0]);
		p = strchr(t,'!');
		if(!p)
			usage();
		*p = 0;
		args[3] = t; /* proto, "tcp" */
		p++;
		p = strchr(p,'!');
		if(!p)
			usage();
		p++;
		args[2] = smprint("%s%s", t, p);
		args[4] = nil; /* to be ndir, "/net/tcp/33" */
		args[5] = nil; /* fix */
		argv = args;
	}
	else{
		if(argc < 2)
			usage();
	}

	if(!verbose){
		close(1);
		fd = open("/dev/null", OWRITE);
		if(fd != 1){
			dup(fd, 1);
			close(fd);
		}
	}

	if(!trusted)
		becomenone();

	syslog(0, listenlog, "started on %s", argv[0]);

	ctl = announce(argv[0], dir);
	if(ctl < 0){
		syslog(0, listenlog, "fatal: announce %s", argv[0]);
		sysfatal("announce %s: %r", argv[0]);
	}

	lbufsize = 0;
	if(maxconnect > 0){
		lbufsize = 2*maxconnect;
		lbuf = calloc(lbufsize, sizeof(laccess));
		if(lbuf == nil)
			sysfatal("lbuf: no memory");
	}

	for(;;){
		/*
		 *  wait for a call (or an error)
		 */
		nctl = listen(dir, ndir);
		if(nctl < 0){
			syslog(0, listenlog, "fatal: listen %s: %r", argv[0]);
			sysfatal("listen %s: %r", argv[0]);
		}
		/* then ndir is, say, "/net/tcp/36"
		*  and we get here remote address using remoteaddr(ndir)
		*  the output is, say, "192.168.1.3"	*/
		//print("%s\n",ndir);
		if(script)
			argv[4] = ndir;

		raddr = remoteaddr(ndir);

		/* reject burst access */
		if(lbufsize){
			/*	check if he did burst access */
			lbuf[lnext].time = now = time(nil);
			strncpy(lbuf[lnext].ip,raddr,16);
			lnext++;
			lnext %= lbufsize;
			n = 0;
			for(i = 0; i < lbufsize; i++)
				if(now - lbuf[i].time < OTIME && strcmp(lbuf[i].ip, raddr) == 0)
					n++;
			if(n > maxconnect){
				reject(nctl, ndir, "rejected");
				close(nctl);
				syslog(0,listenlog,"reject1 %s",raddr);
				sleep(1000);
				continue;
			}
		}

		/* check if he is in accepdir */
		accepted = 0;
		snprint(buf,sizeof(buf),"%s/%s", acceptdir,raddr);
		if(access(buf,0) == 0)
			accepted = 1;

		/* if he is not accepted then check if he is in rejectdir */
		if(!accepted){
			snprint(buf,sizeof(buf),"%s/%s", rejectdir,raddr);
			if(access(buf,0) == 0){
				reject(nctl, ndir, "rejected");
				close(nctl);
				syslog(0,listenlog,"reject2 %s",raddr);
				sleep(1000);
				continue;
			}
		}

		/*
		 *  start a subprocess for the connection
		 */
		switch(rfork(RFFDG|RFPROC|RFNOWAIT|RFENVG|RFNAMEG|RFNOTEG)){
		case -1:
			reject(nctl, ndir, "host overloaded");
			close(nctl);
			continue;
		case 0:
			fd = accept(nctl, ndir);
			if(fd < 0){
				fprint(2, "accept %s: can't open  %s/data: %r\n",
					argv[0], ndir);
				_exits(0);
			}
			syslog(0, listenlog, "accept %s in %s", raddr, ndir);
			fprint(nctl, "keepalive");
			close(ctl);
			close(nctl);
			putenv("net", ndir);
			snprint(data, sizeof data, "%s/data", ndir);
			bind(data, "/dev/cons", MREPL);
			dup(fd, 0);
			dup(fd, 1);
			/* dup(fd, 2); keep stderr */
			close(fd);
			exec(argv[1], argv+1);
			if(argv[1][0] != '/')
				exec(smprint("/bin/%s", argv[1]), argv+1);
			syslog(0, listenlog, "exec: %s %r", argv0);
			exits(nil);
		default:
			close(nctl);
			break;
		}
	}
}