# ------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# -------------------------------------------------------------------------

import datetime
import uuid
import functools
import requests
try:
    from urlparse import urlparse
except ImportError:
    from urllib.parse import urlparse

from uamqp import types

from azure.servicebus.common import mgmt_handlers, mixins
from azure.servicebus.common.constants import (
    ReceiveSettleMode,
    REQUEST_RESPONSE_PEEK_OPERATION,
    REQUEST_RESPONSE_RECEIVE_BY_SEQUENCE_NUMBER,
    REQUEST_RESPONSE_UPDATE_DISPOSTION_OPERATION,
    REQUEST_RESPONSE_GET_MESSAGE_SESSIONS_OPERATION)
from azure.servicebus.common.message import Message
from azure.servicebus.common.utils import parse_conn_str, build_uri
from azure.servicebus.common.errors import ServiceBusConnectionError, ServiceBusResourceNotFound
from azure.servicebus.control_client import ServiceBusService, SERVICE_BUS_HOST_BASE, DEFAULT_HTTP_TIMEOUT
from azure.servicebus.control_client.models import AzureServiceBusResourceNotFound
from azure.servicebus.send_handler import Sender, SessionSender
from azure.servicebus.receive_handler import Receiver, SessionReceiver
from azure.servicebus.base_handler import BaseHandler

class ServiceBusClient(mixins.ServiceBusMixin):
    """A Service Bus client for a namespace with the specified SAS authentication settings.

    :param str service_namespace: Service Bus namespace, required for all operations.
    :param str host_base: Optional. Live host base URL. Defaults to Public Azure.
    :param str shared_access_key_name: SAS authentication key name.
    :param str shared_access_key_value: SAS authentication key value.
    :param int http_request_timeout: Optional. Timeout for the HTTP request, in seconds.
     Default value is 65 seconds.
    :param http_request_session: Optional. Session object to use for HTTP requests.
    :type http_request_session: ~requests.Session
    :param bool debug: Whether to output AMQP network trace to the logger.

    Example:
        .. literalinclude:: ../examples/test_examples.py
            :start-after: [START create_servicebus_client]
            :end-before: [END create_servicebus_client]
            :language: python
            :dedent: 4
            :caption: Create a new instance of the Service Bus client

    """

    def __init__(self, service_namespace=None, host_base=SERVICE_BUS_HOST_BASE,
                 shared_access_key_name=None, shared_access_key_value=None,
                 http_request_timeout=DEFAULT_HTTP_TIMEOUT, http_request_session=None, debug=False):

        self.service_namespace = service_namespace
        self.host_base = host_base
        self.shared_access_key_name = shared_access_key_name
        self.shared_access_key_value = shared_access_key_value
        self.debug = debug
        self.mgmt_client = ServiceBusService(
            service_namespace=service_namespace,
            host_base=host_base,
            shared_access_key_name=shared_access_key_name,
            shared_access_key_value=shared_access_key_value,
            timeout=http_request_timeout,
            request_session=http_request_session)

    @classmethod
    def from_connection_string(cls, conn_str, **kwargs):
        """Create a Service Bus client from a connection string.

        :param conn_str: The connection string.
        :type conn_str: str

        Example:
            .. literalinclude:: ../examples/test_examples.py
                :start-after: [START create_servicebus_client_connstr]
                :end-before: [END create_servicebus_client_connstr]
                :language: python
                :dedent: 4
                :caption: Create a ServiceBusClient via a connection string.

        """
        address, policy, key, _ = parse_conn_str(conn_str)
        parsed_namespace = urlparse(address)
        namespace, _, base = parsed_namespace.hostname.partition('.')
        return cls(
            namespace,
            shared_access_key_name=policy,
            shared_access_key_value=key,
            host_base='.' + base,
            **kwargs)

    def _get_host(self):
        return "sb://" + self.service_namespace + self.host_base

    def get_queue(self, queue_name):
        """Get a client for a queue entity.

        :param queue_name: The name of the queue.
        :type queue_name: str
        :rtype: ~azure.servicebus.servicebus_client.QueueClient
        :raises: ~azure.servicebus.common.errors.ServiceBusConnectionError if the namespace is not found.
        :raises: ~azure.servicebus.common.errors.ServiceBusResourceNotFound if the queue is not found.

        Example:
            .. literalinclude:: ../examples/test_examples.py
                :start-after: [START get_queue_client]
                :end-before: [END get_queue_client]
                :language: python
                :dedent: 8
                :caption: Get the specific queue client from Service Bus client

        """
        try:
            queue = self.mgmt_client.get_queue(queue_name)
        except requests.exceptions.ConnectionError as e:
            raise ServiceBusConnectionError("Namespace: {} not found".format(self.service_namespace), e)
        except AzureServiceBusResourceNotFound:
            raise ServiceBusResourceNotFound("Specificed queue does not exist.")
        return QueueClient.from_entity(
            self._get_host(), queue,
            shared_access_key_name=self.shared_access_key_name,
            shared_access_key_value=self.shared_access_key_value,
            mgmt_client=self.mgmt_client,
            debug=self.debug)

    def list_queues(self):
        """Get clients for all queue entities in the namespace.

        :rtype: list[~azure.servicebus.servicebus_client.QueueClient]
        :raises: ~azure.servicebus.common.errors.ServiceBusConnectionError if the namespace is not found.

        Example:
            .. literalinclude:: ../examples/test_examples.py
                :start-after: [START list_queues]
                :end-before: [END list_queues]
                :language: python
                :dedent: 4
                :caption: List the queues from Service Bus client

        """
        try:
            queues = self.mgmt_client.list_queues()
        except requests.exceptions.ConnectionError as e:
            raise ServiceBusConnectionError("Namespace: {} not found".format(self.service_namespace), e)
        queue_clients = []
        for queue in queues:
            queue_clients.append(QueueClient.from_entity(
                self._get_host(), queue,
                shared_access_key_name=self.shared_access_key_name,
                shared_access_key_value=self.shared_access_key_value,
                mgmt_client=self.mgmt_client,
                debug=self.debug))
        return queue_clients

    def get_topic(self, topic_name):
        """Get a client for a topic entity.

        :param topic_name: The name of the topic.
        :type topic_name: str
        :rtype: ~azure.servicebus.servicebus_client.TopicClient
        :raises: ~azure.servicebus.common.errors.ServiceBusConnectionError if the namespace is not found.
        :raises: ~azure.servicebus.common.errors.ServiceBusResourceNotFound if the topic is not found.

        Example:
            .. literalinclude:: ../examples/test_examples.py
                :start-after: [START get_topic_client]
                :end-before: [END get_topic_client]
                :language: python
                :dedent: 8
                :caption: Get the specific topic client from Service Bus client

        """
        try:
            topic = self.mgmt_client.get_topic(topic_name)
        except requests.exceptions.ConnectionError as e:
            raise ServiceBusConnectionError("Namespace: {} not found".format(self.service_namespace), e)
        except AzureServiceBusResourceNotFound:
            raise ServiceBusResourceNotFound("Specificed topic does not exist.")
        return TopicClient.from_entity(
            self._get_host(), topic,
            shared_access_key_name=self.shared_access_key_name,
            shared_access_key_value=self.shared_access_key_value,
            debug=self.debug)

    def list_topics(self):
        """Get a client for all topic entities in the namespace.

        :rtype: list[~azure.servicebus.servicebus_client.TopicClient]
        :raises: ~azure.servicebus.common.errors.ServiceBusConnectionError if the namespace is not found.

        Example:
            .. literalinclude:: ../examples/test_examples.py
                :start-after: [START list_topics]
                :end-before: [END list_topics]
                :language: python
                :dedent: 4
                :caption: List the topics from Service Bus client

        """
        try:
            topics = self.mgmt_client.list_topics()
        except requests.exceptions.ConnectionError as e:
            raise ServiceBusConnectionError("Namespace: {} not found".format(self.service_namespace), e)
        topic_clients = []
        for topic in topics:
            topic_clients.append(TopicClient.from_entity(
                self._get_host(), topic,
                shared_access_key_name=self.shared_access_key_name,
                shared_access_key_value=self.shared_access_key_value,
                debug=self.debug))
        return topic_clients

    def get_subscription(self, topic_name, subscription_name):
        """Get a client for a subscription entity.

        :param topic_name: The name of the topic.
        :type topic_name: str
        :param subscription_name: The name of the subscription.
        :type subscription_name: str
        :rtype: ~azure.servicebus.servicebus_client.SubscriptionClient
        :raises: ~azure.servicebus.common.errors.ServiceBusConnectionError if the namespace is not found.
        :raises: ~azure.servicebus.common.errors.ServiceBusResourceNotFound if the subscription is not found.

        Example:
            .. literalinclude:: ../examples/test_examples.py
                :start-after: [START get_subscription_client]
                :end-before: [END get_subscription_client]
                :language: python
                :dedent: 8
                :caption: Get the specific subscription client from Service Bus client

        """
        try:
            subscription = self.mgmt_client.get_subscription(topic_name, subscription_name)
        except requests.exceptions.ConnectionError as e:
            raise ServiceBusConnectionError("Namespace: {} not found".format(self.service_namespace), e)
        except AzureServiceBusResourceNotFound:
            raise ServiceBusResourceNotFound("Specificed subscription does not exist.")
        return SubscriptionClient.from_entity(
            self._get_host(), topic_name, subscription,
            shared_access_key_name=self.shared_access_key_name,
            shared_access_key_value=self.shared_access_key_value,
            debug=self.debug)

    def list_subscriptions(self, topic_name):
        """Get a client for all subscription entities in the topic.

        :param topic_name: The topic to list subscriptions for.
        :type topic_name: str
        :rtype: list[~azure.servicebus.servicebus_client.SubscriptionClient]
        :raises: ~azure.servicebus.common.errors.ServiceBusConnectionError if the namespace is not found.
        :raises: ~azure.servicebus.common.errors.ServiceBusResourceNotFound if the topic is not found.

        Example:
            .. literalinclude:: ../examples/test_examples.py
                :start-after: [START list_subscriptions]
                :end-before: [END list_subscriptions]
                :language: python
                :dedent: 4
                :caption: List the subscriptions from Service Bus client

        """
        try:
            subs = self.mgmt_client.list_subscriptions(topic_name)
        except requests.exceptions.ConnectionError as e:
            raise ServiceBusConnectionError("Namespace: {} not found".format(self.service_namespace), e)
        except AzureServiceBusResourceNotFound:
            raise ServiceBusResourceNotFound("Specificed topic does not exist.")
        sub_clients = []
        for sub in subs:
            sub_clients.append(SubscriptionClient.from_entity(
                self._get_host(), topic_name, sub,
                shared_access_key_name=self.shared_access_key_name,
                shared_access_key_value=self.shared_access_key_value,
                debug=self.debug))
        return sub_clients

class SendClientMixin(object):

    def send(self, messages, message_timeout=0, session=None, **kwargs):
        """Send one or more messages to the current entity.

        This operation will open a single-use connection, send the supplied messages, and close
        connection. If the entity requires sessions, a session ID must be either
        provided here, or set on each outgoing message.

        :param messages: One or more messages to be sent.
        :type messages: ~azure.servicebus.common.message.Message or list[~azure.servicebus.common.message.Message]
        :param message_timeout: The period in seconds during which the Message must be
         sent. If the send is not completed in this time it will return a failure result.
        :type message_timeout: int
        :param session: An optional session ID. If supplied this session ID will be
         applied to every outgoing message sent with this Sender.
         If an individual message already has a session ID, that will be
         used instead. If no session ID is supplied here, nor set on an outgoing
         message, a ValueError will be raised if the entity is sessionful.
        :type session: str or ~uuid.Guid
        :raises: ~azure.servicebus.common.errors.MessageSendFailed
        :returns: A list of the send results of all the messages. Each
         send result is a tuple with two values. The first is a boolean, indicating `True`
         if the message sent, or `False` if it failed. The second is an error if the message
         failed, otherwise it will be `None`.
        :rtype: list[tuple[bool, ~azure.servicebus.common.errors.MessageSendFailed]]

        Example:
            .. literalinclude:: ../examples/test_examples.py
                :start-after: [START send_message_service_bus]
                :end-before: [END send_message_service_bus]
                :language: python
                :dedent: 4
                :caption: Send a message to current entity via a single use connection

            .. literalinclude:: ../examples/test_examples.py
                :start-after: [START send_message_service_bus_multiple]
                :end-before: [END send_message_service_bus_multiple]
                :language: python
                :dedent: 4
                :caption: Send messages to current entity via a single use connection

        """
        with self.get_sender(message_timeout=message_timeout, session=session, **kwargs) as sender:
            if isinstance(messages, Message):
                sender.queue_message(messages)
            else:
                try:
                    messages = list(messages)
                except TypeError:
                    raise TypeError("Value of messages must be a 'Message' object or an iterable of 'Message' objects.")

                for m in messages:
                    if not isinstance(m, Message):
                        raise TypeError("Item {} in iterator is not of type 'Message'.".format(m))
                    sender.queue_message(m)

            return sender.send_pending_messages()

    def get_sender(self, message_timeout=0, session=None, **kwargs):
        """Get a Sender for the Service Bus endpoint.

        A Sender represents a single open Connection with which multiple send operations can be made.

        :param message_timeout: The period in seconds during which messages sent with
         this Sender must be sent. If the send is not completed in this time it will fail.
        :type message_timeout: int
        :param session: An optional session ID. If supplied this session ID will be
         applied to every outgoing message sent with this Sender.
         If an individual message already has a session ID, that will be
         used instead. If no session ID is supplied here, nor set on an outgoing
         message, a ValueError will be raised if the entity is sessionful.
        :type session: str or ~uuid.Guid
        :returns: A Sender instance with an unopened connection.
        :rtype: ~azure.servicebus.send_handler.Sender

        Example:
            .. literalinclude:: ../examples/test_examples.py
                :start-after: [START get_sender]
                :end-before: [END get_sender]
                :language: python
                :dedent: 4
                :caption: Get sender client from Service Bus client

        """
        handler_id = str(uuid.uuid4())
        if self.entity and self.requires_session:
            return SessionSender(
                handler_id,
                self.entity_uri,
                self.auth_config,
                session=session,
                debug=self.debug,
                msg_timeout=message_timeout,
                **kwargs)
        return Sender(
            handler_id,
            self.entity_uri,
            self.auth_config,
            session=session,
            debug=self.debug,
            msg_timeout=message_timeout,
            **kwargs)

class ReceiveClientMixin(object):

    def peek(self, count=1, start_from=0, session=None, **kwargs):
        """Browse messages currently pending in the queue.

        Peeked messages are not removed from queue, nor are they locked. They cannot be completed,
        deferred or dead-lettered.

        :param count: The maximum number of messages to try and peek. The default
         value is 1.
        :type count: int
        :param start_from: A message sequence number from which to start browsing messages.
        :type start_from: int
        :param session: If the entity requires sessions, a session ID must be supplied
         in order that only messages from that session will be browsed. If the entity
         does not require sessions this value will be ignored.
        :type session: str
        :rtype: list[~azure.servicebus.common.message.PeekMessage]

        Example:
            .. literalinclude:: ../examples/test_examples.py
                :start-after: [START peek_messages_service_bus]
                :end-before: [END peek_messages_service_bus]
                :language: python
                :dedent: 4
                :caption: Look at specificied number of messages without removing them from queue

        """
        message = {'from-sequence-number': types.AMQPLong(start_from), 'message-count': int(count)}
        if self.entity and self.requires_session:
            if not session:
                raise ValueError("Sessions are required, please set session.")
            message['session-id'] = session

        with BaseHandler(self.entity_uri, self.auth_config, debug=self.debug, **kwargs) as handler:
            return handler._mgmt_request_response(  # pylint: disable=protected-access
                REQUEST_RESPONSE_PEEK_OPERATION,
                message,
                mgmt_handlers.peek_op)

    def list_sessions(self, updated_since=None, max_results=100, skip=0, **kwargs):
        """List session IDs.

        List the Session IDs with pending messages in the queue where the state of the session
        has been updated since the timestamp provided. If no timestamp is provided, all will be returned.
        If the state of a session has never been set, it will not be returned regardless of whether
        there are messages pending.

        :param updated_since: The UTC datetime from which to return updated pending Session IDs.
        :type updated_since: datetime.datetime
        :param max_results: The maximum number of Session IDs to return. Default value is 100.
        :type max_results: int
        :param skip: The page value to jump to. Default value is 0.
        :type skip: int
        :rtype: list[str]

        Example:
            .. literalinclude:: ../examples/test_examples.py
                :start-after: [START list_sessions_service_bus]
                :end-before: [END list_sessions_service_bus]
                :language: python
                :dedent: 4
                :caption: Get the Ids of session which have messages pending in the queue

        """
        if self.entity and not self.requires_session:
            raise ValueError("This is not a sessionful entity.")
        message = {
            'last-updated-time': updated_since or datetime.datetime.utcfromtimestamp(0),
            'skip': types.AMQPInt(skip),
            'top': types.AMQPInt(max_results),
        }
        with BaseHandler(self.entity_uri, self.auth_config, debug=self.debug, **kwargs) as handler:
            return handler._mgmt_request_response(  # pylint: disable=protected-access
                REQUEST_RESPONSE_GET_MESSAGE_SESSIONS_OPERATION,
                message,
                mgmt_handlers.list_sessions_op)

    def receive_deferred_messages(self, sequence_numbers, mode=ReceiveSettleMode.PeekLock, **kwargs):
        """Receive messages by sequence number that have been previously deferred.

        When receiving deferred messages from a partitioned entity, all of the supplied
        sequence numbers must be messages from the same partition.

        :param sequence_numbers: A list of the sequence numbers of messages that have been
         deferred.
        :type sequence_numbers: list[int]
        :param mode: The mode with which messages will be retrieved from the entity. The two options
         are PeekLock and ReceiveAndDelete. Messages received with PeekLock must be settled within a given
         lock period before they will be removed from the queue. Messages received with ReceiveAndDelete
         will be immediately removed from the queue, and cannot be subsequently rejected or re-received if
         the client fails to process the message. The default mode is PeekLock.
        :type mode: ~azure.servicebus.common.constants.ReceiveSettleMode
        :rtype: list[~azure.servicebus.common.message.Message]

        Example:
            .. literalinclude:: ../examples/test_examples.py
                :start-after: [START receive_deferred_messages_service_bus]
                :end-before: [END receive_deferred_messages_service_bus]
                :language: python
                :dedent: 8
                :caption: Get the messages which were deferred using their sequence numbers

        """
        if (self.entity and self.requires_session) or kwargs.get('session'):
            raise ValueError("Sessionful deferred messages can only be received within a locked receive session.")
        if not sequence_numbers:
            raise ValueError("At least one sequence number must be specified.")
        try:
            receive_mode = mode.value.value
        except AttributeError:
            receive_mode = int(mode)
        message = {
            'sequence-numbers': types.AMQPArray([types.AMQPLong(s) for s in sequence_numbers]),
            'receiver-settle-mode': types.AMQPuInt(receive_mode)}
        mgmt_handler = functools.partial(mgmt_handlers.deferred_message_op, mode=receive_mode)
        with BaseHandler(self.entity_uri, self.auth_config, debug=self.debug, **kwargs) as handler:
            return handler._mgmt_request_response(  # pylint: disable=protected-access
                REQUEST_RESPONSE_RECEIVE_BY_SEQUENCE_NUMBER,
                message,
                mgmt_handler)

    def settle_deferred_messages(self, settlement, messages, **kwargs):
        """Settle messages that have been previously deferred.

        :param settlement: How the messages are to be settled. This must be a string
         of one of the following values: 'completed', 'suspended', 'abandoned'.
        :type settlement: str
        :param messages: A list of deferred messages to be settled.
        :type messages: list[~azure.servicebus.common.message.DeferredMessage]

        Example:
            .. literalinclude:: ../examples/test_examples.py
                :start-after: [START settle_deferred_messages_service_bus]
                :end-before: [END settle_deferred_messages_service_bus]
                :language: python
                :dedent: 8
                :caption: Settle deferred messages.

        """
        if (self.entity and self.requires_session) or kwargs.get('session'):
            raise ValueError("Sessionful deferred messages can only be settled within a locked receive session.")
        if settlement.lower() not in ['completed', 'suspended', 'abandoned']:
            raise ValueError("Settlement must be one of: 'completed', 'suspended', 'abandoned'")
        if not messages:
            raise ValueError("At least one message must be specified.")
        message = {
            'disposition-status': settlement.lower(),
            'lock-tokens': types.AMQPArray([m.lock_token for m in messages])}

        with BaseHandler(self.entity_uri, self.auth_config, debug=self.debug, **kwargs) as handler:
            return handler._mgmt_request_response(  # pylint: disable=protected-access
                REQUEST_RESPONSE_UPDATE_DISPOSTION_OPERATION,
                message,
                mgmt_handlers.default)

    def get_receiver(self, session=None, prefetch=0, mode=ReceiveSettleMode.PeekLock, idle_timeout=0, **kwargs):
        """Get a Receiver for the Service Bus endpoint.

        A Receiver represents a single open Connection with which multiple receive operations can be made.

        :param session: A specific session from which to receive. This must be specified for a
         sessionful entity, otherwise it must be None. In order to receive the next available
         session, set this to NEXT_AVAILABLE.
        :type session: str or ~azure.servicebus.common.constants.NEXT_AVAILABLE
        :param prefetch: The maximum number of messages to cache with each request to the service.
         The default value is 0, meaning messages will be received from the service and processed
         one at a time. Increasing this value will improve message throughput performance but increase
         the change that messages will expire while they are cached if they're not processed fast enough.
        :type prefetch: int
        :param mode: The mode with which messages will be retrieved from the entity. The two options
         are PeekLock and ReceiveAndDelete. Messages received with PeekLock must be settled within a given
         lock period before they will be removed from the queue. Messages received with ReceiveAndDelete
         will be immediately removed from the queue, and cannot be subsequently rejected or re-received if
         the client fails to process the message. The default mode is PeekLock.
        :type mode: ~azure.servicebus.common.constants.ReceiveSettleMode
        :param idle_timeout: The timeout in seconds between received messages after which the receiver will
         automatically shutdown. The default value is 0, meaning no timeout.
        :type idle_timeout: int
        :returns: A Receiver instance with an unopened Connection.
        :rtype: ~azure.servicebus.receive_handler.Receiver
        :raises: If the current Service Bus entity requires sessions, a TypeError will
         be raised.

        Example:
            .. literalinclude:: ../examples/test_examples.py
                :start-after: [START get_receiver]
                :end-before: [END get_receiver]
                :language: python
                :dedent: 4
                :caption: Get the receiver client from Service Bus client

        """
        if self.entity and not self.requires_session and session:
            raise ValueError("A session cannot be used with a non-sessionful entitiy.")
        if self.entity and self.requires_session and not session:
            raise ValueError("This entity requires a session.")
        if int(prefetch) < 0 or int(prefetch) > 50000:
            raise ValueError("Prefetch must be an integer between 0 and 50000 inclusive.")

        prefetch += 1
        handler_id = str(uuid.uuid4())
        if session:
            return SessionReceiver(
                handler_id,
                self.entity_uri,
                self.auth_config,
                session=session,
                debug=self.debug,
                timeout=int(idle_timeout * 1000),
                prefetch=prefetch,
                mode=mode,
                **kwargs)
        return Receiver(
            handler_id,
            self.entity_uri,
            self.auth_config,
            debug=self.debug,
            timeout=int(idle_timeout * 1000),
            prefetch=prefetch,
            mode=mode,
            **kwargs)

    def get_deadletter_receiver(
            self, transfer_deadletter=False, prefetch=0,
            mode=ReceiveSettleMode.PeekLock, idle_timeout=0, **kwargs):
        """Get a Receiver for the deadletter endpoint of the queue.

        A Receiver represents a single open Connection with which multiple receive operations can be made.

        :param transfer_deadletter: Whether to connect to the transfer deadletter queue, or the standard
         deadletter queue. Default is False, using the standard deadletter endpoint.
        :type transfer_deadletter: bool
        :param prefetch: The maximum number of messages to cache with each request to the service.
         The default value is 0, meaning messages will be received from the service and processed
         one at a time. Increasing this value will improve message throughput performance but increase
         the change that messages will expire while they are cached if they're not processed fast enough.
        :type prefetch: int
        :param mode: The mode with which messages will be retrieved from the entity. The two options
         are PeekLock and ReceiveAndDelete. Messages received with PeekLock must be settled within a given
         lock period before they will be removed from the queue. Messages received with ReceiveAndDelete
         will be immediately removed from the queue, and cannot be subsequently rejected or re-received if
         the client fails to process the message. The default mode is PeekLock.
        :type mode: ~azure.servicebus.common.constants.ReceiveSettleMode
        :param idle_timeout: The timeout in seconds between received messages after which the receiver will
         automatically shutdown. The default value is 0, meaning no timeout.
        :type idle_timeout: int
        :returns: A Receiver instance with an unopened Connection.
        :rtype: ~azure.servicebus.receive_handler.Receiver

        Example:
            .. literalinclude:: ../examples/test_examples.py
                :start-after: [START get_dead_letter_receiver]
                :end-before: [END get_dead_letter_receiver]
                :language: python
                :dedent: 4
                :caption: Get the dead lettered messages

        """
        if int(prefetch) < 0 or int(prefetch) > 50000:
            raise ValueError("Prefetch must be an integer between 0 and 50000 inclusive.")

        prefetch += 1
        handler_id = str(uuid.uuid4())
        if transfer_deadletter:
            entity_uri = self.mgmt_client.format_transfer_dead_letter_queue_name(self.entity_uri)
        else:
            entity_uri = self.mgmt_client.format_dead_letter_queue_name(self.entity_uri)
        return Receiver(
            handler_id,
            entity_uri,
            self.auth_config,
            debug=self.debug,
            timeout=int(idle_timeout * 1000),
            prefetch=prefetch,
            mode=mode,
            **kwargs)

class QueueClient(SendClientMixin, ReceiveClientMixin, mixins.BaseClient):
    """A queue client.

    The QueueClient class defines a high level interface for sending
    messages to and receiving messages from an Azure Service Bus queue.
    If you do not wish to perform management operations, a QueueClient can be
    instantiated directly to perform send and receive operations to a Queue.
    However if a QueueClient is created directly, a `get_properties` operation will
    need to be completed in order to retrieve the properties of this queue (for example,
    whether it is sessionful).

    :param address: The full URI of the Service Bus namespace. This can optionally
        include URL-encoded access name and key.
    :type address: str
    :param name: The name of the queue to which the Client will connect.
    :type name: str
    :param shared_access_key_name: The name of the shared access policy. This must be supplied
     if not encoded into the address.
    :type shared_access_key_name: str
    :param shared_access_key_value: The shared access key. This must be supplied if not encoded
     into the address.
    :type shared_access_key_value: str
    :param debug: Whether to output network trace logs to the logger. Default is `False`.
    :type debug: bool

    Example:
        .. literalinclude:: ../examples/test_examples.py
            :start-after: [START create_queue_client_directly]
            :end-before: [END create_queue_client_directly]
            :language: python
            :dedent: 8
            :caption: Create a QueueClient.

    """

    def _get_entity(self):
        return self.mgmt_client.get_queue(self.name)

class TopicClient(SendClientMixin, mixins.BaseClient):
    """A topic client.

    The TopicClient class defines a high level interface for sending
    messages to an Azure Service Bus Topic.
    If you do not wish to perform management operations, a TopicClient can be
    instantiated directly to perform send operations to a Topic.

    :param address: The full URI of the Service Bus namespace. This can optionally
        include URL-encoded access name and key.
    :type address: str
    :param name: The name of the topic to which the Client will connect.
    :type name: str
    :param shared_access_key_name: The name of the shared access policy. This must be supplied
     if not encoded into the address.
    :type shared_access_key_name: str
    :param shared_access_key_value: The shared access key. This must be supplied if not encoded
     into the address.
    :type shared_access_key_value: str
    :param debug: Whether to output network trace logs to the logger. Default is `False`.
    :type debug: bool

    Example:
        .. literalinclude:: ../examples/test_examples.py
            :start-after: [START create_topic_client_directly]
            :end-before: [END create_topic_client_directly]
            :language: python
            :dedent: 8
            :caption: Create a TopicClient.

    """

    def _get_entity(self):
        return self.mgmt_client.get_topic(self.name)

class SubscriptionClient(ReceiveClientMixin, mixins.BaseClient):
    """A subscription client.

    The SubscriptionClient class defines a high level interface for receiving
    messages to an Azure Service Bus Subscription.
    If you do not wish to perform management operations, a SubscriptionClient can be
    instantiated directly to perform receive operations from a Subscription.

    :param address: The full URI of the Service Bus namespace. This can optionally
        include URL-encoded access name and key.
    :type address: str
    :param name: The name of the topic to which the Client will connect.
    :type name: str
    :param shared_access_key_name: The name of the shared access policy. This must be supplied
     if not encoded into the address.
    :type shared_access_key_name: str
    :param shared_access_key_value: The shared access key. This must be supplied if not encoded
     into the address.
    :type shared_access_key_value: str
    :param debug: Whether to output network trace logs to the logger. Default is `False`.
    :type debug: bool

    Example:
        .. literalinclude:: ../examples/test_examples.py
            :start-after: [START create_sub_client_directly]
            :end-before: [END create_sub_client_directly]
            :language: python
            :dedent: 8
            :caption: Create a SubscriptionClient.

    """

    def __init__(self, address, name, shared_access_key_name=None,
                 shared_access_key_value=None, debug=False, **kwargs):

        super(SubscriptionClient, self).__init__(
            address, name, shared_access_key_name=shared_access_key_name,
            shared_access_key_value=shared_access_key_value, debug=debug, **kwargs)
        self.topic_name = self.address.path.split("/")[1]

    @classmethod
    def from_connection_string(cls, conn_str, name, topic=None, **kwargs):  # pylint: disable=arguments-differ
        """Create a SubscriptionClient from a connection string.

        :param conn_str: The connection string.
        :type conn_str: str
        :param name: The name of the Subscription.
        :type name: str
        :param topic: The name of the Topic, if the EntityName is
         not included in the connection string.
        :type topic: str
        """
        address, policy, key, entity = parse_conn_str(conn_str)
        entity = topic or entity
        address = build_uri(address, entity)
        address += "/Subscriptions/" + name
        return cls(address, name, shared_access_key_name=policy, shared_access_key_value=key, **kwargs)

    @classmethod
    def from_entity(cls, address, topic, entity, **kwargs):  # pylint: disable=arguments-differ
        client = cls(
            address + "/" + topic + "/Subscriptions/" + entity.name,
            entity.name,
            validated_entity=entity,
            **kwargs)
        return client

    def _get_entity(self):
        return self.mgmt_client.get_subscription(self.topic_name, self.name)


#-------------------------------------------------------------------------
# Copyright (c) Microsoft.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#--------------------------------------------------------------------------
from .constants import (
    DEFAULT_HTTP_TIMEOUT,
    MANAGEMENT_HOST,
)
from .models import (
    MetricDefinitions,
    MetricResponses,
    PublishData,
    Site,
    Sites,
    WebSpace,
    WebSpaces,
)
from .servicemanagementclient import (
    _ServiceManagementClient,
)
from ._common_conversion import (
    _str,
)
from ._serialization import (
    _XmlSerializer,
)

class WebsiteManagementService(_ServiceManagementClient):
    ''' Note that this class is a preliminary work on WebSite
        management. Since it lack a lot a features, final version
        can be slightly different from the current one.
    '''

    def __init__(self, subscription_id=None, cert_file=None,
                 host=MANAGEMENT_HOST, request_session=None,
                 timeout=DEFAULT_HTTP_TIMEOUT):
        '''
        Initializes the website management service.

        subscription_id:
            Subscription to manage.
        cert_file:
            Path to .pem certificate file (httplib), or location of the
            certificate in your Personal certificate store (winhttp) in the
            CURRENT_USER\my\CertificateName format.
            If a request_session is specified, then this is unused.
        host:
            Live ServiceClient URL. Defaults to Azure public cloud.
        request_session:
            Session object to use for http requests. If this is specified, it
            replaces the default use of httplib or winhttp. Also, the cert_file
            parameter is unused when a session is passed in.
            The session object handles authentication, and as such can support
            multiple types of authentication: .pem certificate, oauth.
            For example, you can pass in a Session instance from the requests
            library. To use .pem certificate authentication with requests
            library, set the path to the .pem file on the session.cert
            attribute.
        timeout:
            Optional. Timeout for the http request, in seconds.
        '''
        super(WebsiteManagementService, self).__init__(
            subscription_id, cert_file, host, request_session, timeout)

    #--Operations for web sites ----------------------------------------
    def list_webspaces(self):
        '''
        List the webspaces defined on the account.
        '''
        return self._perform_get(self._get_list_webspaces_path(),
                                 WebSpaces)

    def get_webspace(self, webspace_name):
        '''
        Get details of a specific webspace.

        webspace_name:
            The name of the webspace.
        '''
        return self._perform_get(self._get_webspace_details_path(webspace_name),
                                 WebSpace)

    def list_sites(self, webspace_name):
        '''
        List the web sites defined on this webspace.

        webspace_name:
            The name of the webspace.
        '''
        return self._perform_get(self._get_sites_path(webspace_name),
                                 Sites)

    def get_site(self, webspace_name, website_name):
        '''
        List the web sites defined on this webspace.

        webspace_name:
            The name of the webspace.
        website_name:
            The name of the website.
        '''
        return self._perform_get(self._get_sites_details_path(webspace_name,
                                                              website_name),
                                 Site)

    def create_site(self, webspace_name, website_name, geo_region, host_names,
                    plan='VirtualDedicatedPlan', compute_mode='Shared',
                    server_farm=None, site_mode=None):
        '''
        Create a website.

        webspace_name:
            The name of the webspace.
        website_name:
            The name of the website.
        geo_region:
            The geographical region of the webspace that will be created.
        host_names:
            An array of fully qualified domain names for website. Only one
            hostname can be specified in the azurewebsites.net domain.
            The hostname should match the name of the website. Custom domains
            can only be specified for Shared or Standard websites.
        plan:
            This value must be 'VirtualDedicatedPlan'.
        compute_mode:
            This value should be 'Shared' for the Free or Paid Shared
            offerings, or 'Dedicated' for the Standard offering. The default
            value is 'Shared'. If you set it to 'Dedicated', you must specify
            a value for the server_farm parameter.
        server_farm:
            The name of the Server Farm associated with this website. This is
            a required value for Standard mode.
        site_mode:
            Can be None, 'Limited' or 'Basic'. This value is 'Limited' for the
            Free offering, and 'Basic' for the Paid Shared offering. Standard
            mode does not use the site_mode parameter; it uses the compute_mode
            parameter.
        '''
        xml = _XmlSerializer.create_website_to_xml(webspace_name, website_name, geo_region, plan, host_names, compute_mode, server_farm, site_mode)
        return self._perform_post(
            self._get_sites_path(webspace_name),
            xml,
            Site)

    def delete_site(self, webspace_name, website_name,
                    delete_empty_server_farm=False, delete_metrics=False):
        '''
        Delete a website.

        webspace_name:
            The name of the webspace.
        website_name:
            The name of the website.
        delete_empty_server_farm:
            If the site being deleted is the last web site in a server farm,
            you can delete the server farm by setting this to True.
        delete_metrics:
            To also delete the metrics for the site that you are deleting, you
            can set this to True.
        '''
        path = self._get_sites_details_path(webspace_name, website_name)
        query = ''
        if delete_empty_server_farm:
            query += '&deleteEmptyServerFarm=true'
        if delete_metrics:
            query += '&deleteMetrics=true'
        if query:
            path = path + '?' + query.lstrip('&')
        return self._perform_delete(path)

    def update_site(self, webspace_name, website_name, state=None):
        '''
        Update a web site.

        webspace_name:
            The name of the webspace.
        website_name:
            The name of the website.
        state:
            The wanted state ('Running' or 'Stopped' accepted)
        '''
        xml = _XmlSerializer.update_website_to_xml(state)
        return self._perform_put(
            self._get_sites_details_path(webspace_name, website_name),
            xml, as_async=True)

    def restart_site(self, webspace_name, website_name):
        '''
        Restart a web site.

        webspace_name:
            The name of the webspace.
        website_name:
            The name of the website.
        '''
        return self._perform_post(
            self._get_restart_path(webspace_name, website_name),
            None, as_async=True)

    def get_historical_usage_metrics(self, webspace_name, website_name,
                                     metrics = None, start_time=None, end_time=None, time_grain=None):
        '''
        Get historical usage metrics.

        webspace_name:
            The name of the webspace.
        website_name:
            The name of the website.
        metrics:
            Optional. List of metrics name. Otherwise, all metrics returned.
        start_time:
            Optional. An ISO8601 date. Otherwise, current hour is used.
        end_time:
            Optional. An ISO8601 date. Otherwise, current time is used.
        time_grain:
            Optional. A rollup name, as P1D. OTherwise, default rollup for the metrics is used.
        More information and metrics name at:
        http://msdn.microsoft.com/en-us/library/azure/dn166964.aspx
        '''
        metrics = ('names='+','.join(metrics)) if metrics else ''
        start_time = ('StartTime='+start_time) if start_time else ''
        end_time = ('EndTime='+end_time) if end_time else ''
        time_grain = ('TimeGrain='+time_grain) if time_grain else ''
        parameters = ('&'.join(v for v in (metrics, start_time, end_time, time_grain) if v))
        parameters = '?'+parameters if parameters else ''
        return self._perform_get(self._get_historical_usage_metrics_path(webspace_name, website_name) + parameters,
                                 MetricResponses)

    def get_metric_definitions(self, webspace_name, website_name):
        '''
        Get metric definitions of metrics available of this web site.

        webspace_name:
            The name of the webspace.
        website_name:
            The name of the website.
        '''
        return self._perform_get(self._get_metric_definitions_path(webspace_name, website_name),
                                 MetricDefinitions)

    def get_publish_profile_xml(self, webspace_name, website_name):
        '''
        Get a site's publish profile as a string

        webspace_name:
            The name of the webspace.
        website_name:
            The name of the website.
        '''
        return self._perform_get(self._get_publishxml_path(webspace_name, website_name),
                                 None).body.decode("utf-8")

    def get_publish_profile(self, webspace_name, website_name):
        '''
        Get a site's publish profile as an object

        webspace_name:
            The name of the webspace.
        website_name:
            The name of the website.
        '''
        return self._perform_get(self._get_publishxml_path(webspace_name, website_name),
                                 PublishData)

    #--Helper functions --------------------------------------------------
    def _get_list_webspaces_path(self):
        return self._get_path('services/webspaces', None)

    def _get_webspace_details_path(self, webspace_name):
        return self._get_path('services/webspaces/', webspace_name)

    def _get_sites_path(self, webspace_name):
        return self._get_path('services/webspaces/',
                              webspace_name) + '/sites'

    def _get_sites_details_path(self, webspace_name, website_name):
        return self._get_path('services/webspaces/',
                              webspace_name) + '/sites/' + _str(website_name)

    def _get_restart_path(self, webspace_name, website_name):
        return self._get_path('services/webspaces/',
                              webspace_name) + '/sites/' + _str(website_name) + '/restart/'

    def _get_historical_usage_metrics_path(self, webspace_name, website_name):
        return self._get_path('services/webspaces/',
                              webspace_name) + '/sites/' + _str(website_name) + '/metrics/'

    def _get_metric_definitions_path(self, webspace_name, website_name):
        return self._get_path('services/webspaces/',
                              webspace_name) + '/sites/' + _str(website_name) + '/metricdefinitions/'

    def _get_publishxml_path(self, webspace_name, website_name):
        return self._get_path('services/webspaces/',
                              webspace_name) + '/sites/' + _str(website_name) + '/publishxml/'


#---------------------------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for license information.
#---------------------------------------------------------------------------------------------

import json
import time
import os
from ._internal import _a128cbc_hs256_encrypt, _a128cbc_hs256_decrypt, _JwsHeader, _JwsObject, \
    _JweHeader, _JweObject, _str_to_b64url, _bstr_to_b64url, _b64_to_bstr, _RsaKey

def generate_pop_key():
    """
    Generates a key which can be used for Proof Of Possession token authentication.
    :return:
    """
    return _RsaKey.generate()

class HttpMessageSecurity(object):
    """
    Used for message authorization, encryption and decrtyption.

    This class is intended for internal use only.  Details are subject to non-compatible changes, consumers of the
    azure-keyvault module should not take dependencies on this class or its current implementation.
    """
    def __init__(self, client_security_token=None,
                 client_signature_key=None,
                 client_encryption_key=None,
                 server_signature_key=None,
                 server_encryption_key=None):
        self.client_security_token = client_security_token
        self.client_signature_key = client_signature_key
        self.client_encryption_key = client_encryption_key
        self.server_signature_key = server_signature_key
        self.server_encryption_key = server_encryption_key

    def protect_request(self, request):
        """
        Adds authorization header, and encrypts and signs the request if supported on the specific request.
        :param request: unprotected request to apply security protocol
        :return: protected request with appropriate security protocal applied
        """
        # Setup the auth header on the request
        # Due to limitations in the service we hard code the auth scheme to 'Bearer' as the service will fail with any
        # other scheme or a different casing such as 'bearer', once this is fixed the following line should be replaced:
        # request.headers['Authorization'] = '{} {}'.format(auth[0], auth[1])
        request.headers['Authorization'] = '{} {}'.format('Bearer', self.client_security_token)

        # if the current message security doesn't support message protection, or the body is empty
        # skip protection and return the original request
        if not self.supports_protection() or len(request.body) == 0:
            return request

        plain_text = request.body

        # if the client encryption key is specified add it to the body of the request
        if self.client_encryption_key:
            # note that this assumes that the body is already json and not simple string content
            # this is true for all requests which currently support message encryption, but might
            # need to be revisited when the types of
            body_dict = json.loads(plain_text)
            body_dict['rek'] = {'jwk': self.client_encryption_key.to_jwk().serialize()}
            plain_text = json.dumps(body_dict).encode(encoding='utf8')

        # build the header for the jws body
        jws_header = _JwsHeader()
        jws_header.alg = 'RS256'
        jws_header.kid = self.client_signature_key.kid
        jws_header.at = self.client_security_token
        jws_header.ts = int(time.time())
        jws_header.typ = 'PoP'

        jws = _JwsObject()

        jws.protected = jws_header.to_compact_header()
        jws.payload = self._protect_payload(plain_text)
        data = (jws.protected + '.' + jws.payload).encode('ascii')
        jws.signature = _bstr_to_b64url(self.client_signature_key.sign(data))

        request.headers['Content-Type'] = 'application/jose+json'

        request.prepare_body(data=jws.to_flattened_jws(), files=None)

        return request

    def unprotect_response(self, response, **kwargs):
        """
        Removes protection from the specified response
        :param request: response from the key vault service
        :return: unprotected response with any security protocal encryption removed
        """
        body = response.content
        # if the current message security doesn't support message protection, the body is empty, or the request failed
        # skip protection and return the original response
        if not self.supports_protection() or len(response.content) == 0 or response.status_code != 200:
            return response

        # ensure the content-type is application/jose+json
        if 'application/jose+json' not in response.headers.get('content-type', '').lower():
            raise ValueError('Invalid protected response')

        # deserialize the response into a JwsObject, using response.text so requests handles the encoding
        jws = _JwsObject().deserialize(body)

        # deserialize the protected header
        jws_header = _JwsHeader.from_compact_header(jws.protected)

        # ensure the jws signature kid matches the key from original challenge
        # and the alg matches expected signature alg
        if jws_header.kid != self.server_signature_key.kid \
                or jws_header.alg != 'RS256':
            raise ValueError('Invalid protected response')

        # validate the signature of the jws
        data = (jws.protected + '.' + jws.payload).encode('ascii')
        # verify will raise an InvalidSignature exception if the signature doesn't match
        self.server_signature_key.verify(signature=_b64_to_bstr(jws.signature), data=data)

        # get the unprotected response body
        decrypted = self._unprotect_payload(jws.payload)

        response._content = decrypted
        response.headers['Content-Type'] = 'application/json'

        return response

    def supports_protection(self):
        """
        Determines if the the current HttpMessageSecurity object supports the message protection protocol.
        :return: True if the current object supports protection, otherwise False
        """
        return self.client_signature_key \
               and self.client_encryption_key \
               and self.server_signature_key \
               and self.server_encryption_key

    def _protect_payload(self, plaintext):
        # create the jwe header for the payload
        kek = self.server_encryption_key
        jwe_header = _JweHeader()
        jwe_header.alg = 'RSA-OAEP'
        jwe_header.kid = kek.kid
        jwe_header.enc = 'A128CBC-HS256'

        # create the jwe object
        jwe = _JweObject()
        jwe.protected = jwe_header.to_compact_header()

        # generate the content encryption key and iv
        cek = os.urandom(32)
        iv = os.urandom(16)
        jwe.iv = _bstr_to_b64url(iv)
        # wrap the cek using the server encryption key
        wrapped = _bstr_to_b64url(kek.encrypt(cek))
        jwe.encrypted_key = wrapped

        # encrypt the plaintext body with the cek using the protected header
        # as the authdata to get the ciphertext and the authtag
        ciphertext, tag = _a128cbc_hs256_encrypt(cek, iv, plaintext, jwe.protected.encode('ascii'))

        jwe.ciphertext = _bstr_to_b64url(ciphertext)
        jwe.tag = _bstr_to_b64url(tag)

        # flatten and encode the jwe for the final jws payload content
        flat = jwe.to_flattened_jwe()
        return _str_to_b64url(flat)

    def _unprotect_payload(self, payload):
        # deserialize the payload
        jwe = _JweObject().deserialize_b64(payload)

        # deserialize the payload header
        jwe_header = _JweHeader.from_compact_header(jwe.protected)

        # ensure the kid matches the specified client encryption key
        # and the key wrap alg and the data encryption enc match the expected
        if self.client_encryption_key.kid != jwe_header.kid \
                or jwe_header.alg != 'RSA-OAEP' \
                or jwe_header.enc != 'A128CBC-HS256':
            raise ValueError('Invalid protected response')

        # unwrap the cek using the client encryption key
        cek = self.client_encryption_key.decrypt(_b64_to_bstr(jwe.encrypted_key))

        # decrypt the cipher text to get the unprotected body content
        return _a128cbc_hs256_decrypt(cek,
                                      _b64_to_bstr(jwe.iv),
                                      _b64_to_bstr(jwe.ciphertext),
                                      jwe.protected.encode('ascii'),
                                      _b64_to_bstr(jwe.tag))