DEALER = b'DEALER' ROUTER = b'ROUTER' class CompositeClosableAsyncObject(ClosableAsyncObject):
    """
    Base class for objects that can be asynchronously closed and awaited and
    have ownership of other closable objects.
    """
    def on_open(self):
        self._children = set()

    async def on_close(self):
        await asyncio.gather(
            *[
                self.on_close_child(child)
                for child in self._children
            ],
            loop=self.loop,
        )
        del self._children

    async def on_close_child(self, child):
        """
        Called whenever a child instance must be closed.

        :returns: An awaitable that must only complete when the specified child
            instance is effectively closed.

        May be redefined by child-classes. The default implementation calls
            `close` followed by `wait_closed` on the specified child instance.
        """
        child.close()
        await child.wait_closed()

    def register_child(self, child):
        """
        Register a new child that will be closed whenever the current instance
        closes.

        :param child: The child instance.
        """
        if self.closing:
            child.close()
        else:
            self._children.add(child)
            child.on_closed.connect(self.unregister_child)

    def unregister_child(self, child):
        """
        Unregister an existing child that is no longer to be owned by the
        current instance.

        :param child: The child instance.
        """
        self._children.remove(child)
        child.on_closed.disconnect(self.unregister_child) def cancel_on_closing(func):
    """
    Automatically cancels a function or coroutine when the defining instance
    gets closed.

    :param func: The function to cancel on closing.
    :returns: A decorated function or coroutine.
    """
    @wraps(func)
    async def wrapper(self, *args, **kwargs):
        return await self.await_until_closing(func(self, *args, **kwargs))

    return wrapper class ZAPTemporaryError(ZAPError):
    def __init__(self, text, code=300):
        super().__init__(text, code) class ZAPAuthenticationFailure(ZAPError):
    def __init__(self, text, code=400):
        super().__init__(text, code) class ZAPInternalError(ZAPError):
    def __init__(self, text, code=500):
        super().__init__(text, code) class ZAPError(AZMQError):
    def __init__(self, text, code):
        super().__init__("ZAP error: %s (%s)" % (text, code))
        self.text = text
        self.code = code def buffer_to_metadata(buffer):
    """
    Transform a buffer to a metadata dictionary.

    :param buffer: The buffer, as received in a READY command.
    :returns: A metadata dictionary, with its keys normalized (in
        lowercase).
    """
    offset = 0
    size = len(buffer)
    metadata = {}

    while offset < size:
        name_size = struct.unpack_from('B', buffer, offset)[0]
        offset += 1

        if name_size > size - 4:
            raise ProtocolError(
                "Invalid name size in metadata",
                fatal=True,
            )

        name = buffer[offset:offset + name_size]
        offset += name_size

        value_size = struct.unpack_from('!I', buffer, offset)[0]
        offset += 4

        if value_size > size - name_size - 5:
            raise ProtocolError(
                "Invalid value size in metadata",
                fatal=True,
            )

        value = buffer[offset:offset + value_size]
        offset += value_size
        metadata[name.lower()] = value

    return metadata def metadata_to_buffers(metadata):
    """
    Transform a dict of metadata into a sequence of buffers.

    :param metadata: The metadata, as a dict.
    :returns: A list of buffers.
    """
    results = []

    for key, value in metadata.items():
        assert len(key) < 256
        assert len(value) < 2 ** 32
        results.extend([
            struct.pack('!B', len(key)),
            key,
            struct.pack('!I', len(value)),
            value,
        ])

    return results
import asyncio
from itertools import count
from .constants import (
    DEALER,
    ROUTER,
)
from .common import (
    CompositeClosableAsyncObject,
    cancel_on_closing,
)
from .errors import (
    ZAPTemporaryError,
    ZAPAuthenticationFailure,
    ZAPInternalError,
    ZAPError,
)
from .metadata import (
    buffer_to_metadata,
    metadata_to_buffers,
)
from .log import logger
                metadata = buffer_to_metadata(frames[6]) if frames[6] else {}
            except Exception as ex:
                logger.warning(
                    "Unexpected error while handling ZAP response (%r): %s."
                    " Ignoring.",
                    frames,
                    ex,
                )
                continue

            future = self._requests.get(request_id)

            if not future:
                logger.debug(
                    "Got authentication result for unknown request %r. "
                    "Ignoring.",
                    request_id,
                )
            else:
                if status_code == 200:
                    future.set_result((user_id, metadata))
                elif status_code == 300:
                    future.set_exception(ZAPTemporaryError(status_text))
                elif status_code == 400:
                    future.set_exception(ZAPAuthenticationFailure(status_text))
                elif status_code == 500:
                    future.set_exception(ZAPInternalError(status_text))
                else:
                    future.set_exception(ZAPError(status_text, status_code))

    def _get_next_request_id(self):
        request_id = next(self._request_ids)
        return '%x' % request_id

    @cancel_on_closing
    async def authenticate(
        self,
        domain,
        address,
        identity,
        mechanism,
        credentials,
    ):
        request_id = self._get_next_request_id()
        await self._socket.send_multipart([
            b'',
            b'1.0',
            request_id,
            domain.encode('utf-8'),
            address.encode('ascii'),
            identity,
            mechanism,
        ] + credentials)
        future = asyncio.Future(loop=self.loop)
        self._requests[request_id] = future

        @future.add_done_callback
        def remove_request(_):
            del self._requests[request_id]

        await future
        return future.result()


class BaseZAPAuthenticator(CompositeClosableAsyncObject):
    """
    A base class for ZAP authenticators.
    """
    def on_open(self, context):
        super().on_open()

        self._socket = context.socket(ROUTER)
        self._socket.bind(ZAP_INPROC_ENDPOINT)
        self.register_child(self._socket)
        self.run_task = asyncio.ensure_future(self.run())

    @cancel_on_closing
    async def run(self):
        logger.debug("ZAP authenticator started.")

        try:
            while True:
                frames = await self._socket.recv_multipart()

                try:
                    delimiter_index = frames.index(b'')
                    envelope = frames[:delimiter_index + 1]
                    message = frames[delimiter_index + 1:]
                    assert message[0] == b'1.0'
                    request_id = message[1]
                    domain = message[2].decode('utf-8')
                    address = message[3].decode('ascii')
                    identity = message[4]
                    mechanism = message[5]
                    credentials = message[6:]

                except Exception as ex:
                    logger.warning(
                        "Unexpected error while handling ZAP request (%r): %s."
                        " Ignoring.",
                        frames,
                        ex,
                    )
                    continue

                try:
                    user_id, metadata = await self.on_request(
                        domain=domain,
                        address=address,
                        identity=identity,
                        mechanism=mechanism,
                        credentials=credentials,
                    )

                    await self._socket.send_multipart(envelope + [
                        b'1.0',
                        request_id,
                        b'200',
                        b'',
                        user_id.encode('utf-8') if user_id else b'',

                        b''.join(metadata_to_buffers(metadata))