#include <asm-generic/socket.h>
#include <net/if.h>
#include <netinet/in.h>
#include <netinet/udp.h>
#include <netinet/ip.h>
#include <netpacket/packet.h>
#include <linux/filter.h>
#include <linux/if_ether.h>
#include <arpa/inet.h>
#include <sys/socket.h>
#include <sys/ioctl.h>
#include <err.h>
#include <errno.h>

#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <stdint.h>
#include <stdbool.h>

#define PACKETBUFLEN (1024 * 1024)
#define HDRLEN (sizeof(struct ethhdr) + sizeof(struct iphdr) + sizeof(struct udphdr))

struct dnshdr {
	uint16_t id;
	struct {
		uint16_t qr : 1;
		uint16_t op : 4;
		uint16_t aa : 1;
		uint16_t tc : 1;
		uint16_t rd : 1;
		uint16_t ra : 1;
		uint16_t zr : 3;
		uint16_t rc : 4;
	} flags;
	uint16_t questions;
	uint16_t answers;
	uint16_t authorities;
	uint16_t additionals;
} __attribute__((packed));

struct dnsdec {
	struct dnshdr *hdr;
	size_t pos;

	char *qname;
	uint16_t type;
	uint16_t class;
};

struct packet {
	struct ethhdr *eth;
	struct iphdr *ip;
	struct udphdr *udp;
	uint8_t *data;
	size_t len;
};

struct filter {
	uint32_t src_ip, dst_ip;
	uint16_t dst_port;
};

static struct sock_filter udp_filter[] = {
	{ BPF_LD + BPF_H + BPF_ABS,  0, 0,     12 },
	{ BPF_JMP + BPF_JEQ + BPF_K, 0, 1,  0x800 },
	{ BPF_LD + BPF_B + BPF_ABS,  0, 0,     23 },
	{ BPF_JMP + BPF_JEQ + BPF_K, 0, 1,   0x11 },
	{ BPF_RET + BPF_K,           0, 0, 0xffff },
	{ BPF_RET + BPF_K,           0, 0, 0x0000 },
};

static int sock, ifid;

void
usage(void)
{
	printf("Usage: dnsniff IFACE\n");
	exit(0);
}

void
hexdump(const uint8_t *data, size_t len)
{
	size_t i;

	for (i = 0; i < len; i++) {
		printf("%02X ", data[i]);
		if ((i + 1) % 16 == 0)
			printf("\n");
	}
	if (len % 16 != 0)
		printf("\n");
}

void
sniff_init(const char *ifname)
{
	struct ifreq ifreq;
	struct packet_mreq mreq;
	struct sock_fprog prog;
	struct sockaddr_ll sa;
	int ret;

	sock = socket(AF_PACKET, SOCK_RAW, htons(ETH_P_ALL));
	if (sock < 0) err(1, "failed to open socket");

	/* get interface index */
	strncpy(ifreq.ifr_name, ifname, IFNAMSIZ);
	ret = ioctl(sock, SIOCGIFINDEX, &ifreq);
	if (ret) err(1, "failed to lookup interface index");
	ifid = ifreq.ifr_ifindex;

	/* enable promiscuous mode on interface */
	mreq.mr_ifindex = ifid;
	mreq.mr_type = PACKET_MR_PROMISC;
	ret = setsockopt(sock, SOL_PACKET, PACKET_ADD_MEMBERSHIP,
		&mreq, sizeof(mreq));
	if (ret) err(1, "failed to enable promiscuous mode");

	/* filter for udp */
	prog.len = sizeof(udp_filter) / sizeof(struct sock_filter);
	prog.filter = udp_filter;
	ret = setsockopt(sock, SOL_SOCKET, SO_ATTACH_FILTER,
		&prog, sizeof(prog));
	if (ret) err(1, "failed to attach udp filter");

	/* bind socket to interface */
	sa.sll_family = AF_PACKET;
	sa.sll_ifindex = ifid;
	sa.sll_protocol = htons(ETH_P_ALL);
	ret = bind(sock, (struct sockaddr *)&sa, sizeof(sa));
	if (ret) err(1, "failed to bind interface");
}

bool
sniff(struct packet *pkt, struct filter *flt, uint8_t *buf, size_t len)
{
	char src_tmp[INET_ADDRSTRLEN];
	ssize_t res;

	while (true) {
		res = recvfrom(sock, buf, len, 0, NULL, 0);
		if (res < 0 && errno == EINTR)
			return false;

		if (res < 0) err(1, "failed to read packet");

		if (res < HDRLEN) continue;

		pkt->eth = (struct ethhdr *) buf;
		pkt->ip = (struct iphdr *) (&pkt->eth[1]);
		pkt->udp = (struct udphdr *) (&pkt->ip[1]);
		pkt->data = (uint8_t *) (&pkt->udp[1]);
		pkt->len = htons(pkt->udp->len) - sizeof(struct udphdr);

		if (flt->src_ip && flt->src_ip != pkt->ip->saddr)
			continue;

		if (flt->dst_ip && flt->dst_ip != pkt->ip->daddr)
			continue;

		if (flt->dst_port && flt->dst_port != pkt->udp->dest)
			continue;

		break;
	}

	return true;
}

void
sniff_deinit(void)
{
	struct packet_mreq mreq;
	int ret;

	/* disable promiscuous mode on interface */
	mreq.mr_ifindex = ifid;
	mreq.mr_type = PACKET_MR_PROMISC;
	ret = setsockopt(sock, SOL_PACKET, PACKET_DROP_MEMBERSHIP,
		&mreq, sizeof(mreq));
	if (ret) err(1, "failed to enable promiscuous mode");
}

void
decode_init(struct dnsdec *dec)
{
	memset(dec, 0, sizeof(struct dnsdec));
}

bool
decode_query(struct dnsdec *dec, const uint8_t *data, size_t len)
{
	const uint8_t *end;
	char *c;

	if (!dec->pos) {
		dec->hdr = (struct dnshdr *) data;
		dec->hdr->questions = ntohs(dec->hdr->questions);
		dec->pos = sizeof(struct dnshdr);
	}

	end = memchr(data + dec->pos, '\0', len - dec->pos);
	if (!end || end + 5 > data + len) return false;
	dec->qname = strndup((const char *) data + dec->pos,
		end - data - dec->pos);
	if (!dec->qname) err(1, "strndup");
	for (c = dec->qname; *c; c++) {
		if (*c == 0x0a) *c = '.';
		if (*c == 0x03) *c = '.';
		if (*c == 0x09) *c = '.';
	}

	dec->type = ntohs(*(uint16_t *)(end + 1));
	dec->class = ntohs(*(uint16_t *)(end + 3));
	dec->pos = end + 5 - data;

	return true;
}

void
decode_deinit(struct dnsdec *dec)
{
	free(dec->qname);
}

int
main(int argc, const char **argv)
{
	const char *ifname;
	char src_ip[INET_ADDRSTRLEN];
	char dst_ip[INET_ADDRSTRLEN];
	struct dnsdec dec;
	struct filter flt;
	struct packet pkt;
	uint8_t *buf;

	if (argc != 2) usage();
	ifname = argv[1];

	buf = malloc(PACKETBUFLEN);
	if (!buf) err(1, "malloc");

	sniff_init(ifname);

	flt.dst_ip = 0;
	flt.src_ip = 0;
	flt.dst_port = htons(53);
	while (sniff(&pkt, &flt, buf, PACKETBUFLEN)) {
		inet_ntop(AF_INET, (struct in_addr *) &pkt.ip->saddr,
			src_ip, INET_ADDRSTRLEN);
		inet_ntop(AF_INET, (struct in_addr *) &pkt.ip->daddr,
			dst_ip, INET_ADDRSTRLEN);
		printf("%s -> %s\n", src_ip, dst_ip);
		decode_init(&dec);
		while (decode_query(&dec, pkt.data, pkt.len))
			printf("  %s\n", dec.qname);
		decode_deinit(&dec);
		printf("\n");
	}

	free(buf);
	sniff_deinit();
}