#include <u.h>
#include <libc.h>
#include <bio.h>
#include <ndb.h>

char *
portname(Ndb * db, int n)
{
	Ndbs s;
	Ndbtuple *t, *nt;
	char port[32], *p, *attr = "port";

	snprint(port, sizeof(port), "%d", n);
	t = ndbsearch(db, &s, attr, port);
	while(t){
		for(nt = t; nt; nt = nt->entry){
			if(strcmp(nt->attr, "tcp") == 0){
				p = strdup(nt->val);
				while(t){
					ndbfree(t);
					t = ndbsnext(&s, attr, port);
				}
				return p;
			}
		}
		ndbfree(t);
		t = ndbsnext(&s, attr, port);
	}

	return nil;
}

static void
ding(void *u, char *msg)
{
	USED(u);

	if(strstr(msg, "alarm"))
		noted(NCONT);
	noted(NDFLT);
}

vlong
probe(vlong rtt, char *host, char *net, int n)
{
	int fd;
	vlong now, then;
	char *addr, port[32];

	snprint(port, sizeof(port), "%d", n);
	addr = netmkaddr(host, net, port);

	alarm(rtt*3.1);
	then = nsec();
	fd = dial(addr, 0, 0, 0);
	now = nsec();
	alarm(0);
	if(fd == -1)
		return -1;
	close(fd);

	return (rtt*0.66) + ((now-then)/1000000LL)*0.34;
}

	
void
usage(void)
{
	fprint(2, "usage: %s [-s start-port] [-e end-port] [-t timeout] [net!]host\n", argv0);
	exits("usage");
}

void
main(int argc, char *argv[])
{
	Ndb *db;
	vlong r, rtt;
	int verbose, i, t, x, y, n, s, e;
	char *p, *host, *net;
	int *ptab;

	if((db = ndbopen(nil)) == nil)
		sysfatal("cannot open ndb - %r\n");

	s = 1;
	rtt = 500;
	e = 0xffff;
	verbose = 0;
	ARGBEGIN{
	case 'v':
		verbose++;
		break;
	case 's':
		s = atoi(EARGF(usage()));
		break;
	case 'e':
		e = atoi(EARGF(usage()));
		break;
	case 't':
		rtt = atol(EARGF(usage()));
		break;
	default:
		usage();
	}
	ARGEND;


	if(argc != 1)
		usage();

	net = "tcp";
	host = argv[0];
	if((p = strchr(host, '!')) != nil){
		*p = 0;
		net = host;
		host = p+1;
	}

	notify(ding);
	srand(time(nil) ^ getpid());

	if(s < 1 || s > e)
		s = 1;
	if(e < 0 || e < s || e > 0xffff)
		e = 0xffff;

	ptab = malloc(sizeof(int) * ((e-s)+1));
	if(ptab == nil)
		sysfatal("no memory for port table\n");

	for(i = 0; i < ((e-s)+1); i++)
		ptab[i] = s+i;
	for(i = 0; i < ((e-s)+1)*8; i++){
		x = (rand() % ((e-s)+1));
		y = (rand() % ((e-s)+1));
		t = ptab[x];
		ptab[x] = ptab[y];
		ptab[y] = t;
	}

	for(n = 0; n < ((e-s)+1); n++){
		r = probe(rtt, host, net, ptab[n]);
		if(r > 0 || verbose){
			if(r > 0){
				print("open     %-6d %-12s %lldms\n", ptab[n], portname(db, ptab[n]), r);
				rtt = r;
			}
			else
				print("shut     %-6d %-12s %lldms\n", ptab[n], portname(db, ptab[n]), rtt);
		}
	}

	ndbclose(db);
	exits(0);
}