;/* ==========================================================================
 * ldig - Simple DNS Query Interface
 * --------------------------------------------------------------------------
 * Copyright (c) 2004, 2005, 2006  Barracuda Networks, Inc.
 * Copyright (c) 2006  William Ahern <william@25thandClement.com>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a
 * copy of this software and associated documentation files (the
 * "Software"), to deal in the Software without restriction, including
 * without limitation the rights to use, copy, modify, merge, publish,
 * distribute, sublicense, and/or sell copies of the Software, and to permit
 * persons to whom the Software is furnished to do so, subject to the
 * following conditions:
 *
 * The above copyright notice and this permission notice shall be included
 * in all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
 * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN
 * NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
 * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
 * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
 * USE OR OTHER DEALINGS IN THE SOFTWARE.
 * --------------------------------------------------------------------------
 * History
 *
 * 2006-04-23 (william@25thandClement.com)
 * 	Make work with refactored interface.
 *
 * 2006-02-14 (william@25thandClement.com)
 * 	Published by Barracuda Networks, originally authored by
 * 	employee William Ahern (wahern@barracudanetworks.com).
 * --------------------------------------------------------------------------
 * Description
 *
 * Clone of the dig(1) utility.
 *
 * ==========================================================================
 */
#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>

#include <unistd.h>
#include <sched.h>

#include <netdb.h>		/* NI_MAXHOST */

#include <string.h>

#include <getopt.h>

#include <event.h>

#include <err.h>

#include <assert.h>

#include "lookup.h"


#define USAGE									\
	"Usage: ldig [OPTIONS] [QUERY NAME]\n"					\
	"  -t, --type=TYPE   Query type: PTR, A, AAAA, MX, MX+A+AAAA, etc.\n"	\
	"  -f, --file=PATH   Path to query input: <qtype> <qname>\n"		\
	"  -o, --option=OPT  Set flag: PREFER_A, FALLBACK\n"			\
	"  -l, --list        List supported query types.\n"			\
	"  -v, --verbose     Enable verbose output.\n"				\
	"  -h, --help        Print this usage information.\n\n"			\


static int verbose;
static int pending;


static void split(int *argc, char *argv[], int len, char *line, char *sep) {
	char **ap;
	char *nxt	= line;

	if (!sep)
		sep	= " \t";

	for (*argc = 0, ap = argv; ap < &argv[len - 1];) {
		if ((*ap = strsep(&nxt, sep)) == 0)
			break;

		if (**ap != '\0')
			ap++, (*argc)++;
	}

	*ap	= 0;

	return /* void */;
} /* split() */


void catch_rr(int nanswers, struct lookup_rr *answer, int nadditional, struct lookup_rr *additional, enum lookup_errno status, void *query) {
	char name[128];

	pending--;

	(void)printf("Query %s:",(char *)query);
	
	if (!LOOKUP_SUCCESS(status))
		printf(" FAILED (%s)\n", lookup_errlist[status]);
	else
		printf(" SUCCEEDED\n");

	if (nanswers)
		puts("\t--- ANSWERS ---");
print:
	for (; answer; answer = answer->next) {
		switch (answer->type) {
		case LOOKUP_IN_PTR:
			printf("\t%s",answer->rr.ptr.host);
			break;
		case LOOKUP_IN_TXT:
			printf("\t%s",answer->rr.txt.data);
			break;
		case LOOKUP_IN_SOA:
			printf("\t%s %s %u %d %d %d %u",answer->rr.soa.host,answer->rr.soa.mbox,answer->rr.soa.serial,answer->rr.soa.refresh,answer->rr.soa.retry,answer->rr.soa.expire,answer->rr.soa.minimum);
			break;
		case LOOKUP_IN_NS:
			printf("\t%s",answer->rr.ns.host);
			break;
		case LOOKUP_IN_CNAME:
			printf("\t%s",answer->rr.cname.host);
			break;
		case LOOKUP_IN_MX:
			printf("\t%hu %s",answer->rr.mx.preference,answer->rr.mx.host);
			break;
		case LOOKUP_IN_SRV:
			printf("\t%hu %hu %hu %s",answer->rr.srv.priority,answer->rr.srv.weight,answer->rr.srv.port,answer->rr.srv.host);
			break;
		case LOOKUP_IN_A:
			/* FALL THROUGH */
		case LOOKUP_IN_AAAA:
			if (0 != getnameinfo(&answer->rr.ip.sa.sa,answer->rr.ip.salen,name,sizeof name,NULL,0,NI_NUMERICHOST|NI_NUMERICSERV))
				err(1,"getnameinfo");

			printf("\t%s",name);
			break;
		default:
			assert(0);
		}

		switch (answer->section) {
		case LOOKUP_SECTION_LOCAL:
			puts(" (section: local)");
			break;
		case LOOKUP_SECTION_ANSWERS:
			puts(" (section: answers)");
			break;
		case LOOKUP_SECTION_AUTHORITY:
			puts(" (section: authority)");
			break;
		case LOOKUP_SECTION_ADDITIONAL:
			puts(" (section: additional)");
			break;
		default:
			assert(0);
		}
	}

	if (nadditional) {
		puts("\t--- ADDITIONAL ---");

		answer		= additional;
		nanswers	= nadditional;
		nadditional	= 0;
		
		goto print;
	}

	free(query);

	return /* void */;
} /* catch_rr() */


char *strquery(char *qname, int qtype, int flags) {
	struct {
		const char *name;
		int type;
	} *nxt, map[] = {
		{ "PTR",	LOOKUP_IN_PTR },
		{ "TXT",	LOOKUP_IN_TXT },
		{ "SOA",	LOOKUP_IN_SOA },
		{ "NS",		LOOKUP_IN_NS },
		{ "CNAME",	LOOKUP_IN_CNAME },
		{ "MX",		LOOKUP_IN_MX },
		{ "SRV",	LOOKUP_IN_SRV },
		{ "A",		LOOKUP_IN_A },
		{ "AAAA",	LOOKUP_IN_AAAA },
		{ NULL,		0}
	};
	char type[64];
	unsigned len;
	char *q;

	*type	= '\0';
	len	= 0;

	for (nxt = map; nxt->name; nxt++) {
		if (nxt->type & qtype) {
			len	= strlen(strncat(type, nxt->name, sizeof type - 1));
			len	= strlen(strncat(type, "+", sizeof type - len - 1));
		}
	}

	if (len && type[len - 1] == '+')
		type[--len]	= '\0';

	if (0 >= asprintf(&q,"%s IN %s",qname,type))
		err(EXIT_FAILURE,"asprintf");

	return q;
} /* strquery() */


int doquery(char *qname, int qtype, int flags) {
	struct timeval tv = { 5 , 0 };
	char buf[NI_MAXHOST];
	char ptr[64];
	char *octets[5];
	int noctets;
	char *q;

	if (qtype == LOOKUP_IN_PTR && !strstr(qname,".in-addr.arpa")) {
		strncpy(buf, qname, sizeof buf);
		buf[sizeof buf - 1]	= '\0';

		split(&noctets,octets,sizeof octets / sizeof *octets,buf,".");

		if (noctets != 4)
			return -1;

		(void)snprintf(ptr,sizeof ptr,"%s.%s.%s.%s.in-addr.arpa",octets[3],octets[2],octets[1],octets[0]);
		qname	= ptr;
	}

	q	= strquery(qname,qtype,flags);

	if (verbose)
		fprintf(stderr,"Issuing %s\n",q);

	pending++;

	lookup(qname, strlen(qname), qtype, flags, catch_rr, q, &tv);

	return 0;
} /* doquery() */


int settype(int *type,const char *name) {
	if (0 == strcasecmp("PTR",name))
		*type	|= LOOKUP_IN_PTR;
	else if (0 == strcasecmp("A",name))
		*type	|= LOOKUP_IN_A;
	else if (0 == strcasecmp("AAAA",name))
		*type	|= LOOKUP_IN_AAAA;
	else if (0 == strcasecmp("MX",name))
		*type	|= LOOKUP_IN_MX;
	else if (0 == strcasecmp("SRV",name))
		*type	|= LOOKUP_IN_SRV;
	else if (0 == strcasecmp("TXT",name))
		*type	|= LOOKUP_IN_TXT;
	else if (0 == strcasecmp("SOA",name))
		*type	|= LOOKUP_IN_SOA;
	else if (0 == strcasecmp("NS",name))
		*type	|= LOOKUP_IN_NS;
	else if (0 == strcasecmp("CNAME",name))
		*type	|= LOOKUP_IN_CNAME;
	else
		return 0;

	return 1;
} /* settype() */


int setflag(int *type,const char *name) {
	if (0 == strcasecmp("FALLBACK",name))
		*type	|= LOOKUP_FALLBACK;
	else if (0 == strcasecmp("PREFER_A",name))
		*type	|= LOOKUP_PREFER_A;
	else
		return 0;

	return 1;
} /* setflag() */


int main(int argc, char *argv[]) {
	extern char *optarg;
	extern int optind;
	int c;
	struct option list[] = {
		{"type", 	1, 	NULL,	't'},
		{"file",	1,	NULL,	'f'},
		{"option",	1,	NULL,	'o'},
		{"list",	0,	NULL,	'l'},
		{"verbose",	0,	NULL,	'v'},
		{"help",	0,	NULL,	'h'},
		{NULL,		0,	NULL,	0}
	};
	const char *file	= NULL;
	int type		= 0;
	int flags 		= 0;
	char *types[8];
	int ntypes;
	char *opts[8];
	int nopts;
	FILE *fp	= NULL;
	char *ln;
	int lineno	= 0;
	char linebuf[1024];
	int last;

	event_init();
	lookup_init(0, 0, 0);
	setlinebuf(stdout);

	while (-1 != (c = getopt_long(argc,argv,"t:f:o:lvh",list,NULL))) {
		switch (c) {
		case 't':
			split(&ntypes,types,sizeof types / sizeof *types,optarg,"+");

			while (ntypes--) if (!settype(&type,types[ntypes]))
				warnx("Invalid type: %s",types[ntypes]);

			break;
		case 'f':
			file	= optarg;
			break;
		case 'o':
			split(&nopts,opts,sizeof opts / sizeof *opts,optarg,"+");

			while (nopts--) if (!setflag(&flags,opts[nopts]))
				warnx("Invalid flag: %s",opts[nopts]);

			break;
		case 'l':
			fprintf(stderr,"Query types:\n");
			fprintf(stderr,"	PTR\n");
			fprintf(stderr,"	A\n");
			fprintf(stderr,"	AAAA\n");
			fprintf(stderr,"	MX\n");
			fprintf(stderr,"	MX+A\n");
			fprintf(stderr,"	MX+AAAA\n");
			fprintf(stderr,"	MX+A+AAAA\n");
			fprintf(stderr,"	SRV\n");
			fprintf(stderr,"	SRV+A\n");
			fprintf(stderr,"	SRV+AAAA\n");
			fprintf(stderr,"	SRV+A+AAAA\n");
			fprintf(stderr,"	TXT\n");
			fprintf(stderr,"	SOA\n");
			fprintf(stderr,"	NS\n");
			fprintf(stderr,"	CNAME\n");

			return EXIT_SUCCESS;
		case 'v':
			verbose	= 1;
			break;
/*help:*/
		case 'h':
		default:
			fprintf(stderr,USAGE);
			return (c == 'h')? EXIT_SUCCESS : EXIT_FAILURE;
		}	
	} /* getopt_long() */

	event_init();
	lookup_init(0, 0, 0);
	setlinebuf(stdout);

	if (!type)
		type	|= LOOKUP_IN_A | LOOKUP_IN_AAAA;

	if (optind < argc) {
		while (optind < argc) {
			if (0 != doquery(argv[optind],type,flags))
				warnx("Invalid query: %s",strquery(argv[optind],type,flags));
			optind++;
		}
	} else fp = stdin;
	
	if (file) {
		fp	= fopen(file,"r");

		if (!fp)
			err(EXIT_FAILURE,"Could not open %s",file);
	}

	if (fp) while ((ln = fgets(linebuf, sizeof linebuf, fp))) {
		char *args[3];
		int nargs;

		assert(ln = strdup(ln));

		lineno++;

		split(&nargs, args, sizeof args / sizeof *args, linebuf, " \t");

		type	= LOOKUP_IN_A | LOOKUP_IN_AAAA;

		switch (nargs) {
		case 2:
			type	= 0;

			split(&ntypes, types, sizeof types / sizeof *types, args[0], "+");

			while (ntypes--) if (!settype(&type, types[ntypes]))
				warnx("Invalid type at line %d in %s: %s", lineno, ln, types[ntypes]);

			args[0] = args[1];
		case 1:
			if (0 == doquery(args[0], type, flags))
				break;

			/* FALL THROUGH */
		default:
			warnx("Invalid query: %s", ln);
		}

		if (verbose)
			fprintf(stderr,"%d queries pending\n", pending);

		free(ln);
	}

	for (last = 0;;) {

		if (verbose) {
			if (last != pending) {
				last	= pending;
				fprintf(stderr,"%d queries pending\n",pending);
			}
		}

		event_loop(EVLOOP_NONBLOCK);

		event_loop(EVLOOP_ONCE);

		if (!pending)
			break;
	}

	return 0;
} /* main() */