/*****************************************************
 * This code was compiled and tested on Ubuntu 18.04.1
 * with kernel version 4.15.0
 *****************************************************/

#include <linux/init.h>
#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/netfilter.h>
#include <linux/netfilter_ipv4.h>
#include <linux/ip.h>
#include <linux/tcp.h>
#include <linux/udp.h>


static struct nf_hook_ops *nfho = NULL;

struct last_dns_reply {
	int last_port;
	int match_count;
};

struct last_dns_reply dns_info;
int response_thresh = 10;


static int handle_dns_resp(int cli_port, struct sock *sk) {

	int safe_resp = 1;
	printk(KERN_INFO "response headed for client port: %d, lp: %d, mc: %d\n", cli_port, dns_info.last_port, dns_info.match_count);
	
	if (cli_port == dns_info.last_port) {
		dns_info.match_count = dns_info.match_count + 1;

		if (dns_info.match_count > response_thresh) {
			printk(KERN_INFO "got dns port match count over thresh, so dropping response..\n");
			safe_resp = 0;
		}
	} else {
		dns_info.match_count = 0;
	}
	
	dns_info.last_port = cli_port;
	
	return safe_resp;
}


static unsigned int hfunc(void *priv, struct sk_buff *skb, const struct nf_hook_state *state)
{
	struct iphdr *iph;
	struct udphdr *udph;
	int safe_resp;

	if (!skb)
		return NF_ACCEPT;

	iph = ip_hdr(skb);
	if (iph->protocol == IPPROTO_UDP) {
		udph = udp_hdr(skb);
		if (ntohs(udph->source) == 53) {
			printk(KERN_INFO "handle dns response packet info..\n");
			
			safe_resp = handle_dns_resp(udph->dest, skb->sk);
			if (safe_resp) { 
				return NF_ACCEPT;			
			} else {
				printk(KERN_INFO "dropping unsafe dns resp for client port: %d\n", udph->dest);
			       	return NF_DROP;
			}
			
		}
	}
	else if (iph->protocol == IPPROTO_TCP) {
		return NF_ACCEPT;
	}
	
	return NF_DROP;
}

static int __init LKM_init(void)
{
	nfho = (struct nf_hook_ops*)kcalloc(1, sizeof(struct nf_hook_ops), GFP_KERNEL);
	
	/* Initialize netfilter hook */
	nfho->hook 	= (nf_hookfn*)hfunc;		/* hook function */
	nfho->hooknum 	= NF_INET_PRE_ROUTING;		/* received packets */
	nfho->pf 	= PF_INET;			/* IPv4 */
	nfho->priority 	= NF_IP_PRI_FIRST;		/* max hook priority */
	
	nf_register_net_hook(&init_net, nfho);

	return 1;
}

static void __exit LKM_exit(void)
{
	nf_unregister_net_hook(&init_net, nfho);
	kfree(nfho);
}

module_init(LKM_init);
module_exit(LKM_exit);