/*
 * Sample C code to resolve MX records for an address (on linux).
 *
 * compile with:    gcc getmx.c -o getmx -lresolv
 * usage: ./getmx domain_name.com
 *
 * Copyright (C) 2004 by HL Combrinck
 * This code is distributed under the terms of the GNU General Public License.
 *
 */

#include <stdlib.h>
#include <stdio.h>

#include <unistd.h>
#include <string.h>
#include <netinet/in.h>
#include <arpa/nameser.h>
#include <resolv.h>

struct mx 
{
    int pref;
    char host[1024];
};

#ifndef HFIXEDSZ
# define HFIXEDSZ 12
#endif
#ifndef INT16SZ
# define INT16SZ sizeof(cit_int16_t)
#endif
#ifndef INT32SZ
# define INT32SZ sizeof(cit_int32_t)
#endif

/*
 * Compare the preference of two MX records.  Check the actual
 * number listed in the MX record - if they're the same, randomize.
 */
int mxcomp(int p1, int p2)
{
    if (p1 > p2) return(1);
    else if (p1 < p2) return(0);
    else return(rand() % 2);
}


/*
 * sort_mxrecs()
 *
 * Sort MX records
 *
 */
void sort_mxrecs (struct mx *mxrecs, int nmx)
{
	int a, b;
	struct mx t1, t2;

	if (nmx < 2) return;

	for (a = nmx - 2; a >= 0; --a)
	{
		for (b = 0; b <= a; ++b)
		{
			if (mxcomp(mxrecs[b].pref,mxrecs[b+1].pref))
			{
				memcpy(&t1, &mxrecs[b], sizeof(struct mx));
				memcpy(&t2, &mxrecs[b+1], sizeof(struct mx));
				memcpy(&mxrecs[b], &t2, sizeof(struct mx));
				memcpy(&mxrecs[b+1], &t1, sizeof(struct mx));
			}
		}
	}
}



/* 
 * getmx()
 *
 * Get MX recs for an address.
 *
 * Upon success, it fills 'mxbuff' with one or more MX hosts, delimited by
 * ':' chars, and returns the number of hosts.  0 if none found.
 *
 */
int getmx(char *mxbuff, char *dest, int maxbuffsz)
{
	union
	{
		u_char bytes[1024];
		HEADER header;
    } ans;

	int ret;
	unsigned char *startptr, *endptr, *ptr;
	char expanded_buf[1024];
	unsigned short pref, type;
	int n = 0;
	int qdcount;

	struct mx *mxrecs = NULL;
	int nmx = 0;
	
	ret = res_query (dest, C_IN, T_MX, (unsigned char *)ans.bytes, 
		sizeof(ans));

	if (ret < 0) 
	{
		mxrecs = malloc(sizeof(struct mx));
		mxrecs[0].pref = 0;
		strcpy(mxrecs[0].host, dest);
		nmx = 0;
	}
	else 
	{
		if (ret > sizeof(ans)) ret = sizeof(ans);
	
		startptr = &ans.bytes[0];
		endptr = &ans.bytes[ret];
		ptr = startptr + HFIXEDSZ;	/* skip header */
	
		for (qdcount = ntohs(ans.header.qdcount); qdcount--; 
				ptr += ret + QFIXEDSZ)
		{
			if ((ret = dn_skipname(ptr, endptr)) < 0) return(0);
		}
	
		while(1)
		{
			memset (expanded_buf, 0, sizeof(expanded_buf));
			ret = dn_expand (startptr, endptr, ptr, expanded_buf,
					sizeof(expanded_buf));
			if (ret < 0) break;
			ptr += ret;
	
			GETSHORT (type, ptr);
			ptr += INT16SZ + INT32SZ;
			GETSHORT (n, ptr);
	
			if (type != T_MX) ptr += n;
			else
			{
				GETSHORT(pref, ptr);
				ret = dn_expand(startptr, endptr, ptr, expanded_buf,
						sizeof(expanded_buf));
				ptr += ret;
	
				++nmx;
				if (mxrecs == NULL)
					mxrecs = malloc(sizeof(struct mx));
				else
					mxrecs = realloc (mxrecs, (sizeof(struct mx) * nmx));
	
				mxrecs[nmx - 1].pref = pref;
				strcpy(mxrecs[nmx - 1].host, expanded_buf);
			}
		}
	}

	/* sort by MX pref */
	sort_mxrecs(mxrecs, nmx);

	strcpy(mxbuff, "");
	for (n=0; n<nmx; ++n)
	{
		if (strlen(mxbuff)+strlen(mxrecs[n].host) < maxbuffsz)
			strcat(mxbuff, mxrecs[n].host);
		else
			break;
		strcat(mxbuff, ":");
	}
	/* kill last ':' */
	if (mxbuff[strlen(mxbuff)-1] == ':') mxbuff[strlen(mxbuff)-1] = 0;
	free(mxrecs);
	return(nmx);
}

/*
int main(argc, argv)
	int argc;
	char *argv[];
{
	int n;
	char buff[1024];

	n = getmx (buff, argv[1], sizeof(buff)-1);

	if (!n)
	{
		printf ("error: no MX record(s) found\n");
	}
	else
	{
		printf ("recs=%d (%s)\n", n, argv[1]);
		printf ("mx: %s\n", buff);
	}
}
*/