text: #-------------------------------------------------------------------------
  # Copyright (c) Microsoft Corporation. All rights reserved.
  # Licensed under the MIT License. See License.txt in the project root for
  # license information.
  #--------------------------------------------------------------------------
  
  import os
  import logging
  import sys
  from base64 import b64encode, b64decode
  from hashlib import sha256
  from hmac import HMAC
  from time import time
  try:
      from urllib import quote, quote_plus, urlencode #Py2
  except Exception:
      from urllib.parse import quote, quote_plus, urlencode
  
  import uamqp
  from uamqp import utils, errors
  from uamqp import authentication
  
  
  ### IotHub CLI Extension SAS token implementation
  
  class SasTokenAuthentication():
      """
      Shared Access Signature authorization for Azure IoT Hub.
      Args:
          uri (str): Uri of target resource.
          shared_access_policy_name (str): Name of shared access policy.
          shared_access_key (str): Shared access key.
          expiry (int): Expiry of the token to be generated. Input should
              be seconds since the epoch, in UTC. Default is an hour later from now.
      """
      def __init__(self, uri, shared_access_policy_name, shared_access_key, expiry=None):
          self.uri = uri
          self.policy = shared_access_policy_name
          self.key = shared_access_key
          if expiry is None:
              self.expiry = time() + 3600  # Default expiry is an hour later
          else:
              self.expiry = expiry
  
      def generate_sas_token(self):
          """
          Create a shared access signiture token as a string literal.
          Returns:
              result (str): SAS token as string literal.
          """
          encoded_uri = quote_plus(self.uri)
          ttl = int(self.expiry)
          sign_key = '%s\n%d' % (encoded_uri, ttl)
          signature = b64encode(HMAC(b64decode(self.key), sign_key.encode('utf-8'), sha256).digest())
          result = {
              'sr': self.uri,
              'sig': signature,
              'se': str(ttl)
          }
  
          if self.policy:
              result['skn'] = self.policy
  
          return 'SharedAccessSignature ' + urlencode(result)
  
  
  def get_logger(level):
      uamqp_logger = logging.getLogger("uamqp")
      if not uamqp_logger.handlers:
          handler = logging.StreamHandler(stream=sys.stdout)
          handler.setFormatter(logging.Formatter('%(asctime)s %(name)-12s %(levelname)-8s %(message)s'))
          uamqp_logger.addHandler(handler)
      uamqp_logger.setLevel(level)
      return uamqp_logger
  
  
  log = get_logger(logging.INFO)
  
  
  def _build_iothub_amqp_endpoint_from_target(target):
      hub_name = target['hostname'].split('.')[0]
      endpoint = "{}@sas.root.{}".format(target['key_name'], hub_name)
      endpoint = quote_plus(endpoint)
      sas_token = SasTokenAuthentication(target['hostname'], target['key_name'],
                                         target['access_key'], time() + 360).generate_sas_token()
      endpoint = endpoint + ":{}@{}".format(quote_plus(sas_token), target['hostname'])
      return endpoint
  
  
  def _receive_message(receive_client, target):
      try:
          batch = receive_client.receive_message_batch(max_batch_size=10)
      except errors.LinkRedirect as redirect:
          new_auth = authentication.SASLPlain(redirect.hostname, target['key_name'], target['access_key'])
          #new_auth = authentication.SASTokenAuth.from_shared_access_key(redirect.address.decode('utf-8'), target['key_name'], target['access_key'])
          receive_client.redirect(redirect, new_auth)
          batch = receive_client.receive_message_batch(max_batch_size=10)
  
      while batch:
          log.info("Got batch: {}".format(len(batch)))
          assert len(batch) <= 10
          for message in batch:
              annotations = message.annotations
              log.info("Sequence Number: {}".format(annotations.get(b'x-opt-sequence-number')))
          batch = receive_client.receive_message_batch(max_batch_size=10)
  
  
  def test_iothub_client_receive_sync(live_iothub_config):
      operation = '/messages/events/ConsumerGroups/{}/Partitions/{}'.format(
          live_iothub_config['consumer_group'],
          live_iothub_config['partition'])
      endpoint = _build_iothub_amqp_endpoint_from_target(live_iothub_config)
  
      source = 'amqps://' + endpoint + operation
      log.info("Source: {}".format(source))
      receive_client = uamqp.ReceiveClient(source, debug=False, timeout=5000, prefetch=50)
      try:
          log.info("Created client, receiving...")
          _receive_message(receive_client, live_iothub_config)
      except Exception as e:
          print(e)
          raise
      finally:
          receive_client.close()
      log.info("Finished receiving")
  
  
  if __name__ == '__main__':
      config = {}
      config['hostname'] = os.environ['IOTHUB_HOSTNAME']
      config['device'] = os.environ['IOTHUB_DEVICE']
      config['key_name'] = os.environ['IOTHUB_SAS_POLICY']
      config['access_key'] = os.environ['IOTHUB_SAS_KEY']
      config['consumer_group'] = "$Default"
      config['partition'] = "0"
  
      test_iothub_client_receive_sync(config)
  