/*

Copyright (C) 2000 - 2003 Christian Kreibich <christian@whoop.org>.

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to
deal in the Software without restriction, including without limitation the
rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
sell copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies of the Software and its documentation and acknowledgment shall be
given in the documentation and software packages that this Software was
used.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

*/
#ifdef HAVE_CONFIG_H
#  include <config.h>
#endif

#include <string.h>
#include <unistd.h>
#include <stdlib.h>
#include <stdio.h>
#include <sys/types.h>
#ifdef LINUX
#define __FAVOR_BSD
#endif
#include <netinet/in_systm.h>
#include <netinet/in.h>
#include <netinet/ip.h>

#include <nd.h>
#include <nd_gui.h>
#include <nd_dialog.h>
#include <nd_trace.h>
#include <nd_trace_registry.h>
#include <nd_ip.h>
#include <nd_ipfrag.h>
#include <interface.h>
#include <callbacks.h>
#include <support.h>

void      
nd_ip_frag_show_dialog(LND_Packet *packet)
{
  ND_FragData     *data;
  struct ip       *iphdr;
  char             s[128];
  GtkWidget       *sb1, *sb2;
  GtkLabel        *label;
  GtkObject       *adj;

  if (!packet)
    return;

  iphdr = (struct ip *) libnd_packet_get_data(packet, nd_ip_get(), 0);
  if (!iphdr)
    return;

  data = g_new0(ND_FragData, 1);
  D_ASSERT_PTR(data);

  /* The initial setting is to split the packet
     approximately in half: */
  data->size_orig = ND_IP_PAYLOAD_LENGTH(iphdr);
  data->size1  = data->size_orig/2 - ((data->size_orig/2) % 8);
  data->size2  = data->size_orig - data->size1;
  data->packet = packet;
  data->dialog = create_ip_frag_dialog();

  D_ASSERT_PTR(data->dialog);
  ND_GTK_GET(sb1, data->dialog, "ip_frag1_spinbutton");
  ND_GTK_GET(sb2, data->dialog, "ip_frag2_spinbutton");

  gtk_object_set_data_full(GTK_OBJECT(data->dialog), "data", data,
			   (GtkDestroyNotify) g_free);

  /* Disable callbacks for a moment in order to not trigger the
     callbacks while we're still only initializing settings */
  gtk_signal_handler_block_by_func(GTK_OBJECT(sb1), on_ip_frag1_spinbutton_changed, NULL);
  gtk_signal_handler_block_by_func(GTK_OBJECT(sb2), on_ip_frag2_spinbutton_changed, NULL);

  adj = gtk_adjustment_new(0, 0, data->size_orig, 8, 8, 1);
  gtk_spin_button_set_adjustment(GTK_SPIN_BUTTON(sb1), GTK_ADJUSTMENT(adj));
  gtk_spin_button_set_value(GTK_SPIN_BUTTON(sb1), (gfloat) data->size1);

  adj = gtk_adjustment_new(0, 0, data->size_orig, 8, 8, 1);
  gtk_spin_button_set_adjustment(GTK_SPIN_BUTTON(sb2), GTK_ADJUSTMENT(adj));
  gtk_spin_button_set_value(GTK_SPIN_BUTTON(sb2), (gfloat) data->size2);

  /* Re-enable callbacks ... */
  gtk_signal_handler_unblock_by_func(GTK_OBJECT(sb1), on_ip_frag1_spinbutton_changed, NULL);
  gtk_signal_handler_unblock_by_func(GTK_OBJECT(sb2), on_ip_frag2_spinbutton_changed, NULL);

  /* Get the offsets of our fragments from the fragment offset
     of the packet (maybe it's a fragment already) */
  data->offset1 = (ntohs(iphdr->ip_off) & IP_OFFMASK) * 8;
  data->offset2 = data->offset1 + data->size1;

  D_ASSERT_PTR(data->dialog);
  ND_GTK_GET(label, data->dialog, "ip_frag1_offset_label");
  g_snprintf(s, 128, "%u", data->offset1);
  gtk_label_set_text(label, s);

  ND_GTK_GET(label, data->dialog, "ip_frag2_offset_label");
  g_snprintf(s, 128, "%u", data->offset2);
  gtk_label_set_text(label, s);

  gtk_widget_show(data->dialog);
}


void      
nd_ip_frag_adjust(ND_FragData *data, int fragment_size, int fragment_number)
{
  char             s[128];
  GtkWidget       *sb1, *sb2;
  GtkLabel        *label;
  
  if (!data)
    return;

  D_ASSERT_PTR(data->dialog);

  switch (fragment_number)
    {
    case 0:
      data->size1   = fragment_size;
      data->size2   = data->size_orig - data->size1;
      break;
    case 1:
      data->size2   = fragment_size;
      data->size1   = data->size_orig - data->size2;
      break;
    default:
      D(("Huh? We only split into two fragments!"));
      return;
    }
  
  data->offset2 = data->offset1 + data->size1;

  ND_GTK_GET(sb1, data->dialog, "ip_frag1_spinbutton");
  ND_GTK_GET(sb2, data->dialog, "ip_frag2_spinbutton");

  /* Disable callbacks for a moment to not trigger callbacks */
  gtk_signal_handler_block_by_func(GTK_OBJECT(sb1), on_ip_frag1_spinbutton_changed, NULL);
  gtk_signal_handler_block_by_func(GTK_OBJECT(sb2), on_ip_frag2_spinbutton_changed, NULL);

  if (fragment_number == 0)
    gtk_spin_button_set_value(GTK_SPIN_BUTTON(sb2), (gfloat) data->size2);
  else
    gtk_spin_button_set_value(GTK_SPIN_BUTTON(sb1), (gfloat) data->size1);

  /* Re-enable callbacks ... */
  gtk_signal_handler_unblock_by_func(GTK_OBJECT(sb1), on_ip_frag1_spinbutton_changed, NULL);
  gtk_signal_handler_unblock_by_func(GTK_OBJECT(sb2), on_ip_frag2_spinbutton_changed, NULL);

  ND_GTK_GET(label, data->dialog, "ip_frag2_offset_label");
  g_snprintf(s, 128, "%u", data->offset2);
  gtk_label_set_text(label, s);
}


void
nd_ip_frag_cancel(ND_FragData *data)
{
  if (!data || !data->dialog)
    return;

  gtk_widget_destroy(data->dialog);
}


void      
nd_ip_frag_apply(ND_FragData *data)
{
  if (!data || !data->packet)
    return;

  nd_ip_frag_fragment(data->packet, data->size1, data->size2);
  gtk_widget_destroy(data->dialog);
}


void      
nd_ip_frag_fragment(LND_Packet *packet, guint size1, guint size2)
{
  LND_Protocol  *ip;
  struct ip     *iphdr;
  LND_Packet    *p_copy;
  guchar        *payload_start, *payload_end, *fragment_start, *fragment_end;
  int            remainder_len, offset;

  if (!packet)
    return;

  ip = nd_ip_get();

  /* Create the second fragment: */

  p_copy = libnd_packet_duplicate(packet);
  iphdr = (struct ip *) libnd_packet_get_data(p_copy, ip, 0);
  D_ASSERT_PTR(iphdr);

  if (!iphdr)
    {
      D(("Packet doesn't contain IP\n"));
      return;
    }

  /* Sanity-check input: sum of fragment sizes must
     be equal to packet payload size: */

  if ((size1 + size2) != (guint) (ntohs(iphdr->ip_len) - (iphdr->ip_hl * 4)))
    {
      D(("Fragment sizes don't sum up to payload size\n"));
      return;
    }

  /* First fragment size must be multiple of 8: */
  if (size1 % 8)
    {
      D(("Fragment size is not multiple of 8\n"));
      return;
    }

  offset         = (ntohs(iphdr->ip_off) & IP_OFFMASK) * 8 + size1;
  payload_start  = ((guchar *) iphdr) + (iphdr->ip_hl * 4);
  payload_end    = libnd_packet_get_data_end(p_copy, ip, 0);
  fragment_start = payload_start + size1;  
  remainder_len  = libnd_packet_get_end(p_copy) - fragment_start;

  p_copy->ph.caplen -= size1;
  p_copy->ph.len -= size1;

  memmove(payload_start, fragment_start, remainder_len);

  /* We don't need to adjust IP_MF -- if the original didn't have it set,
     the copy is now the last fragment and doesn't need it. If it did
     have it set, there's going to be another fragment after this one
     anyway. But we need to set the fragment offset:
  */
  iphdr->ip_off = htons((ntohs(iphdr->ip_off) & ~IP_OFFMASK) | ((offset / 8) & IP_OFFMASK));

  /* Adjust IP packet size: */
  iphdr->ip_len = htons(ntohs(iphdr->ip_len) - size1);

  /* .. and fix the checksum. */
  libnd_ip_fix_packet(p_copy);

  /* Realloc the packet's memory. All protocol offsets may become
     invalid at this point, but we re-initialize immediately below: */
  p_copy->data = realloc(p_copy->data, p_copy->ph.caplen);

  /* We've just cut out a part of the packet -- reinitialize the
     packet: */
  libnd_packet_init(p_copy);

  /* Now adjust the original packet: */

  iphdr = (struct ip *) libnd_packet_get_data(packet, ip, 0);
  fragment_start = ((guchar *) iphdr) + (iphdr->ip_hl * 4);
  payload_end    = libnd_packet_get_data_end(packet, ip, 0);
  fragment_end   = fragment_start + size1;
  remainder_len  = libnd_packet_get_end(packet) - payload_end;

  packet->ph.caplen -= size2;
  packet->ph.len -= size2;

  /* Move anything coming after the IP payload to the end of the fragment: */
  if (remainder_len > 0)
    memmove(fragment_end, payload_end, remainder_len);
  packet->data = realloc(packet->data, packet->ph.caplen);

  /* Correct IP fields -- fragmentation flags, IP length, checksum: */
  iphdr->ip_off = htons((ntohs(iphdr->ip_off) | IP_MF) & ~IP_DF);
  iphdr->ip_len = htons(ntohs(iphdr->ip_len) - size2);
  libnd_ip_fix_packet(packet);

  /* And reinitialize the whole thing */
  libnd_packet_init(packet);
  libnd_packet_modified(packet);

  /* Insert new fragment after the original into trace: */
  libnd_tp_insert_packets(packet->part, p_copy,
			  libnd_packet_get_index(packet) + 1);
}



/* Compare function for IP fragment offsets.
   Returns -1/0/1 when the first packet's fragment
   offset is smaller/equal/larger than the second's */
static gint 
ip_off_cmp(gconstpointer d1, gconstpointer d2)
{
  LND_Packet *p1, *p2;
  struct ip *iphdr1, *iphdr2;

  p1 = (LND_Packet *) d1;
  p2 = (LND_Packet *) d2;

  iphdr1 = (struct ip *) libnd_packet_get_data(p1, nd_ip_get(), 0);
  D_ASSERT_PTR(iphdr1);
  iphdr2 = (struct ip *) libnd_packet_get_data(p2, nd_ip_get(), 0);
  D_ASSERT_PTR(iphdr2);

  if (ND_IP_FRAG_OFFSET(iphdr1) < ND_IP_FRAG_OFFSET(iphdr2))
    return (-1);
  if (ND_IP_FRAG_OFFSET(iphdr1) > ND_IP_FRAG_OFFSET(iphdr2))
    return (1);

  return (0);
}


static void 
ip_frag_reassemble_packet(gpointer key, GList *fraglist, gpointer user_data)
{
  struct ip        *base_iphdr, *frag_iphdr;
  GList            *f, *f2, *f3;
  int               payload_increase, trailer_len, covered_offset;
  LND_Packet       *base_packet, *frag_packet, *input_packet;
  guchar           *trailer, *base_payload, *dest;
  

  if (!fraglist || (g_list_length(fraglist) == 1))
    return;

  input_packet = (LND_Packet *) user_data;

  D(("Beginning reassembly.\n"));

  /* Sort fragments in increasing order by offsets. */
  fraglist = g_list_sort(fraglist, ip_off_cmp);

  for (f = fraglist; f; /* Increment depends on fragments, see below */)
    {
      base_packet = (LND_Packet *) f->data;
      base_iphdr = (struct ip *) libnd_packet_get_data(base_packet, nd_ip_get(), 0);
      D_ASSERT_PTR(base_iphdr);

      covered_offset = ND_IP_FRAG_OFFSET(base_iphdr) + ND_IP_PAYLOAD_LENGTH(base_iphdr);
      payload_increase = 0; 
      f2 = g_list_next(f);

      D(("Handling fragment of %u at %u\n", ntohs(base_iphdr->ip_id), ND_IP_FRAG_OFFSET(base_iphdr)));
	
      while (f2)
	{
	  frag_iphdr = (struct ip *) libnd_packet_get_data((LND_Packet *) f2->data, nd_ip_get(), 0);

	  if (ND_IP_FRAG_OFFSET(frag_iphdr) > covered_offset)
	    break;

	  D(("One consecutive fragment.\n"));
	  
	  payload_increase += ND_IP_PAYLOAD_LENGTH(frag_iphdr);
	  covered_offset = ND_IP_FRAG_OFFSET(frag_iphdr) + ND_IP_PAYLOAD_LENGTH(frag_iphdr);
	  f2 = g_list_next(f2);
	}
	
      D(("Payload increase in base packet: %i bytes\n", payload_increase));

      /* f2 is now NULL or the first list item containing a fragment packet
	 that does not consecutively add payload to the previous ones. */

	 
      /* The packet size increases by the amount of all pasted
	 payload sizes: */
      
      base_packet->ph.caplen += payload_increase;
      base_packet->ph.len    += payload_increase;
      base_iphdr->ip_len = htons(ntohs(base_iphdr->ip_len) + payload_increase);

      /* Realloc packet size to new size: */

      base_packet->data =
	realloc(base_packet->data,
		base_packet->ph.caplen);
      
      /* This may have made offsets invalid, so re-init: */
      libnd_packet_init(base_packet);
      base_iphdr = (struct ip *) libnd_packet_get_data(base_packet, nd_ip_get(), 0);

      base_payload = ND_IP_PAYLOAD(base_iphdr);
      trailer = libnd_packet_get_data_end(base_packet, nd_ip_get(), 0);
      trailer_len = libnd_packet_get_end(base_packet) - trailer;

      D(("Reallocd content, adjusted length values.\n"));

      /* Move anything that follows IP to the new trailer location: */
      if (trailer_len)
	{
	  memmove(base_payload + ND_IP_PAYLOAD_LENGTH(base_iphdr) + payload_increase, trailer, trailer_len);
	  D(("Trailer of %i bytes moved to end.\n", trailer_len));
	}
      
      /* Go through list and reassemble (we can skip
	 the first packet since its payload is already there) */
      for (f3 = g_list_next(f); f3 && (f3 != f2) ; f3 = g_list_next(f3))
	{
	  frag_packet = (LND_Packet *) f3->data;
	  frag_iphdr  = (struct ip *) libnd_packet_get_data(frag_packet, nd_ip_get(), 0);

	  D(("Merging one packet.\n"));
	  
	  /* The offset of the new fragment payload must start at
	     the currently covered payload, otherwise we don't have
	     consecutive fragments */

	  dest =
	    base_payload +                  /* Start address of data */
	    ND_IP_FRAG_OFFSET(frag_iphdr) - /* plus fragment offset of the pasted fragment */
	    ND_IP_FRAG_OFFSET(base_iphdr);  /* minus fragment offset ofthe base fragment */

	  /* Copy the payload of the fragment packet to the correct offset
	     in the base fragment */
	  memcpy(dest, ND_IP_PAYLOAD(frag_iphdr), ND_IP_PAYLOAD_LENGTH(frag_iphdr));      

	  /* If we're pasting a last fragment, the base fragment will now
	     be a last fragment as well (since we are only handling consecutive
	     areas). Thus clear MF if necessary: */
	  if ((ntohs(frag_iphdr->ip_off) & IP_MF) == 0)
	    base_iphdr->ip_off = htons(ntohs(base_iphdr->ip_off) & ~IP_MF);
	  
	  /* This fragment is reassembled and can thus be nuked: */
	  if (input_packet == frag_packet)
	    nd_trace_set_current_packet(libnd_packet_get_trace(base_packet), NULL);
	  
	  libnd_packet_remove(frag_packet);
	  libnd_packet_free(frag_packet);
	  f3->data = NULL;
	}

      D(("Fixing checksum in merged packet.\n"));
      libnd_ip_fix_packet(base_packet);
      libnd_packet_modified(base_packet);

      /* If this is a now-complete IP packet, reinitialize it
	 so that payload protocols become visible: */
      if (((ntohs(base_iphdr->ip_off) & IP_OFFMASK) == 0) &&
	  ((ntohs(base_iphdr->ip_off) & IP_MF) == 0))
	{
	  D(("New packet is complete IP datagram, initializing ...\n"));
	  libnd_packet_init(base_packet);
	}

      /* Continue where new consecutive fragments start: */
      f = f2;
    }
    
  return;
  TOUCH(key);
  TOUCH(user_data);
}


static void 
ip_frag_free(gpointer key, GList *fraglist, gpointer user_data)
{
  g_list_free(fraglist);

  return;
  TOUCH(key);
  TOUCH(user_data);
}


void      
nd_ip_frag_reassemble(LND_Packet *input_packet)
{
  struct ip               *iphdr;
  LND_Packet              *packet;
  LND_Trace               *trace;
  GHashTable              *fraggles;
  GList                   *fragslist = NULL;
  LND_PacketIterator       pit;
  LND_PacketIteratorMode   mode;
  int                      fragments = 0, packets = 0;
  char                     message[MAXPATHLEN];

  if (!input_packet || !libnd_packet_get_trace(input_packet))
    return_if_no_current_trace(trace);
  else
    trace = libnd_packet_get_trace(input_packet);

  fraggles = g_hash_table_new(g_direct_hash, g_direct_equal);
  mode = trace->iterator_mode;

  if (mode == LND_PACKET_IT_AREA_R || mode == LND_PACKET_IT_AREA_RW)
    {
      /* We're currently only reassembling packets in memory,
       * so reduce this to apply-to-trace-part if the current
       * mode is apply-to-all:
       */

      mode = LND_PACKET_IT_PART_RW;
    }

  for (libnd_pit_init_mode(&pit, trace, mode); libnd_pit_get(&pit); libnd_pit_next(&pit))
    {
      packet = libnd_pit_get(&pit);

      iphdr = (struct ip *) libnd_packet_get_data(packet, nd_ip_get(), 0);
      if (!iphdr)
	continue;

      /* If this is a fragment, put it into the hashtable. */
      if ((ntohs(iphdr->ip_off) & IP_MF) ||
	  (ntohs(iphdr->ip_off) & IP_OFFMASK))
	{
	  fragslist = g_hash_table_lookup(fraggles, GINT_TO_POINTER((int) (iphdr->ip_id)));

	  if (!fragslist)
	    {
	      packets++;
	      fragslist = g_list_append(fragslist, packet);
	      g_hash_table_insert(fraggles, GINT_TO_POINTER((int) (iphdr->ip_id)), fragslist);
	      D(("Creating fragment list for id %u\n", ntohs(iphdr->ip_id)));
	    }
	  else
	    {
	      fragslist = g_list_append(fragslist, packet);	      
	      D(("Adding fragment to list for id %u\n", ntohs(iphdr->ip_id)));
	    }

	  fragments++;
	}
    }

  if (fragments > 1)
    {
      g_hash_table_foreach(fraggles, (GHFunc) ip_frag_reassemble_packet, input_packet);
      g_snprintf(message, MAXPATHLEN, _("Reassembled %i fragments belonging to %i IP packet(s)."), fragments, packets);
      nd_dialog_message(_("Reassembly results"), message, FALSE);
    }

  g_hash_table_foreach(fraggles, (GHFunc) ip_frag_free, NULL);
  g_hash_table_destroy(fraggles);  
}