# Source code for aioxmpp.ibb.service

    ########################################################################
    # File name: service.py
    # This file is part of: aioxmpp
    #
    # LICENSE
    #
    # This program is free software: you can redistribute it and/or modify
    # it under the terms of the GNU Lesser General Public License as
    # published by the Free Software Foundation, either version 3 of the
    # License, or (at your option) any later version.
    #
    # This program is distributed in the hope that it will be useful, but
    # WITHOUT ANY WARRANTY; without even the implied warranty of
    # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
    # Lesser General Public License for more details.
    #
    # You should have received a copy of the GNU Lesser General Public
    # License along with this program.  If not, see
    # <http://www.gnu.org/licenses/>.
    #
    ########################################################################
    import asyncio
    import random

    from datetime import timedelta

    import aioxmpp
    import aioxmpp.callbacks
    import aioxmpp.errors as errors
    import aioxmpp.service as service
    import aioxmpp.utils as utils

    from . import xso as ibb_xso

    MAX_BLOCK_SIZE = (1 << 16) - 1

    class IBBTransport(asyncio.Transport):
        """
        The transport for IBB sessions.

        .. note:: Never instantiate this class directly, all instances of
                  this class are created the methods
                  :meth:`~aioxmpp.ibb.IBBService.open_session` and
                  :meth:`~aioxmpp.ibb.IBBService.expect_session` of
                  :class:`~aioxmpp.ibb.IBBService`.

        The following keys are supported for
        :meth:`~asyncio.BaseTransport.get_extra_info`:

           `block_size`
               The maximal block size of data in a IBB stanza.

           `peer_jid`
              The JID of the peer.

           `sid`
              The session id of the unerlying IBB session.

           `stanza_type`
              The used stanza type.
        """

        def __init__(self, service, peer_jid, sid, stanza_type, block_size):
            self._protocol = None
            self._service = service
            self._stanza_type = stanza_type
            self._sid = sid
            self._peer_jid = peer_jid
            self._block_size = block_size
            self._incoming_seq = 0
            self._outgoing_seq = 0
            self._closed = False
            self._closing = False

            self.set_write_buffer_limits()
            self._write_buffer = b""
            self._can_write = asyncio.Event()

            self._reading_paused = False
            self._input_buffer = []

            self._write_task = asyncio.ensure_future(self._write_task_main())
            self._write_task.add_done_callback(
                self._handle_close
            )

            self._wait_time = self._service.initial_wait_time.total_seconds()
            self._retries = 0

        def set_write_buffer_limits(self, high=None, low=None):
            if low is None:
                low = 4 * self._block_size

            if high is None:
                high = 8 * self._block_size

            if low < 0:
                raise ValueError("the limits must be positive")

            if high < 0:
                raise ValueError("the limits must be positive")

            if low > high:
                low = high

            self._output_buffer_limit_low = low
            self._output_buffer_limit_high = high

        def get_write_buffer_limits(self):
            return self._output_buffer_limit_low, self._output_buffer_limit_high

        def get_write_buffer_size(self):
            return len(self._write_buffer)

        def set_protocol(self, proto):
            self._protocol = proto
            proto.connection_made(self)

        def get_protocol(self):
            return self._protocol

        def pause_reading(self):
            self._reading_paused = True

        def resume_reading(self):
            self._reading_paused = False
            self._protocol.data_received(b"".join(self._input_buffer))
            self._input_buffer.clear()

        def is_closing(self):
            return self._closing or self._closed

        def get_extra_info(self, key, default=None):
            return {
                "block_size": self._block_size,
                "peer_jid": self._peer_jid,
                "stanza_type": self._stanza_type,
                "sid": self._sid,
            }.get(key, default)

        async def _write_task_main(self):
            e = None
            while True:
                await self._can_write.wait()

                if self._write_buffer:
                    data = self._write_buffer[:self._block_size]

                    if self._stanza_type == ibb_xso.IBBStanzaType.IQ:
                        stanza = aioxmpp.IQ(
                            aioxmpp.IQType.SET,
                            to=self._peer_jid,
                            payload=ibb_xso.Data(
                                self._sid,
                                self._outgoing_seq,
                                data
                            )
                        )
                    elif self._stanza_type == ibb_xso.IBBStanzaType.MESSAGE:
                        # TODO: use some form of tracking for messages
                        stanza = aioxmpp.Message(
                            aioxmpp.MessageType.NORMAL,
                            to=self._peer_jid,
                        )
                        stanza.xep0047_data = ibb_xso.Data(
                            self._sid,
                            self._outgoing_seq,
                            data
                        )

                    try:
                        await self._service.client.send(
                            stanza
                        )
                    except errors.XMPPWaitError:
                        # wait and try again unless max retries have been reached
                        if self._retries < self._service.max_retries:
                            await asyncio.sleep(self._wait_time)
                            self._wait_time *= self._service.wait_backoff_factor
                            self._retries += 1
                            continue
                        else:
                            e = asyncio.TimeoutError()
                            break
                    except errors.StanzaError as _e:
                        # break the loop to close the connection
                        e = _e
                        break
                    # update the internal state after the successful
                    # write: remove the written data from the buffer and
                    # increment the sequence number
                    self._write_buffer = self._write_buffer[len(data):]
                    self._outgoing_seq += 1
                    self._outgoing_seq &= 0xffff

                    # reset the wait time
                    self._wait_time = \
                        self._service.initial_wait_time.total_seconds()
                    self._retries = 0

                if len(self._write_buffer) < self._output_buffer_limit_low:
                    self._protocol.resume_writing()

                if not self._write_buffer:
                    if self._closing:
                        e = None
                        break
                    self._can_write.clear()

            close = ibb_xso.Close()
            close.sid = self._sid
            stanza = aioxmpp.IQ(
                aioxmpp.IQType.SET,
                to=self._peer_jid,
                payload=close,
            )

            try:
                await self._service.client.send(stanza)
            except errors.StanzaError as _e:
                if e is None:
                    e = _e
            finally:
                if e is not None:
                    raise e

        def write(self, data):
            """
            Send `data` over the IBB. If `data` is larger than the block size
            is is chunked and sent in chunks.

            Chunks from one call of :meth:`write` will always be sent in
            series.
            """

            if self.is_closing():
                return

            self._write_buffer += data

            if len(self._write_buffer) >= self._output_buffer_limit_high:
                self._protocol.pause_writing()

            if self._write_buffer:
                self._can_write.set()

        def _connection_closed(self):
            self._write_task.cancel()

        def _handle_close(self, fut):
            e = None
            self._service._remove_session(self._peer_jid, self._sid)
            try:
                e = fut.exception()
            except asyncio.CancelledError:
                pass
            self._protocol.connection_lost(e)
            self._closed = True

        def close(self):
            """
            Close the session.
            """
            if self.is_closing():
                return

            self._closing = True
            # make sure the writer wakes up
            self._can_write.set()

        def abort(self):
            """
            Abort the session.
            """
            if self.is_closing():
                return
            self._connection_closed()

        def _data_received(self, data):
            if self._closed:
                return

            if self._reading_paused:
                self._input_buffer.append(data)
            else:
                self._protocol.data_received(data)

        def _process_iq(self, payload):
            if payload.seq != self._incoming_seq:
                raise errors.XMPPCancelError(
                    condition=errors.ErrorCondition.UNEXPECTED_REQUEST
                )
            self._incoming_seq += 1
            self._incoming_seq &= 0xffff
            self._data_received(payload.content)

        def _process_msg(self, payload):
            if payload.seq != self._incoming_seq:
                return
            self._incoming_seq += 1
            self._incoming_seq &= 0xffff
            self._data_received(payload.content)

    class IBBService(service.Service):
        """
        A service implementing in-band bytestreams.

        Methods for establishing sessions:

        .. automethod:: expect_session

        .. automethod:: open_session

        The following attributes control the establishment of sessions due
        to a received request, that was not announced to the service by
        :meth:`expect_session`:

        .. attribute:: session_limit
           :annotation: = 0

           The maximal number of sessions to be accepted. If there are
           that many or more active sessions, no new sessions are
           accepted, unless they are whitelisted by
           :meth:`expect_session`. (This means, that by default only
           expected sessions are accepted!).

        .. attribute:: default_protocol_factory

           The protocol factory to be used when an unexpected connection
           is established. This *must* be set when changing
           :attr:`session_limit` to a non-zero value.

        .. signal:: on_session_accepted(transport, protocol)

           Fires when a session is established due to a received open
           request that was not expected (compare
           :meth:`expect_session`). This can only happen when
           :attr:`session_limit` is set to another value than its default
           value.

        The following attributes control how the IBB sessions react to
        errors of type wait:

        .. attribute:: max_retries
           :annotation: = 5

           The number of times it is tried to resend a data stanza, when a
           :class:`~aioxmpp.errors.XMPPWaitError` is received. When
           :attr:`max_retries` have been tried, the session is closed.
           `connection_lost` of the protocol receives an
           :class:`asyncio.TimeoutError`.

        .. attribute:: initial_wait_time
           :annotation: = timedelta(seconds=1)

           The time to wait when receiving a
           :class:`~aioxmpp.errors.XMPPWaitError` for the first time.

        .. attribute:: wait_backoff_factor
           :annotation: = 1.2

           The factor by which the wait time is prolonged on each
           successive wait error.
        """

        on_session_accepted = aioxmpp.callbacks.Signal()

        def __init__(self, client, **kwargs):
            super().__init__(client, **kwargs)

            self._sessions = {}

            self.session_limit = 0
            self._expected_sessions = {}

            self.default_protocol_factory = None

            self.client.on_stream_destroyed.connect(
                self._on_stream_destroyed
            )

            self.max_retries = 5
            self.initial_wait_time = timedelta(seconds=1)
            self.wait_backoff_factor = 1.2

        def _on_stream_destroyed(self):
            self._expected_sessions = {}

            # tear down the remaining open sessions
            for session in list(self._sessions.values()):
                session.abort()

        def expect_session(self, protocol_factory, peer_jid, sid):
            """
            Whitelist the session with `peer_jid` and the session id `sid` and
            return it when it is established. This is meant to be used
            with signalling protocols like Jingle and is the counterpart
            to :meth:`open_session`.

            :returns: an awaitable object, whose result is the tuple
                      `(transport, protocol)`
            """
            def on_done(fut):
                del self._expected_sessions[sid, peer_jid]

            _, fut = self._expected_sessions[sid, peer_jid] = (
                protocol_factory, asyncio.Future()
            )
            fut.add_done_callback(on_done)
            return fut

        async def open_session(self, protocol_factory, peer_jid, *,
                               stanza_type=ibb_xso.IBBStanzaType.IQ,
                               block_size=4096, sid=None):
            """
            Establish an in-band bytestream session with `peer_jid` and
            return the transport and protocol.

            :param protocol_factory: the protocol factory
            :type protocol_factory: a nullary callable returning an
                                    :class:`asyncio.Protocol` instance
            :param peer_jid: the JID with which to establish the byte-stream.
            :type peer_jid: :class:`aioxmpp.JID`
            :param stanza_type: the stanza type to use
            :type stanza_type: class:`~aioxmpp.ibb.IBBStanzaType`
            :param block_size: the maximal size of blocks to transfer
            :type block_size: :class:`int`
            :param sid: the session id to use
            :type sid: :class:`str` (must be a valid NMTOKEN)

            :returns: the transport and protocol
            :rtype: a tuple of :class:`aioxmpp.ibb.service.IBBTransport`
                    and :class:`asyncio.Protocol`
            """
            if block_size > MAX_BLOCK_SIZE:
                raise ValueError("block_size too large")

            if sid is None:
                sid = utils.to_nmtoken(random.getrandbits(8*8))
            open_ = ibb_xso.Open()
            open_.stanza = stanza_type
            open_.sid = sid
            open_.block_size = block_size

            # XXX: retry on XMPPModifyError with RESOURCE_CONSTRAINT
            await self.client.send(
                aioxmpp.IQ(
                    aioxmpp.IQType.SET,
                    to=peer_jid,
                    payload=open_,
                )
            )

            handle = self._sessions[sid, peer_jid] = IBBTransport(
                self,
                peer_jid,
                sid,
                stanza_type,
                block_size,
            )

            protocol = protocol_factory()
            handle.set_protocol(protocol)
            return handle, protocol

        @service.iq_handler(
            aioxmpp.IQType.SET,
            ibb_xso.Open)
        async def _handle_open_request(self, iq):
            peer_jid = iq.from_
            sid = iq.payload.sid
            block_size = iq.payload.block_size
            stanza_type = iq.payload.stanza

            if block_size > MAX_BLOCK_SIZE:
                raise errors.XMPPModifyError(
                    condition=errors.ErrorCondition.RESOURCE_CONSTRAINT
                )

            try:
                protocol_factory, expected_future = \
                    self._expected_sessions[sid, peer_jid]
            except KeyError:
                if len(self._sessions) >= self.session_limit:
                    raise errors.XMPPCancelError(
                        condition=errors.ErrorCondition.NOT_ACCEPTABLE
                    )
                expected_future = None
                protocol_factory = self.default_protocol_factory

            if (sid, peer_jid) in self._sessions:
                # disallow opening a session twice
                if expected_future is not None:
                    # is this correct?
                    expected_future.cancel()
                raise errors.XMPPCancelError(
                    condition=errors.ErrorCondition.NOT_ACCEPTABLE
                )

            handle = self._sessions[sid, peer_jid] = IBBTransport(
                self,
                peer_jid,
                sid,
                stanza_type,
                block_size
            )

            protocol = protocol_factory()
            handle.set_protocol(protocol)

            if expected_future is None:
                self.on_session_accepted((handle, protocol))
            else:
                expected_future.set_result((handle, protocol))

        @service.iq_handler(
            aioxmpp.IQType.SET,
            ibb_xso.Close)
        async def _handle_close_request(self, iq):
            peer_jid = iq.from_
            sid = iq.payload.sid

            try:
                session_handle = self._sessions[sid, peer_jid]
            except KeyError:
                raise errors.XMPPCancelError(
                    condition=errors.ErrorCondition.ITEM_NOT_FOUND
                )

            session_handle._connection_closed()

        @service.iq_handler(
            aioxmpp.IQType.SET,
            ibb_xso.Data)
        async def _handle_data(self, iq):
            peer_jid = iq.from_
            sid = iq.payload.sid

            try:
                session_handle = self._sessions[sid, peer_jid]
            except KeyError:
                raise errors.XMPPCancelError(
                    condition=errors.ErrorCondition.ITEM_NOT_FOUND
                )

            session_handle._process_iq(iq.payload)

        @service.inbound_message_filter
        def _handle_message(self, msg):
            if msg.xep0047_data is None:
                return msg

            payload = msg.xep0047_data
            peer_jid = msg.from_
            sid = payload.sid

            try:
                session = self._sessions[sid, peer_jid]
            except KeyError:
                return None

            session._process_msg(payload)

            return None

        def _remove_session(self, peer_jid, sid):
            del self._sessions[sid, peer_jid]