/** WEP/WPA key setting tool for MadWifi driver, version 0.53
 *
 * (C) 2008-03-28 Georg Lukas <georg@madwifi.org>
 *
 * This program can be used to debug the MadWifi Key Cache. Use with caution
 * and without warranty!
 *
 * Instructions:
 *
 *	gcc -I/usr/src/madwifi-ng -Wall wpakey.c -o wpakey
 *	./wpakey -h
 *
 * This code is published under the GNU General Public Licence v2.
 */

#include <iwlib.h>
#include <include/compat.h>
#include <net80211/ieee80211_ioctl.h>

#include <stdio.h>
#include <stdint.h>

#define MACS "%02hhx:%02hhx:%02hhx:%02hhx:%02hhx:%02hhx"
#define MACP(mac) (mac)[0], (mac)[1], (mac)[2], (mac)[3], (mac)[4], (mac)[5]

char *dev = "ath0";
int sock;
int warn_wpa = 1;

int parse_mac(uint8_t *mac, const char *str) {
	if (sscanf(str, MACS,
			&mac[0], &mac[1], &mac[2], &mac[3], &mac[4], &mac[5]
			) < 6) {
		fprintf(stderr, "Invalid MAC address at \"%s\"\n", str);
		return 0;
	} else {
		return 1;
	}
}
int parse_key(uint8_t *key, const char *str) {
	if (sscanf(str, "%02hhx%02hhx%02hhx%02hhx%02hhx%02hhx%02hhx%02hhx%02hhx%02hhx%02hhx%02hhx%02hhx%02hhx%02hhx%02hhx",
			&key[0], &key[1], &key[2], &key[3], &key[4], &key[5], &key[6],
			&key[7], &key[8], &key[9], &key[10], &key[11], &key[12], &key[13],
			&key[14], &key[15]
			) < 16) {
		fprintf(stderr, "Invalid key at \"%s\"\n", str);
		return 0;
	} else {
		return 1;
	}
}

void hexdump(unsigned char *data, ssize_t dlen) {
	//printf("%s: (%i) ", prefix, dlen);
	while (dlen-- > 0) {
		printf("%02hhx", *data++);
	}
}

int set80211param(int op, int arg) {
	struct iwreq iwr;

	memset(&iwr, 0, sizeof(iwr));

	strncpy(iwr.ifr_name, dev, IFNAMSIZ);
	iwr.u.mode = op;
	memcpy(iwr.u.name+4, &arg, 4);

	if (ioctl(sock, IEEE80211_IOCTL_SETPARAM, &iwr) < 0) {
		perror("ioctl(setparam)");
		return -1;
	}
	return 0;
}

int get80211param(int op) {
	struct iwreq iwr;

	memset(&iwr, 0, sizeof(iwr));

	strncpy(iwr.ifr_name, dev, IFNAMSIZ);
	iwr.u.mode = op;

	if (ioctl(sock, IEEE80211_IOCTL_GETPARAM, &iwr) < 0) {
		perror("ioctl(getparam)");
		return -1;
	}
	return iwr.u.mode;
}


int set80211priv(int op, void *data, int len) {
	struct iwreq iwr;

	memset(&iwr, 0, sizeof(iwr));

	strncpy(iwr.ifr_name, dev, IFNAMSIZ);
	iwr.u.data.pointer = data;
	iwr.u.data.length = len;

	if (ioctl(sock, op, &iwr) < 0) {
		perror("ioctl()");
		return -1;
	}
	return iwr.u.data.length;
}

void prep_key(struct ieee80211req_key *wk, int keyidx, uint8_t *mac) {
	memset(wk, 0, sizeof(struct ieee80211req_key));
	wk->ik_keyix = keyidx;

	if (keyidx == IEEE80211_KEYIX_NONE) {
		memcpy(wk->ik_macaddr, mac, 6);
	}
}

char *cipherstrs[] = { "WEP", "TKIP", "OCB", "CCMP", "invalid", "CKIP", "none" };

char *strcipher(int c) {

	if (c > IEEE80211_CIPHER_NONE) return "invalid";
	return cipherstrs[c];
}

char *strflags(int f) {
	static char buf[5];
	char *ff = buf;
	memset(buf, 0, sizeof(buf));

	if (f & IEEE80211_KEY_RECV) *ff++ = 'R';
	if (f & IEEE80211_KEY_XMIT) *ff++ = 'T';
	if (f & IEEE80211_KEY_DEFAULT) *ff++ = 'D';
	if (f & IEEE80211_KEY_GROUP) *ff++ = 'G';
	*ff = '\0';
	return buf;
}

int getkey(int keyidx, uint8_t *mac, int groupkey, int verbose) {
	struct ieee80211req_key wk;
	char idxbuf[6];

	if (warn_wpa && get80211param(IEEE80211_PARAM_WPA) == 0) {
		printf("WARNING: WPA is disabled!\n");
		warn_wpa = 0;
	}
	prep_key(&wk, keyidx, mac);
	if (groupkey)
		wk.ik_flags = IEEE80211_KEY_GROUP;
	if (set80211priv(IEEE80211_IOCTL_GETKEY, &wk, sizeof(wk)) >= 0) {
		if (verbose == 0 && wk.ik_type == IEEE80211_CIPHER_NONE && wk.ik_keylen == 0)
			return 0;
		if (wk.ik_keyix == IEEE80211_KEYIX_NONE) {
			strcpy(idxbuf, groupkey?"MC":"UC");
		} else {
			snprintf(idxbuf, sizeof(idxbuf), "%2x", wk.ik_keyix);
		}
		printf("Key %s: <" MACS "> ", idxbuf, MACP(wk.ik_macaddr));
		if (wk.ik_type == IEEE80211_CIPHER_NONE && wk.ik_flags == 3 && wk.ik_keylen == 0) {
			printf("off\n");
		} else {
			printf("c=%-4s f=%-4s k<%-2i>=", strcipher(wk.ik_type),
				strflags(wk.ik_flags), wk.ik_keylen);
			hexdump(wk.ik_keydata, wk.ik_keylen);
			printf(" rs=%lld ts=%lld\n", wk.ik_keyrsc, wk.ik_keytsc);
		}
		return 0;
	}
	return -1;
}

int delkey(int keyidx, uint8_t *mac) {
	struct ieee80211req_key wk;

	prep_key(&wk, keyidx, mac);
	return set80211priv(IEEE80211_IOCTL_DELKEY, &wk, sizeof(wk));
}

int setkey(int keyidx, uint8_t *mac, int type, int flags, int keylen, uint8_t *key) {
	struct ieee80211req_key wk;

	prep_key(&wk, keyidx, mac);
	wk.ik_type = type;
	wk.ik_flags = flags;
	wk.ik_keylen = keylen;
	memcpy(wk.ik_keydata, key, keylen);

	return set80211priv(IEEE80211_IOCTL_SETKEY, &wk, sizeof(wk));
}


void iter_sta() {
	uint8_t buf[24*1024];
	uint8_t *bufpos;
	ssize_t len;

	if ((len = set80211priv(IEEE80211_IOCTL_STA_INFO, buf, sizeof(buf))) >= 0) {
		struct ieee80211req_sta_info *si;
		bufpos = buf;
		while (len > sizeof(struct ieee80211req_sta_info)) {
			si = (struct ieee80211req_sta_info*)bufpos;

			getkey(IEEE80211_KEYIX_NONE, si->isi_macaddr, 0, 0);
			getkey(IEEE80211_KEYIX_NONE, si->isi_macaddr, 1, 0);

			bufpos += si->isi_len;
			len -= si->isi_len;
		}

	}
}

void set_wpa(int cipher, int wpa, int key) {
	printf("Setting WPA: cipher=%s wpa=%i mgmt=%i\n",
		strcipher(cipher), wpa, key);
	set80211param(IEEE80211_PARAM_MCASTCIPHER, cipher);
	set80211param(IEEE80211_PARAM_UCASTCIPHERS, 1 << cipher);
	set80211param(IEEE80211_PARAM_KEYMGTALGS, key);
	set80211param(IEEE80211_PARAM_WPA, wpa);

	set80211param(IEEE80211_PARAM_PRIVACY, wpa > 0? 1 : 0);
}


void init() {
	sock = socket(PF_INET, SOCK_DGRAM, 0);
	if (sock < 0) {
		perror("socket()");
		exit(32);
	}
}


void help() {
	fprintf(stderr, "Possible options are:\n"
		"	-a		print all group keys\n"
		"	-A		print all keys (default option)\n"
		"	-i <if>		set interface (default: %s)\n"
		"	-w		WPA on\n"
		"	-n		no WPA\n"
		"	-k <idx|mac>	set / read specified key (default: 0)\n"
		"	-c <#>		set ciphers\n"
		"	-f <#|rtgd>	set key flags - integer or: Rx/Tx/Group/Default\n"
		"	[-|<key>]	hex key to be set or '-' to unset.\n"
		"\n"
		"Example:\n"
		"	# activate WPA on ath3, write an RX+TX key to slot #3\n"
		"	./wpakey -i ath3 -w -k 3 -f rt XXXXXXXXXXXXXXXX\n"
		"", dev);
}

int main(int argc, char** argv) {
	int keyidx;
	uint8_t mac[6];
	int cipher = IEEE80211_CIPHER_AES_CCM;
	int flags = IEEE80211_KEY_RECV | IEEE80211_KEY_XMIT;
	int c;

	init();

	if (argc == 1) {
		for (keyidx = 0; keyidx <= 3; keyidx++) {
			getkey(keyidx, NULL, 0, 1);
		}
		iter_sta();
		return 0;
	}
	while ((c = getopt(argc, argv, "aAhi:f:c:k:wn")) != -1) {
		//printf("%c - %s\n", c, optarg);
		switch (c) {
		case 'a':
		case 'A':
			for (keyidx = 0; keyidx <= 3; keyidx++) {
				getkey(keyidx, NULL, 0, 1);
			}
			if (c == 'A')
				iter_sta();
			break;
		case 'h':
			help();
			break;
		case 'i':
			dev = optarg;
			break;
		case 'f':
			if (isalpha(optarg[0])) {
				flags = 0;
				while (*optarg) switch (*optarg++) {
				case 'g': flags |= IEEE80211_KEY_GROUP; break;
				case 'r': flags |= IEEE80211_KEY_RECV; break;
				case 't': flags |= IEEE80211_KEY_XMIT; break;
				case 'd': flags |= IEEE80211_KEY_DEFAULT; break;
				}
			} else
				flags = strtol(optarg, NULL, 0);
			printf("flags = 0x%02x\n", flags);
			break;
		case 'c':
			cipher = strtol(optarg, NULL, 0);
			printf("cipher = %x\n", cipher);
			break;
		case 'w':
			set_wpa(cipher, 3, WPA_ASE_8021X_PSK);
			break;
		case 'n':
			set_wpa(IEEE80211_CIPHER_NONE, 0, WPA_ASE_NONE);
			break;
		case 'k':
			if (strlen(optarg) < 17)
				keyidx = atoi(optarg);
			else {
				keyidx = IEEE80211_KEYIX_NONE;
				parse_mac(mac, optarg);
			}
			getkey(keyidx, mac, 0, 1);
			if (keyidx == IEEE80211_KEYIX_NONE)
				getkey(keyidx, mac, 1, 1);
			break;
		}
	}

	if (optind < argc) {
		uint8_t key[16];
		memset(key, 0, 16);
		if (strcmp(argv[optind], "-") == 0) {
			delkey(keyidx, mac);
			setkey(keyidx, mac, IEEE80211_CIPHER_NONE, 0, 0, NULL);
			getkey(keyidx, mac, 0, 1);
			if (keyidx == IEEE80211_KEYIX_NONE)
				getkey(keyidx, mac, 1, 1);
		} else {
			parse_key(key, argv[optind]);
			setkey(keyidx, mac, cipher,
				flags, 128/8, key);
			getkey(keyidx, mac, 0, 1);
			if (keyidx == IEEE80211_KEYIX_NONE)
				getkey(keyidx, mac, 1, 1);
		}
	}

	return 0;
}