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 base64
  13 import logging
  14 from websocket import create_connection
  15
  16 __author__ = "Radovan Sajben"
  17 __copyright__ = "Copyright(c) 2016, Cisco Systems, Inc."
  18 __license__ = "Eclipse Public License v1.0"
  19 __email__ = "rsajben@cisco.com"
  20
  21
  22 def parse_arguments():
  23     """Use argparse to get arguments,
  24
  25     Returns:
  26         :return: args object
  27     """
  28     parser = argparse.ArgumentParser()
  29     parser.add_argument("--uri", default="ws://127.0.0.1:8185/", help="URI to connect")
  30     parser.add_argument(
  31         "--count", type=int, default=1, help="Number of messages to receive"
  32     )
  33     parser.add_argument(
  34         "--credentials",
  35         default="admin:admin",
  36         help="Basic authorization username:password",
  37     )
  38     parser.add_argument("--logfile", default="wsreceiver.log", help="Log file name")
  39     parser.add_argument(
  40         "--debug",
  41         dest="loglevel",
  42         action="store_const",
  43         const=logging.DEBUG,
  44         default=logging.INFO,
  45         help="Log level",
  46     )
  47     args = parser.parse_args()
  48     return args
  49
  50
  51 class WSReceiver(object):
  52     """Class which receives web socket messages."""
  53
  54     def __init__(self, uri, credentials):
  55         """Initialise and connect to URI.
  56
  57         Arguments:
  58             :uri: uri to connect to
  59             :credentials: user:password used in Basic Auth Header
  60         Returns:
  61             None
  62         """
  63         self.uri = uri
  64         self.credentials = credentials
  65         auth_string = base64.b64encode(credentials.encode("ascii"))
  66         self.headers = {"Authorization": "Basic " + auth_string.decode("ascii")}
  67
  68         logger.info("Connecting to: %s", self.uri)
  69         self.ws = create_connection(self.uri, header=self.headers)
  70
  71     def close(self):
  72         """Close the connection.
  73
  74         Arguments:
  75             None
  76         Returns:
  77             None
  78         """
  79         logger.info("Disconnecting from: %s", self.uri)
  80         self.ws.close()
  81
  82     def receive(self):
  83         """Receive a message.
  84
  85         Arguments:
  86             None
  87         Returns:
  88             :return: received data
  89         """
  90         data = self.ws.recv()
  91         logger.info("Data received:\n%s", data)
  92         return data
  93
  94
  95 if __name__ == "__main__":
  96     args = parse_arguments()
  97     logger = logging.getLogger("logger")
  98     log_formatter = logging.Formatter("%(asctime)s %(levelname)s: %(message)s")
  99     console_handler = logging.StreamHandler()
 100     file_handler = logging.FileHandler(args.logfile, mode="w")
 101     console_handler.setFormatter(log_formatter)
 102     file_handler.setFormatter(log_formatter)
 103     logger.addHandler(console_handler)
 104     logger.addHandler(file_handler)
 105     logger.setLevel(args.loglevel)
 106     receiver = WSReceiver(args.uri, args.credentials)
 107     remains = args.count
 108     logger.info("Expected %d message(s)", remains)
 109     while remains:
 110         logger.info("Waiting for a message ...")
 111         data = receiver.receive()
 112         remains -= 1
 113         logger.info("Remaining messages to receive: %d", remains)
 114     logger.info("Finished ...")
 115     receiver.close()