/*
 *  Paglo Crawler
 *  Copyright (C) 2006-2008 Paglo Labs Inc. All rights reserved.
 *  www.paglo.com
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of version 2 of the GNU General Public License as
 *  published by the Free Software Foundation.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program; if not, write to the Free Software
 *  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
 */
#ifndef ROGUESCANNER
//---------------------------------------------------------------------------
#include <stdexcept>
#include <pcap.h>
#include "ParseSFlow.h"
#include "AnalysisManager.h"
#include "ScanAgentUtils.h"
//---------------------------------------------------------------------------


#define SFLOW_TRACE_FILE "sflow_trace.pcap"

pcap_dumper_t *_SFlowDumpHandle = NULL;

void ParseSFlow(TPacket *Pkt)
{
    /*
   * Don't parse past this!
   */
  const unsigned char *PktEnd = Pkt->GetData() + Pkt->GetCaptureLength();
  
  
  /*
   * Sanity checks.
   */
  if (Pkt->GetIPProto() != IPPROTO_UDP || Pkt->GetDestPort() != 6343 || 
      Pkt->Payload == NULL) {
    return;
  }

  /*
   * More sanity.
   */
  TSFIPv4DatagramHeader *SFlowDatagram = (TSFIPv4DatagramHeader *)Pkt->Payload;
  if (((unsigned char *)SFlowDatagram + sizeof(TSFIPv4DatagramHeader))  > PktEnd) { 
    return;
  }
  
  /*
   * We only support v.5 for now.
   */
  if (ntohl(SFlowDatagram->Version) != 5) {
    return;
  }

  /*
   * We only support IPv4 for now.
   */
  if (ntohl(SFlowDatagram->AddressType) != 1) {
    return;
  }
    
  int Format = 0;
  int Length = 0;
  int NRecords = 0;
  TSFFlowSample *FlowSample = NULL;
  unsigned char *EncapPacketData = NULL;
  int CapLen = 0;
  int Len = 0;
  int Stripped = 0;

  TSFSampleHeader *SampleHdr = (TSFSampleHeader *)((unsigned char *)SFlowDatagram + sizeof(TSFIPv4DatagramHeader));

  for (unsigned int i = 0; i < ntohl(SFlowDatagram->NumSamples); i++) {

    /*
     * More sanity -- "serenity now. insanity later."
     */
    if (((unsigned char *)SampleHdr + sizeof (TSFSampleHeader)) > PktEnd) { 
      return;
    }
    
    Format = htonl(SampleHdr->Format) & 0xFFF;
    Length = htonl(SampleHdr->Length);
    
    switch(Format) {
    case 1:

      FlowSample = (TSFFlowSample *)((unsigned char *)SampleHdr + sizeof(TSFSampleHeader));
      if (((unsigned char *) FlowSample + sizeof(TSFFlowSample)) > PktEnd) { 
        return;
      }

      NRecords = htonl(FlowSample->NRecords);
      for (int j = 0; j < NRecords; j++) {
        TSFFlowRecordHeader *RecordHeader = (TSFFlowRecordHeader *)((unsigned char *)FlowSample + sizeof (TSFFlowSample));
        if (((unsigned char *)RecordHeader + sizeof(TSFFlowRecordHeader)) > PktEnd) { 
          return;
        }

        TSFSampledHeader *Asdf = (TSFSampledHeader *)((unsigned char *)RecordHeader + sizeof(TSFFlowRecordHeader));
        if (((unsigned char *)Asdf + sizeof(TSFSampledHeader)) > PktEnd) { 
          return;
        }
        
        CapLen = ntohl(Asdf->HeaderLength);
        Len = ntohl(Asdf->FrameLength);
        Stripped = ntohl(Asdf->Stripped);
        EncapPacketData = (unsigned char *)Asdf + sizeof (TSFSampledHeader);
        
 
        if ((EncapPacketData + CapLen) < PktEnd) {
          try {
            /*
             * Fudge a pcap_pkthdr structure, so that we can dump the sample's contents to 
             * a .pcap file for easier debugging.
             */
            struct pcap_pkthdr PktHdr;
            PktHdr.ts.tv_sec = 0;
            PktHdr.ts.tv_usec = 0;
            gettimeofday(&PktHdr.ts, NULL);
            PktHdr.caplen = CapLen;
            PktHdr.len = Len;
            
            TPacket *OutPacket = new TPacket(&PktHdr, EncapPacketData);
            OutPacket->IsSFlow = true;
 
#ifdef _DEBUG
           /*
            * Dump the sFlow sample's payload to a .pcap file if we were compiled
            * with _DEBUG and if debug="..." in the config file has DEBUG_SFLOW set.
            */
            if (_DebugConf & DEBUG_SFLOW) {
              /*
               * If we don't have a dump handle then we need to initialize one.
               * We create a dummy pcap_t structure, then truncate/create our trace file.
               * Finally we dump the sample payload to the file if everything else was
               * successful.
               */
              if (_SFlowDumpHandle == NULL) {
                pcap_t *DummyPcapHandle = pcap_open_dead(DLT_EN10MB, 65535);
                if (DummyPcapHandle) {
                  FILE *TraceFileHandle = fopen(SFLOW_TRACE_FILE, "w");
                  if (TraceFileHandle) {
                    fclose(TraceFileHandle);
                  }
                  if ((_SFlowDumpHandle = pcap_dump_open(DummyPcapHandle, SFLOW_TRACE_FILE))) {
                    pcap_dump((unsigned char *)_SFlowDumpHandle, &PktHdr, EncapPacketData);
                  }
                }
              } else {
                pcap_dump((unsigned char *)_SFlowDumpHandle, &PktHdr, EncapPacketData);
              }
            }
#endif            
            DEBUG_MESSAGE(DEBUG_SFLOW, ("Received sFlow sample (CapLen: %d, Len: %d, SMAC: %s, DMAC: %s, SIP: %s, DIP: %s)",
                                        OutPacket->GetCaptureLength(), OutPacket->GetLength(),
                                        OutPacket->GetSrcMac().Print().c_str(),
                                        OutPacket->GetDstMac().Print().c_str(),
                                        OutPacket->GetSrcIP().Valid() ? OutPacket->GetSrcIP().Print().c_str() : "undef",
                                        OutPacket->GetDstIP().Valid() ? OutPacket->GetDstIP().Print().c_str() : "undef"));
            _AnalysisMgr->GetPacketQueue()->Put(OutPacket);
          } catch (exception &E) {
            LogMesg("Malformed packet encountered: %s.", E.what());
          }
        }
      }
      
      /*
       * This will get sanity checked at the top of the loop.
       */
      SampleHdr = (TSFSampleHeader *)(EncapPacketData + CapLen + 2);
      break;
    }
    
  }
}

//---------------------------------------------------------------------------
#endif // !ROGUESCANNER