#!/usr/bin/env python3

import os, re, sys, time
import argparse
import datetime
import queue
import socket
from threading import Thread

scriptName  = os.path.basename( __file__ )

''' udp_cmr.py [-h] [-csv] [-run_seconds SECONDS] [-quiet] UDP_PORT

   This script is used to listen on a UDP port and check whether CMR
   records for a single epoch are properly being sent in a single UDP
   packet.

   The script is terminated either by entering Ctrl-C from a console or
   by specifying a maximum run time. Once terminated, a string of the
   following format is always written.

        NPackets: 5 AvgPerSec: 0.987

   The following option parameters modify behavior:

     -quiet
        Disable logging of CVS data to stderr.

     -csv
        Log data to a CSV file named "udp_cmr-UDPPORT-YYYYMMDD-HHMMSS.csv"
        (for example, "udp_cmr-20001-20220510-013723.csv") containing lines
        as follows.
 
           2022/05/10 00:11:57.001,10.1.107.37,1034,20001,326

        Fields:
           date/time,
           source IP address (receiver)
           source port number (receiver)
           destimation port number (host)
           UDP packet size in bytes

     -run_seconds N_SECONDS
        Run for N_SECONDS then terminate. This might be used for automated
        testing where the only information of interest is the packet count
        and average per second values printed upon termination. An average
        count per second near 1.0 would be a check for correct behavior.

'''

# -------------------------------------------------- monitorUdp()
def monitorUdp( port, q, runSeconds=None ):
  ''' Listen for UDP packets on a port. Put info about packets into the
      queue, including UTC time, src ip, src port, and packet length.
      Optionally terminate after runSeconds when specified.
  '''
  sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
  sock.bind( ("0.0.0.0", port) )
  sock.settimeout(0.2)

  nPacketsReceived = 0

  tStart = datetime.datetime.utcnow()

  exceptionOccurred = False
  timeLimitReached  = False

  while not exceptionOccurred and not timeLimitReached:
    try:
      data, addr = sock.recvfrom(2048)
      nPacketsReceived += 1

      utc = datetime.datetime.utcnow()

      # Put the UDP packet data into the queue as a list.
      #sys.stderr.write("q.put()\n")
      q.put( [utc, addr[0], addr[1], port, len(data)] )
    except socket.timeout:
      pass
    except:
      exceptionOccurred = True
      sys.stderr.write("Caught exception in monitorUdp()\n")

    if not None == runSeconds:
      if (datetime.datetime.utcnow() - tStart).total_seconds() >= runSeconds:
        timeLimitReached = True

  tStop = datetime.datetime.utcnow()
  nPerSec = nPacketsReceived / (tStop - tStart).total_seconds()
  print("NPackets: %d AvgPerSec: %.3f" % ( nPacketsReceived, nPerSec ))

# -------------------------------------------------- queueProcessor()
class queueProcessor():
  ''' Get UDP packet data sent by monitorUdp() through a queue and log
      info stdout and/or a CSV file.
  '''
  def __init__(self, q, name, udpPort, verbose=True, doLogging=False):
    self.q         = q
    self.name      = name
    self.doLogging = doLogging
    self.verbose   = verbose

    self.logFileName = None
    if doLogging:
      utc = datetime.datetime.utcnow()
      fnbase = name.replace( " ", "_" )
      self.logFileName = "%s-%d-%s.csv" % ( fnbase,
                                            udpPort,
                                            utc.strftime("%Y%m%d-%H%M%S") )

    self.enabled = False
    self.running = False
    self.t    = None

  def processor(self):
    self.running = True
    while self.enabled:
      while not self.q.empty():
        # Read item from queue.
        item = q.get()
        #sys.stderr.write("q.get()\n")

        # Unpacket items from list.
        utc     = item[0]
        srcIp   = item[1]
        srcPort = item[2]
        dstPort = item[3]
        dataLen = item[4]

        if self.doLogging or self.verbose:
          # Format UTC time to the millisecond by stripping last 3 chars.
          utcStr = utc.strftime("%Y/%m/%d %H:%M:%S.%f")[:-3]
          logStr = ( "%s,%s,%d,%d,%d" % (
                     utcStr, srcIp, srcPort, dstPort, dataLen ) )

          if self.doLogging:
            with open( self.logFileName, "a" ) as f:
              f.write("%s\n" % ( logStr ))
          if self.verbose:
            sys.stderr.write("%s\n" % ( logStr ))
      time.sleep(1.0)
    self.running = False

  def start(self):
    self.enabled = True

    self.t = Thread( name = self.name, target = self.processor )
    #sys.stderr.write("Start: %s\n" % self.name)
    self.t.start()

  def stop(self):
    self.enabled = False
    sys.stderr.write("Stoping %s ..." % self.name)
    sys.stderr.flush()
    while self.running:
      time.sleep(0.5)
    sys.stderr.write(" stopped\n")

# -------------------------------------------------- main()
parser = argparse.ArgumentParser( scriptName )

parser.add_argument( 'UDP_PORT', type=int,
                     help='UDP destination port at this host' )
parser.add_argument( '-csv', action="store_true", default=False,
                     help='Log all packet times and length to CSV file' )
parser.add_argument( '-run_seconds', type=int, default=None,
                     help='Maximum script runtime in seconds' )
parser.add_argument( '-quiet', action="store_true", default=False,
                     help='Disable packet logging on stderr' )

args = parser.parse_args( sys.argv[1:] )

udpPort    = args.UDP_PORT
verbose    = not args.quiet
doLogging  = args.csv
runSeconds = args.run_seconds

#
# Create a queue and start the logger thread.
#
q = queue.Queue()
logger = queueProcessor( q, scriptName.replace(".py", ""),
                         udpPort,
                         verbose=verbose, doLogging=doLogging )
logger.start()

#
# Monitor UDP packets passing data through the queue to the logger thread.
#
sys.stderr.write("Terminate logging using Ctrl-C\n")
monitorUdp( udpPort, q, runSeconds )

#
# Stop the logger thread.
#
logger.stop()