#include <netinet/tcp.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>

#include "packet_handler.h"

extern struct hash_table* ConnTable; 
extern unsigned int my_dest;
extern unsigned int my_dmask;
extern unsigned int my_source;
extern unsigned int my_smask;
extern double startTime;
extern unsigned int dests[MAX_DESTS];
extern int nd;
extern int batch;
extern char Host[17];

int CHECKID = TRUE;
int transactions = 0;


int isDest(unsigned int address)
{
  int low = -1, high = nd, i;
  do{
    i = (high+low)/2;
    if (dests[i] == address)
      return TRUE;
    if (dests[i] < address)
      low = i;
    else
      high = i;
  }while (high > low+1);
  return FALSE;
}


int compare(unsigned int address)
{
  int result = 0;
  if (!batch)
    {
      if (my_dest)  
	result = ((my_dest & my_dmask) == (address & my_dmask));
      if (my_source)  
	result = ((my_source & my_smask) == (address & my_smask));
    }
  else
    result = isDest(address);
  return result;
}

void closeTransactions(double lasttime)
{
  int i;
  for(i=0;i<ConnTable->size;i++)
    if (ConnTable->entries[i].state != EMPTY)
      {
	int j;
	struct conn_key* key;
	struct conn_stats* data;
	
	key = (struct conn_key*)ConnTable->entries[i].key;
	data = (struct conn_stats*)ConnTable->entries[i].data;

	for(j=0;j<data->nt;j++)
	  closeTransaction(&data->trans[j], lasttime);
	data->nt = 0;
	data->t = 0;
      }
}


void resetTransactions()
{
  int i, j;
  for(i=0;i<ConnTable->size;i++)
    if (ConnTable->entries[i].state != EMPTY)
      {
	struct conn_key* key;
	struct conn_stats* data;
	
	key = (struct conn_key*)ConnTable->entries[i].key;
	data = (struct conn_stats*)ConnTable->entries[i].data;

	for(j=0; j<data->nt;j++)
	  if (data->trans[j]->success != IGNORE)
	    data->trans[j]->success = NOT_FOUND;
      }
}

/* Return TRUE if the application uses first 16 bits or so of the header for ID field */
int hasID(enum app tapp)
{
  switch(tapp)
    {
    case DNS:
      if (CHECKID)
	return TRUE;
      else
      return FALSE;
    case ICMP:
      if (CHECKID)
	return TRUE;
      else
      return FALSE;
    default: 
      return FALSE;
    }
}

/* Return TRUE for applications that can have multiple transactions over the
   same connection, and should follow request/reply paradigm */
int followRequestReply(enum app tapp)
{
  switch(tapp)
    {
    case DNS:
      return TRUE;
    case ICMP:
      return TRUE;
    default: 
      return FALSE;
    }    
}

/* Return minimum inactive delay for applications that can have 
   multiple transactions over the
   same connection, and follow request/reply paradigm but we should 
   use a period
   of inactivity to detect start of new transactions */
double detectInactive(enum app tapp)
{
  switch(tapp)
    {
    case DNS:
    case ICMP:
      return INTERACTIVE_DELAY_THRESHOLD;
    case Telnet:
      return INACTIVE_TIME;
    case FTPData:
      return INACTIVE_TIME;
    case HTTP:
      return INACTIVE_TIME;
    default:
      return MAX_TIME;
    }
}
     
void getTypeAndApp(int dport, enum type* ttype, enum app* tapp)
{
  switch(dport)
    {
    case ICMP_PORT:
      *tapp = ICMP;
      break;
    case HTTP_PORT1:
    case HTTP_PORT2:
    case HTTP_PORT3:
      *tapp = HTTP;
      break;
    case DNS_PORT:
      *tapp = DNS;
      break;
    case SSH_PORT:
      *tapp = FTPTelnet;
      break;
    case SMTP_PORT1:
      *tapp = MailSS;
      break;
    case SMTP_PORT2:
      *tapp = MailUS;
      break;
    case NEWS_PORT:
      *tapp = News;
      break;
    case IRC_PORT:
      *tapp = IRC;
      break;
    case VOIP_PORT:
      *tapp = VoIP;
      break;
    case FTP_PORT1:
      *tapp = FTPData;
      break;
    case FTP_PORT2:
      *tapp = FTPControl;
      break;
    case TELNET_PORT:
      *tapp = Telnet;
      break;
    default:
      *tapp = NOAPP;
      break;
    }
  *ttype = BASIC; 
}

void handle_TCP_packet(double dtime, unsigned int srcIP, unsigned int dstIP, 
		       unsigned short sport, unsigned short dport, 
		       unsigned int seqnum, unsigned int acknum,
		       unsigned short syn, unsigned short ack, 
		       unsigned short fin, unsigned short reset, int len, 
		       int hlen,
		       double owdelay, unsigned short dataport)
{
  int index;
  struct conn_stats* data;
  enum app tapp;
  enum type ttype;
  int ofInterest;
  int primary;
  int toDest;
  int fromSource;

  /* Don't do anything for ports that we don't care about */
  if (!isDstPort(sport) && !isDstPort(dport))
    return;
  /* Get type and application for a given port */
  if (isDstPort(dport))
    getTypeAndApp(dport, &ttype, &tapp);
  else
    getTypeAndApp(sport, &ttype, &tapp);
  /* This is a packet of interest if it is originated by/going to
     a source or a destination we are tracking. For
     FTP data connections we have to do special things because they 
     are initiated in a reverse direction */
  ofInterest = (compare(dstIP) || compare(srcIP));
  if (!ofInterest)
    return;
  /* Check if this is a packet in the primary direction on the connection */
  fromSource = ((compare(dstIP) && sport == FTP_PORT1) || ( compare(srcIP) && sport != FTP_PORT1 && dport != FTP_PORT1 && isDstPort(dport))) && my_source;
  toDest = ((compare(srcIP) && sport == FTP_PORT1) || ( compare(dstIP) && sport != FTP_PORT1 && dport != FTP_PORT1 && isDstPort(dport))) && my_dest;
  primary =  (fromSource || toDest);
  /* For syn packets open a new connection */
  if (syn && !ack && primary)
    {
      int res;
      /* Insert a new connection if it doesn't exist */
      index = find_connection(srcIP, dstIP, sport, dport);
      if (index == NOT_FOUND) 
	{
	  index = insert_connection(srcIP, dstIP, sport, dport);
	  if (index == NOT_FOUND)
	    return;
	  data = (struct conn_stats*)ConnTable->entries[index].data;
	  data->starttime = dtime;
	  data->lasttime = dtime;
	  data->seqnum = seqnum + len + 1;
	  data->loss = 0;
	  res = insertTransaction
	    (data, dtime, ttype, tapp, TCP, srcIP, sport, dstIP, dport, transactions);
	  if(!res)
	    return;
	  transactions++;
	}
      else
	data = (struct conn_stats*)ConnTable->entries[index].data;
      /* Check if we should insert a new transaction */
      if (!setActiveTransaction(data, tapp,0))
	{
	  res=insertTransaction
	    (data, dtime, ttype, tapp, TCP, srcIP, sport, dstIP, dport, transactions);
	  if (!res)
	    return;
	  transactions++;
	}
      data->lasttime = dtime;
    }
  /* If this is a non-syn packet find the appropriate connection direction */
  if (primary)
    index = find_connection(srcIP, dstIP, sport, dport);
  else
    index = find_connection(dstIP, srcIP, dport, sport);


  /* We don't have an appropriate connection */
  if (index == NOT_FOUND)
      return;

  data = (struct conn_stats*)ConnTable->entries[index].data;
  setActiveTransaction(data, tapp,0);
  /* We have no transactions */
  if (!data->t)
    return;

  /* Here create new transaction if needed */
  if (primary && detectNew(data->t, seqnum) && 
      (data->t->usefulBytesSent > 0 && data->t->usefulBytesRcvd > 0) 
      && (dtime - data->t->lastTime > detectInactive(data->t->tapp)))
    {
      removeTransaction(data, dtime, data->t);
      insertTransaction
	(data, dtime, ttype, tapp, TCP, srcIP, sport, dstIP, dport, transactions++);
    }	
 
  /* Calculate jitter */
  if (owdelay != MAX_DELAY)
    {
      if (data->t->ds == 0)
	data->t->ds = owdelay;
      else 
	{
	  data->t->jitter += ((owdelay-data->t->ds) - data->t->jitter)/16;
	  if (data->t->jitter > data->t->maxjitter)
	    data->t->maxjitter = data->t->jitter;
	  data->t->ds = owdelay;
	}
      /* Calculate one-way delay */
      if (owdelay > data->t->oneWayDelay)
	data->t->oneWayDelay = owdelay;
    }
  if (owdelay == MAX_DELAY)
      data->t->lostpackets++;

  /* If this transaction has a control FTP port be the destination, 
     see if we have the data port to assign to it. This is currently
     not used but we may want to combine control and data FTP traffic into
     a single transaction. */
  if (data->t->dport == FTP_PORT2)
    {
      /* Try to find a corresponding FTP data transaction and set its
	 controlport */
      if (data->t->dataport)
	{
	  int index_d = find_connection(data->t->dIP, data->t->sIP, 20, data->t->dataport);
	  if (index_d != NOT_FOUND)
	    {
	      struct conn_stats* data_d = (struct conn_stats*)ConnTable->entries[index_d].data;
	      int res = setActiveTransaction(data_d, FTPTelnet, 0);
	      if (res)
		{
		  if (data_d->t->controlport == 0)
		    {
		      data_d->t->controlport = data->t->sport;
		    }
		}
	    }
	}
      if (dataport > 0)
	{
	  data->t->dataport = dataport;
	}
    }
  /* If this is a FIN update the state, 
     an ack after second FIN closes transaction */
  if ((ack && data->t->finstate == 3) || reset)
    {
      removeTransaction(data, dtime, data->t);
      return;
    }
  if (fin)
    {
      int toadd;
      /* If this is in the original direction, add 1 to state, else add 2 */
      if (primary)
	toadd = 1;
      else 
	toadd = 2;
      if ((data->t->finstate != toadd) && (data->t->finstate < 3))
	data->t->finstate += toadd;
    }
  if (primary)
    addTCPPacket(data->t, dtime, seqnum, len, hlen, index, syn, fin);
  else
    addTCPAck(data->t, dtime, seqnum, acknum, len, hlen, index);
}

void handle_UDP_packet(double dtime, unsigned int srcIP, unsigned int dstIP, 
		       unsigned short sport, unsigned short dport, int len, 
		       int hlen, int identification, double owdelay, int req)
{
  int index;
  struct conn_stats* data;
  enum app tapp;
  enum type ttype;
  int ofInterest;
  int primary;
  int toDest;
  int fromSource;

  /* Don't do anything for ports that we don't care about */
  if (!isDstPort(sport) && !isDstPort(dport))
    return;
  if (isDstPort(sport))
    getTypeAndApp(sport, &ttype, &tapp);
  else
    getTypeAndApp(dport, &ttype, &tapp);
  /* This is a packet of interest if it is originated by/going to
     a source or a destination we are tracking. For
     FTP data connections we have to do special things because they 
     are initiated in a reverse direction */
  ofInterest = (compare(dstIP) || compare(srcIP));
  if (!ofInterest)
    return;
  /* Check if this is a packet in the primary direction on the connection */
  fromSource = compare(srcIP) && isDstPort(dport) && my_source;
  toDest = compare(dstIP) && isDstPort(dport) && my_dest;
  primary = (fromSource || toDest);

  /* Deal with short DNS packets that may not have a complete header */
  if ((dport == DNS_PORT) && (!CHECKID))
    if (len < MAX_DNS_REQ)
      {
	req = TRUE;
      }
    else
      {
	req = FALSE;
      }
  if (primary && req)
  {
    /* First check if connection exists, if not then insert it */
    index = find_connection(srcIP, dstIP, sport, dport);
    /* Check if connection exists in reverse direction, this is only relevant
       if we have packets without ID, to avoid false positives */
    if (!CHECKID && batch)
      {
	int rindex = find_connection(dstIP, srcIP, dport, sport);
	if (rindex != NOT_FOUND)
	  {
	    int res;
	    data = (struct conn_stats*)ConnTable->entries[rindex].data;
	    res = setActiveTransaction(data, tapp, identification);
	    if (res)
		addUDPReply(data->t, dtime, len, hlen, identification);
	    return;
	  }
      }
    if (index == NOT_FOUND)
      {
	int res;
	/* If DNS, only insert connections for requests */
	if (ConnTable->filled < ConnTable->size)
	  index = insert_connection(srcIP, dstIP, sport, dport);
	if (index == NOT_FOUND)
	  return;
	data = (struct conn_stats*)ConnTable->entries[index].data;
	data->starttime = dtime;
	data->lasttime = dtime;
	data->loss = 0;
	/* Here decide about the type etc. */
	res=insertTransaction(data, dtime, ttype, tapp, UDP, srcIP, sport, dstIP, dport, transactions);
	if (!res)
	  return;
	transactions++;
      }
    else
      {
	int res;
	data = (struct conn_stats*)ConnTable->entries[index].data;
	/* ID for VOIP is packet ID so avoid creating new transaction */
	if (dport == VOIP_PORT)
	  identification = 0;
	res = setActiveTransaction(data, tapp, identification);
	if (!res)
	  {
	    res = insertTransaction(data, dtime, ttype, tapp, UDP, srcIP, sport, dstIP, dport, transactions);
	    if (!res)
	      return;
	    transactions++;
	  }
      }
    if (!hasID(data->t->tapp))
      identification = -1;
    /* Calculate jitter */
    if (owdelay != MAX_DELAY)
      {
	if (data->t->ds == 0)
	  data->t->ds = owdelay;
	else 
	  {
	    data->t->jitter += ((owdelay-data->t->ds) - data->t->jitter)/16;
	    if (data->t->jitter > data->t->maxjitter)
	      data->t->maxjitter = data->t->jitter;
	    data->t->ds = owdelay;
	  }
    if(owdelay > data->t->oneWayDelay)
      data->t->oneWayDelay = owdelay;
      }
    if (owdelay == MAX_DELAY)
	data->t->lostpackets++;
    addUDPPacket(data->t, dtime, len, hlen, identification);
  }
  else 
    {
      int res;
      index = find_connection(dstIP, srcIP, dport, sport);
      if (index == NOT_FOUND)
	return;
      data = (struct conn_stats*)ConnTable->entries[index].data;
      res = setActiveTransaction(data, tapp, identification);
      if (res)
	addUDPReply(data->t, dtime, len, hlen, identification);
    }
}

void handle_ICMP_packet(double dtime, unsigned int srcIP, unsigned int dstIP,
			unsigned short dport, unsigned short sport,
			int len, int hlen, 
			int identification, int type, double owdelay)
{
  int index;
  struct conn_stats* data;
  enum app tapp;
  enum type ttype;

  int ofInterest;
  int primary;
  int toDest;
  int fromSource;

  /* Don't do anything for ports that we don't care about */
  if (!isDstPort(sport) && !isDstPort(dport))
    return;
  if (isDstPort(sport))
    getTypeAndApp(sport, &ttype, &tapp);
  else
    getTypeAndApp(dport, &ttype, &tapp);
  /* This is a packet of interest if it is originated by/going to
     a source or a destination we are tracking. For
     FTP data connections we have to do special things because they 
     are initiated in a reverse direction */
  ofInterest = (compare(dstIP) || compare(srcIP));
  if (!ofInterest)
    return;
  /* Check if this is a packet in the primary direction on the connection */
  fromSource = compare(srcIP) && isDstPort(dport) && my_source;
  toDest = compare(dstIP) && isDstPort(dport) && my_dest;
  primary = (fromSource || toDest);
  if (primary && (type == ECHO_REQUEST || !CHECKID))
  {
    /* First check if connection exists, if not then insert it */
    index = find_connection(srcIP, dstIP, sport, dport);
    /* Check if connection exists in reverse direction, this is only relevant
       if we have packets without ID, to avoid false positives */
    if (!CHECKID && batch)
      {
	struct in_addr sin, din;
	int rindex = find_connection(dstIP, srcIP, dport, sport);
	sin.s_addr =htonl(srcIP);
	din.s_addr = htonl(dstIP);
	if (rindex != NOT_FOUND)
	  {
	    int res;
	    data = (struct conn_stats*)ConnTable->entries[rindex].data;
	    res = setActiveTransaction(data, tapp, identification);
	    if (res)
	      {
		addICMPReply(data->t, dtime, len, hlen, identification);
		removeTransaction(data, dtime, data->t);
	      }
	    return;
	  }
      }
    if (index == NOT_FOUND)
      {
	int res;
	if (ConnTable->filled < ConnTable->size)
	  index = insert_connection(srcIP, dstIP, sport, dport);
	if (index == NOT_FOUND)
	  return;
	data = (struct conn_stats*)ConnTable->entries[index].data;
	data->starttime = dtime;
	data->lasttime = dtime;
	data->loss = 0;
	res=insertTransaction(data, dtime, ttype, tapp, ICMP, 
			      srcIP, sport, dstIP, dport, transactions);
	if (!res)
	  return;
	transactions++;
      }
    else 
      {
	int res;
	data = (struct conn_stats*)ConnTable->entries[index].data;
	res = setActiveTransaction(data, tapp, identification);
	if (!res)
	  {
	    res=insertTransaction(data, dtime, ttype, tapp, ICMP, srcIP, sport, dstIP, dport, transactions);
	    if (!res)
	      return;
	    transactions++;
	  }
      }
    /* Calculate jitter */
    if (owdelay != MAX_DELAY)
      {
	if (data->t->ds == 0)
	  data->t->ds = owdelay;
	else 
	  {
	    data->t->jitter += ((owdelay-data->t->ds) - data->t->jitter)/16;
	    if (data->t->jitter > data->t->maxjitter)
	      data->t->maxjitter = data->t->jitter;
	    data->t->ds = owdelay;
	  }
	if (owdelay > data->t->oneWayDelay)
	  data->t->oneWayDelay = owdelay;
      }
    if (owdelay == MAX_DELAY)
      {
	data->t->lostpackets++;
      }
    addICMPPacket(data->t, dtime, len, hlen, identification);
  }
  else if (type == ECHO_REPLY || !CHECKID)
    {
      int res;
      index = find_connection(dstIP, srcIP, dport, sport);
      if (index == NOT_FOUND)
	return;
      data = (struct conn_stats*)ConnTable->entries[index].data;
      res = setActiveTransaction(data, tapp, identification);

      if (res)
	{
	  addICMPReply(data->t, dtime, len, hlen, identification);
	  removeTransaction(data, dtime, data->t);
	}
    }
}