/*
 *
 *	bss_stream: this program reads prismdump output
 *	and generates an output stream of JUST data packets
 *  for a specified SSID or BSSID.  the output format
 *  is the same as the input.
 *
 *	updates: http://www.sublimation.org/802.11
 *
 *	36c9eb5bcddbf9685703423d5c082e8a
 *
 */
#include <stdio.h>
#include <stdlib.h>
#include <ctype.h>
#include <signal.h>
#include <unistd.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <sys/resource.h>
#include <unistd.h>
#include <asm/types.h>
#include <linux/netlink.h>
#include <netinet/in.h>

#define SSID_SIZE	32
#define MAX_MAC		200

__u8 dummybuf[32];
__u8 bcast[6] = { 0xff, 0xff, 0xff, 0xff, 0xff, 0xff };

typedef struct
{
  __u16 frame_control __attribute__ ((packed));
  __u16 duration_id __attribute__ ((packed));
  __u8 mac1[6] __attribute__ ((packed));
  __u8 mac2[6] __attribute__ ((packed));
  __u8 mac3[6] __attribute__ ((packed));
  __u16 sequence __attribute__ ((packed));
  __u8 mac4[6] __attribute__ ((packed));
} wlan_hdr_t;

typedef struct
{
	__u8	bssid[6];
	__u8	ssid[32];
	__u8	addr_list[(6 * MAX_MAC)];
	__u8	num;
	void	*next;
} bss_t;

bss_t *bss_list = NULL;
int	ssid_sort=0;
__u8 targ_id[SSID_SIZE];

typedef struct 
{
  struct timeval ts;
  __u32 caplen;
  __u32 len;
  // int pkt_encap;
} wtap_hdr_t;

typedef struct
{
	wtap_hdr_t wtap;
	wlan_hdr_t wlan;
} packet_t;

typedef struct
{
  __u8 version:2;
  __u8 type:2;
  __u8 subtype:4;
  __u8 toDS:1;
  __u8 fromDS:1;
  __u8 morefrag:1;
  __u8 retry:1;
  __u8 pwr:1;
  __u8 moredata:1;
  __u8 wep:1;
  __u8 rsvd:1;
} frame_control_t;

char *
hexdump (__u8 * x, int y)
{
  int i = -1;
  while (++i < y)
    sprintf (&dummybuf[(i * 2)], "%02x", x[i]);
  return (dummybuf);
}

packet_t *
readpacket (int *size, FILE *stream)
{
	wtap_hdr_t 	wtap_info;
	__u8 		*msgbuf;
	packet_t	*packet;
	
	if (1 != fread(&wtap_info,sizeof(wtap_hdr_t),1,stream))
	{
		fprintf(stderr,"could not read wiretap frame header. aborting...\n");
		return(NULL);
	}
	if (wtap_info.len != wtap_info.caplen)
	{
		fprintf(stderr,"prismdump wiretap header impossibility: caplen!=len (%10u!=%u)\n",
			wtap_info.caplen,wtap_info.len);
		return(NULL);
	}
	*size = (wtap_info.caplen + sizeof(wtap_hdr_t));
	if (NULL == (msgbuf = malloc(*size)))
	{
		perror("malloc");
		exit(-1);
	}
	packet=(packet_t *)msgbuf;
	memcpy((void *)&packet->wtap,(void *)&wtap_info,sizeof(wtap_hdr_t));
	if (1 != fread((void *)&packet->wlan,wtap_info.caplen,1,stream))
	{
		fprintf(stderr,"couldn't read 802.11 frame header! aborting...\n");
		return(NULL);
	}
	return(packet);
}

int bss_find (bss_t *curr, __u8 *addr)
{
	int c=0;
	while (c < curr->num)
	{
		if (0 == memcmp(&curr->addr_list[ 6 * c],addr,6))
			return 1;
		c++;
	}
	return 0;
}

void bss_add (bss_t *curr, __u8 *addr)
{
	if (0 == memcmp (bcast, addr, 6))		// we dont care about broadcasts
		return;
	if (bss_find(curr,addr))			// or addresses we've already seen
		return;
	memcpy(&curr->addr_list[ 6 * (curr->num)], addr, 6);
	curr->num++;
}

bss_t *new_bss(__u8 *bssid, __u8 *da, __u8 *sa, __u8 *ssid)
{
	bss_t *temp;

	if (0 == memcmp (bcast, bssid, 6))		// we dont care about broadcasts
		return NULL;
	if (NULL == (temp = malloc(sizeof(bss_t))))
	{
		perror("malloc");
		exit(-1);
	}
	memset (&temp->addr_list,0xff,MAX_MAC * 6);
	temp->num=0;
	bzero(&temp->ssid,32);
	if (strlen(ssid))
		strncpy(&temp->ssid,ssid,32);
	memcpy(&temp->bssid,bssid,6);
	bss_add(temp,sa);
	bss_add(temp,da);
	return(temp);
}

void
filter_bss (__u8 *bssid, __u8 *da, __u8 *sa, packet_t *packet, int caplen)
{
	bss_t		*temp;
	bss_t		*last;
	__u8 		temp_targ[32];


	for (temp = bss_list; temp;temp=temp->next )
	{
		if (ssid_sort)
		{
			if (((bss_find(temp, sa)) || (bss_find(temp,da))) && 
				(0 == strncmp(targ_id,&temp->ssid,32)))
					fwrite (packet,caplen,1,stdout);
		}
		else
		{
			sprintf(temp_targ,"%s",hexdump((__u8 *)&temp->bssid,6));
			if ((0 == memcmp (&temp->bssid, bssid, 6)) && 
				(0 == strncmp(targ_id,temp_targ,6)))
					fwrite (packet,caplen,1,stdout);
		}
	}
}

void
update_bss(__u8 *bssid, __u8 *da, __u8 *sa, __u8 *ssid, packet_t *packet)
{
	bss_t		*temp;
	bss_t		*last;
	wlan_hdr_t *p80211_hdr;
	frame_control_t *fc;

	p80211_hdr = &packet->wlan;
	fc=(frame_control_t *)&p80211_hdr->frame_control;

#if 0
	printf ("da: %s ",hexdump(da,6));
	printf ("sa: %s ",hexdump(sa,6));
	printf ("bssid: %s ssid:%s\n",hexdump(bssid,6), ssid);
#endif

	if (ssid_sort && (fc->type == 2))
		return;
	if (NULL == bss_list)
	{
		bss_list = new_bss (bssid, da, sa, ssid);
		return;
	}
	else
	{
		for (temp = bss_list; temp;temp=temp->next )
		{
			if (ssid_sort)
			{
				if (0 == strncmp (&temp->ssid, ssid, 32))
				{
					bss_add(temp,sa);
					bss_add(temp,da);
					return;
				}
			}
			else
			{
				if (0 == memcmp (&temp->bssid, bssid, 6))
				{
					if (strlen(ssid))
						strncpy(&temp->ssid,ssid,32);
					bss_add(temp,sa);
					bss_add(temp,da);
					return;
				}
			}
			last=temp;
		}
		last->next=new_bss(bssid, da, sa, ssid);
	}
}

void dump_bsss (void)
{
	bss_t *curr_bss=bss_list;
	int c;

	while (curr_bss)
	{
		fprintf (stderr,"dumping bssid : %s    ssid: %s\n", hexdump((__u8 *)&curr_bss->bssid,6),curr_bss->ssid);
		c=0;
		while (0 != memcmp (&curr_bss->addr_list[c * 6],bcast,6))
		{
			fprintf (stderr,"\taddress %d: %s ", c, hexdump(&curr_bss->addr_list[c * 6],6));
			if (!ssid_sort)
			{
				if (0 == memcmp(&curr_bss->addr_list[c * 6], &curr_bss->bssid, 6))
					fprintf (stderr,"AP\n");
				else
					fprintf (stderr,"STA\n");
			}
			else
				fprintf (stderr,"\n");
			c++;
		}
		curr_bss=curr_bss->next;
	}

}

void
handledata (packet_t *packet, int caplen)
{
	wlan_hdr_t *p80211_hdr;
	frame_control_t *fc;
	__u8		tofrom=0;
	__u8		*da,*sa,*bssid;

	p80211_hdr = &packet->wlan;
	fc = (frame_control_t *) &p80211_hdr->frame_control;
	tofrom = (fc->toDS << 1) | fc->fromDS;
	switch (tofrom) 
	{
		case 0:						// ad hoc
			da=(__u8 *)&p80211_hdr->mac1;
			sa=(__u8 *)&p80211_hdr->mac2;
			bssid=(__u8 *)&p80211_hdr->mac3;
			break;
		case 1:						// AP to STA
			da=(__u8 *)&p80211_hdr->mac1;
			sa=(__u8 *)&p80211_hdr->mac3;
			bssid=(__u8 *)&p80211_hdr->mac2;
			break;
		case 2:						// STA to AP
			da=(__u8 *)&p80211_hdr->mac2;
			sa=(__u8 *)&p80211_hdr->mac3;
			bssid=(__u8 *)&p80211_hdr->mac1;
			break;
		case 3:						// AP to AP
#if 0
			printf ("essid ");
			printf ("ra: %s ",hexdump((__u8 *)&p80211_hdr->mac1,6));
			printf ("ta: %s ",hexdump((__u8 *)&p80211_hdr->mac2,6));
			printf ("da: %s ",hexdump((__u8 *)&p80211_hdr->mac3,6));
			printf ("sa: %s\n",hexdump((__u8 *)&p80211_hdr->mac4,6));
#endif
			return;
	}
	update_bss(bssid,da,sa,"",packet);
	filter_bss(bssid,da,sa,packet,caplen);
	return;
}

void
handlemanagement (packet_t *packet, int caplen)
{
	wlan_hdr_t *p80211_hdr;
	frame_control_t *fc;
	__u8		*da,*sa,*bssid,*msgbuf,temp,ssid[SSID_SIZE];

	bzero (ssid,SSID_SIZE);
	p80211_hdr = &packet->wlan;
	fc=(frame_control_t *)&p80211_hdr->frame_control;
	da=(__u8 *)&p80211_hdr->mac1;
	sa=(__u8 *)&p80211_hdr->mac2;
	bssid=(__u8 *)&p80211_hdr->mac3;
	msgbuf = (__u8 *)p80211_hdr;
	if ((fc->type == 0) && ( fc->subtype == 8)) // just beacon packets
	{
		temp=(msgbuf[37] & 0xff)+1;
		snprintf(ssid,temp,"%s",&msgbuf[38]);
	}   
	update_bss(bssid,da,sa,ssid,packet);
}

int
main (int argc, char **argv)
{
	int caplen, counter, n=0;
	__u8 *msgbuf,fileheader[24];
	wtap_hdr_t *wiretap_hdr;
	wlan_hdr_t *p80211_hdr;
	packet_t	*packet;
	frame_control_t *fc;

	if (argc<2)
		exit(1);
	strncpy(targ_id,argv[1],32);
	counter = argc;
	while (counter-- > 2)
	{
		if (strncmp (argv[counter], "-s", 2) == 0)
			ssid_sort=1;
		if (strncmp (argv[counter], "-?", 2) == 0)
		{
			printf ("\n%s usage:\n\t%s (ssid|bssid) (options) < rawlogfile\n",argv[0],argv[0]);
			printf ("\n\t\tor\n\n");
			printf ("\tprismdump | %s (ssid|bssid) (options)\n",argv[0]);
			printf ("\noptions:\n");
			printf ("\t-?\tfilter by SSID, not BSSID\n");
			printf ("\t-?\tthis help\n\n");
			exit(0);
		}
	}
	if (1!=fread(fileheader,24,1,stdin))
	{
		fprintf (stderr,"could not read the file header\n");
		exit(-1);
	}
	if (1!=fwrite(fileheader,24,1,stdout))
	{
		fprintf (stderr,"could not write the file header\n");
		exit(-1);
	}
	while (1)
	{
		if (NULL == (packet=readpacket(&caplen,stdin)))
			break;
		n++;
		wiretap_hdr= &packet->wtap;
		p80211_hdr = &packet->wlan;
		fc=(frame_control_t *)&p80211_hdr->frame_control;
		if (fc->type == 2)
			handledata(packet,caplen);
		else if (fc->type == 0)
			handlemanagement(packet,caplen);
		free(packet);
	}
//	dump_bsss();
	fprintf (stderr,"%d packets analyzed\n",n);
	fflush(stdout);
	fflush(stderr);
	return(0);
}