/* probe.cpp */
/* Copyright (C) 2005-2010 StarDot Technologies. */
/* Written by Jim Treadway <jim@stardot-tech.com>. */

/* This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <errno.h>
#include <getopt.h>
#include <ifaddrs.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <net/if.h>
#include <net/if_arp.h>
#include <netinet/in.h>
#include <sys/ioctl.h>

#include <algorithm>
#include <iostream>
#include <map>
#include <sstream>
#include <string>
#include <vector>

#define VER_NAME                 "probe"
#define VER_MAJOR                1
#define VER_MINOR                0
#ifndef VER_BUILD
#define VER_BUILD                0
#endif

#undef DEBUG

enum {
	SORT_NONE = 0,
	SORT_HWADDR,
	SORT_IPADDR,
	SORT_TYPE,
	SORT_VERSION,
};

static char *pname;
static unsigned short port = 7364; /* "sd" */
static int xml = 0;
static int sort = SORT_NONE;

static const char shortOpts[] = "hs:vx?";
static const struct option longOpts[] = {
	{ "help",                0, NULL, 'h' },
	{ "sort",                1, NULL, 's' },
	{ "version",             0, NULL, 'v' },
	{ "xml",                 0, NULL, 'x' },
	{ NULL,                  0, NULL,  0  }
};

static void usage(void)
{
	printf("Usage: %s [OPTION...]\n", pname);
	printf("Probe LAN for StarDot network devices.\n");
	printf("\nOptions:\n");
	printf("  -s, --sort         Sort the output.\n");
	printf("  -x, --xml          Display output in XML.\n");
	printf("\nHelp and version options:\n");
	printf("  -h, --help         Display this help and exit.\n");
	printf("  -v, --version      Display version information and exit.\n");
	printf("\nReport bugs to <bugs@stardot-tech.com>.\n");
}

static void version(void)
{
	printf("%s %d.%d.%d\n", VER_NAME, VER_MAJOR, VER_MINOR, VER_BUILD);
	printf("Copyright (C) 2005-2010 StarDot Technologies."
	       "  All rights reserved.\n");
	printf("Written by Jim Treadway <jim@stardot-tech.com>.\n");
}

static int parse_cmd_line(int *argc, char **argv[])
{
	int c, index = 0, len;

	while ((c = getopt_long(*argc, *argv, shortOpts, longOpts, &index))
	       != EOF) {
		switch (c) {
		case 's':
			len = strlen(optarg);
			if (!strncasecmp(optarg, "hwaddr", len))
				sort = SORT_HWADDR;
			else if (!strncasecmp(optarg, "ipaddr", len))
				sort = SORT_IPADDR;
			else if (!strncasecmp(optarg, "type", len))
				sort = SORT_TYPE;
			else if (!strncasecmp(optarg, "version", len))
				sort = SORT_VERSION;
			else {
				printf("%s: invalid argument `%s' for `%s'\n",
				       pname, optarg, "sort");
				printf("Try '%s --help' for more information.\n",
				       pname);
				exit(1);
			}
			break;
		case 'x':
			xml = 1;
			break;

		case 'h':
		case '?':
			usage();
			exit(0);
			break;
		case 'v':
			version();
			exit(0);
			break;
		default:
			printf("Try '%s --help' for more information.\n",
			       pname);
			exit(1);
		}
	}

	*argc -= optind;
	*argv += optind;
	return optind;
}

static void set_pname(char *argv0)
{
#ifdef _WIN32
	char *p;

	if ((pname = strrchr(argv0, '\\')) != NULL)
		pname++;
	else
		pname = argv0;
	if ((p = strchr(pname, '.')) != NULL)
		*p = 0;
	strlwr(pname);
	strcpy(argv0, pname);  /* fix up for getopt_long error messages */
#else
	pname = argv0;
#endif
}

static const char *hwaddr(const char *dev, int fd, struct sockaddr_in addr)
{
	static char buf[32];
	struct arpreq arp;

	strcpy(buf, "Unknown");

	memset(&arp, 0, sizeof(arp));
	memcpy(&arp.arp_pa, &addr, sizeof(arp.arp_pa));
	arp.arp_ha.sa_family = ARPHRD_ETHER;
	strcpy(arp.arp_dev, dev);

	if (!ioctl(fd, SIOCGARP, &arp)) {
		sprintf(buf, "%02X%02X%02X-%02X%02X%02X",
			arp.arp_ha.sa_data[0] & 0xff,
			arp.arp_ha.sa_data[1] & 0xff,
			arp.arp_ha.sa_data[2] & 0xff,
			arp.arp_ha.sa_data[3] & 0xff,
			arp.arp_ha.sa_data[4] & 0xff,
			arp.arp_ha.sa_data[5] & 0xff);
	}

	return buf;
}

static void split(const std::string &str, char delim,
		  std::vector<std::string> &out)
{
	std::istringstream iss(str);
	std::string token;

	out.clear();
	while (std::getline(iss, token, delim))
		out.push_back(token);
}

typedef std::less<std::string> lstr;
struct device {
	unsigned int ipaddr;
	std::string hwaddr;
	std::string type;
	std::string version;
	std::string info;	/* the entire probe string read from device */
	std::map<std::string, std::string, lstr> attr;

	device() : ipaddr(0), hwaddr(), type(), version(), info(), attr() {}
};

static std::vector<struct device> devices;

static bool sort_fn(const struct device &lhs, const struct device &rhs)
{
	switch (sort) {
	case SORT_HWADDR:
		return lhs.hwaddr < rhs.hwaddr;
	case SORT_IPADDR:
		return lhs.ipaddr < rhs.ipaddr;
	case SORT_TYPE:
		return lhs.type < rhs.type;
	case SORT_VERSION:
		return lhs.version < rhs.version;
	default:
		return false;
	}
}

static void add_device(const char *dev, int fd, struct sockaddr_in addr,
		       const char *const buf)
{
	struct device d;
	d.ipaddr = (unsigned int) addr.sin_addr.s_addr;
	d.version = "Unknown";
	d.type = "Unknown";
	d.info = buf;

	/* StarDot/NetCam3-5407 Version 1.1.51-pre60 0030f4-000000 name=val.. */
	std::string str(buf);
	std::vector<std::string> list;
	split(str, ' ', list);
	if (list.size() >= 1) {
		d.type = list[0];
		int n = d.type.find('/');
		if (n >= 0) {
			d.type = d.type.substr(n + 1);
			n = d.type.find('-');
			if (n >= 0)
				d.type.erase(n);
		}
	}

	if (list.size() >= 3)
		d.version = list[2].c_str();

	for (int i = 3; i < (int) list.size(); i++) {
		std::vector<std::string> args;
		split(list[i], '=', args);

		if (args.size() < 2)
			continue;

		if (args[0] == "id")
			d.hwaddr = args[1].c_str();
		else if (args[0] == "subtype")
			d.type += " " + args[1];
		else
			d.attr[args[0]] = args[1];
	}

	if (!d.hwaddr.length())
		d.hwaddr = hwaddr(dev, fd, addr);

	devices.push_back(d);
}

static void probe(const char *dev, struct sockaddr_in *saddr, int bcast)
{
#ifdef _WIN32
	SOCKET fd;
#else
	int fd;
#endif

	if ((fd = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP)) < 0) {
		perror("socket");
		exit(1);
	}

	if (bcast) {
		int n = 1;
		if (setsockopt(fd, SOL_SOCKET, SO_BROADCAST, (void *) &n,
			       sizeof(n))) {
			perror("setsockopt(SO_BROADCAST)");
			exit(1);
		}
	}

	saddr->sin_family = AF_INET;
	saddr->sin_port = htons(port);
#ifdef DEBUG
	printf("sending packet to %s:%d\n", inet_ntoa(saddr->sin_addr),
	       ntohs(saddr->sin_port));
#endif
	/* send a dummy msg */
	const char buf[1] = { 0 };
	if (sendto(fd, buf, sizeof(buf), 0, (struct sockaddr *) saddr,
		   (int) sizeof(struct sockaddr)) < 0) {
		perror("sendto");
		exit(1);
	}

	for ( ; ; ) {
		char buf[256];
		fd_set fds;
		struct timeval tv;
		int retval;

		FD_ZERO(&fds);
		FD_SET(fd, &fds);
		tv.tv_sec = 3;
		tv.tv_usec = 0;
		retval = select(fd + 1, &fds, NULL, NULL, &tv);

		if (retval > 0) {
			struct sockaddr_in caddr = { 0 };
			socklen_t size = sizeof(caddr);
			int n = recvfrom(fd, buf, sizeof(buf), 0,
					 (struct sockaddr *) &caddr, &size);


			if (n >= 0) {
#ifdef DEBUG
				printf("received %d byte%s from %s:%d\n",
				       n, n == 1 ? "" : "s",
				       inet_ntoa(caddr.sin_addr),
				       ntohs(caddr.sin_port));
#endif
				add_device(dev, fd, caddr, buf);
			} else
				fprintf(stderr, "%s: %s\n", "recvfrom",
					strerror(errno));
		} else if (retval == 0) {
#ifdef DEBUG
			printf("timed out\n");
#endif
			break;
		} else {
			fprintf(stderr, "%s: %s\n", "select", strerror(errno));
			break;
		}

	}

#ifdef _WIN32
	closesocket(fd);
#else
	close(fd);
#endif
}

extern int main(int argc, char *argv[])
{
	struct ifaddrs *ifs, *p;

	set_pname(argv[0]);	/* Fix up for getopt_long error messages */
	parse_cmd_line(&argc, &argv);

	if (getifaddrs(&ifs) != 0) {
		perror("getifaddrs");
		exit(1);
	}

	if (xml)
		printf("<?xml version=\"1.0\" encoding=\"ISO-8859-1\"?>\n");

	for (p = ifs; p != NULL; p = p->ifa_next) {
		const char *ifname = p->ifa_name;
		struct sockaddr_in *saddr;
		char addr[32], bcast[32];

		if (p->ifa_addr->sa_family != AF_INET)
			continue;
		if (!(p->ifa_flags & IFF_UP))
			continue;
		if (!(p->ifa_flags & IFF_BROADCAST))
			continue;
		if (strchr(ifname, ':')) /* alias? */
			continue;

		saddr = (struct sockaddr_in *) p->ifa_addr;
		strncpy(addr, inet_ntoa(saddr->sin_addr), sizeof(addr));
		saddr = (struct sockaddr_in *) p->ifa_broadaddr;
		strncpy(bcast, inet_ntoa(saddr->sin_addr), sizeof(bcast));

		printf("%sProbing %s [%s/%s]%s\n", xml ? "<!-- " : "",
		       ifname, addr, bcast, xml ? " -->" : "");

		probe(ifname, saddr, 1);
	}

	freeifaddrs(ifs);

	if (sort)
		std::sort(devices.begin(), devices.end(), sort_fn);

	if (xml)
		printf("<probe>\n");

	for (std::vector<struct device>::const_iterator i = devices.begin();
	     i != devices.end(); ++i) {
		const struct device &d = (*i);
		in_addr addr;
		addr.s_addr = d.ipaddr;

		if (!xml) {
			printf("\t%16s %s %s Version %s\n",
			       inet_ntoa(addr), d.hwaddr.c_str(),
			       d.type.c_str(), d.version.c_str());
		} else {
			printf("\t<device>\n");
			printf("\t\t<ip>%s</ip>\n", inet_ntoa(addr));
			printf("\t\t<info>%s</info>\n", d.info.c_str());
			printf("\t\t<type>%s</type>\n", d.type.c_str());
			printf("\t\t<version>%s</version>\n",
			       d.version.c_str());
			printf("\t\t<id>%s</id>\n", d.hwaddr.c_str());

			for (std::map<std::string, std::string, lstr>
				     ::const_iterator j = d.attr.begin();
			     j != d.attr.end(); ++j) {
				const char *name = (*j).first.c_str();
				const char *value = (*j).second.c_str();
				printf("\t\t<%s>%s</%s>\n", name, value, name);
			}

			printf("\t</device>\n");
		}
	}

	if (xml)
		printf("</probe>\n");

	exit(0);
}