UDP TCP Echo

udptcpecho.py

#! /usr/bin/python

from __future__ import absolute_import
import sys, threading, argparse, pprint, logging, logging.handlers, os, __main__, platform
from socket import *

# Setup format for pprint.
pp = pprint.PrettyPrinter(indent=4)

BUFSIZE = 1024
args = None
logger = None

def main():
    global args
    global logger

    config_setup()
    logging_setup()

    if args.listen:
      server()
    else:
      client()

def config_setup():
  global args

  # Define command line arguments
  parser = argparse.ArgumentParser(prog=u'PROG')
  parser.add_argument(u'--verbose', u'-v', action=u'count', help=u'verbose multi level', default=1)

  group1 = parser.add_mutually_exclusive_group(required=True)
  group1.add_argument(u'--tcp', u'-t', action=u'store_true', help=u'use TCP')
  group1.add_argument(u'--udp', u'-u', action=u'store_true', help=u'use UDP')

  group2 = parser.add_mutually_exclusive_group(required=True)
  group2.add_argument(u'--client', u'-c', default=u'localhost', help=u'Client Address to connect to')
  group2.add_argument(u'--listen', u'-l', action=u'store_true', help=u'run as listening server')

  parser.add_argument(u'--port', u'-p', help=u'Listen or Destination port', type=int, default=u"50000")

  # Read in and parse the command line arguments
  args = parser.parse_args()

def logging_setup():
    global logger

    fn = os.path.splitext(os.path.basename(__main__.__file__))[0]

    # Setup display and file logging with level support.
    logFormatter = logging.Formatter(u"%(asctime)s [%(threadName)-12.12s] [%(levelname)-5.5s]  %(message)s")
    logger = logging.getLogger()
    logfn = "{0}/{1}.log".format('/var/log/'+ fn +'/', fn)
    if not os.path.isfile(logfn)    :
        logfn = u"{0}.log".format(fn)
    fileHandler = logging.handlers.RotatingFileHandler(logfn, maxBytes=2*1024*1024, backupCount=2)

    fileHandler.setFormatter(logFormatter)
    #fileHandler.setLevel(logging.DEBUG)
    logger.addHandler(fileHandler)

    consoleHandler = logging.StreamHandler()
    #consoleHandler.setLevel(logging.DEBUG)
    consoleHandler.setFormatter(logFormatter)
    logger.addHandler(consoleHandler)

    # Dictionary to translate Count of -v's to logging level
    verb = { 0 : logging.WARN,
             1 : logging.INFO,
             2 : logging.DEBUG,
           }

    # zero adjust for zero offset to make it easier to understand
    args.verbose = int(args.verbose) - 1
    try:
      # set logging level from command line arg.
      logger.setLevel(verb[(args.verbose)])
    except:
      # if out of range use levels, and account for standard levels
      if (args.verbose - 2) > 0:
        logger.setLevel(logging.DEBUG - (args.verbose - 2))
      else:
        logger.setLevel(logging.DEBUG)

    logger.info(u'Starting script ' + os.path.join(os.path.dirname(os.path.realpath(__file__)), __file__))

    # log which levels of debug are enabled.
    logger.log(logging.DEBUG-9, "discrete log level = " + str(logging.DEBUG-9))
    logger.log(logging.DEBUG-8, "discrete log level = " + str(logging.DEBUG-8))
    logger.log(logging.DEBUG-7, "discrete log level = " + str(logging.DEBUG-7))
    logger.log(logging.DEBUG-6, "discrete log level = " + str(logging.DEBUG-6))
    logger.log(logging.DEBUG-5, "discrete log level = " + str(logging.DEBUG-5))
    logger.log(logging.DEBUG-4, "discrete log level = " + str(logging.DEBUG-4))
    logger.log(logging.DEBUG-3, "discrete log level = " + str(logging.DEBUG-3))
    logger.log(logging.DEBUG-2, "discrete log level = " + str(logging.DEBUG-2))
    logger.log(logging.DEBUG-1, "discrete log level = " + str(logging.DEBUG-1))
    logger.log(logging.DEBUG,   "discrete log level = " + str(logging.DEBUG  ))
    logger.info(u'verbose = ' + str(args.verbose) + ", logger level = " + str(logger.getEffectiveLevel()))
    logger.debug(u'debug level enabled')
    logger.info(u'info  level enabled')
    #logger.warn(u'warn  level enabled')
    #logger.error(u'error  level enabled')
    #logger.critical(u'critical  level enabled')

    logger.info(u'pytho version ' + platform.python_version())

def handle_client_connection(client_socket):
    request = client_socket.recv(BUFSIZE)
    logger.debug(u'Received {}'.format(request))
    client_socket.send(str('ACK!').encode('UTF-8'))
    client_socket.close()

def server():
    global args
    global logger
    logger.debug(u"listen mode")
    logger.debug(args.tcp)
    logger.debug(args.udp)
    if (args.tcp and (not args.udp)):
      logger.debug(u"tcp mode")

      tcp = socket(AF_INET, SOCK_STREAM)
      tcp.bind((u'', args.port))
      tcp.listen(5)  # max backlog of connections

      while True:
          client_sock, address = tcp.accept()
          logger.debug(u'Accepted connection from {}:{}'.format(address[0], address[1]))
          client_handler = threading.Thread(
              target=handle_client_connection,
              args=(client_sock,)  # without comma you'd get a... TypeError: handle_client_connection() argument after * must be a sequence, not _socketobject
          )
          client_handler.start()

      logger.debug(u'Listening on 0.0.0.0:{}'.format(args.port))

    elif ((not args.tcp) and args.udp):
      logger.debug(u"udp mode")
      udp = socket(AF_INET, SOCK_DGRAM)
      udp.bind((u'', args.port))

      while 1:
          data, addr = udp.recvfrom(BUFSIZE)
          logger.debug(u'received {} bytes from {}'.format(len(data), addr))
          if data:
            udp.sendto(data, addr)
            logger.debug(u'sent {} bytes back to {}'.format(data, addr))

def client():
    global args
    global logger

    addr = (args.client, args.port)
    udp = socket(AF_INET, SOCK_DGRAM)
    logger.debug(u'udp echo client ready, reading stdin')
    while 1:
        line = sys.stdin.readline()
        if not line:
            break
        sent = udp.sendto(str(line).encode('UTF-8'), addr)
        logger.debug(u'sent msg')
        logger.debug(u'string = ' + line)
        logger.debug(str(line).encode('UTF-8'))

        data, server = udp.recvfrom(1024)
        logger.debug(u'received {!r}'.format(data))
        #quit()

main()