#include <u.h>
#include <libc.h>
#include <thread.h>
#include <ip.h>
#include <libsec.h>
#include "dat.h"
#include "fns.h"


static void
setKey(int fd, int idx, char*k, int l, int isunicast)
{
	char *s;

	USED(l);
	s = isunicast ? "unicast" : "broadcast";
	appendlog(getKeysbuf(), 0, "%s\tkey%1d %s", s, idx, k);
	loglog("%s key%1d %s", s, idx, k);
	if (fprint(fd, "key%1d %s\n", idx, k) < 0)
		loglog("error setting %s key %d: %r", s, idx);
}

static void
decode(uchar *buf, int blen, uchar *data, int dlen, uchar *iv, int ilen)
{
	uchar *p;
	RC4state rc4s;
	uchar seed[256];

	assert(dlen <= blen);
	memcpy(buf, data, dlen);
	assert(ilen + sizeof(theSessionKey) <= sizeof(seed));
	p = seed;
	memcpy(p, iv, ilen);
	p += ilen;
	memcpy(p, theSessionKey, sizeof(theSessionKey));
	p += sizeof(theSessionKey);
	setupRC4state(&rc4s, seed, p-seed);
	rc4(&rc4s, buf, dlen);
}

void
handleRc4Key(int fd, RC4KeyDesc* m, int mlen)
{
	int len, idx, isunicast, sz;
	char b[40], *p;
	uchar *k, dec[256];
	int i;

	len = nhgets(m->ln);
	idx = (m->idx & RC4KeyIdx) + 1;
	isunicast = m->idx & RC4KeyUnicast;
	sz = mlen - RC4KeyDescHlen;

	if (debug) print("rc4 key len=%ud idx=%ud sz=%d isunicast=%d\n", len, idx, sz, isunicast);

	if (sz < len) {
		if (debug) print("use sessionkey\n");
		k = theSessionKey;
	} else {
		if (debug) print("use included key\n");
		decode(dec, sizeof(dec), m->data, len, m->iv, 2);
		k = dec;
	}

	memset(b, 0, sizeof(b));
	p = b;
	for (i=0; i<len; i++) {
		sprint(p, "%02x", k[i]);
		p+=2;
	}
	setKey(fd, idx, b, p-b, isunicast);
}

void
handleKey(int fd, Eapol* e, int len)
{
	KeyDesc *k;
	int eapollen;

	eapollen=nhgets(e->ln);
	if (eapollen > len) {
		loglog("ignoring short key msg (eapollen=%d > len=%d)", eapollen, len);
		return;
	}
	k = (KeyDesc*)e->data;
	if (debug) print("handling key frame len=%ud type=%d\n", len, k->tp);
	switch (k->tp) {
	case KeyTpRC4:
		handleRc4Key(fd,  (RC4KeyDesc*)k->data, eapollen - KeyDescHlen);
		break;
	default:
		loglog("no support for key type %d", k->tp);
	}
}