#include <stdio.h>
#include <stdlib.h>
#include "util.h"
#include "slitheen.h"

static fragment_table *frag_table;

/**
 * This function takes as argument a packet fragment. It checks
 * to see if there are other pending packets. If so, it
 * looks for matching identification numbers and attempts
 * to reconstruct a packet. If the packet reconstruction was
 * complete, returns a 1 and the input packet will point to
 * the completely reassembled ip packet. If the reconstruction
 * requires more fragments, it returns a 0. Returns -1 on error
 */
int report_fragment(const u_char *data, u_char *packet, u_int length){
	int i, ret;
	unsigned char *partial_packet;
	const struct ip_header *ip_hdr;

	ip_hdr = (struct ip_header*) data;
	printf("fragment reported.\n");
	/* First check to see if this belongs to a partial packet */
	for (i=0; i<frag_table->len; i++){
		printf("checking fragments.\n");
		fragment *candidate = frag_table->table+i;
		if (candidate->id == ip_hdr->id){
			/*update fragment in table with new data */
			printf("Matches existing packet.\n");
			ret = update_fragment(i, data, length);
			if(ret == 1){
				/* Complete, copy into packet ptr */
				printf("Successfully reconstructed packet!\n");
				for(i=0; i<ip_hdr->len; i++){
					packet[i] = candidate->data[i];
				}
			}
			return ret;
		}
	}

	printf("Creating new partial packet.\n");
	/* New packet, create a new table entry */
	fragment new_fragment;
	new_fragment.id = ip_hdr->id;
	new_fragment.ip_len = ip_hdr->len;
	partial_packet = malloc(65535);
	new_fragment.data = partial_packet;
	if(add_fragment(new_fragment))
		return -1;

	/* Now add data into fragment */
	printf("New partial packet created.\n");
	return update_fragment(frag_table->len -1, data, length);

}

/**
 * Initializes the table of all partialy reconstructed
 * packets. Sets the maximum length to defined MAX_FRAGMENTS
 */
int init_fragment_table(void) {

	frag_table = malloc(sizeof(fragment_table));
	
	frag_table->table = (fragment *) malloc(sizeof(fragment)*MAX_FRAGMENTS);
	if( frag_table->table == NULL){
		fprintf(stderr, "malloc failed.\n");
		return 1;
	}
	frag_table->len = 0;
	frag_table->max_len = MAX_FRAGMENTS;

	return 0;
}

/**
 * Adds a new fragment to the table of partial packets.
 */
int add_fragment(fragment new_fragment){
	fragment *ptr;

	if(frag_table->len == frag_table->max_len){
		//grow_table();
		return(1);
	}

	ptr = frag_table->table + frag_table->len;
	*ptr = new_fragment;

	frag_table->len ++;

	return 0;
}

/**
 * Updates an existing partial packet by copying data into
 * packet. returns a 1 if reconstruction is complete, 0 if not,
 * and -1 on error.
 */
int update_fragment(int index, const u_char *data, u_int length){
	int i, data_length;
	unsigned char *src, *dest;
	const struct ip_header *ip_hdr;
	u_short offset;
	fragment *ptr;
	u_int size_ip_hdr;

	ip_hdr = (struct ip_header*) data;
	offset = ip_hdr->flagsoff & 0x1fff;
	size_ip_hdr = IP_HEADER_LEN(ip_hdr);

	ptr = (fragment *) frag_table->table+index;
	dest = ptr->data + offset;
	src = (unsigned char *) data + size_ip_hdr;

	/* Copy all but the last 4 bytes to partial packet
	 * these belong to the ethernet footer */
	data_length = length - size_ip_hdr - 4;

	for(i=0; i< data_length; i++){
		dest[i] = src[i];
	}

	if(i== ip_hdr->len){
		/* return a 1! */
		printf("Filled fragment %d from %d to %d.\n", ip_hdr->id, offset, offset + data_length);
		return 1;
	}

	printf("Filled fragment %d from %d to %d.\n", ip_hdr->id, offset, offset+data_length);
	printf("Still need %d to %d.\n", offset + data_length, ip_hdr->len);

	return 0;
}

void remove_fragment(int index){
	/* free ptr to data */

	/* remove reference in table to fragment */
	int i;
	fragment *ptr;

	ptr = frag_table->table + index;
	
	for(i=0; i< frag_table->len - index-1; i++){
		*ptr = *(ptr + 1);
		ptr += i;
	}
	frag_table->len --;

}