#define _GNU_SOURCE

#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <fcntl.h>
#include <stdlib.h>
#include <unistd.h>
#include <limits.h>
#include <errno.h>

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

#include <netinet/in.h>
#include <netinet/ip.h>
#include <netinet/ip6.h>
#include <netinet/tcp.h>
#include <netinet/udp.h>
#include <netinet/ip_icmp.h>

#include <arpa/inet.h>
#include <net/if.h>

#include <linux/if_packet.h>
#include <linux/if_ether.h>
#include <linux/filter.h>

// for mapping DNS query -> uid
/* #define BUF_SIZE 8192 */
/* #include <errno.h> */
/* #include <linux/inet_diag.h> */
/* #include <linux/netlink.h> */
/* #include <linux/rtnetlink.h> */
/* #include <linux/sock_diag.h> */
/* static void dump_udp_sockets(); */


#include "agent_context.h"
#include "output.h"
#include "sniffer.h"
#include "error.h"
#include "time_common.h"
#include "net.h"
#include "dns.h"

// apply a BPF filter to the socket
// TODO user-specified patterns
void sniffer_apply_filter(sock_t s) {
	// tcpdump -dd '(ip or ip6) and (udp port 53 or (tcp[tcpflags] & (tcp-syn|tcp-fin|tcp-rst) != 0))'
	struct sock_filter bpf[] = {
		{ 0x28, 0, 0, 0x0000000c },
		{ 0x15, 0, 15, 0x00000800 },
		{ 0x30, 0, 0, 0x00000017 },
		{ 0x15, 0, 7, 0x00000011 },
		{ 0x28, 0, 0, 0x00000014 },
		{ 0x45, 19, 0, 0x00001fff },
		{ 0xb1, 0, 0, 0x0000000e },
		{ 0x48, 0, 0, 0x0000000e },
		{ 0x15, 15, 0, 0x00000035 },
		{ 0x48, 0, 0, 0x00000010 },
		{ 0x15, 13, 14, 0x00000035 },
		{ 0x15, 0, 13, 0x00000006 },
		{ 0x28, 0, 0, 0x00000014 },
		{ 0x45, 11, 0, 0x00001fff },
		{ 0xb1, 0, 0, 0x0000000e },
		{ 0x50, 0, 0, 0x0000001b },
		{ 0x45, 7, 8, 0x00000007 },
		{ 0x15, 0, 7, 0x000086dd },
		{ 0x30, 0, 0, 0x00000014 },
		{ 0x15, 0, 5, 0x00000011 },
		{ 0x28, 0, 0, 0x00000036 },
		{ 0x15, 2, 0, 0x00000035 },
		{ 0x28, 0, 0, 0x00000038 },
		{ 0x15, 0, 1, 0x00000035 },
		{ 0x6, 0, 0, 0x00040000 },
		{ 0x6, 0, 0, 0x00000000 },
	};

	struct sock_fprog program = {
		.len = sizeof(bpf) / sizeof(bpf[0]),
		.filter = bpf,
	};

	if (setsockopt(s, SOL_SOCKET, SO_ATTACH_FILTER, &program, sizeof(program)) == -1) {
		error_fatal("FATAL: setsockopt SO_ATTACH_FILTER: %s", strerror(errno));
	}
}

sock_t sniffer_init_interface(const char *interface, bool promisc) {
	if (!interface) {
		error("setup_sniffer: NULL interface");
		return -1;
	}

	sock_t sniffer = socket(AF_PACKET, SOCK_RAW, htons(ETH_P_ALL));
	if (sniffer < 0) {
		error("socket: %s", strerror(errno));
		return -1;
	}

	// set non-blocking
	int flags = fcntl(sniffer, F_GETFL, 0);
	if (flags == -1) {
		error("fcntl F_GETFL: %s", strerror(errno));
		close(sniffer);
		return -1;
	}

	if (fcntl(sniffer, F_SETFL, flags | O_NONBLOCK) == -1) {
		error("fcntl F_SETFL: %s", strerror(errno));
		close(sniffer);
		return -1;
	}

	struct sockaddr_ll sll;
	struct ifreq ifr;

	memset(&ifr, 0, sizeof(ifr));
	strncpy(ifr.ifr_name, interface, IFNAMSIZ - 1);
	if (ioctl(sniffer, SIOCGIFINDEX, &ifr) == -1) {
		error("ioctl SIOCGIFINDEX: %s", strerror(errno));
		close(sniffer);
		return -1;
	}

	memset(&sll, 0, sizeof(sll));
	sll.sll_family = AF_PACKET;
	sll.sll_ifindex = ifr.ifr_ifindex;
	sll.sll_protocol = htons(ETH_P_ALL);
	if (bind(sniffer, (struct sockaddr *)&sll, sizeof(sll)) == -1) {
		error("bind: %s", strerror(errno));
		close(sniffer);
		return -1;
	}

	// TODO Enable promiscuous mode if desired
	if (promisc) {
		struct packet_mreq mreq;
		memset(&mreq, 0, sizeof(mreq));
		mreq.mr_ifindex = ifr.ifr_ifindex;
		mreq.mr_type = PACKET_MR_PROMISC;
		if (setsockopt(sniffer, SOL_PACKET, PACKET_ADD_MEMBERSHIP, &mreq, sizeof(mreq)) == -1) {
			error("setsockopt PACKET_ADD_MEMBERSHIP: %s", strerror(errno));
			close(sniffer);
			return -1;
		}
	}

	sniffer_apply_filter(sniffer);

	return sniffer;
}

static void json_add_tcp_flag_array(json_t *buf, uint8_t flags) {
	json_add_int(buf, "flags_int", flags);
	json_add_array_start(buf, "flags");

	if (flags & TH_SYN)  json_array_add_string(buf, "SYN");
	if (flags & TH_ACK)  json_array_add_string(buf, "ACK");
	if (flags & TH_FIN)  json_array_add_string(buf, "FIN");
	if (flags & TH_RST)  json_array_add_string(buf, "RST");
	if (flags & TH_PUSH) json_array_add_string(buf, "PSH");
	if (flags & TH_URG)  json_array_add_string(buf, "URG");
	if (flags & TH_ECE)  json_array_add_string(buf, "ECE");
	if (flags & TH_CWR)  json_array_add_string(buf, "CWR");

	json_end_array(buf);
}

static void sniffer_log_tcp_event(agent_context_t *ctx,
								  const char *ifname,
								  const char *proto,
								  const char *src_ip, uint16_t src_port,
								  const char *dst_ip, uint16_t dst_port,
								  int ttl_or_hop,
								  uint8_t flags) {
	json_t buf;
	json_init(&buf);
	json_start_object(&buf);

	json_add_double(&buf, "timestamp", timestamp());
	json_add_string(&buf, "hostname", ctx->hostname);
	json_add_string(&buf, "event_type", "tcp");
	json_add_string(&buf, "interface", ifname);
	json_add_string(&buf, "protocol", proto);
	json_add_string(&buf, "source_ip", src_ip);
	json_add_int(&buf, "source_port", src_port);
	json_add_string(&buf, "dest_ip", dst_ip);
	json_add_int(&buf, "dest_port", dst_port);
	json_add_int(&buf, (strcmp(proto, "ipv6") == 0) ? "hop_limit" : "ttl", ttl_or_hop);
	json_add_tcp_flag_array(&buf, flags);

	json_end_object(&buf);
	output(json_get(&buf));
	json_free(&buf);
}

static void sniffer_log_dns_event(agent_context_t *ctx, const char *ifname, const struct dns_question *q) {
	json_t buf;
	json_init(&buf);
	json_start_object(&buf);
	json_add_double(&buf, "timestamp", timestamp());
	json_add_string(&buf, "hostname", ctx->hostname);
	json_add_string(&buf, "event_type", "dns_query");
	json_add_string(&buf, "interface", ifname);
	json_add_string(&buf, "question", q->name);
	json_add_int(&buf, "qtype", q->qtype);
	json_add_string(&buf, "qtype_str", dns_type_to_string(q->qtype));
	json_end_object(&buf);
	output(json_get(&buf));
	json_free(&buf);
}

static void sniffer_parse_ipv4(agent_context_t *ctx, const char *ifname, uint8_t *data, size_t len) {
	if (len < sizeof(struct iphdr)) {
		return;
	}

	struct iphdr *ip_header = (struct iphdr *)data;

	// TODO handle fragmentation better.
	if (ntohs(ip_header->frag_off) & 0x1FFF) {
		//printf("Skipping fragmented IPv4 packet\n");
		return;
	}

	struct sockaddr_in source, dest;
	memset(&source, 0, sizeof(source));
	source.sin_addr.s_addr = ip_header->saddr;
	memset(&dest, 0, sizeof(dest));
	dest.sin_addr.s_addr = ip_header->daddr;

	char src_ip[INET_ADDRSTRLEN];
	char dst_ip[INET_ADDRSTRLEN];
	inet_ntop(AF_INET, &source.sin_addr, src_ip, sizeof(src_ip));
	inet_ntop(AF_INET, &dest.sin_addr, dst_ip, sizeof(dst_ip));

	uint8_t *transport = data + ip_header->ihl * 4;

	switch (ip_header->protocol) {
	case IPPROTO_TCP:
		if (len < (transport - data) + sizeof(struct tcphdr)) {
			// not big enough for TCP header
			return;
		}

		struct tcphdr *tcp = (struct tcphdr *)transport;

		size_t tcp_header_len = tcp->doff * 4;
		if (len < (transport - data) + tcp_header_len) {
			// incomplete TCP header
			return;
		}

		uint16_t tcp_sport = ntohs(tcp->source);
		uint16_t tcp_dport = ntohs(tcp->dest);

		// TODO map socket to PID
		sniffer_log_tcp_event(ctx, ifname, "ipv4", src_ip, tcp_sport, dst_ip, tcp_dport, ip_header->ttl, tcp->th_flags);
		break;

	case IPPROTO_UDP:
		if (len < (transport - data) + sizeof(struct udphdr)) {
			// not big enough for UDP header
			return;
		}

		struct udphdr *udp = (struct udphdr *)transport;
		uint16_t udp_sport = ntohs(udp->source);
		uint16_t udp_dport = ntohs(udp->dest);

		if (udp_sport == 53 || udp_dport == 53) { // DNS
			size_t dns_len = len - (transport - data) - sizeof(struct udphdr);
			const uint8_t *dns_payload = (const uint8_t *)(udp + 1);

			//pid_t pid = match_udp_inode(ip_header->saddr, sport);
			struct dns_question qs[MAX_DNS_QUESTIONS];
			size_t count = parse_dns_udp(dns_payload, dns_len, qs, MAX_DNS_QUESTIONS);
			for (int i = 0; i < count; i++) {
				sniffer_log_dns_event(ctx, ifname, &qs[i]);
			}
		}
		break;

	case IPPROTO_ICMP:
		//printf("ICMP: %s -> %s\n", src_ip, dst_ip);
		break;

	default:
		// shouldn't get here due to bpf filter
		//printf("Unknown IPv4 protocol: %d\n", ip_header->protocol);
		break;
	}
}

static void sniffer_parse_ipv6(agent_context_t *ctx, const char *ifname, uint8_t *data, size_t len) {
	if (len < sizeof(struct ip6_hdr))
		return;

	struct ip6_hdr *ip6h = (struct ip6_hdr *)data;

	char src_ip[INET6_ADDRSTRLEN];
	char dst_ip[INET6_ADDRSTRLEN];
	inet_ntop(AF_INET6, &ip6h->ip6_src, src_ip, sizeof(src_ip));
	inet_ntop(AF_INET6, &ip6h->ip6_dst, dst_ip, sizeof(dst_ip));

	uint8_t *transport = data + sizeof(struct ip6_hdr);
	size_t transport_len = len - sizeof(struct ip6_hdr);

	switch (ip6h->ip6_nxt) {
	case IPPROTO_TCP:
		if (transport_len < sizeof(struct tcphdr)) {
			return;
		}
		struct tcphdr *tcp = (struct tcphdr *)transport;
		uint16_t tcp_sport = ntohs(tcp->source);
		uint16_t tcp_dport = ntohs(tcp->dest);

		// TODO map socket to PID
		sniffer_log_tcp_event(ctx, ifname, "ipv6", src_ip, tcp_sport, dst_ip, tcp_dport, ip6h->ip6_hlim, tcp->th_flags);
		break;

	case IPPROTO_UDP:
		if (transport_len < sizeof(struct udphdr)) {
			return;
		}
		struct udphdr *udp = (struct udphdr *)transport;
		uint16_t udp_sport = ntohs(udp->source);
		uint16_t udp_dport = ntohs(udp->dest);

		if (udp_sport == 53 || udp_dport == 53) { // DNS
			size_t dns_len = len - (transport - data) - sizeof(struct udphdr);
			const uint8_t *dns_payload = (const uint8_t *)(udp + 1);

			//pid_t pid = match_udp_inode(ip_header->saddr, sport);
			struct dns_question qs[MAX_DNS_QUESTIONS];
			int count = parse_dns_udp(dns_payload, dns_len, qs, MAX_DNS_QUESTIONS);
			for (int i = 0; i < count; i++) {
				sniffer_log_dns_event(ctx, ifname, &qs[i]);
			}
		}
		break;

	case IPPROTO_ICMPV6:
		//printf("ICMPv6: %s -> %s\n", src_ip, dst_ip);
		break;

	default:
		// shouldn't get here due to bpf filter
		//printf("Unknown IPv6 NextHeader: %d\n", ip6h->nexthdr);
		break;
	}
}

static void sniffer_parse_ethernet(agent_context_t *ctx, const char *ifname, uint8_t *buffer, size_t len) {
	if (len < sizeof(struct ethhdr)) {
		return;
	}

	struct ethhdr *eth = (struct ethhdr *)buffer;
	uint16_t       ethertype = ntohs(eth->h_proto);
	size_t         offset = sizeof(struct ethhdr);

	// TODO log vlan tags
	while (ethertype == ETH_P_8021Q) {
		if (len < offset + 4)
			return;
		ethertype = ntohs(*(uint16_t *)(buffer + offset + 2));
		offset += 4;
	}

	switch(ethertype) {
	case ETH_P_IP:
		sniffer_parse_ipv4(ctx, ifname, buffer + offset, len - offset);
		break;
	case ETH_P_IPV6:
		sniffer_parse_ipv6(ctx, ifname, buffer + offset, len - offset);
		break;
	default:
		// unknown/unsupported Ethertype
		break;
	}
}

void sniffer_handle_packet(sock_t sniffer, agent_context_t *ctx) {
	uint8_t buffer[2048];
	struct sockaddr_ll sll;

	struct iovec iov = {
		.iov_base = buffer,
		.iov_len = sizeof(buffer)
	};

	struct msghdr msg = {
		.msg_name = &sll,
		.msg_namelen = sizeof(sll),
		.msg_iov = &iov,
		.msg_iovlen = 1
	};

	ssize_t data_size = recvmsg(sniffer, &msg, 0);
	if (data_size < 0) {
		error("recvmsg: %s", strerror(errno));
		return;
	}

	char ifname[IFNAMSIZ] = "unknown";
	if_indextoname(sll.sll_ifindex, ifname);

	sniffer_parse_ethernet(ctx, ifname, buffer, data_size);
}

// This attempts to parse /proc/net/udp to find the inode of the
// socket, then walk /proc/PID/fd/* to correlate it to a process. this
// is too slow to give meaningful results most of the time
/* pid_t match_udp_inode(uint32_t ip_be, uint16_t port_be) { */
/* 	FILE *fp = fopen("/proc/net/udp", "r"); */
/* 	if (!fp) { */
/* 		return 0; */
/* 	} */

/* 	char line[512]; */
/* 	fgets(line, sizeof(line), fp); // skip header */

/* 	char local_addr[64]; */
/* 	uint32_t ip_hex; */
/* 	uint32_t port_hex; */
/* 	unsigned long inode = 0; */

/* 	while (fgets(line, sizeof(line), fp)) { */
/* 		sscanf(line, "%*d: %64[0-9A-Fa-f]:%x %*s %*s %*s %*s %*s %lu", */
/* 			   local_addr, */
/* 			   &port_hex, */
/* 			   &inode); */

/* 		sscanf(local_addr, "%x", &ip_hex); */

/* 		if (ip_hex == ntohl(ip_be) && port_hex == ntohs(port_be)) { */
/* 			break; */
/* 		} else { */
/* 			inode = 0; */
/* 		} */
/* 	} */

/* 	fclose(fp); */
/* 	if (!inode) { */
/* 		return 0; */
/* 	} */

/* 	DIR *proc = opendir("/proc"); */
/* 	if (!proc) { */
/* 		return 0; */
/* 	} */

/* 	struct dirent *entry; */
/* 	while ((entry = readdir(proc))) { */
/* 		if (!isdigit(entry->d_name[0])) { */
/* 			continue; */
/* 		} */

/* 		char fd_dir[PATH_MAX]; */
/* 		snprintf(fd_dir, sizeof(fd_dir), "/proc/%s/fd", entry->d_name); */

/* 		DIR *fd = opendir(fd_dir); */
/* 		if (!fd) { */
/* 			continue; */
/* 		} */

/* 		struct dirent *fd_entry; */
/* 		while((fd_entry = readdir(fd))) { */
/* 			char link_path[PATH_MAX + 1024]; */
/* 			char target[PATH_MAX]; */

/* 			snprintf(link_path, sizeof(link_path), "%s/%s", fd_dir, fd_entry->d_name); */
/* 			ssize_t len = readlink(link_path, target, sizeof(target) - 1); */
/* 			if (len == -1) { */
/* 				continue; */
/* 			} */
/* 			target[len] = '\0'; */

/* 			char socket_str[64] = {0}; */
/* 			snprintf(socket_str, sizeof(socket_str), "socket:[%lu]", inode); */
/* 			if (strcmp(target, socket_str) == 0) { */
/* 				closedir(fd); */
/* 				closedir(proc); */
/* 				return atoi(entry->d_name); */
/* 			} */
/* 		} */

/* 		closedir(fd); */
/* 	} */

/* 	closedir(proc); */

/* 	return 0; */
/* } */



// TODO investigate this further. Would be nice to correlate DNS
// requests to a process id, but this search method is too slow.

// This is able to get the UID initiating the DNS query and inode, but
// traversing /proc to match inode->PID is too slow.
/* static void dump_udp_sockets() { */
/*     int sock = socket(AF_NETLINK, SOCK_DGRAM, NETLINK_INET_DIAG); */
/*     if (sock < 0) { */
/*         perror("socket"); */
/*         exit(1); */
/*     } */

/*     char req_buf[BUF_SIZE]; */
/*     memset(req_buf, 0, sizeof(req_buf)); */

/*     struct nlmsghdr *nlh = (struct nlmsghdr *)req_buf; */
/*     struct inet_diag_req_v2 *req = (struct inet_diag_req_v2 *)(nlh + 1); */

/*     nlh->nlmsg_len = NLMSG_LENGTH(sizeof(*req)); */
/*     nlh->nlmsg_type = SOCK_DIAG_BY_FAMILY; */
/*     nlh->nlmsg_flags = NLM_F_REQUEST | NLM_F_DUMP; */
/*     nlh->nlmsg_seq = 123456; */

/*     req->sdiag_family = AF_INET; */
/*     req->sdiag_protocol = IPPROTO_UDP; */
/*     req->idiag_ext = 0; */
/*     req->idiag_states = -1; */

/*     struct sockaddr_nl sa = {0}; */
/*     sa.nl_family = AF_NETLINK; */

/*     if (sendto(sock, req_buf, nlh->nlmsg_len, 0, (struct sockaddr *)&sa, sizeof(sa)) < 0) { */
/*         perror("sendto"); */
/*         close(sock); */
/*         exit(1); */
/*     } */

/*     char buffer[BUF_SIZE]; */
/*     while (1) { */
/*         int len = recv(sock, buffer, sizeof(buffer), 0); */
/*         if (len < 0) { */
/*             perror("recv"); */
/*             break; */
/*         } */

/*         struct nlmsghdr *nlh = (struct nlmsghdr *)buffer; */
/*         for (; NLMSG_OK(nlh, len); nlh = NLMSG_NEXT(nlh, len)) { */
/*             if (nlh->nlmsg_type == NLMSG_DONE) */
/*                 return; */
/*             if (nlh->nlmsg_type == NLMSG_ERROR) { */
/*                 fprintf(stderr, "netlink error\n"); */
/*                 return; */
/*             } */

/*             struct inet_diag_msg *diag = (struct inet_diag_msg *)NLMSG_DATA(nlh); */
/*             struct sockaddr_in src = { */
/*                 .sin_family = AF_INET, */
/*                 .sin_port = diag->id.idiag_sport, */
/*                 .sin_addr.s_addr = diag->id.idiag_src[0], */
/*             }; */
/*             struct sockaddr_in dst = { */
/*                 .sin_family = AF_INET, */
/*                 .sin_port = diag->id.idiag_dport, */
/*                 .sin_addr.s_addr = diag->id.idiag_dst[0], */
/*             }; */

/*             char src_ip[INET_ADDRSTRLEN]; */
/*             char dst_ip[INET_ADDRSTRLEN]; */
/*             inet_ntop(AF_INET, &src.sin_addr, src_ip, sizeof(src_ip)); */
/*             inet_ntop(AF_INET, &dst.sin_addr, dst_ip, sizeof(dst_ip)); */

/*             printf("UDP %s:%d -> %s:%d  inode=%u uid=%u\n", */
/*                    src_ip, ntohs(src.sin_port), */
/*                    dst_ip, ntohs(dst.sin_port), */
/*                    diag->idiag_inode, */
/*                    diag->idiag_uid); */
/*         } */
/*     } */

/*     close(sock); */
/* } */