#include <unistd.h>
#include <string.h>
#include <stdint.h>
#include <sys/socket.h>
#include <linux/if_packet.h>
#include <net/ethernet.h>
#include <netinet/ip.h>
#include <linux/bpf.h>
#include <linux/filter.h>
#include <arpa/inet.h>
#include <errno.h>

/* Remove later */
#include <time.h>

#include "main.h"
#include "uring.h"
#include "igmp.h"

struct igmp {
	struct uring_task task;
	struct uring_task_buf tbuf;
};

#define ETH_HDR_LEN		14
#define IPV4_MIN_HDR_LEN	20
#define IGMP_MIN_LEN		8

struct __attribute__((packed, scalar_storage_order("big-endian"))) ipv4_hdr {
	unsigned version:4;
	unsigned ihl:4;
	unsigned dscp:6;
	unsigned ecn:2;
	unsigned length:16;
	unsigned id:16;
	unsigned flags:3;
	unsigned fragment_offset:13;
	unsigned ttl:8;
	unsigned protocol:8;
	unsigned checksum:16;
	unsigned src:32;
	unsigned dst:32;
	unsigned options[];
};

enum igmp_type {
	IGMP_MEMBERSHIP_QUERY		= 0x11,
	IGMP_V1_MEMBERSHIP_REPORT	= 0x12,
	IGMP_V2_MEMBERSHIP_REPORT	= 0x16,
	IGMP_V3_MEMBERSHIP_REPORT	= 0x22,
	IGMP_V2_LEAVE_GROUP		= 0x17
};

union igmp_msg {
	struct __attribute__((packed, scalar_storage_order("big-endian"))) {
		enum igmp_type type:8;
		unsigned unknown:8;
		unsigned checksum:16;
	} common;

	struct __attribute__((packed, scalar_storage_order("big-endian"))) {
		enum igmp_type type:8;
		unsigned resptime:8;
		unsigned checksum:16;
		unsigned addr:32;
	} v2;

	struct __attribute__((packed, scalar_storage_order("big-endian"))) {
		enum igmp_type type:8;
		unsigned reserved1:8;
		unsigned checksum:16;
		unsigned reserved2:16;
		unsigned nrecs:16;
		uint8_t records[];
	} v3;
};

enum igmp_v3_record_type {
	IGMP_V3_REC_MODE_IS_INCL	= 1,
	IGMP_V3_REC_MODE_IS_EXCL	= 2,
	IGMP_V3_REC_MODE_CH_INCL	= 3,
	IGMP_V3_REC_MODE_CH_EXCL	= 4
};

union igmp_v3_record {
	struct __attribute__((packed, scalar_storage_order("big-endian"))) {
		enum igmp_v3_record_type type:8;
		unsigned auxlen:8;
		unsigned nsrcs:16;
		unsigned addr:32;
		uint32_t saddr[];
	};
};

static inline unsigned short
from32to16(unsigned int x)
{
	/* add up 16-bit and 16-bit for 16+c bit */
	x = (x & 0xffff) + (x >> 16);
	/* add up carry.. */
	x = (x & 0xffff) + (x >> 16);
	return x;
}

static unsigned int
do_csum(const unsigned char *buff, int len)
{
	int odd;
	unsigned int result = 0;

	if (len <= 0)
		goto out;
	odd = 1 & (unsigned long) buff;
	if (odd) {
#if __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__
		result += (*buff << 8);
#else
		result = *buff;
#endif
		len--;
		buff++;
	}
	if (len >= 2) {
		if (2 & (unsigned long) buff) {
			result += *(unsigned short *) buff;
			len -= 2;
			buff += 2;
		}
		if (len >= 4) {
			const unsigned char *end = buff + ((unsigned)len & ~3);
			unsigned int carry = 0;
			do {
				unsigned int w = *(unsigned int *) buff;
				buff += 4;
				result += carry;
				result += w;
				carry = (w > result);
			} while (buff < end);
			result += carry;
			result = (result & 0xffff) + (result >> 16);
		}
		if (len & 2) {
			result += *(unsigned short *) buff;
			buff += 2;
		}
	}
	if (len & 1)
#if __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__
		result += *buff;
#else
		result += (*buff << 8);
#endif
	result = from32to16(result);
	if (odd)
		result = ((result >> 8) & 0xff) | ((result & 0xff) << 8);
out:
	return result;
}

static inline bool
csum_valid(const char *buf, size_t len)
{
	return do_csum((unsigned const char *)buf, len) == 0xffff;
}

static void
igmp_match()
{
	/* Sent with approx 120-130 sec intervals */
	time_t t = time(NULL);
	struct tm *tm = localtime(&t);
	char s[64];
	strftime(s, sizeof(s), "%c", tm);
	debug(DBG_IGMP, "multicast request discovered at: %s\n", s);
	//start announce
}

static void
igmp_parse(struct cfg *cfg, struct igmp *igmp)
{
	char *buf = igmp->task.tbuf->buf;
	size_t len = igmp->task.tbuf->len;
	struct ipv4_hdr *hdr = (struct ipv4_hdr *)buf;
	size_t body_len;
	union igmp_msg *igmp_msg;

	if (len <= IPV4_MIN_HDR_LEN)
		return;

	if (hdr->version != 4)
		return;

	if (hdr->ihl * 4 < IPV4_MIN_HDR_LEN)
		return;

	if (hdr->length < hdr->ihl * 4)
		return;

	if (hdr->length != len)
		return;

	if (hdr->fragment_offset > 0)
		return;

	if (hdr->protocol != IPPROTO_IGMP)
		return;

	if (!csum_valid(buf, hdr->ihl * 4))
		return;

	body_len = hdr->length - hdr->ihl * 4;
	igmp_msg = (union igmp_msg *)(buf + hdr->ihl * 4);

	if (body_len < IGMP_MIN_LEN)
		return;

	switch (igmp_msg->common.type) {
	case IGMP_V1_MEMBERSHIP_REPORT:
		debug(DBG_IGMP, "igmp_v1_membership_report\n");
		/* fall through */

	case IGMP_V2_MEMBERSHIP_REPORT: {
		struct in_addr src;
		char src_str[INET_ADDRSTRLEN];
		struct in_addr dst;
		char dst_str[INET_ADDRSTRLEN];
		struct in_addr grp;
		char grp_str[INET_ADDRSTRLEN];

		src.s_addr = htonl(hdr->src);
		inet_ntop(AF_INET, &src, src_str, sizeof(src_str));
		dst.s_addr = htonl(hdr->dst);
		inet_ntop(AF_INET, &dst, dst_str, sizeof(dst_str));
		grp.s_addr = htonl(igmp_msg->v2.addr);
		inet_ntop(AF_INET, &grp, grp_str, sizeof(grp_str));

		debug(DBG_IGMP, "igmp_v2_membership_report %s -> %s (%s)\n",
		      src_str, dst_str, grp_str);

		if (body_len != IGMP_MIN_LEN) {
			error("IGMPv2 invalid size\n");
			break;
		}

		if (!csum_valid((char *)igmp_msg, body_len)) {
			error("IGMPv2 invalid checksum\n");
			break;
		}

		debug(DBG_IGMP, "Inet addr: 0x%x\n", inet_addr("224.0.2.60"));
		debug(DBG_IGMP, "Inet addr: 0x%x\n", cinet_addr(224,0,2,60));
		debug(DBG_IGMP, "Inet addr: 0x%x\n", chtobe32(cinet_addr(224,0,2,60)));
		if (htonl(hdr->dst) != cinet_addr(224,0,2,60)) {
			debug(DBG_IGMP, "IGMPv2 invalid dst addr\n");
			break;
		}

		if (htonl(igmp_msg->v2.addr) != cinet_addr(224,0,2,60)) {
			debug(DBG_IGMP, "IGMPv2 invalid grp addr\n");
			break;
		}

		igmp_match();
		break;
	}

	case IGMP_V3_MEMBERSHIP_REPORT: {
		char *pos = (char *)igmp_msg;
		struct in_addr src;
		char src_str[INET_ADDRSTRLEN];
		struct in_addr dst;
		char dst_str[INET_ADDRSTRLEN];

		src.s_addr = htonl(hdr->src);
		inet_ntop(AF_INET, &src, src_str, sizeof(src_str));
		dst.s_addr = htonl(hdr->dst);
		inet_ntop(AF_INET, &dst, dst_str, sizeof(dst_str));

		debug(DBG_IGMP, "igmp_v3_membership_report %s -> %s\n",
		      src_str, dst_str);

		debug(DBG_IGMP, "IGMPv3\n"
		      " * Type: %x\n"
		      " * Reserved: %u\n"
		      " * Csum: %u\n"
		      " * Reserved: %u\n"
		      " * NRecs: %u\n"
		      " * Size: %zu bytes\n",
		      igmp_msg->v3.type,
		      igmp_msg->v3.reserved1,
		      igmp_msg->v3.checksum,
		      igmp_msg->v3.reserved2,
		      igmp_msg->v3.nrecs,
		      sizeof(igmp_msg->v3));

		if (!csum_valid(pos, body_len)) {
			error("IGMPv3 csum invalid\n");
			break;
		}

		if (htonl(hdr->dst) != cinet_addr(224,0,0,22)) {
			debug(DBG_IGMP, "IGMPv2 invalid dst addr\n");
			break;
		}

		body_len -= sizeof(igmp_msg->v3);
		pos += sizeof(igmp_msg->v3);

		for (unsigned rec = 0; rec < igmp_msg->v3.nrecs; rec++) {
			union igmp_v3_record *record = (union igmp_v3_record *)pos;
			struct in_addr grp;
			char grp_str[INET_ADDRSTRLEN];

			if (body_len < sizeof(*record)) {
				error("IGMPv3 too short\n");
				break;
			}

			grp.s_addr = htonl(record->addr);
			inet_ntop(AF_INET, &grp, grp_str, sizeof(grp_str));
			debug(DBG_IGMP, "received IGMPv3 record to %s\n", grp_str);

			debug(DBG_IGMP, "IGMPv3 rec\n"
			      " * Type: %u\n"
			      " * Auxlen: %u\n"
			      " * NSrcs: %u\n"
			      " * Addr: %s\n"
			      " * Size: %zu bytes\n",
			      record->type,
			      record->auxlen,
			      record->nsrcs,
			      grp_str,
			      sizeof(*record));

			body_len -= sizeof(*record);
			pos += sizeof(*record);

			if (body_len < record->nsrcs * sizeof(uint32_t) + record->auxlen) {
				error("IGMPv3 too short\n");
				break;
			}

			for (unsigned addr = 0; addr < record->nsrcs; addr++) {
				struct in_addr grp_src;
				char grp_src_str[INET_ADDRSTRLEN];

				grp_src.s_addr = htonl(record->saddr[addr]);
				inet_ntop(AF_INET, &grp_src, grp_src_str, sizeof(grp_src_str));
				debug(DBG_IGMP, "received IGMPv3 record src %s\n",
				      grp_src_str);

				body_len -= sizeof(record->saddr[addr]);
				pos += sizeof(record->saddr[addr]);
			}

			if ((htonl(record->addr) == cinet_addr(224,0,2,60)) &&
			    ((record->type == IGMP_V3_REC_MODE_IS_INCL) ||
			     (record->type == IGMP_V3_REC_MODE_CH_INCL)))
				igmp_match();

			body_len -= record->auxlen;
			pos += record->auxlen;
		}

		break;
	}

	case IGMP_MEMBERSHIP_QUERY:
		debug(DBG_IGMP, "igmp_membership_query\n");
		break;

	case IGMP_V2_LEAVE_GROUP:
		debug(DBG_IGMP, "igmp_v2_leave_group\n");
		break;

	default:
		debug(DBG_IGMP, "IGMP msg type %02hhx\n", igmp_msg->common.type);
		break;
	}

	buf += hdr->length;
	len -= hdr->length;
}

static void
igmp_read_cb(struct cfg *cfg, struct uring_task *task, int res)
{
	struct igmp *igmp = container_of(task, struct igmp, task);

	debug(DBG_IGMP, "task %p, igmp %p, res %i\n", task, igmp, res);

	if (res < 0 || task->dead)
		return;

	task->tbuf->len = res;

	if (task->addr.storage.ss_family == AF_PACKET ||
	    task->addr.ll.sll_protocol == htons(ETH_P_IP))
		igmp_parse(cfg, igmp);
	else
		debug(DBG_IGMP, "invalid packet type received\n");

	uring_tbuf_read(cfg, &igmp->task, igmp_read_cb);
}

static void
igmp_free(struct uring_task *task)
{
	struct igmp *igmp = container_of(task, struct igmp, task);

	debug(DBG_IGMP, "task %p, igmp %p\n", task, igmp);
	xfree(igmp);
}

void
igmp_refdump(struct igmp *igmp)
{
	if (!igmp)
		return;

	uring_task_refdump(&igmp->task);
}

void
igmp_delete(struct cfg *cfg)
{
	struct igmp *igmp = cfg->igmp;

	if (!igmp)
		return;

	debug(DBG_IGMP, "closing fd %i\n", igmp->task.fd);
	uring_task_destroy(cfg, &igmp->task);
	cfg->igmp = NULL;
}

void
igmp_init(struct cfg *cfg)
{
	static const struct sock_filter filter[] = {
		BPF_STMT(BPF_LD + BPF_W + BPF_LEN, 0),					/* A <- packet length */
		BPF_JUMP(BPF_JMP + BPF_JGE + BPF_K, sizeof(struct iphdr), 1, 0),	/* A < sizeof(iphdr) */
                BPF_STMT(BPF_RET + BPF_K, 0),						/* drop packet */

		BPF_STMT(BPF_LD + BPF_B + BPF_ABS, 0 /* iphdr[0] */),			/* A <- version + ihl */
		BPF_STMT(BPF_ALU + BPF_RSH + BPF_K, 4),					/* A <- A >> 4 (version) */
                BPF_JUMP(BPF_JMP + BPF_JEQ + BPF_K, 0x04, 1, 0),			/* A != 4 */
                BPF_STMT(BPF_RET + BPF_K, 0),						/* drop packet */

		BPF_STMT(BPF_LD + BPF_B + BPF_ABS, offsetof(struct iphdr, protocol)),	/* A <- ip protocol */
                BPF_JUMP(BPF_JMP + BPF_JEQ + BPF_K, IPPROTO_IGMP, 1, 0),		/* A != IPPROTO_IGMP */
                BPF_STMT(BPF_RET + BPF_K, 0),						/* drop packet */

		BPF_STMT(BPF_LD + BPF_W + BPF_ABS, offsetof(struct iphdr, daddr)),	/* A <- ip dst addr */
                BPF_JUMP(BPF_JMP + BPF_JEQ + BPF_K, chtobe32(cinet_addr(224,0,2,60)), 2, 0),	/* A != 224.0.2.60 */
                //BPF_JUMP(BPF_JMP + BPF_JEQ + BPF_K, 0xe000023c, 2, 0),			/* A != 224.0.2.60 */
                BPF_JUMP(BPF_JMP + BPF_JEQ + BPF_K, chtobe32(cinet_addr(224,0,0,22)), 1, 0),	/* A != 224.0.0.22 */
                //BPF_JUMP(BPF_JMP + BPF_JEQ + BPF_K, 0xe0000016, 1, 0),			/* A != 224.0.0.22 */
                BPF_STMT(BPF_RET + BPF_K, 0),						/* drop packet */

		BPF_STMT(BPF_LDX + BPF_B + BPF_MSH, 0 /* iphdr[0] */),			/* X <- pkt->ihl * 4 */
		BPF_STMT(BPF_LD + BPF_IMM, 20),						/* A <- 20 */
		BPF_JUMP(BPF_JMP + BPF_JGT + BPF_X, 0, 0, 1),				/* A > X */
                BPF_STMT(BPF_RET + BPF_K, 0),						/* drop packet */

		BPF_STMT(BPF_LD + BPF_H + BPF_ABS, offsetof(struct iphdr, tot_len)),	/* A <- ip tot_len */
		BPF_JUMP(BPF_JMP + BPF_JGT + BPF_X, 0, 1, 0),				/* A <= ip->ihl * 4 */ 
                BPF_STMT(BPF_RET + BPF_K, 0),						/* drop packet */
		BPF_STMT(BPF_ALU + BPF_SUB + BPF_X, 0),					/* A <- A - X (bodylen) */
		BPF_JUMP(BPF_JMP + BPF_JGE + BPF_K, 8, 1, 0),				/* A < 8 */
                BPF_STMT(BPF_RET + BPF_K, 0),						/* drop packet */

		BPF_STMT(BPF_LD + BPF_H + BPF_ABS, offsetof(struct iphdr, tot_len)),	/* A <- ip->tot_len */
		BPF_STMT(BPF_MISC + BPF_TAX, 0),					/* X <- A */
		BPF_STMT(BPF_LD + BPF_W + BPF_LEN, 0),					/* A <- packet length */
		BPF_JUMP(BPF_JMP + BPF_JEQ + BPF_X, 0, 1, 0),				/* A != ip->tot_len */
                BPF_STMT(BPF_RET + BPF_K, 0),						/* drop packet */

                BPF_STMT(BPF_RET + BPF_K, (uint32_t) -1),				/* accept packet */
        };
        static const struct sock_fprog fprog = {
                .len = ARRAY_SIZE(filter),
                .filter = (struct sock_filter*) filter,
        };
	struct sockaddr_ll addr = {
		.sll_family = AF_PACKET,
		.sll_ifindex = 0,
		.sll_pkttype = PACKET_MULTICAST,
	};
	struct igmp *igmp;
	int sfd;
	int opt;

	igmp = zmalloc(sizeof(*igmp));
	if (!igmp)
		return;

	/*
	 * Kernel limitation, must be ETH_P_ALL, not ETH_P_IP or we won't get
	 * outgoing packets, https://lkml.org/lkml/1999/12/23/112
	 */
	sfd = socket(AF_PACKET, SOCK_DGRAM | SOCK_CLOEXEC, htons(ETH_P_ALL));
	if (sfd < 0) {
		if (errno == EACCES)
			verbose("igmp socket: permission denied\n");
		else
			error("igmp socket: %m");
		goto out_free;
	}

	if (setsockopt(sfd, SOL_SOCKET, SO_ATTACH_FILTER, &fprog, sizeof(fprog)) < 0)
		perrordie("igmp setsockopt(SO_ATTACH_FILTER): %m");

	if (setsockopt(sfd, SOL_SOCKET, SO_LOCK_FILTER, &opt, sizeof(opt)) < 0)
		perrordie("igmp setsockopt(SO_LOCK_FILTER): %m");

	/*
	struct ifreq ifreq = {
		.ifr_name = "enp8s0",
		.ifr_ifindex = 0,
	};

	strcpy(ifreq.ifr_name, "enp8s0");

	if (ioctl(sfd, SIOCGIFINDEX, &ifreq) < 0) {
		perror("ioctl");
		exit(EXIT_FAILURE);
	}

	fprintf(stderr, "Ifindex = %i\n", ifreq.ifr_ifindex);

	struct packet_mreq mreq = {
		.mr_ifindex = ifreq.ifr_ifindex,
		.mr_type = PACKET_MR_ALLMULTI
	};

	r = setsockopt(sfd, SOL_PACKET, PACKET_ADD_MEMBERSHIP, &mreq, sizeof(mreq));
	if (r < 0) {
		perror("setsockopt");
		exit(EXIT_FAILURE);
	}
	*/

	/* can't set .sll_protocol to htons(ETH_P_IP), see comment above */
	if (bind(sfd, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
		error("bind: %m");
		goto out_fd;
	}

	uring_task_init(&igmp->task, "igmp", uring_parent(cfg), igmp_free);
	uring_task_set_fd(&igmp->task, sfd);
	uring_task_set_buf(&igmp->task, &igmp->tbuf);
	igmp->task.addr.addrlen = sizeof(igmp->task.addr.ll);
	uring_tbuf_recvmsg(cfg, &igmp->task, igmp_read_cb);
	
	cfg->igmp = igmp;

	return;

out_fd:
	close(sfd);
out_free:
	xfree(igmp);
}