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

int nosnmp = 0;
int nonames = 0;
int nobootp = 0;
int nodomain = 0;
int nonetbios = 0;
int norip = 0;
char *netroot = "/net";

#define BGET16(p)	((p)[0]<<8 | (p)[1])
#define BPUT16(p, x)	((p)[1] = (x), (p)[0] = ((x)>>8))

enum {
	NBstat		= 0x21,
	NBin		= 1,
};

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

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

static int
nbname(char *rv, char *name, char pad)
{
	char *p, c;
	int i;
	int done=0;

	p = rv;
	*p++ = 0x20;
	for(i=0; i<16; i++) {
		c = pad;
		if(!done && name[i] == '\0')
			done = 1;
		if(!done)
			c = toupper(name[i]);
		*p++ = (c >> 4) + 'A';
		*p++ = (c & 0xf) + 'A';
	}
	*p++ = 0x0;
	return (p-rv);
}

char *
strlwr(char *str)
{
	char *p = str;
	for (; *p; p++)
		if (isupper(*p))
			*p = tolower(*p);
	return str;
}

char *
nbhost(char *host)
{
	int num, n, i, fd, trn;
	char *addr, buf[520], *p;
	static char name[20];
	char *wildcard = " CKAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA";

	trn = (getpid() ^ time(0)) & 0xffff;

	if ((addr = netmkaddr(host, "udp", "netbios-ns")) == nil)
		sysfatal("%s can't make address - %r", host);

	if((fd = dial(addr, "netbios-ns", 0, 0)) < 0)
		sysfatal("%s can't dial - %r", addr);
	n = 0;
	p = buf;
	memset(buf, 0, sizeof(buf));

	BPUT16(p+n, trn);	n+= 2;	// TRNid
	BPUT16(p+n, 0);		n+= 2;	// flags/type
	BPUT16(p+n, 0);		n+= 2;	// # questions
	BPUT16(p+n, 0);		n+= 2;	// # answers
	BPUT16(p+n, 0);		n+= 2;	// # authority RRs
	BPUT16(p+n, 0);		n+= 2;	// # Aditional RRs
	strcpy(p+n, wildcard);	n+= strlen(wildcard) +1;
	BPUT16(p+n, NBstat);	n+= 2;
	BPUT16(p+n, NBin);	n+= 2;

	if (write(fd, buf, n) != n)
		return nil;

	notify(ding);

	memset(buf, 0, sizeof(buf));

	for (i = 0; i < 3; i++){
		alarm(300);
		n = read(fd, buf, sizeof(buf));
		alarm(0);
		if (BGET16(p) == trn)
			break;
	}
	close(fd);
	if (n < 64)
		return "?";

	p = buf +56;
	num = *p++;			// number of names

	memset(name, 0, sizeof(name));
	for (i = 0; i < num; i++){
		memcpy(name, p, 16);
		p += 18;
		if (name[15] == 0){
			name[15] = 0;
			return strlwr(name);
		}
	}

	return "?";
}

void
usage(void)
{
	fprint(2, "usage: %s [-b] [-d] [-n] [-r] [-N] [-s] [-x mntpoint]\n", argv0);
	fprint(2, "  -N     don't resolve DNS names\n");
	fprint(2, "  -d     ignore DNS packets\n");
	fprint(2, "  -b     ignore BootP packets\n");
	fprint(2, "  -n     ignore NetBios packets\n");
	fprint(2, "  -r     ignore RIP packets\n");
	fprint(2, "  -s     ignore snmp packets\n");
	fprint(2, "  -x xx  alternative network mountpoint\n");
	exits("usage");
}

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;
		if ((p = strchr(buf, '\n')) != nil)
			*p = 0;
		return buf;
	}
	return "";
}

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

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

int skip(char c, int l, int r)
{
	if (nobootp && c == 'u' && l == 68 && r == 67)
		return 1;
	if (nobootp && c == 'u' && l == 67 && r == 68)
		return 1;
	if (nodomain && c == 'u' && r == 53)
		return 1;
	if (norip && c == 'u' && l == 520 && r == 520)
		return 1;
	if (nonetbios && c == 'u' && l == 137 && r == 137)
		return 1;
	if (nonetbios && c == 'u' && l == 138 && r == 138)
		return 1;
	if (nosnmp && c == 'u' && l == 161)
		return 1;
	return 0;
}


void
guard(char *proto)
{
	long when;
	int ctl, nctl;
	char *p, *ip, *tstr, *lprt, *rprt, *sys, *rem, *loc;
	char addr[16], dir[40], ndir[40];

	snprint(addr, sizeof(addr), "%s!*!*", proto);
	if ((ctl = announce(addr, dir)) < 0)
		sysfatal("announce %s: %r", addr);

	while(1){
		sys = lprt = rprt = "?";
		if((nctl = listen(dir, ndir)) < 0)
			sysfatal("listen %s: %r", addr);
		when = time(0);
		ip = remoteaddr(ndir);
		loc = localport(ndir);
		reject(ctl, ndir, "unexpected");
		close(nctl);

		if ((rem = strchr(ip, '!')) != nil)
			*rem++ = 0;

		if (skip(*proto, atoi(loc), atoi(rem)))
			continue;

		tstr = ctime(when);
		tstr[strlen(tstr) -1] = 0;

		print("%s %4s %-16s %-4s -> %-4s", tstr +4, proto, ip, rem, loc);

		if (! nonames){
			if ((p = csgetvalue(netroot, "ip", ip, "dom", nil)) != nil)
				sys = p;
			else
				switch(atoi(loc)){
				case 137:
				case 138:
				case 139:
					sys = nbhost(ip);
					break;
				default:
					break;
				}

			if ((p = csgetvalue(netroot, "port", loc, proto, nil)) != nil)
				lprt = p;
			if ((p = csgetvalue(netroot, "port", rem, proto, nil)) != nil)
				rprt = p;
			print("   %s %s -> %s", sys, rprt, lprt);
		}
		print("\n");
	}
}

void
main(int argc, char *argv[])
{
	int pid;

	ARGBEGIN{
	case 's':
		nosnmp++;
	case 'r':
		norip++;
		break;
	case 'n':
		nonetbios++;
		break;
	case 'N':
		nonames++;
		break;
	case 'x':
		netroot = EARGF(usage());
		break;
	case 'b':
		nobootp++;
		break;
	case 'd':
		nodomain++;
		break;
	default:
		usage();
		break;
	}ARGEND;

	if (argc)
		usage();

	if ((pid = fork()) == -1)
		sysfatal("fork failed: %r");
	if (pid)
		guard("tcp");
	else
		guard("udp");
	exits(nil);
}