#include <sys/param.h>
#include <sys/ioctl.h>
#include <sys/socket.h>
#include <sys/sysctl.h>
#include <sys/time.h>
#include <sys/sockio.h>

#include <net/ethernet.h>
#include <net/if.h>
#include <net/if_dl.h>
#include <net/if_types.h>
#include <net/if_media.h>
#include <net/route.h>

#include <net80211/ieee80211.h>
#include <net80211/ieee80211_crypto.h>
#include <net80211/ieee80211_ioctl.h>

#include <ifaddrs.h>
#include <ctype.h>
#include <err.h>
#include <errno.h>
#include <fcntl.h>
#include <inttypes.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

static char name[IFNAMSIZ];	/* name of interface */
static int verbose = 0;
static int only_remote = 0;
static uint8_t localaddr[6] = { 0, };
static struct ether_addr remaddr;

static u_int
ieee80211_mhz2ieee(u_int freq)
{
	if (freq == 2484)
		return 14;
	if (freq < 2484)
		return (freq - 2407) / 5;
	if (freq < 5000)
		return 15 + ((freq - 2512) / 20);
	return (freq - 5000) / 5;
}

static const char *
getcaps(int capinfo)
{
	static char capstring[32];
	char *cp = capstring;

	if (capinfo & IEEE80211_CAPINFO_ESS)
		*cp++ = 'E';
	if (capinfo & IEEE80211_CAPINFO_IBSS)
		*cp++ = 'I';
	if (capinfo & IEEE80211_CAPINFO_CF_POLLABLE)
		*cp++ = 'c';
	if (capinfo & IEEE80211_CAPINFO_CF_POLLREQ)
		*cp++ = 'C';
	if (capinfo & IEEE80211_CAPINFO_PRIVACY)
		*cp++ = 'P';
	if (capinfo & IEEE80211_CAPINFO_SHORT_PREAMBLE)
		*cp++ = 'S';
	if (capinfo & IEEE80211_CAPINFO_PBCC)
		*cp++ = 'B';
	if (capinfo & IEEE80211_CAPINFO_CHNL_AGILITY)
		*cp++ = 'A';
	if (capinfo & IEEE80211_CAPINFO_SHORT_SLOTTIME)
		*cp++ = 's';
	if (capinfo & IEEE80211_CAPINFO_RSN)
		*cp++ = 'R';
	if (capinfo & IEEE80211_CAPINFO_DSSSOFDM)
		*cp++ = 'D';
	*cp = '\0';
	return capstring;
}

static void
printie(const char* tag, const uint8_t *ie, size_t ielen, int maxlen)
{
	printf("%s", tag);
	if (verbose) {
		maxlen -= strlen(tag)+2;
		if (2*ielen > maxlen)
			maxlen--;
		printf("<");
		for (; ielen > 0; ie++, ielen--) {
			if (maxlen-- <= 0)
				break;
			printf("%02x", *ie);
		}
		if (ielen != 0)
			printf("-");
		printf(">");
	}
}

/* unaligned little endian access */
#define LE_READ_4(p)					\
	((u_int32_t)					\
	 ((((const u_int8_t *)(p))[0]      ) |		\
	  (((const u_int8_t *)(p))[1] <<  8) |		\
	  (((const u_int8_t *)(p))[2] << 16) |		\
	  (((const u_int8_t *)(p))[3] << 24)))

static int __inline
iswpaoui(const u_int8_t *frm)
{
	return frm[1] > 3 && LE_READ_4(frm+2) == ((WPA_OUI_TYPE<<24)|WPA_OUI);
}

static int __inline
iswmeoui(const u_int8_t *frm)
{
	return frm[1] > 3 && LE_READ_4(frm+2) == ((WME_OUI_TYPE<<24)|WME_OUI);
}

static int __inline
isatherosoui(const u_int8_t *frm)
{
	return frm[1] > 3 && LE_READ_4(frm+2) == ((ATH_OUI_TYPE<<24)|ATH_OUI);
}

static void
printies(const u_int8_t *vp, int ielen, int maxcols)
{
	while (ielen > 0) {
		switch (vp[0]) {
		case IEEE80211_ELEMID_VENDOR:
			if (iswpaoui(vp))
				printie(" WPA", vp, 2+vp[1], maxcols);
			else if (iswmeoui(vp))
				printie(" WME", vp, 2+vp[1], maxcols);
			else if (isatherosoui(vp))
				printie(" ATH", vp, 2+vp[1], maxcols);
			else
				printie(" VEN", vp, 2+vp[1], maxcols);
			break;
		case IEEE80211_ELEMID_RSN:
			printie(" RSN", vp, 2+vp[1], maxcols);
			break;
		default:
			printie(" ???", vp, 2+vp[1], maxcols);
			break;
		}
		ielen -= 2+vp[1];
		vp += 2+vp[1];
	}
}

static void
list_stations(int s, int head)
{
	uint8_t buf[24*1024];
	struct ieee80211req ireq;
	struct ieee80211req_sta_info *si, *ni;
	uint8_t *vp;
	uint8_t *cp;
	int len, print;

	(void) memset(&ireq, 0, sizeof(ireq));
	(void) strncpy(ireq.i_name, name, sizeof(ireq.i_name));
	ireq.i_type = IEEE80211_IOC_STA_INFO;
	ireq.i_data = buf;
	ireq.i_len = sizeof(buf);
	if (ioctl(s, SIOCG80211, &ireq) < 0)
		err(1, "unable to get station information");
	len = ireq.i_len;
	if (len < sizeof(struct ieee80211req_sta_info))
		return;

	if (head)
		printf("%-17.17s %4s %4s %4s %4s %4s %6s %6s %4s %3s\n"
			, "ADDR"
			, "AID"
			, "CHAN"
			, "RATE"
			, "RSSI"
			, "IDLE"
			, "TXSEQ"
			, "RXSEQ"
			, "CAPS"
			, "ERP"
		);

	cp = buf;
	ni = (struct ieee80211req_sta_info *)cp;

	do {

		si = ni;
		vp = (u_int8_t *)(si + 1);

		cp += si->isi_len;
		len -= si->isi_len;
		ni = (struct ieee80211req_sta_info *)cp;

		if (len < sizeof(struct ieee80211req_sta_info))
			ni = NULL;

		print = 0;

		if (memcmp (localaddr, si->isi_macaddr, sizeof (localaddr)) != 0 &&
		    (!ni || memcmp (ni->isi_macaddr, si->isi_macaddr, sizeof (ni->isi_macaddr)) != 0))
			print = 1;

		if (only_remote)
			print = (memcmp(&remaddr, si->isi_macaddr, sizeof (remaddr)) == 0);

		if (print) {
			printf("%s %4u %4d %3dM %4d %4d %6d %6d %-4.4s %3x"
				, ether_ntoa((const struct ether_addr*) si->isi_macaddr)
				, IEEE80211_AID(si->isi_associd)
				, ieee80211_mhz2ieee(si->isi_freq)
				, (si->isi_rates[si->isi_txrate] & IEEE80211_RATE_VAL)/2
				, si->isi_rssi
				, si->isi_inact
				, si->isi_txseqs[0]
				, si->isi_rxseqs[0]
				, getcaps(si->isi_capinfo)
				, si->isi_erp
			);
			printies(vp, si->isi_ie_len, 24);
			printf("\n");
		}

	} while (ni);
}

static void
load_lladdr (int s)
{
	struct ifaddrs *ifap, *p;

    if (getifaddrs(&ifap) != 0)
		err(1, "couldn't get interface addresses");

    for (p = ifap; p; p = p->ifa_next) {
		if (p->ifa_addr->sa_family == AF_LINK &&
 			strcmp(name, p->ifa_name) == 0) {
			struct sockaddr_dl* sdp = (struct sockaddr_dl*) p->ifa_addr;
			memcpy(localaddr, sdp->sdl_data + sdp->sdl_nlen, 6);
			break;
		}
	}

	freeifaddrs(ifap);
}

int
main (int argc, char* argv[])
{
	struct ether_addr *eth;
	int s, i;

	if (argc < 2 || argc > 3) {
		fprintf (stderr, "usage: athsta <interface> [macaddr]\n");
		exit (2);
	}

	strncpy (name, argv[1], sizeof (name));
	name[sizeof (name) - 1] = 0;

	if (argc == 3) {
		eth = ether_aton (argv[2]);
		if (!eth)
			err (1, "invalid mac address: %s", argv[2]);
		memcpy (&remaddr, eth, sizeof (remaddr));
		only_remote = 1;
	}

	if ((s = socket (AF_INET, SOCK_DGRAM, 0)) < 0)
		err(1, "creating socket failed");

	load_lladdr (s);

	for (i = 0; ; i++) {
		list_stations (s, !(i % 10));
		fflush(stdout);
		sleep (1);
	}

	return 0;
}