/*
 *  capoptim 0.1 - Analyze and optimize wireless captures
 *
 *  Copyright (C) 2006 SuD
 *  Copyright (C) 2006 Thomas d'Otreppe
 *  Copyright (C) 2004,2005  Christophe Devine (pcap2ivs and mergeivs)
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; either version 2 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program; if not, write to the Free Software
 *  Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 */
 
#include <string.h>
#include <stdio.h>
#include <time.h>
#include <stdlib.h>
#include <unistd.h>

#define TCPDUMP_MAGIC           0xA1B2C3D4
#define TCPDUMP_CIGAM           0xD4C3B2A1

#define LINKTYPE_ETHERNET       1
#define LINKTYPE_IEEE802_11     105
#define LINKTYPE_PRISM_HEADER   119
#define LINKTYPE_RADIOTAP_HDR   127

#define SWAP32(x)       \
    x = ( ( ( x >> 24 ) & 0x000000FF ) | \
          ( ( x >>  8 ) & 0x0000FF00 ) | \
          ( ( x <<  8 ) & 0x00FF0000 ) | \
          ( ( x << 24 ) & 0xFF000000 ) );

struct pcap_file_header
{
    uint magic;
    ushort version_major;
    ushort version_minor;
    int thiszone;
    uint sigfigs;
    uint snaplen;
    uint linktype;
};

struct pcap_pkthdr
{
    int tv_sec;
    int tv_usec;
    uint caplen;
    uint len;
};

typedef struct {
	int n;
	unsigned char * p;
} up;

#define UBLOCK 1024
up * u;
int sizu=0;
int nu=0;

#define DBLOCK 1024
up * du;
int ndu=0;
int pd;

int unique(unsigned char * p, int n) {
	int i, idx;
	unsigned char * tmp;

	if (n==0)
		return 0;
	
	// data packets
	if ( (p[0]&0x0C) == 0x08) {
		for (i=0; i<ndu; i++) {
			idx=(pd+i)%DBLOCK;
			if (du[idx].n!=n)
				continue;
			tmp=du[idx].p;

			// search in data buffer
			if (n>=0x18) {
				if ( tmp[0]==p[0] && ((tmp[1]&0xF7)==(p[1]&0xF7)) && memcmp(tmp+4, p+4, n-4) == 0) 
					return i;
			} else if ( memcmp(tmp, p, n) == 0)
				return i;
		}

		// not found in buffer, assume it's unique
		
		if (ndu<DBLOCK) {
			// grow buffer until limit
			idx=(pd+ndu)%DBLOCK;
			du[idx].n=n;
			du[idx].p=(unsigned char*)malloc(n);
			memcpy(du[idx].p, p, n);
			ndu++;
		} else {
			// overwrite circular buffer
			free(du[pd].p);
			du[pd].n=n;
			du[pd].p=(unsigned char*)malloc(n);
			memcpy(du[pd].p, p, n);
			pd = (pd+1)%DBLOCK;
		}
		
		return -1;
	}
	
	// ctrl, mgt and others
	for (i=0; i<nu; i++) {
		if (u[i].n == n) {
			tmp=u[i].p;

			// ignore retry retry
			if (n>=2 && ( (tmp[0]!=p[0]) || ((tmp[1]&0xF7)!=(p[1]&0xF7)) ) )
				continue;
			// beacons and probe-resp have a timestamp	
			if ( ((p[0]&0xFC)==0x80 || (p[0]&0xFC)==0x50) && (n>=0x20)) {
				if ( (memcmp(tmp+4, p+4, 0x16-4) == 0) &&
				     (memcmp(tmp+0x20, p+0x20, n-0x20) == 0) )
					     return i;
			// standard mgt have seq 
			} else if ( ((p[0]&0x0C)==0) && (n>=0x18) ) {
				if ( (memcmp(tmp+4, p+4, 0x16-4) == 0) &&
				     (memcmp(tmp+0x18, p+0x18, n-0x18) == 0) )
					     return i;
			} else {
				if ( memcmp(tmp+2, p+2, n-2) == 0)
				return i;
			}
		}
	}

	// not found, it's unique
	
	// realloc buffer if needed
	if (nu+1>sizu) {
		u=realloc(u, (sizu+UBLOCK) * sizeof(up));
		sizu+=UBLOCK;
	}
	// append frame to buffer
	u[nu].n=n;
	u[nu].p=(unsigned char*)malloc(n);
	memcpy(u[nu].p, p, n);

	nu++;

	return -1;
}


void usage(char * progname) {
	if (strrchr(progname, '/') != NULL) {
		progname=strrchr(progname, '/')+1;
	}
       printf("\n %s 0.1 - (C) 2006 SuD."
              "  Original work: Thomas d\'Otreppe, Christophe Devine\n"
              "\n   usage: ", progname);
       printf( "%s [-vv] <pcap file> [<output file>]\n"
                  "        Analyze and optimize wireless pcap capture by removing duplicate or equivalent packets\n", progname);
       puts("");	  
}

#define fwrite2(ptr, tam, nmiemb, flujo) if (fwrite(ptr, tam, nmiemb, flujo) != nmiemb) { perror("Write error"); remove(fn_out); exit(-1); }

int main( int argc, char *argv[] )
{
    time_t tt;
    int n;
    FILE *f_in, *f_out;
    char *fn_in, *fn_out;
    unsigned char *h80211;
    unsigned char buffer[65536];
    struct pcap_file_header pfh;
    struct pcap_pkthdr pkh;    
    int plen;
    typedef enum { D,  MB, MAS, MDS, MAU, MDU, MQ, MP, MO,  CR, CC, CA, CO,  O,  T } tipo;
    int b[T+1];
    int bs[T+1];
    int a[T+1];
    int as[T+1];
    int *pt, *ps;
    int verb=0;
   
    while ((n=getopt(argc, argv, "v")) != -1) {
	switch(n) {
		case 'v':
			verb++;
			break;
		case '?':
			usage(argv[0]);
			return -1;
		default:
			break;
	}
    }

    if (optind >= argc || optind < argc-2) {
	usage(argv[0]);
	return -1;
    } else {
        fn_in=argv[optind];
	if (optind<argc-1) {
		fn_out=argv[optind+1];
		if (strcmp(fn_out, fn_in)==0) {
			fprintf(stderr, "Input and output file MUST NOT be the same!\n");
			return -1;
		}
	} else
		fn_out=NULL;
    }
    
    for (n=0; n<=T; n++) {
    	b[n]=bs[n]=a[n]=as[n]=0;
    }

    sizu=UBLOCK;
    u=(up*)malloc(sizu * sizeof (up));

    ndu=0;
    du=(up*)malloc(DBLOCK * sizeof (up));
    pd=0;

    /* check the input pcap file */

    printf( "Opening %s\n", fn_in );

    if( ( f_in = fopen( fn_in, "rb" ) ) == NULL )
    {
        perror( "fopen failed" );
        return( 1 );
    }

    n = sizeof( pfh );

    if( fread( &pfh, 1, n, f_in ) != (size_t) n )
    {
        perror( "fread(pcap file header) failed" );
        return( 1 );
    }

    if( pfh.magic != TCPDUMP_MAGIC &&
        pfh.magic != TCPDUMP_CIGAM )
    {
        printf( "\"%s\" isn't a pcap file (expected "
                "TCPDUMP_MAGIC).\n", fn_in );
        return( 1 );
    }

    if( pfh.magic == TCPDUMP_CIGAM )
        SWAP32( pfh.linktype );

    if( pfh.linktype != LINKTYPE_IEEE802_11 &&
        pfh.linktype != LINKTYPE_PRISM_HEADER &&
        pfh.linktype != LINKTYPE_RADIOTAP_HDR )
    {
        printf( "\"%s\" isn't a regular 802.11 "
                "(wireless) capture.\n", fn_in );
        return( 1 );
    }

    /* create the output ivs file */

    if (fn_out != NULL) {
	    printf( "Creating %s\n", fn_out );

	    if( ( f_out = fopen( fn_out, "wb+" ) ) == NULL )
	    {
	        perror( "fopen failed" );
	        return( 1 );
	    }

    	    fwrite2( &pfh, 1, sizeof(pfh), f_out );
    } else f_out=NULL;


    tt = time( NULL ) - 1;

	a[T] = 0;

    while( 1 )
    {
        if( time( NULL ) - tt > 0 )
        {
            printf( "\33[KRead %d packets...\r", b[T] );
            fflush( stdout );
            tt = time( NULL );
        }

        /* read one packet */

        n = sizeof( pkh );

        if( fread( &pkh, 1, n, f_in ) != (size_t) n )
            break;

        if( pfh.magic == TCPDUMP_CIGAM )
            SWAP32( pkh.caplen );

        plen = pkh.caplen;

        if( plen <= 0 || plen > 65535 )
        {
            printf( "Corrupted file? Invalid packet length %d.\n", plen );
            return( 1 );
        }

        if( fread( buffer, 1, plen, f_in ) != (size_t) plen )
            break;

	b[T]++;
	bs[T]+=pkh.caplen;

        h80211 = buffer;

        /* remove any prism/radiotap header */

        if( pfh.linktype == LINKTYPE_PRISM_HEADER )
        {
            if( h80211[7] == 0x40 )
                n = 64;
            else
            {
                n = *(int *)( h80211 + 4 );

                if( pfh.magic == TCPDUMP_CIGAM )
                    SWAP32( n );
            }

            if( n < 8 || n >= (int) plen )
		    goto ok;

            h80211 += n; plen -= n;
        }

        if( pfh.linktype == LINKTYPE_RADIOTAP_HDR )
        {
            n = *(unsigned short *)( h80211 + 2 );

            if( n <= 0 || n >= (int) plen )
                goto ok;

            h80211 += n; plen -= n;
        }

        /* check if WEP data, min. length & bssid */

	#define COMP(VAL, MASC)  ( (h80211[0]&MASC) == (VAL&MASC) )
	#define SUMA(VAR)  { b[VAR]++ ; bs[VAR] += pkh.caplen; pt=&a[VAR] ; ps=&as[VAR]; }

	if      (COMP(0x08,0x0C)) SUMA(D)
	
	else if (COMP(0x80,0xFC)) SUMA(MB)
	else if (COMP(0x00,0xCC)) SUMA(MAS)
	else if (COMP(0xA0,0xFC)) SUMA(MDS)
	else if (COMP(0xB0,0xFC)) SUMA(MAU)
	else if (COMP(0xC0,0xFC)) SUMA(MDU)
	else if (COMP(0x40,0xFC)) SUMA(MQ)
	else if (COMP(0x50,0xFC)) SUMA(MP)
	else if (COMP(0x00,0x0C)) SUMA(MO)
	
	else if (COMP(0xB4,0xFC)) SUMA(CR)
	else if (COMP(0xC4,0xFC)) SUMA(CC)
	else if (COMP(0xD4,0xFC)) SUMA(CA)
	else if (COMP(0x04,0x0C)) SUMA(CO)

	else SUMA(O)

	// check for uniqueness of frame
	if (f_out != NULL)
	    if (unique(h80211, pkh.caplen) != -1)
		continue;

	(*pt)++;
	(*ps) += pkh.caplen;

     ok:
        ++a[T];
	as[T] += pkh.caplen;
	if (f_out != NULL) {
		fwrite2(&pkh, 1, sizeof(pkh), f_out);
		fwrite2(buffer, 1, pkh.caplen, f_out);
	}
    }
    fclose( f_in );
    if (f_out != NULL)
    	fclose( f_out );
	
    printf( "\33[2KRead %d packets.\n", b[T] );
	
    if (f_out != NULL) {
	    if ( a[T] > 0 ) 
		printf( "Written %d packets.\n", a[T]);
	    else {
		remove ( fn_out );	
		puts("No packets written");
	    }
    }

    if (verb>0) {
	    
     puts("");

     printf("Stats:\n"
           " Data       (08&C): % 8d (% 9d bytes)\n"
           " Management (00&C): % 8d (% 9d bytes)\n",
	   b[D],bs[D],
	   b[MB]+b[MAS]+b[MDS]+b[MAU]+b[MDU]+b[MQ]+b[MP]+b[MO],
	   bs[MB]+bs[MAS]+bs[MDS]+bs[MAU]+bs[MDU]+bs[MQ]+bs[MP]+bs[MO]);

     if (verb>1) printf(
           "  Beacon      (80): % 8d (% 9d bytes)\n"
           "  (Re)Assoc(00-30): % 8d (% 9d bytes)\n"
           "  Disassoc    (A0): % 8d (% 9d bytes)\n"
           "  Auth        (B0): % 8d (% 9d bytes)\n"
           "  Deauth      (C0): % 8d (% 9d bytes)\n"
           "  ProbeReq    (40): % 8d (% 9d bytes)\n"
           "  ProbeResp   (50): % 8d (% 9d bytes)\n"
           "  Other       (..): % 8d (% 9d bytes)\n",
	   b[MB],bs[MB], b[MAS],bs[MAS], b[MDS],bs[MDS], b[MAU],bs[MAU], b[MDU],bs[MDU], b[MQ],bs[MQ], b[MP],bs[MP], b[MO],bs[MO]);

     printf(
           " Control    (04&C): % 8d (% 9d bytes)\n",
	   b[CR]+b[CC]+b[CA]+b[CO], bs[CR]+bs[CC]+bs[CA]+bs[CO]);

     if (verb>1) printf(
           "  RTS         (B4): % 8d (% 9d bytes)\n"
           "  CTS         (C4): % 8d (% 9d bytes)\n"
           "  Ack         (D4): % 8d (% 9d bytes)\n"
           "  Other       (..): % 8d (% 9d bytes)\n",
	   b[CR], bs[CR], b[CC], bs[CC], b[CA], bs[CA], b[CO], bs[CO]);

     printf(
           " Other        (..): % 8d (% 9d bytes)\n"
           " Total            : % 8d (% 9d bytes)\n",
	   b[O], bs[O],
	   b[T], bs[T]);
    
     puts("");
     
     if (f_out != NULL) {

      printf("Stats after optimization:\n"
           " Data       (08&C): % 8d (% 9d bytes)\n"
           " Management (00&C): % 8d (% 9d bytes)\n",
	   a[D],as[D],
	   a[MB]+a[MAS]+a[MDS]+a[MAU]+a[MDU]+a[MQ]+a[MP]+a[MO],
	   as[MB]+as[MAS]+as[MDS]+as[MAU]+as[MDU]+as[MQ]+as[MP]+as[MO]);
	   
      if (verb>1) printf(
           "  Beacon      (80): % 8d (% 9d bytes)\n"
           "  (Re)Assoc(00-30): % 8d (% 9d bytes)\n"
           "  Disassoc    (A0): % 8d (% 9d bytes)\n"
           "  Auth        (B0): % 8d (% 9d bytes)\n"
           "  Deauth      (C0): % 8d (% 9d bytes)\n"
           "  ProbeReq    (40): % 8d (% 9d bytes)\n"
           "  ProbeResp   (50): % 8d (% 9d bytes)\n"
           "  Other       (..): % 8d (% 9d bytes)\n",
	   a[MB],as[MB], a[MAS],as[MAS], a[MDS],as[MDS], a[MAU],as[MAU], a[MDU],as[MDU], a[MQ],as[MQ], a[MP],as[MP], a[MO],as[MO]);
	   
      printf(
           " Control    (04&C): % 8d (% 9d bytes)\n",
	   a[CR]+a[CC]+a[CA]+a[CO], as[CR]+as[CC]+as[CA]+as[CO]);
	   
      if (verb>1) printf(
           "  RTS         (B4): % 8d (% 9d bytes)\n"
           "  CTS         (C4): % 8d (% 9d bytes)\n"
           "  Ack         (D4): % 8d (% 9d bytes)\n"
           "  Other       (..): % 8d (% 9d bytes)\n",
	   a[CR], as[CR], a[CC], as[CC], a[CA], as[CA], a[CO], as[CO]);

      printf(
           " Other        (..): % 8d (% 9d bytes)\n"
           " Total            : % 8d (% 9d bytes)\n",
	   a[O], as[O],
	   a[T], as[T]);

      puts("");

      printf("Optimized size: %0.2f%%\n", (100.0*as[T])/bs[T]);
     }
     
    }
	   
    return( 0 );
}
