/*
 * spp_defrag v1.0b17 -  24 July 2000
 * 
 * Purpose: IP defragmentation preprocessor
 * Author:  Dragos Ruiu (dr@dursec.com/dr@v-wave.com)
 * Acknowledgements:    Code skeleton from diphen@agitation.net
 *			There's a good paper about splay trees
 *			by the developer and they rock.
 *			Marty Roesch(roesch@hiverworld.com 
 *			and Ron Gula(rgula@networksecuritywizards.com)
 *			helped and did something few in the security
 *			community do, shared knowledge.  Thanks.
 *
 * notes:
 *		This defragger implementation differs from the usual
 *		hash table and linked list technique in that it uses
 *		self caching splay trees. My hypothesis is that this
 *		provides substantial performance improvements.  I
 *		hope this module will be able to demonstrate or debunk
 *		this hypothesis. 
 * 
 * Splay Tree:
 * `A self-organizing data structure which uses rotations to move an
 *  accessed key to the root. This leaves recently accessed nodes near
 *  the top of the tree, making them very quickly searchable (Skiena 1997, p. 177). 
 * 
 * All splay tree operations run in O(log n) time _on_average_, where n is the
 * number of items in the tree, assuming you start with an empty tree.  Any single
 * operation can take Theta(n) time in the worst-case, but operations slower than
 * O(log n) time happen rarely enough that they don't affect the average.
 * 
 * Although 2-3-4 trees make a stronger guarantee (_every_ operation on a 2-3-4
 * tree takes O(log n) time), splay trees have several advantages.  Splay trees
 * are simpler and easier to program, and they can take advantage of circumstances
 * in which lots of find operations occur on a small number of items.  Because of
 * their simplicity, splay tree insertions and deletions are typically faster in
 * practice (sometimes by a constant factor, sometimes asymptotically).  Find
 * operations can be faster or slower, depending on circumstances.  Splay trees
 * really excel in applications where a fraction of the items are the targets of
 * most of the find operations, because they're designed to give especially fast
 * access to items that have been accessed recently.
 *
 * a good reference on splay trees is:
 * http://www.cs.umbc.edu/courses/undergraduate/341/fall98/frey/ClassNotes/Class17/splay.html
 * 
 * References 
 * Skiena, S. S. The Algorithm Design Manual. New York: Springer-Verlag, pp. 177 and 179, 1997. 
 * Sleator, D. and Tarjan, R. "Self-Adjusting Binary Search Trees." J. ACM 32, 652-686, 1985. 
 * Tarjan, R. Data Structures and Network Algorithms. Philadelphia, PA: SIAM Press, 1983. 
 * Wood, D. Data Structures, Algorithms, and Performance. Reading, MA: Addison-Wesley, 1993. 
 * 
 */


#include <pcap.h>
#include "decode.h"

/* Warning these are only valid for IPv4
   also note that these are still 
   in raw network order (see htons)  
   and these may fail on 64 bit machines 
   but I'm not worried about Unobtanium yet
*/
#define ID(x)		*((short*)((char*)x->iph+4))
#define PROTO(x)	*((char*)x->iph+9)
#define SADDR(x)        (x->iph->ip_src.s_addr)
#define DADDR(x)        (x->iph->ip_dst.s_addr)
#define	DATA(x)		((char*)x->iph+20)

/* Uh-oh hope this wierdness is right :-) */
#define FOFF(x)		(int)((x->frag_offset)<<3)
#define DF(x)		(x->df)
#define MF(x)		(x->mf)


/* fragment ID structure  */
typedef Packet *frag;

typedef struct tree_node Tree;
struct tree_node {
    Tree * left, * right;
    frag key;
    int size;   /* maintained to be the number of nodes rooted here */
};

Tree *froot;


/*  These next declarations are for the fragment timeout and 
    and cleanup/sweeping process... time math routines are from
    an obscure piece of old code for a defunct video camera product
*/

#define FRAGTIMEOUTSEC	10	/* 10 seconds let's play safe for now */
#define FRAGTIMEOUTUSEC	0	/* 0 micro seconds		      */
#define FASTSWEEPLIM	16000000	/* memory use threhold for fast sweep */

long int fragmemuse;
int fragsweep;  /* the packet timeout / garbage collection stuff  */

typedef struct _timestruct
{
        u_int32_t tv_sec;
        u_int32_t tv_usec;
} time_struct;

time_struct fragtimeout;
 


/******Timestamp Routines******/

#define TIME_LT(x,y) (x tv_sec<y tv_sec||(x tv_sec==y tv_sec&&x tv_usec<y tv_usec))

void addtime(time_struct *op1, time_struct *op2, time_struct *result)
{
    result->tv_usec = op1->tv_usec+op2->tv_usec;
    if (result->tv_usec > 999999)
    {
        result->tv_usec -= 1000000;
        op1->tv_sec++;
     }
     result->tv_sec = op1->tv_sec+op2->tv_sec;
}



/******Splay Tree Stuff******/

/* Function: fragcompare(i,j)				    */
/* This is the splay tree comparison.                       */
/* Returns 1 if i>j; 0 if i=j; -1 if i<j;                   */

int fragcompare(i,j)
frag i,j;
{
	if(!j)
	{
		if(!i)
			return(0);
		else
			return(1);
	} else {
		 if(!i)
			return(-1);
	}
             if( SADDR(i) > SADDR(j) )
		 { return (1); }
	else if( SADDR(i) < SADDR(j) )
		 { return (-1); }
	else if( DADDR(i) > DADDR(j) )
		 { return (1); }
	else if( DADDR(i) < DADDR(j) )
		 { return (-1); }
	else if( PROTO(i) > PROTO(j) )
		 { return (1); }
	else if( PROTO(i) < PROTO(j) )
		 { return (-1); }
	else if( ID(i) > ID(j) )
		 { return (1); }
	else if( ID(i) < ID(j) )
		 { return (-1); }
	else if( FOFF(i) > FOFF(j) )
		 { return (1); }
	else if( FOFF(i) < FOFF(j) )
		 { return (-1); }
	else if( i->dsize > j->dsize )
		 { return (1); }
	else if( i->dsize < j->dsize )
		 { return (-1); }
	return (0);
}
 
/* This macro returns the size of a node.  Unlike "x->size",     */
/* it works even if x=NULL.  The test could be avoided by using  */
/* a special version of NULL which was a real node with size 0.  */

#define node_size(x) ((!x) ? 0 : ((x)->size))
 
/* Function: fragsplay(i, t)					*/
/* Splay using the key i (which may or may not be in the tree.) */
/* The starting root is t, size fields are maintained		*/

Tree *fragsplay(frag i, Tree *t) 
{
    Tree N, *l, *r, *y;
    int comp, root_size, l_size, r_size;
    if(!t) return t;
    N.left = N.right = NULL;
    l = r = &N;
    root_size = node_size(t);
    l_size = r_size = 0;
 
    for (;;) {
        comp = fragcompare(i, t->key);
        if(comp < 0) {
            if(!t->left) break;
            if(fragcompare(i, t->left->key) < 0) {
                y = t->left;                           /* rotate right */
                t->left = y->right;
                y->right = t;
                t->size = node_size(t->left) + node_size(t->right) + 1;
                t = y;
                if(!t->left) break;
            }
            r->left = t;                               /* link right */
            r = t;
            t = t->left;
            r_size += 1+node_size(r->right);
        } else if (comp > 0) {
            if(!t->right) break;
            if(fragcompare(i, t->right->key) > 0) {
                y = t->right;                          /* rotate left */
                t->right = y->left;
                y->left = t;
		t->size = node_size(t->left) + node_size(t->right) + 1;
                t = y;
                if(!t->right) break;
            }
            l->right = t;                              /* link left */
            l = t;
            t = t->right;
            l_size += 1+node_size(l->left);
        } else {
            break;
        }
    }
    l_size += node_size(t->left);  /* Now l_size and r_size are the sizes of */
    r_size += node_size(t->right); /* the left and right trees we just built.*/
    t->size = l_size + r_size + 1;

    l->right = r->left = NULL;

    /* The following two loops correct the size fields of the right path  */
    /* from the left child of the root and the right path from the left   */
    /* child of the root.                                                 */
    for(y = N.right; y; y = y->right) {
        y->size = l_size;
        l_size -= 1+node_size(y->left);
    }
    for(y = N.left; y; y = y->left) {
        y->size = r_size;
        r_size -= 1+node_size(y->right);
    }
 
    l->right = t->left;                                /* assemble */
    r->left = t->right;
    t->left = N.right;
    t->right = N.left;

    return t;
}

/* Function: Tree * fraginsert(frag i, Tree * t) 		*/
/* Insert frag i into the tree t, if it is not already there. 	*/
/* Return a pointer to the resulting tree.                   	*/

Tree *fraginsert(frag i, Tree * t) {
    Tree * new;
    if(t) {
	t = fragsplay(i,t);
	if (fragcompare(i, t->key)==0) {
	    return t;  /* it's already there */
	}
    }
    new = (Tree *) malloc (sizeof (Tree));
    if(!new) { ErrorMessage("Ran out of space\n"); return(t);}
    if(!t) {
	new->left = new->right = NULL;
    } else if(fragcompare(i, t->key) < 0) {
	new->left = t->left;
	new->right = t;
	t->left = NULL;
	t->size = 1+node_size(t->right);
    } else {
	new->right = t->right;
	new->left = t;
	t->right = NULL;
	t->size = 1+node_size(t->left);
    }
    new->key = i;
    new->size = 1 + node_size(new->left) + node_size(new->right);
    return new;
}

/* Function: Tree * fragdelete(frag i, Tree *t) 	*/
/* Deletes i from the tree if it's there.               */
/* Return a pointer to the resulting tree.              */

Tree *fragdelete(frag i, Tree *t) {
    Tree * x;
    int tsize;

    if (!t) return NULL;
    tsize = t->size;
    t = fragsplay(i,t);
    if(fragcompare(i, t->key) == 0) {               /* found it */
	if(!t->left) {
	    x = t->right;
	} else {
	    x = fragsplay(i, t->left);
	    x->right = t->right;
	}
	
	free(t);
	if(x) {
	    x->size = tsize-1;
	}
	return x;
    } else {
	return t;                         /* It wasn't there */
    }
}

/* Function: Tree *fragfind_rank(int r, Tree *t)		   */ 
/* Returns a pointer to the node in the tree with the given rank.  */
/* Returns NULL if there is no such node.                          */
/* Does not change the tree.  To guarantee logarithmic behavior,   */
/* the node found here should be splayed to the root.              */

Tree *fragfind_rank(int r, Tree *t) {
    int lsize;
    if ((r < 0) || (r >= node_size(t))) return NULL;
    for (; t != 0 ;) {
	lsize = node_size(t->left);
	if (r < lsize) {
	    t = t->left;
	} else if (r > lsize) {
	    r = r - lsize -1;
	    t = t->right;
	} else {
	    return t;
	}
    }
    return NULL;
}

/* Function: fragget_rank(t)					   */
/* returns the rank of a tree element				   */

int fragget_rank(Tree *t) {
	if(!t)
		return(0);
	return(node_size(t->left)+1);
}




/******Snort Stuff******/

/*
 * Function: SetupDefrag()
 * Purpose:
 * Registers the preprocessor keyword and initialization function
 * into the preprocessor list.  This is the function that gets called from
 * InitPreprocessors() in plugbase.c.
 * Arguments: None.
 * Returns: void function
 */
void SetupDefrag()
{
	RegisterPreprocessor("defrag", DefragInit);


}

/*
 * Function: DefragInit(u_char *)
 * Purpose:
 * Calls the argument parsing function, performs final setup on data
 * structs, links the preproc function into the function list.
 * Arguments: args => ptr to argument string
 * Returns: void function
 */
void DefragInit(u_char *args)
{
	AddFuncToPreprocList(PreprocDefrag);

	froot = NULL;  /* initialize empty fragment tree */
	fragmemuse = 0;	/* No memory yet */
	fragtimeout.tv_sec = FRAGTIMEOUTSEC;
	fragtimeout.tv_usec = FRAGTIMEOUTUSEC;
	fragsweep = 0;
}



/******Fragmentation Stuff******/

/* Function: fragaddrmatch(i,j) 				*/
/* returns true if the fragments belong to the same reassembly	*/

int fragaddrmatch(i,j)
frag i,j;
/* This is the comparison.                                   */
/* Returns 1 if i j have matching addresses else 0           */
{
 
	if(( SADDR(i) == SADDR(j) )
		&& ( DADDR(i) == DADDR(j) )
		&& ( ID(i) == ID(j) )
		&& ( PROTO(i) == PROTO(j) ))

		return (1);
	return (0);
}
 
/*
 * Function: Packet *ReassembleIP(frag *froot)
 * Purpose: Generate a Packet * and pass it to ProcessPacket, then deallocate packet
 * Arguments: froot - root of the tree containing the last packet of the frame
 * Returns: tree with reassembled fragments deleted
 */
Tree *ReassembleIP(Tree *froot)
{
	Packet *p;
	int writecount = 0;
	char *np;
        u_char *tmp;
	u_int psize;
	int fragrank;
	int overhead, moreoverhead;

	psize = (froot->key)->dsize + FOFF(froot->key); /* last frag is at top of tree */
	/* hopefully moreoverhead here finds out about all the wierd MAC sizes like FDDI and ATM */
	moreoverhead = (char*)froot->key->iph - (char*)froot->key->pkt;
	overhead = (char*)froot->key->pkt - (char*)froot->key->pkth; /* +4 mystery points */
	/* Linux fix code */
	if(overhead < 1 || overhead > sizeof( struct pcap_pkthdr) + 4)
		overhead = sizeof(struct pcap_pkthdr) + 4;
 
	p = (Packet *)malloc(sizeof(Packet));
	if (!p)
		return NULL;
        memcpy(p, froot->key, sizeof(Packet));

	/* SEMI BOGUS - setting up ethernet time header - required by one of the
	 * decode routines - we set the timestamp equal to froot (last fragment)
	 * note also thet our bogus reassembledgrams may not have the bpf header
	 * contiguous and datagram access through this pointer is suspect. Been warned.
	 */

	p->pkth = (struct pcap_pkthdr *)malloc(psize + overhead + moreoverhead + sizeof(IPHdr) + 20);
	if(!p->pkth)
	{
		free(p);
		return NULL;
	}

	p->iph = (IPHdr *)((u_char*)p->pkth + overhead + moreoverhead);
	p->pkt = (u_char*)p->iph - moreoverhead;

	/*
	 * Now copy the header and fragments into the newly-allocated buffer,
	 * reconstructing the packet chunk by chunk. Current algorithm is first 
	 * write to an area wins, play with splay order to change.
	 * we start from the last fragment and work back....
	 */

	/* fill in packet header bpf header first*/
	memcpy(p->pkth, froot->key->pkth, overhead); 

	/* then the mac junk, split into two copies for Linux non-contiguous headers */
	memcpy((char*)p->pkth + overhead, froot->key->pkt, moreoverhead); 

	/*
	 * then the IP header just to be paranoid for debugging because in the
         * real world we would do these in one copy
 	 */
        tmp = (u_char *) froot->key->iph;
	memcpy(p->iph, tmp, sizeof(IPHdr)); 
  
	p->pkth->caplen = psize + overhead + moreoverhead + sizeof(IPHdr);
	p->pkth->len = p->pkth->caplen;

	/*
	 * Clear the more fragments bit, and set the length in the ip header
	 * (in network byte order).
	 */
	p->iph->ip_len = htons(psize);
	p->iph->ip_off = 0;
        p->frag_flag = 0;

	fragrank = fragget_rank(froot);
	while(fragrank > 0 && froot && fragaddrmatch(p,froot->key))
	{
		if(FOFF(froot->key) + froot->key->dsize <= psize)
		{

			memcpy((char*)DATA(p)+FOFF(froot->key), DATA(froot->key), froot->key->dsize);
			writecount += froot->key->dsize;

		} else
			(*AlertFunc)(p, "Fragmentation Overflow Attack");

		if(froot)
		{
			fragmemuse -= froot->key->pkth->caplen + sizeof(Packet);
			np = (char *)froot->key->pkth;  /* address for free later */
			fragrank = fragget_rank(froot);
			froot = fragdelete(froot->key, froot);
			if(np)
				free(np);
		}
	}

	/* and now some jolt2proofing */
	if(psize > 8192)  /* only kick in for monstergrams */
	{
		if(writecount > psize/2)	/*packets have to be half full for us to look at them */
		{	
			if(p)
				ProcessPacket(NULL, p->pkth, p->pkt);
		} else {
			(*AlertFunc)(p, "Mostly Empty Fragmented Packet Discarded!");
		}
	} else {
		if(p)
			ProcessPacket(NULL, p->pkth, p->pkt);
	}
	
	free(p->pkth);	/* set the frankensteingram free.... */
	free(p);
	return(froot);
}

/*
 * Function: PreprocDefrag(Packet *)
 * Purpose:
 * Driver function for the IP defragmentation preprocessor.
 * Arguments: p => pointer to the current packet data struct
 * Returns: void function
 */
void PreprocDefrag(Packet *p)
{
        Packet *new;
	time_struct timecheck;
	Tree *found;
	frag freetemp;
	int overhead;
        int cap;
        u_char *tmp;

	if(!p || !p->pkth || !p->pkt)
	{
		ErrorMessage("%s\n","Garbage Packet with Null Pointer discarded!");
		return;
	}
	if(!p->iph)
		return;	/* return quietly for arps and such with no ip headers */
	if(p->frag_flag)
	{  /* heads up, live fragments inbound  */
		overhead = (char*)p->pkt - (char*)p->pkth; /* +4 mystery points */
		/* Linux fix code */
		if( overhead < 1 || overhead > sizeof(struct pcap_pkthdr) + 4)
			overhead = sizeof(struct pcap_pkthdr) + 4;

		new = malloc(sizeof(Packet));
		if(!new)
		{
			ErrorMessage("Cannot allocate fragment buffer(usage 0x%X)\n",fragmemuse);
			return;
		}
		fragmemuse += sizeof(Packet);

		memcpy(new, p, sizeof(Packet));
		cap = p->pkth->caplen + overhead;
		tmp = calloc(cap + 20, sizeof(char));
                new->pkth = (struct pcap_pkthdr *) tmp;
		if(!new->pkth)
		{
			ErrorMessage("Cannot allocate fragment buffer(usage %X)\n",fragmemuse);
			return;
		}

		new->pkt = (u_char*)new->pkth + overhead;
		new->iph = (IPHdr*)((u_char*)new->pkt + ((u_char*)p->iph - (u_char*)p->pkt));
  
		fragmemuse += p->pkth->caplen + overhead;
		/* we do this with two memcopies to cope with Linux non-contiguous bpf headers */
		memcpy(new->pkth, p->pkth, overhead);
		memcpy(new->pkt, p->pkt, p->pkth->caplen);
		froot = fraginsert(new, froot); 

		/* now check if we have to reassemble anything... */
		if(!MF(p))
			froot = ReassembleIP(froot);

		/* OK now check for a fragment timeout - sweeping the clutter away :-) */
		if(froot)
		{
			if(++fragsweep > node_size(froot)+1)
				fragsweep = 1;
			found = fragfind_rank(fragsweep,froot);
			if(found)
			{
				froot = fragsplay(found->key, froot);
	                        addtime((time_struct *)&(froot->key->pkth->ts), &fragtimeout, &timecheck);
                                if(TIME_LT(timecheck. , p->pkth->ts.))
                                {
					fragmemuse -= froot->key->pkth->caplen + sizeof(Packet);
					freetemp = froot->key;
					froot = fragdelete(froot->key, froot);
					free(freetemp->pkth);  /* free packet copy */
					free(freetemp);
					fragsweep--;
				}
			}

			/* and now do the whole thing again if we are being a memory hog */
			if(froot && fragmemuse > FASTSWEEPLIM)
			{
				if(++fragsweep > node_size(froot))
					fragsweep = 1;
				found = fragfind_rank(fragsweep,froot);
				if(found)
				{
					froot = fragsplay(found->key, froot);
					addtime((time_struct *)&(froot->key->pkth->ts), &fragtimeout, &timecheck);
					if(TIME_LT(timecheck. , p->pkth->ts. ))
					{
						fragmemuse -= froot->key->pkth->caplen + sizeof(Packet);
						freetemp = froot->key;
						froot = fragdelete(froot->key, froot);
						free(freetemp->pkth);  /* free packet copy */
						free(freetemp);
						fragsweep--;
					}
				}
			}
		}
	}
	return;
}


