#include <u.h>
#include <libc.h>
#include <bio.h>
#include <ctype.h>
#include <auth.h>
#include <mp.h>
#include <libsec.h>
#include <ip.h>
#include "/sys/src/cmd/ndb/dns.h"

enum {
	Cnone = 254,	/* rfc 2136 § 1.3 - add to dns.h */
};

enum {
	MaxRRs = 255,	/* arbitrary and can be raised up to 65536 if needed - rfc 2136 § 2.2 */
};

typedef struct Record {
	char	*name;
	u16int	type;
	u16int	class;
	u16int	ttl;
	u16int	rdlength;
	uchar	*rdata;
} Record;

typedef struct Signature {
	char	*kname;
	char	*aname;
	long	time;
	u16int	fudge;
	uchar	digest[MD5dlen];
} Signature;

typedef struct Msg {
	char	*zone;
	Record	**prereqs;	/* nil terminated */
	Record	**updates;	/* nil terminated */
	Signature	sig;
} Msg;

enum {
	Fi16 = 1,
	Fi32,
	Fip,
	Fipv6,
	Fstring,
	Ftxt,
	Fname,
};

typedef struct RdataFormat {
	int	rrtype;
	int	ftypes[8];	/* zero terminated */
} RdataFormat;

RdataFormat rdfmttab[] = {
	{Ta,	{Fip}},	/* rfc 1035 § 3.4.1 */
	{Taaaa,	{Fipv6}},	/* rfc 1886 */
	{Tafsdb,	{Fi16, Fname}},	/* rfc 1183 § 1 */
	{Tcname,	{Fname}},	/* rfc 1035 § 3.3.1 */
	{Thinfo,	{Fstring, Fstring}},	/* rfc 1035 § 3.3.2 */
	{Tmb,	{Fname}},	/* rfc 1035 § 3.3.3 */
	{Tmd,	{Fname}},	/* rfc 1035 § 3.3.4 */
	{Tmf,	{Fname}},	/* rfc 1035 § 3.3.5 */
	{Tmg,	{Fname}},	/* rfc 1035 § 3.3.6 */
	{Tminfo,	{Fname, Fname}},	/* rfc 1035 § 3.3.7 */
	{Tmr,	{Fname}},	/* rfc 1035 § 3.3.8 */
	{Tmx,	{Fi16, Fname}},	/* rfc 1035 § 3.3.9 */
	{Tns,	{Fname}},	/* rfc 1035 § 3.3.11 */
	{Tptr,	{Fname}},	/* rfc 1035 § 3.3.12 */
	{Trp,	{Fname, Fname}},	/* rfc 1183 § 2 */
	{Tsoa,	{Fname, Fname, Fi32, Fi32, Fi32, Fi32, Fi32}},	/* rfc 1035 § 3.3.13 */
	{Tspf,	{Ftxt}},	/* rfc 4408 § 3.1.1 */
	{Tsrv,	{Fi16, Fi16, Fi16, Fname}},	/* rfc 2782 */
	{Ttxt,	{Ftxt}},	/* rfc 1035 § 3.3.14 */
	{Tx25,	{Fstring}},	/* rfc 1183 § 3.1 */
};

enum {
	Sprereq,
	Supdate,
};

typedef struct VerbType {
	char	*name;
	int	needsrdata;
	int	section;
	int	class;
} VerbType;

VerbType verbtab[] = {
	{"yx",	0,	Sprereq,	Call},	/* rfc 2136 §§ 2.4.1, 2.4.4 */
	{"yx",	1,	Sprereq,	Cin},	/* rfc 2136 § 2.4.2 */
	{"nx",	0,	Sprereq,	Cnone},	/* rfc 2136 §§ 2.4.3, 2.4.5 */
	{"a",	1,	Supdate,	Cin},	/* rfc 2136 § 2.5.1 */
	{"d",	0,	Supdate,	Call},	/* rfc 2136 §§ 2.5.2, 2.5.3 */
	{"d",	1,	Supdate,	Cnone},	/* rfc 2136 §§ 2.5.4 */
};

/* copied from /sys/src/cmd/ndb/dn.c */
char *rrtname[] =
{
[Ta]		"ip",
[Tns]		"ns",
[Tmd]		"md",
[Tmf]		"mf",
[Tcname]	"cname",
[Tsoa]		"soa",
[Tmb]		"mb",
[Tmg]		"mg",
[Tmr]		"mr",
[Tnull]		"null",
[Twks]		"wks",
[Tptr]		"ptr",
[Thinfo]	"hinfo",
[Tminfo]	"minfo",
[Tmx]		"mx",
[Ttxt]		"txt",
[Trp]		"rp",
[Tafsdb]	"afsdb",
[Tx25]		"x.25",
[Tisdn]		"isdn",
[Trt]		"rt",
[Tnsap]		"nsap",
[Tnsapptr]	"nsap-ptr",
[Tsig]		"sig",
[Tkey]		"key",
[Tpx]		"px",
[Tgpos]		"gpos",
[Taaaa]		"ipv6",
[Tloc]		"loc",
[Tnxt]		"nxt",
[Teid]		"eid",
[Tnimloc]	"nimrod",
[Tsrv]		"srv",
[Tatma]		"atma",
[Tnaptr]	"naptr",
[Tkx]		"kx",
[Tcert]		"cert",
[Ta6]		"a6",
[Tdname]	"dname",
[Tsink]		"sink",
[Topt]		"opt",
[Tapl]		"apl",
[Tds]		"ds",
[Tsshfp]	"sshfp",
[Tipseckey]	"ipseckey",
[Trrsig]	"rrsig",
[Tnsec]		"nsec",
[Tdnskey]	"dnskey",
[Tspf]		"spf",
[Tuinfo]	"uinfo",
[Tuid]		"uid",
[Tgid]		"gid",
[Tunspec]	"unspec",
[Ttkey]		"tkey",
[Ttsig]		"tsig",
[Tixfr]		"ixfr",
[Taxfr]		"axfr",
[Tmailb]	"mailb",
[Tmaila]	"maila",
[Tall]		"all",
		0,
};

/* copied from /sys/src/cmd/ndb/dn.c */
char *rname[Rmask+1] =
{
[Rok]			"ok",
[Rformat]		"format error",
[Rserver]		"server failure",
[Rname]			"bad name",
[Runimplimented]	"unimplemented",
[Rrefused]		"we don't like you",
[Ryxdomain]		"name should not exist",
[Ryxrrset]		"rr set should not exist",
[Rnxrrset]		"rr set should exist",
[Rnotauth]		"not authorative",
[Rnotzone]		"not in zone",
[Rbadvers]		"bad opt version",
/* [Rbadsig]		"bad signature", */
[Rbadkey]		"bad key",
[Rbadtime]		"bad signature time",
[Rbadmode]		"bad mode",
[Rbadname]		"duplicate key name",
[Rbadalg]		"bad algorithm",
};

/* copied from /sys/src/cmd/ndb/dn.c:/^rrtype */
int
rrtype(char *atype)
{
	int i;

	for(i = 0; i <= Tall; i++)
		if(rrtname[i] && strcmp(rrtname[i], atype) == 0)
			return i;

	/* make any a synonym for all */
	if(strcmp(atype, "any") == 0)
		return Tall;
	else if(isascii(atype[0]) && isdigit(atype[0]))
		return atoi(atype);
	else
		return -1;
}

RdataFormat *
findrdfmt(int type)
{
	RdataFormat *p;

	for(p=rdfmttab; p < rdfmttab + nelem(rdfmttab); p++) {
		if(p->rrtype == type)
			return p;
	}
	return nil;
}

VerbType *
findverbtype(char *name, int hasrdata)
{
	VerbType *v;
	int errtype;

	errtype = 0;
	for(v=verbtab; v < verbtab + nelem(verbtab); v++) {
		if(strcmp(v->name, name) == 0) {
			errtype = 1;
			if(v->needsrdata == hasrdata)
				return v;
		}
	}

	switch(errtype) {
	case 0:
		werrstr("unknown verb: %s", name);
		break;
	case 1:
		werrstr("arity mismatch");
		break;
	default:
		abort();
	}
	return nil;
}

char *
estrdup(char *s)
{
	char *t;

	if((t = strdup(s)) == nil)
		sysfatal("estrdup: %r");
	return t;
}

void
gname(char *name, uchar **p)
{
	int n, i;

	while((n = *(*p)++) != 0) {
		for(i=0; i < n; i++)
			*name++ = *(*p)++;
		*name++ = '.';
	}
	*name = '\0';
}

int
g16(uchar *p)
{
	int n;

	n  = *p++ << 8;
	n |= *p;
	return n;
}

void
p16(uchar **p, int n)
{
	*(*p)++ = n >> 8;
	*(*p)++ = n;
}

void
p32(uchar **p, int n)
{
	*(*p)++ = n >> 24;
	*(*p)++ = n >> 16;
	*(*p)++ = n >> 8;
	*(*p)++ = n;
}

void
p48(uchar **p, u64int n)
{
	*(*p)++ = n >> 40;
	*(*p)++ = n >> 32;
	*(*p)++ = n >> 24;
	*(*p)++ = n >> 16;
	*(*p)++ = n >> 8;
	*(*p)++ = n;
}

void
pmem(uchar **p, void *v, int len)
{
	memmove(*p, v, len);
	*p += len;
}

void
pname(uchar **p, char *s)
{
	uchar *len;

	while (*s){
		len = (*p)++;
		while(*s && *s != '.')
			*(*p)++ = *s++;
		*len = *p - len - 1;	/* FIXME - check for a valid length (rfc 1034 § 3.1) */
		if(*s == '.')
			s++;
	}
	*(*p)++ = 0;
}

void
precord(uchar **p, Record *r)
{
	pname(p, r->name);
	p16(p, r->type);
	p16(p, r->class);
	p32(p, r->ttl);
	p16(p, r->rdlength);
	pmem(p, r->rdata, r->rdlength);
}

void
psig(uchar **p, Signature *s, int forreal)
{
	uchar buf[1024], *q;

	q = buf;
	pname(&q, s->aname);
	p48(&q, s->time);
	p16(&q, s->fudge);
	if(forreal) {
		p16(&q, sizeof s->digest);
		pmem(&q, s->digest, sizeof s->digest);
		p16(&q, 0);
	}
	p16(&q, 0);
	p16(&q, 0);

	pname(p, s->kname);
	if(forreal)
		p16(p, Ttsig);
	p16(p, Call);
	p32(p, 0);
	if(forreal)
		p16(p, q-buf);
	pmem(p, buf, q-buf);
}

int
rcount(void **p)
{
	int n;
	for(n=0; p[n] != nil; n++)
		;
	return n;
}

void
pmsg(uchar **p, Msg *m, int forreal)
{
	Record **r;

	/* header section (rfc 2136 § 2.2, rfc 1035 § 4.1.1) */
	p16(p, 0);
	p16(p, 5<<11);
	p16(p, 1);
	p16(p, rcount(m->prereqs));
	p16(p, rcount(m->updates));
	if(forreal)
		p16(p, 1);
	else
		p16(p, 0);

	/* zone section (rfc 2136 § 2.3, rfc 1035 § 4.1.2) */
	pname(p, m->zone);
	p16(p, Tsoa);
	p16(p, Cin);

	/* prerequisite section (rfc 2136 § 2.4, rfc 1035 § 4.1.3) */
	for(r=m->prereqs; *r != nil; r++)
		precord(p, *r);

	/* update section (rfc 2136 § 2.5, rfc 1035 § 4.1.3) */
	for(r=m->updates; *r != nil; r++)
		precord(p, *r);

	/* tsig section (rfc 2845 § 2.3, rfc 1035 § 4.1.3) */
	psig(p, &m->sig, forreal);
}

int
fieldsize(int ftype, char *value)
{
		switch(ftype) {
		case Fi16:
			return 2;
		case Fi32:
			return 4;
		case Fip:
			return 4;
		case Fipv6:
			return 16;
		case Fstring:
		case Ftxt:
			return 1+strlen(value);
		case Fname:
			/* account for a possible omission in the input of a final . */
			return strlen(value)+1;
		}

		abort();
		return 69105;
}

int
prdata(Record *r, int type, char **fields)
{
	RdataFormat *fmt;
	int *ftypes, *ft;	/* 0 terminated */
	char **f;	/* nil terminated */
	int ftlen, flen;
	int rdsize;
	uchar *p;

	/* find the record data's prototype */
	if((fmt = findrdfmt(type)) == nil) {
		werrstr("rrtype %s isn't implemented", rrtname[type]);
		return -1;
	}
	ftypes = fmt->ftypes;

	/* count the fields in the prototype */
	for(ft=ftypes; *ft != 0; ft++)
		;
	ftlen = ft - ftypes;

	/* count the fields in the input */
	for(f=fields; *f != nil; f++)
		;
	flen = f - fields;

	if(flen != ftlen) {
		werrstr("rrtype %s expects %d fields, saw %d", rrtname[type], ftlen, flen);
		return -1;
	}

	/* allocate a buffer large enough for the record data */
	rdsize = 0;
	for(ft=ftypes, f=fields; *ft != 0; ft++, f++)
		rdsize += fieldsize(*ft, *f);
	if((r->rdata = malloc(rdsize)) == nil)
		sysfatal("prdata: malloc: %r");
	p = r->rdata;

	for(ft=ftypes, f=fields; *ft != 0; ft++, f++) {
		switch(*ft) {
		case Fi16: {
				int x;
				x = atoi(*f);
				if(x > 1<<16 - 1) {
					werrstr("expected i16, saw %s", *f);
					return -1;
				}
				p16(&p, atoi(*f));
			}
			break;
		case Fi32:
			p32(&p, strtoul(*f, nil, 0));
			break;
		case Fip: {
				vlong x;
				uchar addr[IPaddrlen];
				x = parseip(addr, *f);
				if(x < 0) {
					werrstr("malformed ip: %s", *f);
					return -1;
				}
				if(x == 6) {
					werrstr("expected ip, saw ipv6: %s", *f);
					return -1;
				}
				memcpy(p, addr+IPv4off, IPv4addrlen);
			}
			p += 4;
			break;
		case Fipv6: {
				vlong x;
				uchar addr[IPaddrlen];
				x = parseip(addr, *f);
				if(x < 0) {
					werrstr("malformed ipv6: %s", *f);
					return -1;
				}
				if(x != 6) {
					werrstr("expected ipv6, saw ip: %s", *f);
					return -1;
				}
				memcpy(p, addr, IPaddrlen);
			}
			p += 16;
			break;
		case Fstring:
		case Ftxt: {
				/* FIXME - allow more than one string - rfc 1035 § 3.3.14 */
				/* FIXME - distinguish from an Fstring as used in hinfo - rfc § 3.3.2 */
				int n;
				n = strlen(*f);
				if(n > 255) {
					werrstr("txt fields are limited to 255 bytes, saw %d", n);
					return -1;
				}
				*p = n;
				memmove(p+1, *f, n);
				p += 1+n;
			}
			break;
		case Fname:
			pname(&p, *f);
			break;
		default:
			abort();
		}
	}
	r->rdlength = p - r->rdata;

	return 0;
}

void
sign(Msg *m, char *keyname, char *key)
{
	uchar buf[4096], *p;

	m->sig.kname = keyname;
	m->sig.aname = "hmac-md5.sig-alg.reg.int";
	m->sig.time = time(nil);
	m->sig.fudge = 300;

	p = buf;
	pmsg(&p, m, 0);

	hmac_md5(buf, p-buf, (uchar *)key, strlen(key), m->sig.digest, nil);
}

char *
walkup(char *name)
{
	char *cp;

	cp = strchr(name, '.');
	if(cp)
		return cp+1;
	else if(*name)
		return "";
	else
		return 0;
}

char *
walktosoa(char *name)
{
	int fd, n;
	char *query;

	if((fd = open("/net/dns", ORDWR)) < 0)
		sysfatal("open /dev/dns: %r");

	for(; name != nil; name = walkup(name)) {
		query = smprint("%s soa", name);
		if(query == nil)
			sysfatal("walktosoa: smprint query: %r");
		n = pwrite(fd, query, strlen(query), 0);
		free(query);
		if(n >= 0)
			break;
	}

	close(fd);
	return name;
}

char *
lookupns(char *zone)
{
	int fd, n;
	char *query, buf[1024], *ns;

	if((fd = open("/net/dns", ORDWR)) < 0)
		sysfatal("open /dev/dns: %r");

	if((query = smprint("%s ns", zone)) == nil)
		sysfatal("lookupns: smprint query: %r");
	n = write(fd, query, strlen(query));
	free(query);
	if(n < 0)
		return nil;

	alarm(3000);
	n = read(fd, buf, sizeof buf);
	alarm(0);
	if(n < 0)
		return nil;
	if(n < 1) {
		werrstr("unusual answer");
		return nil;
	}
	buf[n] = '\0';
	ns = buf+1;

	close(fd);
	return estrdup(ns);
}

void
timeout(void *, char *msg)
{
	if(strncmp(msg, "alarm", 5) == 0) {
		werrstr("timeout");
		noted(NCONT);
	}
	noted(NDFLT);
}

void
usage(void)
{
	fprint(2, "usage: %s\n", argv0);
	exits("usage");
}

#define inputerr(fmt, ...) { \
	fprint(2, "%s:%d " fmt "\n", filename, lineno, __VA_ARGS__); \
	badinput++; \
}

void
main(int argc, char **argv)
{
	char *filename;
	int lineno;
	int badinput;
	Biobuf bin;
	char *line;
	Msg m;
	Record *prereqs[MaxRRs+1], *updates[MaxRRs+1];	/* nil terminated */
	Record **p, **u, **r;
	char *name, *ns, *addr;
	UserPasswd *userpass;
	uchar buf[4096], *q;	/* FIXME - size properly */
	int fd;
	int n;
	u16int txid, rxid, rcode;

	ARGBEGIN {
	default:
		usage();
	} ARGEND

	if(argc != 0)
		usage();

	filename = "<stdin>";
	lineno = 0;
	badinput = 0;
	Binit(&bin, 0, OREAD);
	notify(timeout);

	p = prereqs;
	u = updates;
	name = nil;
	while((line = Brdline(&bin, '\n')) != nil) {
		char *toks[16];	/* nil terminated */
		int ntoks;
		int type;
		VerbType *v;

		line[Blinelen(&bin)-1] = '\0';
		lineno++;
		ntoks = tokenize(line, toks, nelem(toks)-1);
		toks[ntoks] = nil;
		if(ntoks < 3) {
			if(ntoks != 0)
				inputerr("%s", "syntax error");
			continue;
		}

		if(p - prereqs + 1 > MaxRRs)
			sysfatal("too many prereqs; limit is %d", MaxRRs);
		if(u - updates + 1 > MaxRRs)
			sysfatal("too many updates; limit is %d", MaxRRs);

		if(name == nil)
			name = estrdup(toks[1]);

		if(strlen(toks[0]) > sizeof((*r)->name) - 1) {
			inputerr("%s", "name too long");
			continue;
		}

		if((type = rrtype(toks[2])) < 0) {
			inputerr("unknown rrtype: %s", toks[2]);
			continue;
		}

		if((v = findverbtype(toks[0], ntoks > 3)) == nil) {
			inputerr("%s%r", "");
			continue;
		}

		switch(v->section) {
		case Sprereq:
			r = p++;
			break;
		case Supdate:
			r = u++;
			break;
		default:
			abort();
			SET(r);
		}
		if((*r = malloc(sizeof(**r))) == nil)
			sysfatal("malloc: %r");

		/* ttl is defined as 0 except for add updates where
		 * the server will use the soa minimum - rfc 1035 § 6.2
		 */
		(*r)->name = estrdup(toks[1]);
		(*r)->type = type;
		(*r)->class = v->class;
		(*r)->ttl = 0;
		if(v->needsrdata) {
			if(prdata(*r, type, toks+3) < 0)
				inputerr("%s%r", "");
		} else {
			(*r)->rdlength = 0;
		}
	}
	*p = nil;
	*u = nil;

	if(badinput)
		sysfatal("bad input");
	if(lineno == 0)
		exits(nil);

	if((m.zone = walktosoa(name)) == nil)
		sysfatal("couldn't find soa: %s: %r", name);
	m.prereqs = prereqs;
	m.updates = updates;

	if((ns = lookupns(m.zone)) == nil)
		sysfatal("couldn't find ns: %r");
	addr = netmkaddr(ns, "udp", "dns");

	if((userpass = auth_getuserpasswd(amount_getkey, "proto=pass zone=%s", m.zone)) == nil)
		sysfatal("getpass: %r");
	sign(&m, userpass->user, userpass->passwd);

	/* rfc 1035 § 4.2.2 */
	if(strncmp(addr, "tcp", 3) == 0) {
		q = buf+2;
		pmsg(&q, &m, 1);
		n = q - buf;
		q = buf;
		p16(&q, n);
	} else {
		q = buf;
		pmsg(&q, &m, 1);
		n = q - buf;
	}

	if((fd = dial(addr, nil, nil, nil)) < 0)
		sysfatal("dial: %r");
	if(write(fd, buf, n) < 0)
		sysfatal("write fd: %r");

	txid = g16(buf);	/* save the message id - rfc 1035 § 4.1.1 */
	alarm(3000);
	for(;;) {
		n = read(fd, buf, sizeof buf);
		if(n < 0)
			sysfatal("read fd: %r");
		if(n < 2)
			sysfatal("couldn't read response id");
		if((rxid = g16(buf)) == txid)
			break;
		fprint(2, "received a non sequitur %ud, expected %ud\n", rxid, txid);
	}
	alarm(0);

	if(n < 4)
		sysfatal("truncated response");
	if((rcode = g16(buf+2) & Rmask) != Rok) {
		if(rcode < nelem(rname) && rname[rcode] != nil)
			sysfatal("rcode %d: %s", rcode, rname[rcode]);
		else
			sysfatal("rcode %d", rcode);
	}

	exits(nil);
}