#include <config.h>
#include <stdio.h>
#include <stdio.h>
#include <assert.h>
#include <unistd.h>
#include <inttypes.h>
#include <mailutils/mailutils.h>
#include "libmf.h"
#include "dns.h"

static int quiet = 0;
static char const *suffix;
static int suffix_len;

void
status_print(dns_status status)
{
	if (quiet > 1)
		return;
	switch (status) {
	case dns_success:
		if (quiet)
			return;
		printf("OK");
		break;
	case dns_not_found:
		printf("NOTFOUND");
		break;
	case dns_failure:
		printf("FAILURE");
		break;
	case dns_temp_failure:
		printf("TEMPFAIL");
		break;
	case dns_too_many:
		printf("TOOMANY");
		break;
	default:
		mu_error("%s:%d: unrecognized status", __FILE__, __LINE__);
		abort();
	}
	putchar('\n');
}

static int
hostname_cmp(const void *a, const void *b)
{
	return strcmp(*(const char**) a, *(const char**) b);
}

void
reply_print_str(struct dns_reply *reply, int sorted)
{
	int i;

	if (sorted)
		qsort(reply->data.str, reply->count, sizeof reply->data.str[0],
		      hostname_cmp);
	for (i = 0; i < reply->count; i++) {
		char const *str = reply->data.str[i];
		int len = strlen(str);
		if (len > suffix_len
		    && memcmp(str + len - suffix_len, suffix, suffix_len) == 0)
			len -= suffix_len;
		printf("%*.*s\n", len, len, reply->data.str[i]);
	}
}

static int
ipaddr_cmp(const void *a, const void *b)
{
	uint32_t ipa = ntohl(((struct in_addr*)a)->s_addr);
	uint32_t ipb = ntohl(((struct in_addr*)b)->s_addr);
	if (ipa < ipb)
		return -1;
	if (ipa > ipb)
		return 1;
	return 0;
}

static int
ip6addr_cmp(const void *a, const void *b)
{
	return memcmp(a, b, sizeof(struct in6_addr));
}

void
reply_print_ip(struct dns_reply *reply, int sorted)
{
	int i;

	if (sorted)
		qsort(reply->data.ip, reply->count, sizeof reply->data.ip[0],
		      ipaddr_cmp);
	for (i = 0; i < reply->count; i++) {
		printf("%s\n", inet_ntoa(reply->data.ip[i]));
	}
}

void
reply_print_ip6(struct dns_reply *reply, int sorted)
{
	int i;

	if (sorted)
		qsort(reply->data.ip, reply->count, sizeof reply->data.ip6[0],
		      ip6addr_cmp);
	for (i = 0; i < reply->count; i++) {
		char ipstr[IPMAX_DOTTED_BUFSIZE];
		printf("%s\n", inet_ntop(AF_INET6, &reply->data.ip6[i],
					 ipstr, sizeof(ipstr)));
	}
}

void
reply_print(struct dns_reply *reply, int sorted)
{
	if (quiet > 2)
		return;
	switch (reply->type) {
	case dns_reply_str:
		reply_print_str(reply, sorted);
		break;
	case dns_reply_ip:
		reply_print_ip(reply, sorted);
		break;
	case dns_reply_ip6:
		reply_print_ip6(reply, sorted);
		break;
	default:
		mu_error("%s:%d: unrecognized reply type", __FILE__, __LINE__);
		abort();
	}
}

static int
res_soa(int argc, char **argv)
{
	int resolve_family = resolve_ip4;
	struct dns_reply reply;
	dns_status status;

	if (argc == 2) {
		if (strcmp(argv[0], "-ip") == 0)
			resolve_family = resolve_ip4;
		else if (strcmp(argv[0], "-ip6") == 0)
			resolve_family = resolve_ip6;
		else {
			mu_error("bad arguments");
			abort();
		}
		argc--;
		argv++;
	}
	assert(argc == 1);
	status = soa_check(argv[0], resolve_family, &reply);
	status_print(status);
	if (status == dns_success) {
		reply_print(&reply, 1);
		dns_reply_free(&reply);
	}
	return status;
}

static int
res_a(int argc, char **argv)
{
	dns_status status;
	struct dns_reply reply;

	assert(argc == 1);
	status = a_lookup(argv[0], &reply);
	status_print(status);
	if (status == dns_success) {
		reply_print(&reply, 1);
		dns_reply_free(&reply);
	}
	return status;
}

static int
res_aaaa(int argc, char **argv)
{
	dns_status status;
	struct dns_reply reply;

	assert(argc == 1);
	status = aaaa_lookup(argv[0], &reply);
	status_print(status);
	if (status == dns_success) {
		reply_print(&reply, 1);
		dns_reply_free(&reply);
	}
	return status;
}

static int
res_ptr(int argc, char **argv)
{
	dns_status status;
	struct dns_reply reply;

	assert(argc == 1);
	status = dns_resolve_ipstr(argv[0], &reply);
	status_print(status);
	if (status == dns_success) {
		reply_print(&reply, 1);
		dns_reply_free(&reply);
	}
	return status;
}

static int
res_ptr_val(int argc, char **argv)
{
	dns_status status;
	struct dns_reply reply;

	assert(argc == 1);
	status = ptr_validate(argv[0], &reply);
	status_print(status);
	if (status == dns_success) {
		reply_print(&reply, 1);
		dns_reply_free(&reply);
	}
	return status;
}

static int
res_txt(int argc, char **argv)
{
	dns_status status;
	struct dns_reply reply;

	assert(argc == 1);
	status = txt_lookup(argv[0], &reply);
	status_print(status);
	if (status == dns_success) {
		reply_print(&reply, 1);
		dns_reply_free(&reply);
	}
	return status;
}

static int
res_mx(int argc, char **argv)
{
	dns_status status;
	struct dns_reply reply;
	int ip = resolve_none;

	if (argc == 2) {
		if (strcmp(argv[0], "-ip") == 0)
			ip = resolve_ip4;
		else if (strcmp(argv[0], "-ip6") == 0)
			ip = resolve_ip6;
		else {
			mu_error("bad arguments");
			abort();
		}
		argc--;
		argv++;
	}
	assert(argc == 1);
	status = mx_lookup(argv[0], ip, &reply);
	status_print(status);
	if (status == dns_success) {
		reply_print(&reply, 0);
		dns_reply_free(&reply);
	}
	return status;
}

static int
res_ns(int argc, char **argv)
{
	dns_status status;
	struct dns_reply reply;
	int ip = resolve_none;

	if (argc == 2) {
		if (strcmp(argv[0], "-ip") == 0)
			ip = resolve_ip4;
		else if (strcmp(argv[0], "-ip6") == 0)
			ip = resolve_ip6;
		else {
			mu_error("bad arguments");
			abort();
		}
		argc--;
		argv++;
	}
	assert(argc == 1);
	status = ns_lookup(argv[0], ip, &reply);
	status_print(status);
	if (status == dns_success) {
		reply_print(&reply, 0);
		dns_reply_free(&reply);
	}
	return status;
}

static int
res_spf(int argc, char **argv)
{
	dns_status status;
	char *record;

	status = spf_lookup(argv[0], &record);
	status_print(status);
	if (status == dns_success) {
		printf("%s\n", record);
		free(record);
	}
	return status;
}

static int
res_host(int argc, char **argv)
{
	dns_status status;
	struct dns_reply reply;

	assert(argc == 1);
	status = dns_resolve_ipstr(argv[0], &reply);
	status_print(status);
	if (status == dns_success) {
		reply_print_str(&reply, 0);
		dns_reply_free(&reply);
	}
	return status;
}

static int
res_ip(int argc, char **argv)
{
	dns_status status;
	char *str;
	int resolve_family = resolve_ip4;

	if (argc == 2) {
		if (strcmp(argv[0], "-ip") == 0)
			resolve_family = resolve_ip4;
		else if (strcmp(argv[0], "-ip6") == 0)
			resolve_family = resolve_ip6;
		else {
			mu_error("bad arguments");
			abort();
		}
		argc--;
		argv++;
	}
	assert(argc == 1);
	status = dns_resolve_hostname(argv[0], resolve_family, &str);
	status_print(status);
	if (status == dns_success) {
		printf("%s\n", str);
		free(str);
	}
	return status;
}

struct mode {
	char const *name;
	int (*func)(int, char **);
};

struct mode mode[] = {
	{ "soa", res_soa },
	{ "host", res_host },
	{ "ip", res_ip },
	{ "a", res_a },
	{ "aaaa", res_aaaa },
	{ "ptr", res_ptr },
	{ "txt", res_txt },
	{ "mx", res_mx },
	{ "ns", res_ns },
	{ "ptr_val", res_ptr_val },
	{ "spf", res_spf },
	{ NULL }
};

int
main(int argc, char **argv)
{
	int i;
	mu_opool_t op = NULL;
	char *stmt;
	char *resolv_conf = NULL;

	mu_set_program_name(argv[0]);
	mu_stdstream_setup(MU_STDSTREAM_RESET_NONE);
	dnsbase_init();
	while ((i = getopt(argc, argv, "+dc:f:h:qs:S:")) != EOF) {
		switch (i) {
		case 'c':{
			unsigned long n;
			char *p;

			errno = 0;
			n = strtoul(optarg, &p, 10);
			if (errno || *p) {
				mu_error("bad argument to -c");
				exit(1);
			}
			dns_max_cname_chain = n;
			break;
		}
		case 'd':
			mu_debug_enable_category("dns", 3,
					 MU_DEBUG_LEVEL_UPTO(MU_DEBUG_PROT));
			break;
		case 'f':
			resolv_conf = optarg;
			break;
		case 'h':
			mu_asprintf(&stmt, "nameserver %s\n", optarg);
			if (!op)
				mu_opool_create(&op, MU_OPOOL_ENOMEMABRT);
			mu_opool_appendz(op, stmt);
			free(stmt);
			break;
		case 'q':
			quiet++;
			break;
		case 's':
			if (!op)
				mu_opool_create(&op, MU_OPOOL_ENOMEMABRT);
			mu_opool_appendz(op, optarg);
			mu_opool_append_char(op, '\n');
			break;
		case 'S':
			suffix = optarg;
			suffix_len = strlen(suffix);
			break;
		default:
			exit(1);
		}
	}

	assert(!(resolv_conf && op));

	if (resolv_conf)
		dnsbase_file_init(resolv_conf);

	if (op) {
		mu_opool_append_char(op, 0);
		stmt = mu_opool_finish(op, NULL);
		dnsbase_real_init(stmt);
		mu_opool_destroy(&op);
	}

	argc -= optind;
	argv += optind;
	assert(argc > 1);
	for (i = 0; mode[i].name; i++) {
		if (strcmp(mode[i].name, argv[0]) == 0)
			exit(mode[i].func(argc - 1, argv + 1));
	}
	mu_error("%s:%d: unrecognized mode", __FILE__, __LINE__);
	abort();
}
