/*
 * -------------
 * bwmeas
 *   version 0.9
 *   dated 20050906
 * -------------
 *
 * A program to measure the bandwidth used
 * on the network with the possibility of 
 * specifying a bpf filter to select a specific
 * type of traffic.
 *
 * Copyright (C) 2005 - SISMS vof
 *
 * This program is copyrighted software. 
 * It must be purchased if it is to be used in commercial endeavors, 
 * but is free for non-commercial use.
 */

#include <stdlib.h>  // for exit()
#include <signal.h>  // for signal()
#include <unistd.h>  // for getopt()
#include <assert.h>  // for assert()

extern "C" 
{ 
#include <pcap.h> 
}

#include <iostream>
#include <fstream>  // for ofstream
using namespace std;

// some default values
char *pcDevice         = "eth0"; // default capture device is first Ethernet interface
char *pcFilterString   =  NULL;   // no default bpf filter
char *pcOutputFilename = "/dev/stdout";  // show results on stdout
int   iSnapLength      =  1514;  // max Ethernet size (without FCS which is not captured anyway)
int   iPromiscuous     =  1;     // place interface in promiscuous mode
int   iTimeout         =  1000;  // (milliseconds)
int   uiAlarmTimeout   =  0;
double dTimeStep       =  1e6;   // (in microseconds !!)
int   iShowBandwidth   =  0;
char  acSeparator[2]   = { ' ', 0 };

// some global variables
pthread_t sThreadPacketRecorder;
pthread_t sThreadStatisticsDumper;
ofstream fout;

unsigned long ulPacketCount = 0;
unsigned long ulByteCount   = 0;
pthread_mutex_t mutex = PTHREAD_MUTEX_INITIALIZER;

void ProcessPacket(u_char *pucUserData, const struct pcap_pkthdr *psPacketHeader, const u_char *pucData)
{
  // enter the protection region 
  // (where we modify variables that are also used by the other thread)
  int iResult = pthread_mutex_lock(&mutex);
  // assert(iResult);

  // process this packet
  ulPacketCount++;
  ulByteCount += psPacketHeader->len;
  // (note: we pass "->len" instead of "->caplen" since we 
  //  want to take into account the length of the full packet on
  //  the wire; we don't care about the fact that we didn't
  //  capture the complete length of the packet)

  // leave the protection region 
  iResult = pthread_mutex_unlock(&mutex);
  // assert(iResult);
}

void *threadPacketRecorder(void *pvArgument)
{
  // start the capture
  int  iResult;
  char acErrorMessage [PCAP_ERRBUF_SIZE] ={0};

  // select the device to sniff on
  if (pcDevice == NULL)
  { // find a valid device to sniff on
    pcDevice = pcap_lookupdev (acErrorMessage);
    assert(pcDevice != NULL);
  }

  // get the network address and the network mask for this device
  bpf_u_int32 uiNetworkAddress;  
  bpf_u_int32 uiNetworkMask;  
  iResult = pcap_lookupnet (pcDevice, &uiNetworkAddress, &uiNetworkMask, acErrorMessage);
  assert(iResult>=0);

  // open the device for sniffing
  pcap_t *psDescriptor = pcap_open_live (pcDevice, iSnapLength, iPromiscuous, iTimeout, acErrorMessage);
  assert(psDescriptor != NULL);

  // parse the capture filter
  struct bpf_program sFilterProgram;
  int                iFlagOptimize  =1;
  iResult = pcap_compile (psDescriptor, &sFilterProgram, pcFilterString, iFlagOptimize, uiNetworkMask);
  assert(iResult>=0);
  iResult = pcap_setfilter (psDescriptor, &sFilterProgram);
  assert(iResult>=0);

  // start the capture
  pcap_handler   printer      = ProcessPacket;
  unsigned char *pucUserData  = (unsigned char *) NULL;
  iResult = pcap_loop (psDescriptor, 0, printer, pucUserData);
  assert(iResult>=0);

  // ready
  return NULL;
}

void *threadStatisticsDumper(void *pvArgument)
{
  // get the current time
  struct timeval sTimeInit;  
  int iResult = gettimeofday(&sTimeInit, NULL);
  assert(iResult==0);
  double dTimeNextDeadline = sTimeInit.tv_sec * 1e6 + sTimeInit.tv_usec;

  while (1)
  { 
    // get the current time
    struct timeval sTimeNow;  
    int iResult = gettimeofday(&sTimeNow, NULL);
    assert(iResult==0);

    // compute current time in usec
    double dCurrentTime = sTimeNow.tv_sec * 1e6 + sTimeNow.tv_usec; 

    // enter the protection region 
    // (where we modify variables that are also used by the other thread)
    iResult = pthread_mutex_lock(&mutex);
    // assert(iResult);

    // copy the counters and reset them
    unsigned long ulCopyOfPacketCount = ulPacketCount;
    unsigned long ulCopyOfByteCount   = ulByteCount;
    ulPacketCount = 0;
    ulByteCount   = 0;

    // leave the protection region 
    iResult = pthread_mutex_unlock(&mutex);
    // assert(iResult);
  
    // write the results to the output file
    fout   << dCurrentTime
    << acSeparator << ulCopyOfPacketCount
    << acSeparator << ulCopyOfByteCount;

    // compute and print the bandwidth (optional)
    if (iShowBandwidth)
    {
      double dBandwidth = ulCopyOfByteCount * 8e6 / dTimeStep;  // (timestep is in us)
      fout << acSeparator << dBandwidth << acSeparator;
      
      if      (dBandwidth < 1e3)
      { fout << "(" << dBandwidth << "bps)"; }
      else if (dBandwidth < 1e6)
      { fout << "(" << dBandwidth/1e3 << "kbps)"; }
      else
      { fout << "(" << dBandwidth/1e6 << "Mbps)"; }
    }

    // end the line on the output file
    fout << endl;

    // compute the next deadline
    dTimeNextDeadline += dTimeStep;

    // get the current time again
    iResult = gettimeofday(&sTimeNow, NULL);
    assert(iResult==0);

    // compute the time to wait until the next deadline
    double dTimeToSleep = dTimeNextDeadline - (sTimeNow.tv_sec * 1e6 + sTimeNow.tv_usec);

    // sleep until the next deadline
    if (dTimeToSleep>0)
    { usleep((unsigned int) dTimeToSleep); }
  }
}

void terminateHandler(int iSignal)
{
  // ignore new signals of this type
  signal(SIGINT,  SIG_IGN);      
  signal(SIGTERM, SIG_IGN);      
  signal(SIGALRM, SIG_IGN);      

  // cancel the recorder and statistics threads
  int iResult;
  iResult = pthread_cancel(sThreadPacketRecorder);
  assert(iResult==0);
  iResult = pthread_cancel(sThreadStatisticsDumper);
  assert(iResult==0);

  // ready
  return;
}

int main(int argc, char **argv)
{
  // parse the command-line options
  int c;
  extern char *optarg;
  while	((c = getopt(argc, argv, "i:f:o:s:w:bch")) != -1)
  {
    switch (c) 
    {
      case 'i':
      { pcDevice = optarg;
      } break;
      case 'f':
      { pcFilterString = optarg;
      } break;
      case 'o':
      { pcOutputFilename = optarg;
      } break;
      case 's':
      { dTimeStep = atof(optarg) * 1e6;
      } break;
      case 'w':
      { uiAlarmTimeout = atoi(optarg);
      } break;
      case 'b':
      { iShowBandwidth = 1;
      } break;
      case 'c':
      { acSeparator[0] = ',';
      } break;
      default:
      { cout << "Usage: " << argv[0] << " [options]" << endl;
        cout << "Options:" << endl;
        cout << "  -i device          (default is eth0)" << endl;
        cout << "  -f filter-string   (default is none)" << endl;
        cout << "  -o output-file     (default is stdout)" << endl;
        cout << "  -s time-step       (in seconds, default is 1s)" << endl;
        cout << "  -w duration        (in seconds, default is to never stop)" << endl;
        cout << "  -b                 (default is not to show the bandwidth)" << endl;
        cout << "  -c                 (default is not to use CVS format)" << endl;
        exit(0);
      } break;
    }
  }

  // install signal handlers
  sighandler_t pResult;
  pResult = signal(SIGINT,  terminateHandler);
  assert(pResult!=SIG_ERR);
  pResult = signal(SIGTERM, terminateHandler);
  assert(pResult!=SIG_ERR);
  pResult = signal(SIGALRM, terminateHandler);
  assert(pResult!=SIG_ERR);
  
  // set an alarm to stop us after a user-defined timeout
  if (uiAlarmTimeout >0)
  { alarm(uiAlarmTimeout); }

  // open the output stream
  fout.open(pcOutputFilename, ios::out);
  assert(! fout.bad());

  // set precision (important !)
  // (so we store the timestamps up to ms precision)
  fout.precision(14);

  // set the attributes for the threads
  pthread_attr_t sThreadAttr;
  int iResult;
  iResult = pthread_attr_init(&sThreadAttr);
  assert(iResult==0);
  iResult = pthread_attr_setschedpolicy(&sThreadAttr, SCHED_FIFO); 
  assert(iResult==0);

  // initialize the mutex
  iResult = pthread_mutex_init(&mutex, NULL);
  // assert(iResult);

  // start the recording thread
  sched_param sSchedParam;
  sSchedParam.sched_priority = 50;
  iResult = pthread_attr_setschedparam(&sThreadAttr, &sSchedParam);
  assert(iResult==0);
  iResult = pthread_create(&sThreadPacketRecorder, &sThreadAttr, threadPacketRecorder, NULL);
  assert(iResult==0);

  // start the statistics thread
  sSchedParam.sched_priority = 51;
  iResult = pthread_attr_setschedparam(&sThreadAttr, &sSchedParam);
  assert(iResult==0);
  iResult = pthread_create(&sThreadStatisticsDumper, &sThreadAttr, threadStatisticsDumper, NULL);
  assert(iResult==0);

  // cleanup
  iResult = pthread_attr_destroy(&sThreadAttr);
  assert(iResult==0);

  // wait for the threads to finish
  iResult = pthread_join(sThreadPacketRecorder, NULL);
  assert(iResult==0);
  iResult = pthread_join(sThreadStatisticsDumper, NULL);
  assert(iResult==0);

  // the capture was stopped, clean up 
  fout.close();

  iResult = pthread_mutex_destroy(&mutex);
  // assert(iResult);

  // the end
  return 0;
}