# Source code for aioxmpp.roster.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 logging

    import aioxmpp.service

    import aioxmpp.callbacks as callbacks
    import aioxmpp.errors as errors
    import aioxmpp.stanza as stanza
    import aioxmpp.structs as structs

    from aioxmpp.utils import namespaces

    from . import xso as roster_xso

    logger = logging.getLogger(__name__)

    _Sentinel = object()

    class Item:
        """
        Represent an entry in the roster. These entries are mutable, see the
        documentation of :class:`Service` for details on the lifetime of
        :class:`Item` instances within a :class:`Service` instance.

        .. attribute:: jid

           The :class:`~aioxmpp.JID` of the entry. This is always a bare
           JID.

        .. attribute:: name

           The display name of the entry, if any.

        .. attribute:: groups

           A :class:`set` of names of groups in which the roster entry is.

        .. attribute:: subscription

           The subscription status of the entry. One of ``"none"``, ``"to"``,
           ``"from"`` and ``"both"`` (in contrast to :class:`.xso.Item`,
           ``"remove"`` cannot occur here).

        .. attribute:: ask

           The ``ask`` attribute of the roster entry.

        .. attribute:: approved

           The ``approved`` attribute of the roster entry.

        The data of a roster entry can conveniently be exported to JSON:

        .. automethod:: export_as_json

        To mutate the roster entry, some handy methods are provided:

        .. automethod:: update_from_json

        .. automethod:: update_from_xso_item

        To create a roster entry from a :class:`.xso.Item`, use the
        :meth:`from_xso_item` class method.

        .. automethod:: from_xso_item

        .. note::

           Do not confuse this with the XSO :class:`.xso.Item`.

        """
        def __init__(self, jid, *,
                     approved=False,
                     ask=None,
                     subscription="none",
                     name=None,
                     groups=()):
            super().__init__()
            self.jid = jid
            self.subscription = subscription
            self.approved = approved
            self.ask = ask
            self.name = name
            self.groups = set(groups)

        def update_from_xso_item(self, xso_item):
            """
            Update the attributes (except :attr:`jid`) with the values obtained
            from the gixen `xso_item`.

            `xso_item` must be a valid :class:`.xso.Item` instance.
            """
            self.subscription = xso_item.subscription
            self.approved = xso_item.approved
            self.ask = xso_item.ask
            self.name = xso_item.name
            self.groups = {group.name for group in xso_item.groups}

        @classmethod
        def from_xso_item(cls, xso_item):
            """
            Create a :class:`Item` with the :attr:`jid` set to the
            :attr:`.xso.Item.jid` obtained from `xso_item`. Then update that
            instance with `xso_item` using :meth:`update_from_xso_item` and return
            it.
            """
            item = cls(xso_item.jid)
            item.update_from_xso_item(xso_item)
            return item

        def export_as_json(self):
            """
            Return a :mod:`json`-compatible dictionary which contains the
            attributes of this :class:`Item` except its JID.
            """
            result = {
                "subscription": self.subscription,
            }

            if self.name:
                result["name"] = self.name

            if self.ask is not None:
                result["ask"] = self.ask

            if self.approved:
                result["approved"] = self.approved

            if self.groups:
                result["groups"] = sorted(self.groups)

            return result

        def update_from_json(self, data):
            """
            Update the attributes of this :class:`Item` using the values obtained
            from the dictionary `data`.

            The format of `data` should be the same as the format returned by
            :meth:`export_as_json`.
            """
            self.subscription = data.get("subscription", "none")
            self.approved = bool(data.get("approved", False))
            self.ask = data.get("ask", None)
            self.name = data.get("name", None)
            self.groups = set(data.get("groups", []))

    class RosterClient(aioxmpp.service.Service):
        """
        A roster client :class:`aioxmpp.service.Service`.

        The interaction with a roster service happens mainly by accessing the
        attributes holding the state and using the events to be notified of state
        changes:

        Attributes for accessing the roster:

        .. attribute:: items

           A dictionary mapping :class:`~aioxmpp.JID` instances to corresponding
           :class:`Item` instances.

        .. attribute:: groups

           A dictionary which allows group-based access to :class:`Item`
           instances. The dictionaries keys are the names of the groups, the values
           are :class:`set` instances, which hold the :class:`Item` instances in
           that group.

           At no point one can observe empty :class:`set` instances in this
           dictionary.

        The :class:`Item` instances stay the same, as long as they represent the
        identical roster entry on the remote side. That is, if the name or
        subscription state are changed in the server side roster, the :class:`Item`
        instance stays the same, but the attributes are mutated. However, if the
        entry is removed from the server roster and re-added later for the same
        JID, it will be a different :class:`Item` instance.

        Signals:

        .. signal:: on_initial_roster_received()

           Fires when the initial roster has been received. Note that if roster
           versioning is used, the initial roster may not be up-to-date. The server
           is allowed to tell the client to re-use its local state and deliver
           changes using roster pushes. In that case, the
           :meth:`on_initial_roster_received` event fires immediately, so that the
           user sees whatever roster has been set up for versioning before the
           stream was established; updates pushed by the server are delivered using
           the normal events.

           The roster data has already been imported at the time the callback is
           fired.

           Note that the initial roster is diffed against whatever is in the local
           store and events are fired just like for normal push updates. Thus, in
           general, you won’t need this signal; it might be better to listen for
           the events below.

        .. signal:: on_entry_added(item)

           Fires when an `item` has been added to the roster. The attributes of the
           `item` are up-to-date when this callback fires.

           When the event fires, the bookkeeping structures are already updated.
           This implies that :meth:`on_group_added` is called before
           :meth:`on_entry_added` if the entry adds a new group.

        .. signal:: on_entry_name_changed(item)

           Fires when a roster update changed the name of the `item`. The new name
           is already applied to the `item`.

        .. signal:: on_entry_subscription_state_changed(item)

           Fires when a roster update changes any of the :attr:`Item.subscription`,
           :attr:`Item.ask` or :attr:`Item.approved` attributes. The new values are
           already applied to `item`.

           The event always fires once per update, even if the update changes
           more than one of the above attributes.

        .. signal:: on_entry_added_to_group(item, group_name)

           Fires when an update adds an `item` to a group. The :attr:`Item.groups`
           attribute is already updated (not only with this, but also other group
           updates, including removals) when this event is fired.

           The event fires for each added group in an update, thus it may fire more
           than once per update.

           The name of the new group is in `group_name`.

           At the time the event fires, the bookkeeping structures for the group
           are already updated; this implies that :meth:`on_group_added` fires
           *before* :meth:`on_entry_added_to_group` if the entry added a new group.

        .. signal:: on_entry_removed_from_group(item, group_name)

           Fires when an update removes an `item` from a group. The
           :attr:`Item.groups` attribute is already updated (not only with this,
           but also other group updates, including additions) when this event is
           fired.

           The event fires for each removed group in an update, thus it may fire
           more than once per update.

           The name of the new group is in `group_name`.

           At the time the event fires, the bookkeeping structures are already
           updated; this implies that :meth:`on_group_removed` fires *before*
           :meth:`on_entry_removed_from_group` if the removal of an entry from a
           group causes the group to vanish.

        .. signal:: on_entry_removed(item)

           Fires after an entry has been removed from the roster. The entry is
           already removed from all bookkeeping structures, but the values on the
           `item` object are the same as right before the removal.

           This implies that :meth:`on_group_removed` fires *before*
           :meth:`on_entry_removed` if the removal of an entry causes a group to
           vanish.

        .. signal:: on_group_added(group)

            Fires after a new group has been added to the bookkeeping structures.

            :param group: Name of the new group.
            :type group: :class:`str`

            At the time the event fires, the group is empty.

            .. versionadded:: 0.9

        .. signal:: on_group_removed(group)

            Fires after a new group has been removed from the bookkeeping
            structures.

            :param group: Name of the old group.
            :type group: :class:`str`

            At the time the event fires, the group is empty.

            .. versionadded:: 0.9

        Modifying roster contents:

        .. automethod:: set_entry

        .. automethod:: remove_entry

        Managing presence subscriptions:

        .. automethod:: approve

        .. automethod:: subscribe

        .. signal:: on_subscribe(stanza)

           Fires when a peer requested a subscription. The whole stanza recevied is
           included as `stanza`.

           .. seealso::

              To approve a subscription request, use :meth:`approve`.

        .. signal:: on_subscribed(stanza)

           Fires when a peer has confirmed a previous subscription request. The
           ``"subscribed"`` stanza is included as `stanza`.

        .. signal:: on_unsubscribe(stanza)

           Fires when a peer cancelled their subscription for our presence. As per
           :rfc:`6121`, the server forwards the ``"unsubscribe"`` presence stanza
           (which is included as `stanza` argument) *before* sending the roster
           push.

           Unless your application is interested in the specific cause of a
           subscription state change, it is not neccessary to use this signal; the
           subscription state change will be covered by
           :meth:`on_entry_subscription_state_changed`.

        .. signal:: on_unsubscribed(stanza)

           Fires when a peer cancelled our subscription. As per :rfc:`6121`, the
           server forwards the ``"unsubscribed"`` presence stanza (which is
           included as `stanza` argument) *before* sending the roster push.

           Unless your application is interested in the specific cause of a
           subscription state change, it is not neccessary to use this signal; the
           subscription state change will be covered by
           :meth:`on_entry_subscription_state_changed`.

        Import/Export of roster data:

        .. automethod:: export_as_json

        .. automethod:: import_from_json

        To make use of roster versioning, use the above two methods. The general
        workflow is to :meth:`export_as_json` the roster after disconnecting and
        storing it for the next connection attempt. **Before** connecting, the
        stored data needs to be loaded using :meth:`import_from_json`. This only
        needs to happen after a new :class:`Service` has been created, as roster
        services won’t delete roster contents between two connections on the same
        :class:`.Client` instance.

        .. versionchanged:: 0.8

           This class was formerly known as :class:`aioxmpp.roster.Service`. It
           is still available under that name, but the alias will be removed in
           1.0.
        """

        ORDER_AFTER = [
            aioxmpp.dispatcher.SimplePresenceDispatcher,
        ]

        on_initial_roster_received = callbacks.Signal()
        on_entry_name_changed = callbacks.Signal()
        on_entry_subscription_state_changed = callbacks.Signal()
        on_entry_removed = callbacks.Signal()
        on_entry_added = callbacks.Signal()
        on_entry_added_to_group = callbacks.Signal()
        on_entry_removed_from_group = callbacks.Signal()

        on_group_added = callbacks.Signal()
        on_group_removed = callbacks.Signal()

        on_subscribed = callbacks.Signal()
        on_subscribe = callbacks.Signal()
        on_unsubscribed = callbacks.Signal()
        on_unsubscribe = callbacks.Signal()

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

            self._bse_token = client.before_stream_established.connect(
                self._request_initial_roster
            )

            self.__roster_lock = asyncio.Lock()

            self.items = {}
            self.groups = {}
            self.version = None

        def _update_entry(self, xso_item):
            try:
                stored_item = self.items[xso_item.jid]
            except KeyError:
                stored_item = Item.from_xso_item(xso_item)
                self.items[xso_item.jid] = stored_item
                for group in stored_item.groups:
                    try:
                        group_members = self.groups[group]
                    except KeyError:
                        group_members = self.groups.setdefault(group, set())
                        self.on_group_added(group)
                    group_members.add(stored_item)
                self.on_entry_added(stored_item)
                return

            to_call = []

            if stored_item.name != xso_item.name:
                to_call.append(self.on_entry_name_changed)

            if (stored_item.subscription != xso_item.subscription or
                    stored_item.approved != xso_item.approved or
                    stored_item.ask != xso_item.ask):
                to_call.append(self.on_entry_subscription_state_changed)

            old_groups = set(stored_item.groups)

            stored_item.update_from_xso_item(xso_item)

            new_groups = set(stored_item.groups)

            removed_from_groups = old_groups - new_groups
            added_to_groups = new_groups - old_groups

            for cb in to_call:
                cb(stored_item)

            for group in added_to_groups:
                try:
                    group_members = self.groups[group]
                except KeyError:
                    group_members = self.groups.setdefault(group, set())
                    self.on_group_added(group)
                group_members.add(stored_item)
                self.on_entry_added_to_group(stored_item, group)

            for group in removed_from_groups:
                groupset = self.groups[group]
                groupset.remove(stored_item)
                if not groupset:
                    del self.groups[group]
                    self.on_group_removed(group)
                self.on_entry_removed_from_group(stored_item, group)

        @aioxmpp.service.iq_handler(
            aioxmpp.structs.IQType.SET,
            roster_xso.Query)
        @asyncio.coroutine
        def handle_roster_push(self, iq):
            if iq.from_ and iq.from_ != self.client.local_jid.bare():
                raise errors.XMPPAuthError(errors.ErrorCondition.FORBIDDEN)

            request = iq.payload

            with (yield from self.__roster_lock):
                for item in request.items:
                    if item.subscription == "remove":
                        try:
                            old_item = self.items.pop(item.jid)
                        except KeyError:
                            pass
                        else:
                            self._remove_from_groups(old_item, old_item.groups)
                            self.on_entry_removed(old_item)
                    else:
                        self._update_entry(item)

                self.version = request.ver

        @aioxmpp.dispatcher.presence_handler(
            aioxmpp.structs.PresenceType.SUBSCRIBE,
            None)
        def handle_subscribe(self, stanza):
            self.on_subscribe(stanza)

        @aioxmpp.dispatcher.presence_handler(
            aioxmpp.structs.PresenceType.SUBSCRIBED,
            None)
        def handle_subscribed(self, stanza):
            self.on_subscribed(stanza)

        @aioxmpp.dispatcher.presence_handler(
            aioxmpp.structs.PresenceType.UNSUBSCRIBED,
            None)
        def handle_unsubscribed(self, stanza):
            self.on_unsubscribed(stanza)

        @aioxmpp.dispatcher.presence_handler(
            aioxmpp.structs.PresenceType.UNSUBSCRIBE,
            None)
        def handle_unsubscribe(self, stanza):
            self.on_unsubscribe(stanza)

        def _remove_from_groups(self, item_to_remove, groups):
            for group in groups:
                try:
                    group_members = self.groups[group]
                except KeyError:
                    continue
                group_members.remove(item_to_remove)
                if not group_members:
                    del self.groups[group]
                    self.on_group_removed(group)

        @asyncio.coroutine
        def _request_initial_roster(self):
            iq = stanza.IQ(type_=structs.IQType.GET)
            iq.payload = roster_xso.Query()

            with (yield from self.__roster_lock):
                logger.debug("requesting initial roster")
                if self.client.stream_features.has_feature(
                        roster_xso.RosterVersioningFeature):
                    logger.debug("requesting incremental updates (old ver = %s)",
                                 self.version)
                    iq.payload.ver = self.version

                response = yield from self.client.send(
                    iq,
                    timeout=self.client.negotiation_timeout.total_seconds()
                )

                if response is None:
                    logger.debug("roster will be updated incrementally")
                    self.on_initial_roster_received()
                    return True

                self.version = response.ver
                logger.debug("roster update received (new ver = %s)", self.version)

                actual_jids = {item.jid for item in response.items}
                known_jids = set(self.items.keys())

                removed_jids = known_jids - actual_jids
                logger.debug("jids dropped: %r", removed_jids)

                for removed_jid in removed_jids:
                    old_item = self.items.pop(removed_jid)
                    self._remove_from_groups(old_item, old_item.groups)
                    self.on_entry_removed(old_item)

                logger.debug("jids updated: %r", actual_jids - removed_jids)
                for item in response.items:
                    self._update_entry(item)

                self.on_initial_roster_received()
                return True

        def export_as_json(self):
            """
            Export the whole roster as currently stored on the client side into a
            JSON-compatible dictionary and return that dictionary.
            """
            return {
                "items": {
                    str(jid): item.export_as_json()
                    for jid, item in self.items.items()
                },
                "ver": self.version
            }

        def import_from_json(self, data):
            """
            Replace the current roster with the :meth:`export_as_json`-compatible
            dictionary in `data`.

            No events are fired during this activity. After this method completes,
            the whole roster contents are exchanged with the contents from `data`.

            Also, no data is transferred to the server; this method is intended to
            be used for roster versioning. See below (in the docs of
            :class:`Service`).
            """
            self.version = data.get("ver", None)

            self.items.clear()
            self.groups.clear()
            for jid, data in data.get("items", {}).items():
                jid = structs.JID.fromstr(jid)
                item = Item(jid)
                item.update_from_json(data)
                self.items[jid] = item
                for group in item.groups:
                    self.groups.setdefault(group, set()).add(item)

        @asyncio.coroutine
        def set_entry(self, jid, *,
                      name=_Sentinel,
                      add_to_groups=frozenset(),
                      remove_from_groups=frozenset(),
                      timeout=None):
            """
            Set properties of a roster entry or add a new roster entry. The roster
            entry is identified by its bare `jid`.

            If an entry already exists, all values default to those stored in the
            existing entry. For example, if no `name` is given, the current name of
            the entry is re-used, if any.

            If the entry does not exist, it will be created on the server side.

            The `remove_from_groups` and `add_to_groups` arguments have to be based
            on the locally cached state, as XMPP does not support sending
            diffs. `remove_from_groups` takes precedence over `add_to_groups`.

            `timeout` is the time in seconds to wait for a confirmation by the
            server.

            Note that the changes may not be visible immediately after his
            coroutine returns in the :attr:`items` and :attr:`groups`
            attributes. The :class:`Service` waits for the "official" roster push
            from the server for updating the data structures and firing events, to
            ensure that consistent state with other clients is achieved.

            This may raise arbitrary :class:`.errors.XMPPError` exceptions if the
            server replies with an error and also any kind of connection error if
            the connection gets fatally terminated while waiting for a response.
            """

            existing = self.items.get(jid, Item(jid))

            post_groups = (existing.groups | add_to_groups) - remove_from_groups
            post_name = existing.name
            if name is not _Sentinel:
                post_name = name

            item = roster_xso.Item(
                jid=jid,
                name=post_name,
                groups=[
                    roster_xso.Group(name=group_name)
                    for group_name in post_groups
                ])

            yield from self.client.send(
                stanza.IQ(
                    structs.IQType.SET,
                    payload=roster_xso.Query(items=[
                        item
                    ])
                ),
                timeout=timeout
            )

        @asyncio.coroutine
        def remove_entry(self, jid, *, timeout=None):
            """
            Request removal of the roster entry identified by the given bare
            `jid`. If the entry currently has any subscription state, the server
            will send the corresponding unsubscribing presence stanzas.

            `timeout` is the maximum time in seconds to wait for a reply from the
            server.

            This may raise arbitrary :class:`.errors.XMPPError` exceptions if the
            server replies with an error and also any kind of connection error if
            the connection gets fatally terminated while waiting for a response.
            """
            yield from self.client.send(
                stanza.IQ(
                    structs.IQType.SET,
                    payload=roster_xso.Query(items=[
                        roster_xso.Item(
                            jid=jid,
                            subscription="remove"
                        )
                    ])
                ),
                timeout=timeout
            )

        def approve(self, peer_jid):
            """
            (Pre-)approve a subscription request from `peer_jid`.

            :param peer_jid: The peer to (pre-)approve.

            This sends a ``"subscribed"`` presence to the peer; if the peer has
            previously asked for a subscription, this will seal the deal and create
            the subscription.

            If the peer has not requested a subscription (yet), it is marked as
            pre-approved by the server. A future subscription request by the peer
            will then be confirmed by the server automatically.

            .. note::

                Pre-approval is an OPTIONAL feature in :rfc:`6121`. It is announced
                as a stream feature.
            """
            self.client.enqueue(
                stanza.Presence(type_=structs.PresenceType.SUBSCRIBED,
                                to=peer_jid)
            )

        def subscribe(self, peer_jid):
            """
            Request presence subscription with the given `peer_jid`.

            This is deliberately not a coroutine; we don’t know whether the peer is
            online (usually) and they may defer the confirmation very long, if they
            confirm at all. Use :meth:`on_subscribed` to get notified when a peer
            accepted a subscription request.
            """
            self.client.enqueue(
                stanza.Presence(type_=structs.PresenceType.SUBSCRIBE,
                                to=peer_jid)
            )

        def unsubscribe(self, peer_jid):
            """
            Unsubscribe from the presence of the given `peer_jid`.
            """
            self.client.enqueue(
                stanza.Presence(type_=structs.PresenceType.UNSUBSCRIBE,
                                to=peer_jid)
            )