//===========================================================================
//  $Name: cflowd-2-1-b1 $
//  $Id: dnswatch.cc,v 1.5 2000/10/24 16:18:54 dwm Exp $
//===========================================================================
//  CAIDA Copyright Notice
//
//  By accessing this software, cflowd++, you are duly informed
//  of and agree to be bound by the conditions described below in this
//  notice:
//
//  This software product, cflowd++, is developed by Daniel W. McRobb, and
//  copyrighted(C) 1998 by the University of California, San Diego
//  (UCSD), with all rights reserved.  UCSD administers the CAIDA grant,
//  NCR-9711092, under which part of this code was developed.
//
//  There is no charge for cflowd++ software. You can redistribute it
//  and/or modify it under the terms of the GNU General Public License,
//  v.  2 dated June 1991 which is incorporated by reference herein.
//  cflowd++ is distributed WITHOUT ANY WARRANTY, IMPLIED OR EXPRESS, OF
//  MERCHANTABILITY OR FITNESS FOR A PARTICULAR PURPOSE or that the use
//  of it will not infringe on any third party's intellectual property
//  rights.
//
//  You should have received a copy of the GNU GPL along with cflowd++.
//  Copies can also be obtained from:
//
//    http://www.gnu.org/copyleft/gpl.html
//
//  or by writing to:
//
//    University of California, San Diego
//
//    SDSC/CAIDA
//    9500 Gilman Dr., MS-0505
//    La Jolla, CA 92093 - 0505  USA
//
//  Or contact:
//
//    info@caida.org
//===========================================================================
//
//  This was a quick hack program to collect data on the traffic coming
//  from root DNS server F (replies to DNS queries).  We keep track of
//  the number of flows, packets and bytes of traffic sent from port 53
//  (regardless of source IP address) to each destination host.  When we
//  receive a SIGUSR1, we dump our data to stdout or to a specified
//  output file (using the -o command-line option).  When we receive a
//  SIGHUP, we dump our data *and* clear our data (restart data
//  collection).
//
//  This program hooks into cflowdmux in the same manner as cflowd
//  and flowwatch.
//
//  Daniel W. McRobb
//  CAIDA
//  October 5, 1999
//
//===========================================================================

extern "C" {
#include <sys/types.h>
#include <errno.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <stdio.h>
#include <string.h>
  
#include "caida_t.h"
}

#include "artslocal.h"

#ifdef HAVE_FSTREAM
  #include <fstream>
#else
  #include <fstream.h>
#endif

#include <stdiostream>

#ifdef HAVE_IOMANIP
  #include <iomanip>
#else
  #include <iomanip.h>
#endif

#include "CflowdFlowFilter.hh"
#include "CflowdConfig.hh"
#include "CflowdConfigLex.hh"
#include "CflowdRawFlowConverter.hh"
#include "CflowdPacketQueue.hh"
#include "CflowdVersion.hh"
#include "Signal.hh"
#include "DnsReceiver.hh"

typedef map<ipv4addr_t,DnsReceiver,less<ipv4addr_t> >  DnsReceiverMap_t;

extern int errno;

static const string          rcsid = "@(#) $Name: cflowd-2-1-b1 $ $Id: dnswatch.cc,v 1.5 2000/10/24 16:18:54 dwm Exp $";
static const CflowdVersion   g_cflowdVersion(rcsid);
static CflowdConfig          g_cflowdConfig;
static CflowdPacketQueue     g_packetQueue;
static DnsReceiverMap_t      g_dnsReceiverMap;
static string                g_outFilename = "";
static ostream              *g_outStream = (ostream *)0;

//-------------------------------------------------------------------------
//                      void Usage(const char *argv0)                      
//.........................................................................
//  
//-------------------------------------------------------------------------
void Usage(const char *argv0)
{
  cerr << "usage: " << argv0
       << " [-v] [-c configFile] [-f filterExpression] [-o outFile]"
       << endl;
  return;
}

//----------------------------------------------------------------------------
//     static void AddFlowToDnsReceiverMap(const CflowdRawFlow & flow) 
//............................................................................
//  Adds data from a flow to our DnsReceiverMap.
//----------------------------------------------------------------------------
static void AddFlowToDnsReceiverMap(const CflowdRawFlow & flow)
{
  //  Easy since g_dnsReceiverMap is an instance of a class which
  //  inherits from an STL map.
  g_dnsReceiverMap[flow.DstIpAddr()].AddFlow(flow);

  return;
}

//----------------------------------------------------------------------------
//               static void DumpDnsReceiverMap() 
//............................................................................
//  Dumps the contents of g_dnsReceiverMap to *g_outStream.  The format
//  is 1 line per DNS receiver (host), 4 columns per line 
//  (whitespace-delimited):
//  hostIpAddress  numFlows  numPkts  numBytes
//----------------------------------------------------------------------------
static void DumpDnsReceiverMap()
{
  DnsReceiverMap_t::const_iterator  dnsIter;
  struct in_addr                    inAddr;
  ostream                          *outStream = (ostream *)0;

  if (! g_outFilename.empty()) {
    outStream = new ofstream(g_outFilename.c_str());
    if ((! outStream) || (! (*outStream))) {
      cerr << "failed to open output file '" << g_outFilename.c_str() << "': "
           << strerror(errno) << " {" << __FILE__ << ":"
           << __LINE__ << "}" << endl;
      return;
    }
  }
  else {
     outStream = new ostdiostream(stdout);
  }
  
  outStream->seekp(0,ios::beg);
  
  //  Print a simple 2-line header.
  (*outStream) << setiosflags(ios::left)
               << setw(16) << "Destination Host" 
               << resetiosflags(ios::left)
               << setw(16) << "NumFlows" 
               << setw(16) << "NumPackets"
               << setw(16) << "NumBytes"
               << endl
               << "--------------------------------"
               << "--------------------------------" 
               << endl;
  
  //  Print the data for each DNS receiver.
  for (dnsIter = g_dnsReceiverMap.begin(); dnsIter != g_dnsReceiverMap.end();
       dnsIter++) {
    inAddr.s_addr = dnsIter->first;
    (*outStream) << setiosflags(ios::left)
                 << setw(16) << inet_ntoa(inAddr)
                 << resetiosflags(ios::left)
                 << setw(16) << dnsIter->second.NumFlows()
                 << setw(16) << dnsIter->second.NumPackets()
                 << setw(16) << dnsIter->second.NumBytes() << endl;
  }

  delete(outStream);
  
  return;
}

//----------------------------------------------------------------------------
//                  static void HandleSigTerm(int signum) 
//............................................................................
//  
//----------------------------------------------------------------------------
static void HandleSigTerm(int signum)
{
  return;
}

//----------------------------------------------------------------------------
//                  static void HandleSigUsr1(int signum) 
//............................................................................
//  
//----------------------------------------------------------------------------
static void HandleSigUsr1(int signum)
{
  return;
}

//-------------------------------------------------------------------------
//                   static void HandleSigHup(int signum)                  
//.........................................................................
//  
//-------------------------------------------------------------------------
static void HandleSigHup(int signum)
{
  return;
}

//-------------------------------------------------------------------------
//                     int main(int argc, char *argv[])                    
//.........................................................................
//  
//-------------------------------------------------------------------------
int main(int argc, char *argv[])
{
  struct sockaddr_un   servSockAddr;
  int                  servSockAddrLen;
  int                  sockFd;
  char                *configFilename = NULL;
  extern int           optind;
  extern char         *optarg;
  int                  optChar;
  CflowdFlowFilter     flowFilter;
  string               flowFilterString = "srcport == 53";
  Signal               sigTerm(SIGTERM);
  Signal               sigUsr1(SIGUSR1);
  Signal               sigHup(SIGHUP);
  
  while ((optChar = getopt(argc,argv,"c:f:o:v")) != EOF) {
    switch (optChar) {
      case 'c':
        configFilename = (char *)strdup(optarg);
        break;
      case 'f':
        flowFilterString = optarg;
        break;
      case 'o':
        g_outFilename = optarg;
        break;
      case 'v':
        cerr << g_cflowdVersion.Name() << "  " << g_cflowdVersion.Id() << endl;
        exit(0);
        break;
      case '?':
      default:
        Usage(argv[0]);
        exit(1);
        break;
    }
  }

  if (configFilename == NULL) {
    configFilename =
      (char *)strdup(CflowdConfig::k_defaultCflowdConfigFile.c_str());
  }
  
  //  Load the configuration file to get information about the
  //  packet queue.
  if (LoadConfigFile(configFilename,g_cflowdConfig) < 0) {
    cerr << "failed to load config file '" << configFilename << "': "
         << strerror(errno) << endl;
    Usage(argv[0]);
    exit(1);
  }

  //  Open the packet queue.
  if (g_packetQueue.Open(configFilename,g_cflowdConfig.PacketBufSize()) < 0) {
    cerr << "[A] failed to open packet queue!  Exiting." << endl;
    exit(1);
  }

  //  Compile the flow filter.
  if (flowFilter.Compile(flowFilterString.c_str()) < 0) {
    Usage(argv[0]);
    exit(1);
  }

  //  Install signal handlers.
  sigUsr1.InstallHandler(HandleSigUsr1);
  sigTerm.InstallHandler(HandleSigTerm);
  sigHup.InstallHandler(HandleSigHup);
  
  //  Wait on packet queue semaphore.
  g_packetQueue.GetLock();
  
  for (;;) {
    uint16_t      exportVersion;
    ipv4addr_t    ciscoAddr;
    const char   *pktPtr;
    int           pktNum, flowNum, flowArrNum;
    int           numPackets, numFlows;

    CflowdCiscoMap::iterator    ciscoMapIter;
    CiscoFlowHeaderV1_t        *hdrPtrV1;
    CiscoFlowEntryV1_t         *entryPtrV1;
    CiscoFlowHeaderV5_t        *hdrPtrV5;
    CiscoFlowEntryV5_t         *entryPtrV5;
    CiscoFlowHeaderV6_t        *hdrPtrV6;
    CiscoFlowEntryV6_t         *entryPtrV6;

    numPackets = g_packetQueue.NumPackets();
    CflowdRawFlow flows[numPackets * 30];
    flowArrNum = 0;

    for (pktNum = 0; pktNum < numPackets; pktNum++) {
      //  get reference to next packet.
      pktPtr = g_packetQueue.GetPacket(ciscoAddr);
      ciscoMapIter = g_cflowdConfig.CiscoMap().find(ciscoAddr);
      if (ciscoMapIter == g_cflowdConfig.CiscoMap().end()) {
        struct in_addr  addrIn;
        addrIn.s_addr = ciscoAddr;
        cerr << "[E] received packet from a source for which"
             << " we're not configured (" << inet_ntoa(addrIn) << ")" << endl;
        continue;
      }
    
      //  Get the flow-export version of the packet.
      exportVersion = ntohs(*(uint16_t *)pktPtr);
      
      //  Deal with the packet.
      switch (exportVersion) {
        case 1:
          hdrPtrV1 = (CiscoFlowHeaderV1_t *)(&pktPtr[0]);
          entryPtrV1 = (CiscoFlowEntryV1_t *)(&pktPtr[0] +
                                              sizeof(CiscoFlowHeaderV1_t));
          for (flowNum = 0; flowNum < ntohs(hdrPtrV1->count); flowNum++) {
            flows[flowArrNum].Init(ciscoAddr,hdrPtrV1,entryPtrV1);
            entryPtrV5++;
            flowArrNum++;
          }
          break;
        case 5:
          hdrPtrV5 = (CiscoFlowHeaderV5_t *)(&pktPtr[0]);
          entryPtrV5 = (CiscoFlowEntryV5_t *)(&pktPtr[0] +
                                              sizeof(CiscoFlowHeaderV5_t));
          for (flowNum = 0; flowNum < ntohs(hdrPtrV5->count); flowNum++) {
            flows[flowArrNum].Init(ciscoAddr,hdrPtrV5,entryPtrV5);
            entryPtrV5++;
            flowArrNum++;
          }
          break;
        case 6:
          hdrPtrV6 = (CiscoFlowHeaderV6_t *)(&pktPtr[0]);
          entryPtrV6 = (CiscoFlowEntryV6_t *)(&pktPtr[0] +
                                              sizeof(CiscoFlowHeaderV6_t));
          for (flowNum = 0; flowNum < ntohs(hdrPtrV6->count); flowNum++) {
            flows[flowArrNum].Init(ciscoAddr,hdrPtrV6,entryPtrV6);
            entryPtrV5++;
            flowArrNum++;
          }
          break;
        default:
          break;
      }
    }

    //  Release semaphore for packet queue shared memory.
    g_packetQueue.ReleaseLock();

    if (sigUsr1.Caught()) {
      sigUsr1.Catch();
      DumpDnsReceiverMap();
    }

    if (sigHup.Caught()) {
      sigHup.Catch();
      DumpDnsReceiverMap();
      g_dnsReceiverMap.clear();
    }

    if (sigTerm.Caught()) {
      DumpDnsReceiverMap();
      exit(0);
    }
      
    //  Walk the flows we received and add data to the DnsReceiverMap.
    numFlows = flowArrNum;
    for (flowNum = 0; flowNum < numFlows; flowNum++) {
      if (flowFilter.Matches(&flows[flowNum])) {
        AddFlowToDnsReceiverMap(flows[flowNum]);
      }
    }

    //  Wait on next buffer.
    g_packetQueue.ToggleBuffers(false);
  }
  
}