/*-
 * ----------------------------------------------------------------------------
 * "THE BEER-WARE LICENSE" (Revision 69):
 * <syrinx@FreeBSD.org> wrote this file.  As long as you retain this notice you
 * can do whatever you want with this stuff. If we meet some day, and you think
 * this stuff is worth it, you can buy me a beer in return.   -Shteryana Shopova
 * ----------------------------------------------------------------------------
 */

#include <sys/types.h>
#ifdef	LINUX
#include <limits.h>
#else
#include <sys/limits.h>
#endif /* LINUX */
#include <sys/socket.h>

#include <arpa/inet.h>
#include <netinet/in.h>

#include <err.h>
#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <signal.h>
#include <unistd.h>

/* signal handling */
static int work;
#define WORK_INPUT_SOURCE	0x0001
#define	WORK_MRECV		0x0002
#define	WORK_CLEANUP		0x0004

/* Filter modes */
#define FILTER_INCLUDE		1
#define	FILTER_EXCLUDE		2

/* Default IGMP version */
#define	IGMP_V3			0
#define	IGMP_V2			2
#define	IGMP_V1			1

/* Source address list and structure. */
struct mcsource {
	struct in_addr msaddr;
	struct {
		struct mcsource *sle_next;  /* next element */;
	} mlink;
};

struct mcsource_list {
	int count;
	struct mcsource *slh_first;
} mslist = { 0, NULL };

static void
mslist_free(struct mcsource_list *headp)
{
	struct mcsource *ms;

	while ((ms = headp->slh_first) != NULL) {
		headp->slh_first = (headp->slh_first->mlink).sle_next;

		free(ms);
	}

	headp->count = 0;
}

static int
mslist_add(struct in_addr *msource)
{
	struct mcsource *ms;

	if ((ms = (struct mcsource *)malloc(sizeof(*ms))) == NULL)
		err(1, "malloc failed");

	memset(ms, 0, sizeof(*ms));
	ms->msaddr.s_addr = msource->s_addr;

	ms->mlink.sle_next = mslist.slh_first;
	mslist.slh_first = ms;

	return (0);
}

/* static */ void
mslist_remove(struct mcsource *mentry)
{
	if (mslist.slh_first == mentry)
		mslist.slh_first = (mslist.slh_first)->mlink.sle_next;
	else {
		struct mcsource *current = mslist.slh_first;

		while (current->mlink.sle_next != mentry)
			current = current->mlink.sle_next;

		current->mlink.sle_next = (current->mlink.sle_next)->mlink.sle_next;
	}

	free(mentry);
}

static struct mcsource *
mslist_find(struct in_addr *msource)
{
	struct mcsource *ms;

	for (ms = mslist.slh_first; ms != NULL; ms = (ms->mlink).sle_next)
		if (ms->msaddr.s_addr == msource->s_addr)
			break;

	return (ms);
}

static void
#ifdef LINUX
mrecv_onstop(int s)
#else
mrecv_onstop(int s __unused)
#endif
{
	work |= WORK_INPUT_SOURCE;
}

static void
#ifdef LINUX
mrecv_oncont(int s)
#else
mrecv_oncont(int s __unused)
#endif
{
	work |= WORK_MRECV;
}

static void
#ifdef LINUX
mrecv_onint(int s)
#else
mrecv_onint(int s __unused)
#endif
{
	work |= WORK_CLEANUP;
}

static void
init_signals(void)
{
	struct sigaction sa;

	sa.sa_handler = mrecv_onstop;
	sa.sa_flags = 0;
	sigemptyset(&sa.sa_mask);
	if (sigaction(SIGTSTP, &sa, NULL))
		err(1, "sigaction(%d)", SIGTSTP);

	sa.sa_handler = mrecv_oncont;
	sa.sa_flags = 0;
	sigemptyset(&sa.sa_mask);
	if (sigaction(SIGCONT, &sa, NULL))
		err(1, "sigaction(%d)", SIGCONT);

	sa.sa_handler = mrecv_onint;
	sa.sa_flags = 0;
	sigemptyset(&sa.sa_mask);
	if (sigaction(SIGINT, &sa, NULL))
		err(1, "sigaction(%d)", SIGINT);
}

static int
add_socket_membership(int sock, struct in_addr *ifa, struct in_addr *ga,
	struct mcsource_list *sa_list, int mode)
{
	int sop, ssop, sd_size;
	void *sdata;
	struct ip_mreq mreq;
	struct ip_mreq_source smreq;
	struct mcsource *sa;

	if (ifa == NULL || ga == NULL || sa_list == NULL)
		return (-1);

	memset(&mreq, 0, sizeof(mreq));
	memset(&smreq, 0, sizeof(smreq));

	if (mode == FILTER_INCLUDE) {
		sop = ssop = IP_ADD_SOURCE_MEMBERSHIP;
		smreq.imr_multiaddr = *ga;
		smreq.imr_interface = *ifa;
		if ((sa = sa_list->slh_first) == NULL) {
			fprintf(stderr, "INCLUDE mode - must specify at least "
			    "one source\n");
			return (-1);
		}
		smreq.imr_sourceaddr = sa->msaddr;
		sa = (sa)->mlink.sle_next;
		sdata = &smreq;
		sd_size = sizeof(struct ip_mreq_source);
	} else if (mode == FILTER_EXCLUDE) {
		sop = IP_ADD_MEMBERSHIP;
		ssop = IP_BLOCK_SOURCE;
		mreq.imr_multiaddr = *ga;
		mreq.imr_interface = *ifa;
		sa = sa_list->slh_first;
		sdata = &mreq;
		sd_size = sizeof(struct ip_mreq);
	} else
		return (-1);


	if ((setsockopt(sock, IPPROTO_IP, sop, sdata, sd_size))) {
		fprintf(stderr, "setsockopt(s=%d, o=%d) failed - %s", sock, sop,
		    strerror(errno));
		return (-1);
	}


	/* Join/Prune the additional sources acording to filter type */
	for ( ; sa; sa = (sa)->mlink.sle_next) {
		memset(&smreq, 0, sizeof(smreq));
		smreq.imr_multiaddr = *ga;
		smreq.imr_interface = *ifa;
		smreq.imr_sourceaddr = sa->msaddr;

		if ((setsockopt(sock, IPPROTO_IP, ssop, &smreq, sizeof(smreq)))) {
			fprintf(stderr, "setsockopt(s=%d, o=%d) failed - %s",
			    sock, ssop, strerror(errno));
			return (-1);
		}
	}

	return (0);
}

static int
drop_socket_membership(int sock, struct in_addr *ifa, struct in_addr *ga)
{
	struct ip_mreq mreq;

	if (ifa == NULL || ga == NULL)
		return (-1);

	memset(&mreq, 0, sizeof(mreq));
	mreq.imr_multiaddr = *ga;
	mreq.imr_interface = *ifa;

	if (setsockopt (sock, IPPROTO_IP, IP_DROP_MEMBERSHIP, &mreq,
	    sizeof(mreq)) < 0) {
		fprintf(stderr, "setsockopt(s=%d, o=%d) failed - %s",
		    sock, IP_DROP_MEMBERSHIP, strerror(errno));
		return (-1);
	}

	return (0);
}

void
usage(void)
{
	(void)fprintf(stderr,
	    "usage: mrecv [-e | -f | -i] [-g group-ip] [-I if-ip] [-p port]\n"
		 "[-s source-ip] [-v version]\n");
	(void)fprintf(stderr, "SIGTSTP inputs new addresses\n");
	(void)fprintf(stderr, "SIGCONT reads socket\n");

	exit(1);
}

int
main (int argc, char **argv)
{
	char data[1500];
	unsigned long suint;
	uint8_t mode, version;
	int32_t count;
	int ch, sock;
	struct in_addr group, if_addr, ms_addr;
	struct sockaddr_in bind_addr, from_addr;
	socklen_t fromlen;

	init_signals();

	work = 0;
	mode = FILTER_EXCLUDE;
	version = IGMP_V3;
	count = 0;
	group.s_addr = htonl(INADDR_ALLHOSTS_GROUP);
	if_addr.s_addr = htonl(INADDR_LOOPBACK);
	memset(&bind_addr, 0, sizeof(struct sockaddr_in));
	bind_addr.sin_port = (in_port_t) random() % 0xFFFF;

	while ((ch = getopt(argc, argv, "eifg:I:p:s:v:")) != EOF && count < argc) {
		switch (ch) {
		case 'e':
			break;

		case 'i':
			mode = FILTER_INCLUDE;
			break;

		case 'f':
			work = WORK_MRECV;
			break;

		case 'g':
			if (inet_aton(optarg, &group) != 1 ||
			    !IN_MULTICAST(ntohl(group.s_addr)))
				err(1, "Invalid group address - %s)", optarg);
			++count;
			break;

		case 'I':
			if (inet_aton (optarg, &if_addr) != 1)
				err(1, "Invalid interface address - %s)", optarg);
			++count;
			break;

		case 'p':
			if ((suint = strtoul(optarg, NULL, 10)) == ULONG_MAX)
				err(1, "Invalid port - %s)", optarg);
			else
				bind_addr.sin_port = (in_port_t) suint;
			++count;
			break;

		case 's':
			if (inet_aton (optarg, &ms_addr) != 1)
				err(1, "Invalid source address - %s)", optarg);
			if (mslist_find(&ms_addr) == NULL)
				(void)mslist_add(&ms_addr);
			++count;
			break;

		case 'v':
			if ((suint = strtoul(optarg, NULL, 10)) == ULONG_MAX ||
			    suint > 3)
				err(1, "Invalid version - %s)", optarg);
			else if (suint < 3)
				version = (uint8_t) suint;
			++count;
			break;

		default:
			usage();
		}
		++count;
	}

	/* XXX: set needed sysctl's here */

	bind_addr.sin_family = AF_INET;
	bind_addr.sin_addr = group;

	if ((sock = socket (AF_INET, SOCK_DGRAM, IPPROTO_IP)) == -1)
		err(1, "socket failed");


	if ((add_socket_membership(sock, &if_addr, &group, &mslist, mode)))
		errx(1, "Failed to add groupe membership");

	if (bind (sock, (struct sockaddr *) &bind_addr, sizeof(bind_addr)) < 0)
		err(1, "bind failed");

	for (;;) {
		if (work != 0 ) {
			if ((work & WORK_CLEANUP) != 0)
				break;
			else if ((work & WORK_INPUT_SOURCE) != 0) {
				work = 0;
				printf("SIGTSTP received\n");
			} else if ((work & WORK_MRECV) != 0) {
				/* XXX: FIXME */
				recvfrom (sock, data, 1500, 0,
				    (struct sockaddr *) &from_addr, &fromlen);
			}
		}
	}

	printf("SIGINT received\n");
	/* usage(); */

	(void)drop_socket_membership(sock, &if_addr, &group);

	mslist_free(&mslist);
	exit (0);
}