#include "IpDissector.h"
#ifdef WIN32
#include <Winsock2.h>
#else
#include <netinet/in.h>
#endif

IpDissector::IpDissector() :
	m_ipHdr(0)
{
}

IpDissector::~IpDissector() {
  delete m_ipHdr;
}

IpDissector * IpDissector::Clone() const {
  return new IpDissector();
}

bool IpDissector::Dissect(const unsigned char * pData, ssize_t nDataLen) {
  if (!m_ipHdr)
	  m_ipHdr = new IpHeader;
  memcpy(m_ipHdr, pData, sizeof(IpHeader));
  pData += sizeof(IpHeader);
#if BYTE_ORDER == LITTLE_ENDIAN
  m_ipHdr->totalLength = ntohs(m_ipHdr->totalLength);
#endif
  SavePayload(pData, m_ipHdr->totalLength - m_ipHdr->headerLength*4);
  SaveLayerInfo();
  return true;
}

bool IpDissector::NeedReassembly() const {
  if(m_ipHdr->moreFragment || m_ipHdr->fragOffset || m_ipHdr->fragOffset1)
    return true;
  return false;
}

bool IpDissector::Reassemble(Dissector * pDissector, ProtocolInfo * pProtocolInfo) {
  if(!NeedReassembly())
    return false;
  
  if(m_holes.size() == 0) 
    {
      // No hole defined even thougth reassembly is needed
      // It means that Reassemble is called for first time
      // Initialize reassembly structure and add self
      m_holes.push_back(std::pair<unsigned short, unsigned short>(0, (unsigned short)(- 1)));
      Reassemble(this, pProtocolInfo);
    }
  
  // Inspect potential fragment
  IpDissector * pIpDissector = dynamic_cast<IpDissector *>(pDissector);
  if((pIpDissector->m_ipHdr) && (pIpDissector->m_ipHdr->id == m_ipHdr->id))
    {
      // Build First and Last values
      unsigned short nFirstFragmentByte = pIpDissector->m_ipHdr->fragOffset;
      nFirstFragmentByte = nFirstFragmentByte << 8;
      nFirstFragmentByte += pIpDissector->m_ipHdr->fragOffset1;
      nFirstFragmentByte *= 8;

      unsigned short nLastFragmentByte = nFirstFragmentByte 
	+ pIpDissector->m_ipHdr->totalLength
	- pIpDissector->m_ipHdr->headerLength * 4;

      // Compare to missing data
      std::list<std::pair<unsigned short, unsigned short> >::iterator it;
      int nPayloadSize = GetPayloadSize();
      for(it = m_holes.begin(); it != m_holes.end(); ++it) 
	{
	  if(nFirstFragmentByte > (*it).second) 
	    continue;
	  if(nLastFragmentByte < (*it).first) 
	    continue;
	  if(!pIpDissector->m_ipHdr->moreFragment) 
	    {
	      // Last fragment received
	      (*it).second = nLastFragmentByte;
	      nPayloadSize = nLastFragmentByte;
	    }
	
	  if(nFirstFragmentByte > (*it).first)
	    m_holes.push_back(std::pair<unsigned short, unsigned short>((*it).first, nFirstFragmentByte - 1));
	  if(nLastFragmentByte < (*it).second)
	    m_holes.push_back(std::pair<unsigned short, unsigned short>(nLastFragmentByte + 1, (*it).second));
	  m_holes.erase(it);
	  std::pair<const unsigned char *, const ssize_t> data = pIpDissector->DetachData();
	  m_fragments[nFirstFragmentByte] = data.first;
	  m_fragmentSizes[nFirstFragmentByte] = data.second;

      // received fragment is not needed anymore
      if(pIpDissector != this) 
	    delete pIpDissector;

      // check if packet is now complete
      if(m_holes.size() == 0) 
	  {
        m_ipHdr->moreFragment = 0;
	    m_ipHdr->fragOffset = 0;
	    m_ipHdr->fragOffset1 = 0;	
	  
	    // Reassembly !
	    ssize_t nOffset = 0;
	    AllocPayload(nPayloadSize); // payloadSize is valid since last fragment has been received
	    while((nOffset != nPayloadSize) && (m_fragments.find(nOffset) != m_fragments.end())) 
	      {
		SetPayloadAt(m_fragments[nOffset], m_fragmentSizes[nOffset], nOffset);
		// Clean up
		delete m_fragments[nOffset];
		nOffset += m_fragmentSizes[nOffset];
	      }
	    m_fragments.clear();
	    m_fragmentSizes.clear();
	  }
	  return true;
	}
    }
  return false;
}

const EProtocolType IpDissector::GetUpperLayerType() const 
{
	EProtocolType res = EProtocolType_Unknown;
	switch(m_ipHdr->protocol)
	{
	case 6:
		res = EProtocolType_Tcp;
		break;
	case 17:
		res = EProtocolType_Udp;
		break;
	}
	return res;
}

ProtocolInfoElement * IpDissector::CreateLayerInfo()
{
	IPv4Info * pRes = new IPv4Info();
	unsigned int nSrcAddr, nDstAddr;
	memcpy(&nSrcAddr, m_ipHdr->srcAddr, sizeof(int));
	memcpy(&nDstAddr, m_ipHdr->destAddr, sizeof(int));
	pRes->SetSourceAddress(nSrcAddr);
	pRes->SetDestinationAddress(nDstAddr);
	return pRes;
}
