tools/wstools/wsreceiver.py

   1 """WebSocket data receiver.
   2
   3 The tool receives and logs data from specified URI"""
   4
   5 # Copyright (c) 2016 Cisco Systems, Inc. and others.  All rights reserved.
   6 #
   7 # This program and the accompanying materials are made available under the
   8 # terms of the Eclipse Public License v1.0 which accompanies this distribution,
   9 # and is available at http://www.eclipse.org/legal/epl-v10.html
  10
  11 import argparse
  12 import logging
  13 from websocket import create_connection
  14
  15 __author__ = "Radovan Sajben"
  16 __copyright__ = "Copyright(c) 2016, Cisco Systems, Inc."
  17 __license__ = "Eclipse Public License v1.0"
  18 __email__ = "rsajben@cisco.com"
  19
  20
  21 def parse_arguments():
  22     """Use argparse to get arguments,
  23
  24     Returns:
  25         :return: args object
  26     """
  27     parser = argparse.ArgumentParser()
  28     parser.add_argument("--uri", default="ws://127.0.0.1:8185/", help="URI to connect")
  29     parser.add_argument("--count", type=int, default=1, help="Number of messages to receive")
  30     parser.add_argument("--logfile", default="wsreceiver.log", help="Log file name")
  31     parser.add_argument("--debug", dest="loglevel", action="store_const",
  32                         const=logging.DEBUG, default=logging.INFO, help="Log level")
  33     args = parser.parse_args()
  34     return args
  35
  36
  37 class WSReceiver(object):
  38     """Class which receives web socket messages."""
  39
  40     def __init__(self, uri):
  41         """Initialise and connect to URI.
  42
  43         Arguments:
  44             :uri: uri to connect to
  45         Returns:
  46             None
  47         """
  48         self.uri = uri
  49         logger.info("Connecting to: %s", self.uri)
  50         self.ws = create_connection(self.uri)
  51
  52     def close(self):
  53         """Close the connection.
  54
  55         Arguments:
  56             None
  57         Returns:
  58             None
  59         """
  60         logger.info("Disconnecting from: %s", self.uri)
  61         self.ws.close()
  62
  63     def receive(self):
  64         """Receive a message.
  65
  66         Arguments:
  67             None
  68         Returns:
  69             :return: received data
  70         """
  71         data = self.ws.recv()
  72         logger.info("Data received:\n%s", data)
  73         return data
  74
  75 if __name__ == "__main__":
  76     args = parse_arguments()
  77     logger = logging.getLogger("logger")
  78     log_formatter = logging.Formatter("%(asctime)s %(levelname)s: %(message)s")
  79     console_handler = logging.StreamHandler()
  80     file_handler = logging.FileHandler(args.logfile, mode="w")
  81     console_handler.setFormatter(log_formatter)
  82     file_handler.setFormatter(log_formatter)
  83     logger.addHandler(console_handler)
  84     logger.addHandler(file_handler)
  85     logger.setLevel(args.loglevel)
  86     receiver = WSReceiver(args.uri)
  87     remains = args.count
  88     logger.info("Expected %d message(s)", remains)
  89     while remains:
  90         logger.info("Waiting for a message ...")
  91         data = receiver.receive()
  92         remains -= 1
  93         logger.info("Remaining messages to receive: %d", remains)
  94     logger.info("Finished ...")
  95     receiver.close()