/*
 * Quick copy paste combo of
 * wolc: https://github.com/timofurrer/WOL
 * ether-wake.c: v1.09 11/12/2003 Donald Becker, http://www.scyld.com/
 * so we can have one tool that supports RAW mode (ethernet frame type 0x0842)
 * and encapsulated UDP mode (like the wol or wakeonlan cmdline utils do).
 *
 * GPL2
 */

static const char usage_msg[] =
"Usage:\n"
"   %s [-b] [-i <ifname>] [-p aa:bb:cc:dd[:ee:ff]] \\\n"
"      00:11:22:33:44:55 [, 11:22:33:44:55:66 [, ...]]\n"
"   %s -d <dest_ip> [-p aa:bb:cc:dd[:ee:ff]] \\\n"
"      00:11:22:33:44:55 [, 11:22:33:44:55:66 [, ...]]\n"
"\n"
"The first variant sends a raw ethernet frame with type 0x0842 (WoL) to the\n"
"specified destination(s). In this mode, the source interface must be known\n"
"(-i), which will default to eth0 if omitted.\n"
"\n"
"The second variant sends the UDP version of a WoL packet to port 9 of the\n"
"given destination(s). This required the destination parameter (-d) to be set\n"
"to either the limited broadcast address, or the directed broadcast address of\n"
"the destination network, i.e. 10.1.2.255 if the destination host is\n"
"10.1.2.3/24.\n"
"This requires the final router connected to the destination subnet to be\n"
"configured accordingly.\n"
"\n"
"Options for raw mode (etherwake):\n"
"   --broadcast\n"
"   -b          Send wake-up packet to the broadcast MAC-address, instead of\n"
"               destination MAC address.\n"
"   --interface IFNAME\n"
"   -i IFNAME   Use interface IFNAME instead of the default 'eth0'.\n"
"\n"
"Options for UDP mode (wakeonlan):\n"
"   --destination IPADDR\n"
"   -d IPADDR   Destination (broadcast) IP address.\n"
"\n"
"Options for listening for WOL packets:\n"
"   --listen\n"
"   -l          Listen for incoming WOL packets and print their contents.\n"
"               Pass twice to force listening for UDP and RAW packets and\n"
"               bail out if this fails. Otherwise, continue as long as at\n"
"               least one of the sockets could be bound successfully.\n"
"\n"
"Common options:\n"
"   --verbose\n"
"   -v          Verbose mode, print some additional information.\n"
"   --debug\n"
"   -D          Increase the debug level, print much more information.\n"
"   --password PASSWD\n"
"   -p PASSWD   Append the four or six byte password PW to the packet.\n"
"               A password is only required for a few adapter types.\n"
"               The password may be specified in ethernet hex format\n"
"               or dotted decimal (Internet address),\n"
"               f.i. 00:22:44:66:88:aa or 192.168.1.1\n"
"\n"
"UDP mode will be selected if -d is passed on the command line. Any raw mode\n"
"options (-b, -i) won't work in this mode so program execution is aborted.\n"
;

#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <errno.h>
#include <netpacket/packet.h>
#include <netinet/ether.h>
#include <net/ethernet.h>
#include <linux/if.h>
#include <sys/ioctl.h>
#include <getopt.h>
#include <poll.h>
#include <fcntl.h>

#define PACKET_SIZE (200)
#define CHECK_PACKET_SIZE(x) do { \
		if ( (x) >= PACKET_SIZE ) { fprintf( stderr, "Packet size overflow.\n" ); abort(); } \
	} while(0)

static struct {
	const char *ifname;
	struct in_addr dest;
	char wolpw[6];
	int wolpwlen;
	bool broadcast;
	bool verbose; // Spam
	bool debug; // Even more spam
} options;

static int parse_errors = 0;

static int get_dest_addr(const char *arg, struct ether_addr *dest_ether);
static int get_fill(unsigned char *pkt, const uint8_t *dest_ether, bool with_header);
static void get_wol_pw(const char *optarg);
static int get_source_addr(int s, const char *ifname, unsigned char *output);
static struct in_addr get_ip(const char *arg);
static int rawsocket();
static int udpsocket();
static int etherwake(int count, char **macs);
static int wakeonlan(int count, char **maclist);
static int wolisten(bool force);
static int rawrecv(int sock);
static int udprecv(int sock);
static int find_wol(const char *addr, const char *buffer, size_t buflen, const char *type);
static void formatpw(const uint8_t *pw, int len, char *output, size_t outlen);
static void sock_nonblock(int sock);

static void print_usage(const char *a0)
{
	printf( usage_msg, a0, a0 );
}

int main(int argc, char **argv)
{
	int opt, dolisten = 0;
	bool gotip = false, gotif = false;
	static const char *optString = "bDi:p:vhd:l";
	static const struct option longOpts[] = {
		{ "broadcast", no_argument, NULL, 'b' },
		{ "debug", no_argument, NULL, 'D' },
		{ "interface", required_argument, NULL, 'i' },
		{ "password", required_argument, NULL, 'p' },
		{ "verbose", no_argument, NULL, 'v' },
		{ "help", no_argument, NULL, 'h' },
		{ "destination", required_argument, NULL, 'd' },
		{ "listen", no_argument, NULL, 'l' },
		{ 0, 0, 0, 0 }
	};
	// Conveinience: Support MAC addr with dashes
	for ( int i = 1; i < argc; ++i ) {
		int j = 0;
		while ( argv[i][j] != '\0' ) {
			if ( argv[i][j] == '-' && j > 1 ) argv[i][j] = ':';
			j++;
		}
	}
	while ( ( opt = getopt_long( argc, argv, optString, longOpts, NULL ) ) != -1 ) {
		switch ( opt ) {
		case 'b': options.broadcast = true; gotif = true;  break;
		case 'D': options.debug = true;                    break;
		case 'i': options.ifname = optarg; gotif = true;   break;
   	case 'p': get_wol_pw( optarg );                    break;
		case 'v': options.verbose = true;                  break;
		case 'h': print_usage( argv[0] );                  return 0;
		case 'd': options.dest = get_ip( optarg ); gotip = true; break;
		case 'l': dolisten++;                              break;
		case '?': parse_errors++;                          break;
		}
	}
	if ( dolisten && ( gotip || gotif || optind < argc ) ) {
		fprintf( stderr, "Cannot combine -l with any other options.\n" );
		return 1;
	}
	if ( dolisten ) {
		return wolisten( dolisten > 1 );
	}
	if ( gotip && gotif ) {
		fprintf( stderr, "Cannot use -d and -b/-i at the same time.\n" );
		parse_errors++;
	}
	if ( parse_errors == 0 && optind >= argc ) {
		fprintf( stderr, "No MAC address(es) given.\n" );
		parse_errors++;
	}
	if ( parse_errors != 0 ) {
		fprintf( stderr, "Try %s --help.\n", argv[0] );
		return 1;
	}
	// Seems we're somewhat prepared to go
	if ( !gotip ) {
		// raw (etherwake) mode
		if ( options.verbose ) {
			printf( "Using RAW mode.\n" );
		}
		return etherwake( argc - optind, argv + optind );
	}
	// UDP (wakeonlan) mode
	if ( options.verbose ) {
		printf( "Using UDP mode.\n" );
	}
	return wakeonlan( argc - optind, argv + optind );
}

static int etherwake(int count, char **maclist)
{
	unsigned char outpack[PACKET_SIZE];
	int s = rawsocket();
	if ( s == -1 )
		return 2;
	if ( options.ifname == NULL ) {
		options.ifname = "eth0";
	}
#if defined(PF_PACKET)
	struct sockaddr_ll whereto = {0};
	do {
		struct ifreq ifr;
		snprintf( ifr.ifr_name, IFNAMSIZ, "%s", options.ifname );
		if (ioctl(s, SIOCGIFINDEX, &ifr) == -1) {
			perror( "ioctl SIOCGIFINDEX on raw socket" );
			fprintf( stderr, "Are you sure %s exists?\n", options.ifname );
			return 3;
		}
		whereto.sll_family = AF_PACKET;
		whereto.sll_ifindex = ifr.ifr_ifindex;
	} while (0);
#else
	struct sockaddr whereto = {0};	/* who to wake up */
#endif
	if ( get_source_addr( s, options.ifname, outpack + 6 ) != 0 ) {
		fprintf( stderr, "Cannot determine source MAC address for outgoing packet, using bogus one\n" );
		outpack[6] = 2;
		outpack[7] = outpack[8] = 34;
		outpack[9] = outpack[10] = outpack[11] = 56;
	}
	struct ether_addr dest_ether;
	for ( int idx = 0; idx < count; ++idx ) {
		if ( get_dest_addr( maclist[idx], &dest_ether ) != 0 ) {
			fprintf( stderr, "Cannot parse mac address '%s', skipping...\n", maclist[idx] );
			continue;
		}
		int packet_size = get_fill( outpack, dest_ether.ether_addr_octet, true );
		if ( options.wolpwlen > 0 ) {
			memcpy( outpack + packet_size, options.wolpw, options.wolpwlen );
			packet_size += options.wolpwlen;
		}
		CHECK_PACKET_SIZE( packet_size );
		if ( options.debug ) {
			printf( "The final packet for %s is:", maclist[idx] );
			for ( int i = 0; i < packet_size; i++ ) printf( " %2.2x", outpack[i] );
			printf( ".\n" );
		}
#if defined(PF_PACKET)
		/* The manual page incorrectly claims the address must be filled.
			We do so because the code may change to match the docs. */
		whereto.sll_halen = ETH_ALEN;
		memcpy(whereto.sll_addr, outpack, ETH_ALEN);
#else
		whereto.sa_family = 0;
		snprintf( whereto.sa_data, sizeof(whereto.sa_data), "%s", options.ifname );
#endif
		printf( "Sending magic packet to %s...\n", maclist[idx] );
		int bytes = sendto( s, outpack, packet_size, 0, (struct sockaddr *)&whereto,
						sizeof(whereto) );
		if ( bytes == -1 ) {
			perror( "sendto" );
		} else if ( bytes != packet_size ) {
			fprintf( stderr, "Packet truncated, sent %d of %d bytes!\n", bytes, packet_size );
		} else if ( options.debug ) {
			printf( "Sendto worked! %d bytes.\n", bytes );
		}
	}
	close( s );
	return 0;
}

static int rawsocket()
{
	int s;
	int one = 1;
	/* Note: PF_INET, SOCK_DGRAM, IPPROTO_UDP would allow SIOCGIFHWADDR to
	   work as non-root, but we need SOCK_PACKET to specify the Ethernet
	   destination address. */
#if defined(PF_PACKET)
	s = socket(PF_PACKET, SOCK_RAW, 0);
#else
	s = socket(AF_INET, SOCK_PACKET, SOCK_PACKET);
#endif
	if (s == -1) {
		if (errno == EPERM) {
			fprintf(stderr, "This program must be run as root for raw mode to work.\n");
		} else {
			perror("RAW socket()");
		}
		return -1;
	}
	/* This is necessary for broadcasts to work */
	if ( setsockopt( s, SOL_SOCKET, SO_BROADCAST, (char *)&one, sizeof(one) ) == -1 ) {
		perror( "RAW setsockopt: SO_BROADCAST" );
		return -1;
	}
	return s;
}

int udpsocket()
{
	int sock;
	int optval = 1;

	sock = socket( AF_INET, SOCK_DGRAM, IPPROTO_UDP );
	if ( sock == -1 ) {
		perror( "UDP socket()" );
		return -1;
	}
	if ( setsockopt( sock, SOL_SOCKET, SO_BROADCAST, (char *) &optval, sizeof( optval ) ) == -1 ) {
		perror( "UDP setsockopt: SO_BROADCAST" );
		return -1;
	}

	return sock;
}

static int wakeonlan(int count, char **maclist)
{
	struct ether_addr dest_ether;
	unsigned char outpack[PACKET_SIZE];
	int sock = udpsocket();
	if ( sock == -1 )
		return 2;
	struct sockaddr_in addr = {
		.sin_family = AF_INET,
		.sin_port = htons( 9 ),
		.sin_addr = options.dest,
	};
	for ( int idx = 0; idx < count; ++idx ) {
		if ( get_dest_addr( maclist[idx], &dest_ether ) != 0 ) {
			fprintf( stderr, "Cannot parse mac address '%s', skipping...\n", maclist[idx] );
			continue;
		}
		int packet_size = get_fill( outpack, dest_ether.ether_addr_octet, false );
		if ( options.wolpwlen > 0 ) {
			memcpy( outpack + packet_size, options.wolpw, options.wolpwlen );
			packet_size += options.wolpwlen;
		}
		CHECK_PACKET_SIZE( packet_size );
		if ( options.debug ) {
			printf( "The final packet for %s is:", maclist[idx] );
			for ( int i = 0; i < packet_size; i++ ) printf( " %2.2x", outpack[i] );
			printf( ".\n" );
		}
		printf( "Sending magic packet to %s...\n", maclist[idx] );
		int bytes = sendto( sock, outpack, packet_size, 0, (struct sockaddr *)&addr,
						sizeof(addr) );
		if ( bytes == -1 ) {
			perror( "sendto" );
		} else if ( bytes != packet_size ) {
			fprintf( stderr, "Packet truncated, sent %d of %d bytes!\n", bytes, packet_size );
		} else if ( options.debug ) {
			printf( "Sendto worked! %d bytes.\n", bytes );
		}
	}
	close( sock );
	return 0;
}

static int get_fill(unsigned char *pkt, const uint8_t *bindestmac, bool with_header)
{
	int offset = 0;
	int i;

	if ( with_header ) {
		if ( options.broadcast ) {
			memset( pkt, 0xff, 6 );
		} else {
			memcpy( pkt, bindestmac, 6 );
		}
		// Leave byte 6-11 alone, already set to source MAC addr
		// Next, type
		pkt[12] = 0x08;
		pkt[13] = 0x42;
		offset = 14;
	}

	// Write sync bytes payload
	memset(pkt+offset, 0xff, 6);
	// Repeat destination address 16 times
	offset += 6;
	for (i = 0; i < 16; i++) {
		memcpy(pkt + offset, bindestmac, 6);
		offset += 6;
	}
	CHECK_PACKET_SIZE( offset );
	return offset;
}

static void get_wol_pw(const char *optarg)
{
	unsigned int passwd[6];
	int byte_cnt;
	int i;

	byte_cnt = sscanf(optarg, "%2x:%2x:%2x:%2x:%2x:%2x",
					  &passwd[0], &passwd[1], &passwd[2],
					  &passwd[3], &passwd[4], &passwd[5]);
	if (byte_cnt < 4) { // Support 4 chars in this notion
		byte_cnt = sscanf(optarg, "%u.%u.%u.%u",
						  &passwd[0], &passwd[1], &passwd[2], &passwd[3]);
	}
	if (byte_cnt < 4) {
		fprintf(stderr, "Unable to read the Wake-On-LAN password.\n");
		parse_errors++;
		return;
	}
	if (byte_cnt > 6) {
		fprintf(stderr, "The impossible happened.\n");
		abort();
	}
	options.wolpwlen = byte_cnt;
	for (i = 0; i < byte_cnt; i++) {
		if ( passwd[i] > 255 ) {
			fprintf( stderr, "Password has invalid value for byte %d: %u.\n",
					i, passwd[i] );
			parse_errors++;
		}
		options.wolpw[i] = passwd[i];
	}
	if ( options.verbose ) {
		char pw[20];
		formatpw( (uint8_t*)options.wolpw + i, byte_cnt, pw, sizeof(pw) );
		printf("The Magic packet password is %s.\n", pw);
	}
}

/**
 * Get MAC address of given interface (binary 6 byte)
 * returns 0 on success
 */
static int get_source_addr(int s, const char *ifname, unsigned char *output)
{
	/* Fill in the source address, if possible.
	   The code to retrieve the local station address is Linux specific. */
	struct ifreq if_hwaddr;

	snprintf(if_hwaddr.ifr_name, IFNAMSIZ, "%s", ifname);
	if (ioctl(s, SIOCGIFHWADDR, &if_hwaddr) < 0) {
		fprintf(stderr, "SIOCGIFHWADDR on %s failed: %s\n", ifname,
				strerror(errno));
		/* Magic packets still work if our source address is bogus, but
			we fail just to be anal. */
		return 1;
	}
	memcpy(output, if_hwaddr.ifr_hwaddr.sa_data, 6);

	if ( options.verbose ) {
		char mac[20];
		formatpw( (uint8_t*)if_hwaddr.ifr_hwaddr.sa_data, 6, mac, sizeof(mac) );
		printf("The hardware address (SIOCGIFHWADDR) of %s is type %d"
				" (%s).\n", ifname, if_hwaddr.ifr_hwaddr.sa_family, mac);
	}
	return 0;
}

/**
 * Convert the host ID string to a MAC address.
 * The string may be a
 * Host name
 * IP address string
 * MAC address string
*/
static int get_dest_addr(const char *hostid, struct ether_addr *dest_ether)
{
	struct ether_addr *eap;

	eap = ether_aton(hostid);
	if (eap != NULL) {
		*dest_ether = *eap;
		if (options.debug) {
			printf("The target station address is %s.\n",
					ether_ntoa(dest_ether));
		}
	} else if (ether_hostton(hostid, dest_ether) == 0) {
		if (options.debug) {
			printf("Station address for hostname %s is %s.\n",
					hostid, ether_ntoa(dest_ether));
		}
	} else {
		fprintf(stderr,
					  "The Magic Packet host address must be "
					  "specified as\n"
					  "  - a station address, 00:11:22:33:44:55, or\n"
					  "  - a station address, 00-11-22-33-44-55, or\n"
					  "  - a hostname with a known 'ethers' entry.\n");
		return -1;
	}
	return 0;
}

static struct in_addr get_ip(const char *arg)
{
	struct in_addr tmp;
	if ( inet_aton( arg, &tmp ) == 0 ) {
		fprintf( stderr, "Converting destination address via inet_aton failed.\n" );
		parse_errors++;
	}
	return tmp;
}

static int wolisten(bool force)
{
	// TODO: RAW
	struct pollfd socks[2] = {{.events = POLLIN},{.events = POLLIN}};
	int num = 0;
	// UDP setup
	int su = udpsocket();
	if ( su == -1 ) {
		if ( force )
			return 1;
	} else {
		struct sockaddr_in addr = {
			.sin_family = AF_INET,
			.sin_port = htons( 9 ),
			.sin_addr = { INADDR_ANY },
		};
		socklen_t addrlen = sizeof(addr);
		sock_nonblock( su );
		if ( bind( su, (struct sockaddr*)&addr, addrlen ) == -1 ) {
			perror( "bind on UDP socket" );
			close( su );
			if ( force )
				return 2;
		} else {
			socks[num++].fd = su;
			if ( options.verbose ) {
				printf( "Listening for UDP WOL packets\n" );
			}
		}
	}
	// RAW setup
	int sr = rawsocket();
	if ( sr == -1 ) {
		if ( force )
			return 1;
	} else {
		struct sockaddr_ll addr = {
			.sll_protocol = htons( 0x0842 ),
			.sll_family = AF_PACKET,
		};
		socklen_t addrlen = sizeof(addr);
		sock_nonblock( sr );
		if ( bind( sr, (struct sockaddr*)&addr, addrlen ) == -1 ) {
			perror( "bind on RAW socket" );
			close( sr );
			if ( force )
				return 2;
		} else {
			socks[num++].fd = sr;
			if ( options.verbose ) {
				printf( "Listening for RAW WOL packets\n" );
			}
		}
	}
	// Ready to go?
	if ( num == 0 )
		return 1;
	for ( ;; ) {
		if ( poll( socks, num, -1 ) == -1 ) {
			if ( errno == EINTR || errno == EAGAIN )
				continue;
			perror( "poll" );
			break;
		}
		for ( int i = 0; i < num; ++i ) {
			if ( socks[i].revents == 0 )
				continue;
			if ( socks[i].fd == su ) {
				udprecv( su );
			} else if ( socks[i].fd == sr ) {
				rawrecv( sr );
			} else {
				abort();
			}
		}
	}
	close( su );
	return 1;
}

static const uint8_t ETHER_HEADER_LEN = 6 + 6 + 2; // 2x MAC + type

static int rawrecv(int sock)
{
	char buffer[PACKET_SIZE], addrstr[50];
	struct sockaddr_ll src = {0};
	socklen_t srclen = sizeof(src);
	ssize_t ret = recvfrom( sock, buffer, sizeof(buffer), 0, (struct sockaddr*)&src, &srclen );
	if ( ret < 0 ) {
		if ( errno == EINTR || errno == EAGAIN )
			return 0;
		perror( "recvfrom for RAW socket" );
		return 1;
	}
	if ( ret < ETHER_HEADER_LEN ) {
		printf( "Short raw packet received (%d bytes)\n", (int)ret );
		return 0;
	}
	formatpw( (uint8_t*)src.sll_addr, src.sll_halen, addrstr, sizeof(addrstr) );
	if ( options.debug ) {
		printf( "Received RAW packet of size %d from %s!\n", (int)ret, addrstr );
	}
	return find_wol( addrstr, buffer + ETHER_HEADER_LEN, (size_t)ret - ETHER_HEADER_LEN, "RAW" );
}

static int udprecv(int sock)
{
	char buffer[PACKET_SIZE], addrstr[50];
	struct sockaddr_in src = {0};
	socklen_t srclen = sizeof(src);
	ssize_t ret = recvfrom( sock, buffer, sizeof(buffer), 0, (struct sockaddr*)&src, &srclen );
	if ( ret < 0 ) {
		if ( errno == EINTR || errno == EAGAIN )
			return 0;
		perror( "recvfrom for UDP socket" );
		return 1;
	}
	if ( inet_ntop( AF_INET, &src.sin_addr, addrstr, sizeof(addrstr) ) == NULL ) {
		perror( "Cannot determine source address of received UDP packet" );
		return 0;
	}
	if ( options.debug ) {
		printf( "Received UDP datagram of size %d from %s!\n", (int)ret, addrstr );
	}
	return find_wol( addrstr, buffer, (size_t)ret, "UDP" );
}

static int find_wol(const char *addrstr, const char *buffer, size_t buflen, const char *type)
{
	const uint8_t *sync = (uint8_t*)memmem( buffer, buflen, "\xff\xff\xff\xff\xff\xff", 6 );
	if ( sync == NULL ) {
		if ( options.debug ) {
			printf( "... no sync bytes found\n" );
		}
		return 0;
	}
	sync += 6;
	const uint8_t *end = (uint8_t*)buffer + buflen;
	const uint8_t *pos = sync;
	int count = 0;
	while ( count < 16 && pos + 6 <= end ) {
		if ( memcmp( sync, pos, 6 ) != 0 )
			break;
		pos += 6;
		count++;
	}
	if ( count == 16 ) {
		char passwd[20];
		char mac[20];
		formatpw( sync, 6, mac, sizeof(mac) );
		formatpw( pos, end - pos, passwd, sizeof(passwd) );
		printf( "MAGIC source=%s dest=%s type=%s passwd=%s\n", addrstr, mac, type, passwd );
	} else if ( options.verbose ) {
		printf( "Sync bytes found, with %d repetitions of the following 6 bytes.\n", count );
	}
	return 0;
}

static void formatpw(const uint8_t *pw, int len, char *output, size_t outlen)
{
	if ( outlen <= 0 )
		return;
	if ( len == 4 ) {
		snprintf( output, outlen, "%u.%u.%u.%u", pw[0], pw[1], pw[2], pw[3] );
	} else if ( len == 6 ) {
		snprintf( output, outlen, "%02x:%02x:%02x:%02x:%02x:%02x", pw[0], pw[1], pw[2], pw[3], pw[4], pw[5] );
	} else {
		*output = '\0';
	}
}

static void sock_nonblock(int sock)
{
	errno = 0;
	int f = fcntl( sock, F_GETFL );
	if ( errno != 0 ) {
		perror( "fcntl(F_GETFL)" );
		return;
	}
	if ( fcntl( sock, F_SETFL, f | O_NONBLOCK ) == -1 ) {
		perror( "fcntl(F_SETFL, ... | O_NONBLOCK)" );
	}
}