/*
// AUTHOR
// N. Nielsen
//
// VERSION
// 0.2
//
// LICENSE
// This software is in the public domain.
//
// 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 author(s) 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.
//
// SUPPORT
// Send bug reports to: <nielsen@memberwebs.com>
*/

#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <errno.h>
#include <stdarg.h>
#include <unistd.h>

#include <sys/types.h>
#include <sys/socket.h>
#include <sys/sockio.h>
#include <sys/sysctl.h>
#include <sys/ioctl.h>

#include <net/if.h>
#include <net/if_dl.h>
#include <net/if_types.h>
#include <net/route.h>

#include <netinet/in.h>
#include <netinet/if_ether.h>

#include <arpa/nameser.h>
#include <arpa/inet.h>

#include <resolv.h>

#if HAVE_CONFIG_H
#include <config.h>
#endif

/* The main processing functions */
static void processFile(FILE* fIn, FILE* fOut);
static void processReverse(const char* address, FILE* fOut);
static void processNone(const char* address, FILE* fOut);
static void processAddress(const char* address, FILE* fOut);
static void processHost(const char* address, FILE* fOut);
static void processMAC(const char* address, FILE* fOut);

static void atExit();
static void usage();

/* Are we in reverse resolve mode? */
int g_doReverse = 0;
int g_doVerbose = 0;

/* These are globals for efficiency */
u_char* g_packetBuf = NULL;
char* g_domainBuf = NULL;
u_char* g_routeBuf = NULL;
size_t g_routeLen = 0;


int main(int argc, char** argv)
{
	int useFiles = 0;
	int ch;

	/* Register clean up function */
	atexit(atExit);

	/* Parse command line options here */
	while((ch = getopt(argc, argv, "rfv")) != -1)
	{
		switch(ch)
		{
		/* Treat arguments as files */
		case 'f':
			useFiles = 1;
			break;

		case 'v':
			g_doVerbose = 1;
			break;

		/* Do reverse resolve */
		case 'r':
			g_doReverse = 1;
			break;

		case '?':
		default:
			usage();
		}
	}

	argc -= optind;
	argv += optind;

	if(argc)
	{
		while(argc > 0)
		{
			/* Process the argument as a file full of addresses */
			if(useFiles)
			{
				FILE* file = fopen(argv[0], "r");
				if(file == NULL)
					err(1, "can't open file: %s", argv[0]);

				processFile(file, stdout);

				fclose(file);
			}

			/* Process each argument as an address */
			else
			{
				processAddress(argv[0], stdout);
			}

			argc--;
			argv++;
		}
	}
	else
	{
		/* No args, just process stdin etc... */
		processFile(stdin, stdout);
	}

	exit(0);
}


/**
 * Eats comments and spaces
 */
static void eatSpace(FILE* file)
{
	char ch;

	do
	{
		ch = getc(file);

		if(ch == '#')
		{
			while(ch != '\n')
				ch = getc(file);
		}
	}
	while(isspace(ch));

	ungetc(ch, file);
}


/**
 * Read addresses from a file and process
 */
static void processFile(FILE* fIn, FILE* fOut)
{
	/* Allocate memory, cache for efficiency */
	if(g_domainBuf == NULL)
	{
		g_domainBuf = (char*)malloc(sizeof(char) * MAXDNAME);
		if(g_domainBuf == NULL)
			errx(1, "out of memory.");
	}


	/* And process rest of file */
	eatSpace(fIn);

	while(!feof(fIn) && !ferror(fIn))
	{
		char ch;
		int i = 0;

		while(1)
		{
			ch = fgetc(fIn);

			if(isspace(ch))
				break;

			if(i < MAXDNAME)
				g_domainBuf[i] = ch;

			i++;
		}

		g_domainBuf[i] = 0;

		processAddress(g_domainBuf, fOut);

		eatSpace(fIn);
	}

	/* TODO: Should this be a warning or quit? */
	if(ferror(fIn))
		err(1, "error reading file");
}


/**
 * Main address processing function
 */
static void processAddress(const char* address, FILE* fOut)
{
	if(g_doVerbose)
		fprintf(stderr, "resolve: resolving %s\n", address);

	/* If it's an IP then .... */
	if(strlen(address) == strspn(address, "0123456789."))
	{
		if(g_doReverse)
			processReverse(address, fOut);

		else
			processNone(address, fOut);
	}

	/* If it's a MAC address then .... */
	else if(strlen(address) == strspn(address, "0123456789abcdefABCDEF-:"))
	{
		if(g_doReverse)
			processNone(address, fOut);
		else
			processMAC(address, fOut);
	}

	/* If it's a net block then .... */
	if(strlen(address) == strspn(address, "0123456789./"))
	{
		if(g_doReverse)
			warnx("can't reverse resolve net block: %s", address);
		else
			processNone(address, fOut);
	}

	/* Otherwise it should be a domain name. try to resolve it. */
	else
	{
		if(g_doReverse)
			processNone(address, fOut);
		else
			processHost(address, fOut);
	}
}


/**
 * Used if address is aleady (hopefully) in correct form
 */
static void processNone(const char* address, FILE* fOut)
{
	fprintf(fOut, address);
	fputc('\n', fOut);
}


/**
 * Forward resolve a host name into an IP
 */
static void processHost(const char* address, FILE* fOut)
{
	int ret;
	ns_msg handle;
	int rrnum;
	ns_rr rr;

	/* Allocation cached for efficiency */
	if(g_packetBuf == NULL)
	{
		g_packetBuf = (u_char*)malloc(sizeof(u_char) * PACKETSZ);
		if(g_packetBuf == NULL)
			errx(1, "out of memory.");
	}

	/* Do a DNS lookup */
	ret = res_search(address, C_IN, ns_t_a, g_packetBuf, PACKETSZ);
	if(ret == -1)
		warnx("couldn't resolve: %s\n", address);
	else
	{
		/* Initialize a handle to this response. */
		if(ns_initparse(g_packetBuf, ret, &handle) < 0)
			warn("couldn't parse dns response: %s\n", strerror(errno));
		else
		{
			/* Parse the packet */
			for(rrnum = 0; rrnum < ns_msg_count(handle, ns_s_an); rrnum++)
			{
				if(!ns_parserr(&handle, ns_s_an, rrnum, &rr))
				{
					if(ns_rr_type(rr) == ns_t_a &&	/* It's a host adress */
					   ns_rr_rdlen(rr) == 0x04)		/* And has an IP */
					{
						struct in_addr addr;
						memcpy(&(addr.s_addr), ns_rr_rdata(rr), 0x04);
						fprintf(fOut, inet_ntoa(addr));
						fputc('\n', fOut);
					}
				}
			}
		}
	}
}


/**
 * Reverse resolve an IP into a host name
 */
static void processReverse(const char* address, FILE* fOut)
{
	int ret;
	ns_msg handle;
	int rrnum;
	ns_rr rr;
	u_int32_t ha;
	char name[NS_MAXDNAME];
	struct in_addr addr;

	/* Is it a valid IP address first of all */
	if(!inet_aton(address, &addr))
	{
		warnx("invalid ip address: %s\n", address);
		return;
	}

	/* Format query */
	ha = ntohl(addr.s_addr);

	sprintf(name, "%u.%u.%u.%u.IN-ADDR.ARPA.",
			(ha) & 0xff,
			(ha >> 8) & 0xff,
			(ha >> 16) & 0xff,
			(ha >> 24) & 0xff);

	/* Allocation cached for efficiency */
	if(g_packetBuf == NULL)
	{
		g_packetBuf = (u_char*)malloc(sizeof(u_char) * PACKETSZ);
		if(g_packetBuf == NULL)
			errx("out of memory.");
	}

	/* Do that lookup */
	ret = res_search(name, C_IN, ns_t_ptr, g_packetBuf, PACKETSZ);
	if(ret == -1)
		warnx("couldn't resolve: %s\n", name);
	else
	{
		/* Initialize a handle to this response. */
		if(ns_initparse(g_packetBuf, ret, &handle) < 0)
			warn("couldn't parse dns response");
		else
		{
			/* Parse the packet */
			for(rrnum = 0; rrnum < ns_msg_count(handle, ns_s_an); rrnum++)
			{
				if(!ns_parserr(&handle, ns_s_an, rrnum, &rr))
				{
					if(ns_rr_type(rr) == ns_t_ptr) /* It's a domain pointer */
					{
			            /* Expand the host's name */
						if (ns_name_uncompress(
							ns_msg_base(handle),/* Start of the packet   */
							ns_msg_end(handle), /* End of the packet     */
							ns_rr_rdata(rr),    /* Position in the packet*/
							name,				/* Result                */
							MAXDNAME)           /* Size of nsList buffer */
                                  < 0) /* Negative: error       */
						{
							warn("couldn't parse dns response");
						}
						else
						{
							fprintf(fOut, "%s\n", name);
						}
					}
				}
			}
		}
	}
}


/**
 * Parse components of a MAC address
 */
static int parseMAC(const char* address, u_char* mac)
{
	int i;

	for(i = 0; i < 6; i++)
	{
		char* end;
		int part;

		part = strtol(address, &end, 16);

		if(i != 5 && *end != ':' && *end != '-')
			return 0;

		if(part < 0 || part > 255)
			return 0;

		mac[i] = (u_char)part;

		address = end + 1;
	}

	return 1;
}


#define ROUNDUP(a) \
	((a) > 0 ? (1 + (((a) - 1) | (sizeof(long) - 1))) : sizeof(long))

/**
 * Reverse resolve MAC address if in the ARP table
 */
static void processMAC(const char* address, FILE* fOut)
{
	u_char* next;
	struct rt_msghdr* rtm;
	struct sockaddr_inarp* sin;
	struct sockaddr_dl* sdl;
	u_char mac[6];

	/* Load the kernel routing table */
	if(g_routeBuf == NULL)
	{
		int mib[6];

		mib[0] = CTL_NET;
		mib[1] = PF_ROUTE;
		mib[2] = 0;
		mib[3] = AF_INET;
		mib[4] = NET_RT_FLAGS;
		mib[5] = RTF_LLINFO;

		if(sysctl(mib, 6, NULL, &g_routeLen, NULL, 0) < 0)
			errx(1, "can't load routing table");

		g_routeBuf = malloc(g_routeLen);

		if(g_routeBuf == NULL)
			errx(1, "out of memory");

		if(sysctl(mib, 6, g_routeBuf, &g_routeLen, NULL, 0) < 0)
			errx(1, "can't load routing table");
	}

	/* Get MAC bytes */
	if(!parseMAC(address, mac))
	{
		warn("invalid MAC address: %s", address);
		return;
	}

	/* Look for it in the routing table */
	for(next = g_routeBuf; next < (g_routeBuf + g_routeLen); next += rtm->rtm_msglen)
	{
		rtm = (struct rt_msghdr*)next;
		sin = (struct sockaddr_inarp*)(rtm + 1);
		(char*)sdl = (char*)sin + ROUNDUP(sin->sin_len);

		if(!memcmp(LLADDR(sdl), mac, sizeof(mac)))
		{
			fprintf(fOut, inet_ntoa(sin->sin_addr));
			fputc('\n', fOut);
			return;
		}
	}

	/* oops */
	warnx("unlisted MAC address: %s", address);
}



static void usage()
{
	fprintf(stderr, "usage: resolve [-rv] -f file ...\n");
	fprintf(stderr, "       resolve [-rv] address ...\n");
	exit(2);
}


/*
 * Clean up function.
 */
static void atExit()
{
	if(g_packetBuf != NULL)
		free(g_packetBuf);
	if(g_domainBuf != NULL)
		free(g_domainBuf);
}