Coverage for slidge/contact/roster.py: 55%
149 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-10 12:26 +0000
1from __future__ import annotations
3import asyncio
4import logging
5import warnings
6from collections.abc import Mapping
7from typing import TYPE_CHECKING, ClassVar, NotRequired, TypedDict, Unpack
9import sqlalchemy as sa
10import sqlalchemy.orm
11from slixmpp import JID
12from slixmpp.exceptions import IqError, IqTimeout, XMPPError
14from ..core.collection import BaseLocalPart
15from ..db.models import Contact
16from ..util.types import HoleBound
17from ..util.util import derive_wired_class, timeit
18from .contact import LegacyContact, _ContactProps
20if TYPE_CHECKING:
21 from ..core.session import BaseSession
24class _RosterItemDict(TypedDict):
25 subscription: str
26 groups: NotRequired[list[str]]
27 name: NotRequired[str]
30class LegacyRoster[T: LegacyContact = LegacyContact](BaseLocalPart[Contact, T]):
31 """
32 Virtual roster ("contact list") of a gateway user.
34 Every :class:`.BaseSession` instance will have its own :class:`.LegacyRoster` instance
35 accessible via the :attr:`.BaseSession.contacts` attribute.
37 Typically, you will mostly use the :meth:`.LegacyRoster.get_or_create` function to
38 retrieve a contact instance by its legacy ID.
40 You might need to override :meth:`slidge.BaseGateway.legacy_id_to_localpart`
41 and :meth:`slidge.BaseGateway.localpart_to_legacy_id` to customize or escape
42 the mapping between legacy IDs and localparts (which are **case-insensitive**).
43 """
45 # References to `LegacyRoster` need for this to be defined.
46 _contact_cls: ClassVar[type[LegacyContact]] = LegacyContact # type:ignore[type-abstract]
47 """
48 The concrete :class:`.LegacyContact` subclass this roster produces.
50 Derived automatically from the generic parameter, e.g.,
51 ``class Roster(LegacyRoster[Contact])`` produces ``Contact`` instances.
52 """
54 def __init_subclass__(cls, **kwargs: object) -> None:
55 super().__init_subclass__(**kwargs)
56 derive_wired_class(cls, LegacyRoster, "_contact_cls")
58 def __init__(self, session: BaseSession) -> None:
59 super().__init__(
60 session=session,
61 model=Contact,
62 child_cls=self._contact_cls, # type:ignore[arg-type]
63 store=session.xmpp.store.contacts,
64 legacy_type="contact",
65 )
67 self.log = logging.getLogger(f"{session.user_jid.bare}:roster")
68 self.user_legacy_id: str | None = None
69 self.ready: asyncio.Future[bool] = session.xmpp.loop.create_future()
71 def known_contacts(self, only_friends: bool = True) -> Mapping[str, T]:
72 if only_friends:
73 return {c.jid.bare: c for c in self if c.is_friend}
74 return {c.jid.bare: c for c in self}
76 def _create(self, legacy_id: str, **kwargs: Unpack[_ContactProps]) -> T: # type:ignore[override]
77 if name := kwargs.pop("name", None):
78 kwargs["nick"] = name # type:ignore[typeddict-unknown-key]
79 return super()._create(legacy_id, **kwargs)
81 @property
82 def __privileged(self) -> bool:
83 priv = (
84 self.xmpp.plugin["xep_0356"]
85 .granted_privileges[self._user.jid.server]
86 .roster
87 == "both"
88 )
89 if not priv:
90 warnings.warn(
91 f"Slidge is unprivileged for {self._user.jid.server}, roster push is not possible",
92 UserWarning,
93 )
94 return priv
96 @timeit
97 async def _fill(self) -> None:
98 user_roster = await self.__get_user_roster() if self.__privileged else None
100 updates = []
101 presences = []
103 client_types = set[str]()
104 async for contact in self._batch_fill():
105 await self.__backfill(contact)
106 if not (
107 contact.is_friend
108 and self._session.user.preferences.get("roster_push", True)
109 ):
110 continue
112 self._session.create_task(
113 contact._broadcast_pubsub_items(),
114 f"startup pep broadcast for {contact}",
115 )
117 if user_roster is None:
118 if contact.is_friend:
119 self.__add_to_roster_fallback(contact)
120 continue
122 contact._stored.ensure_loaded("presence", self._orm)
123 if contact._stored.client_type not in client_types:
124 await self.xmpp._populate_caps_cache(contact)
125 client_types.add(contact._stored.client_type)
126 presences.append(contact._make_last_presence())
127 updates.append({"id": contact._stored.id, "added_to_roster": True})
128 item = user_roster.get(contact.jid.bare)
129 await self.__roster_push_contact(contact, item)
131 for presence in presences:
132 presence["to"] = self._user.jid
133 self._session.xmpp.send(presence)
135 def __add_to_roster_fallback(self, contact: LegacyContact) -> None:
136 contact.send_friend_request(
137 f"I'm already your friend on {self.xmpp.COMPONENT_TYPE}, but "
138 "slidge is not allowed to manage your roster."
139 )
141 async def __backfill(self, contact: LegacyContact) -> None:
142 contact._stored.ensure_loaded("last_msg", self._orm)
143 since = (
144 HoleBound(
145 id=contact._stored.last_sent_msg_legacy_id,
146 timestamp=contact._stored.last_sent_msg_date,
147 )
148 if contact._stored.last_sent_msg_legacy_id is not None
149 and contact._stored.last_sent_msg_date is not None
150 else None
151 )
152 try:
153 await contact.backfill(since)
154 except Exception:
155 contact.log.exception("Error during backfill")
157 async def _add_to_roster(self, contact: LegacyContact) -> None:
158 """
159 Add this contact to the user roster using :xep:`0356`.
160 (slidge private method)
161 """
162 if (
163 self.__privileged
164 and self._session.user.preferences.get("roster_push", True)
165 and not self.is_filling
166 and not contact.added_to_roster
167 ):
168 # Roster get queries always return the full roster so probably very
169 # wasteful to do here. If needed, we need to try to cache the user
170 # roster somehow.
171 # item = (await self.__get_user_roster()).get(contact.jid.bare)
172 if await self.__roster_push_contact(contact):
173 contact.added_to_roster = True
175 # we only broadcast pubsub events for contacts added to the roster
176 # so if something was set before, we need to push it now
177 contact.send_last_presence(force=True)
179 def __get_roster_item(self, contact: Contact) -> _RosterItemDict:
180 item: _RosterItemDict = {
181 "subscription": "both" if contact.is_friend else "none",
182 "groups": [self.xmpp.ROSTER_GROUP],
183 }
184 return item
186 async def _sync(self) -> tuple[int, int, int]:
187 if not self.__privileged:
188 raise XMPPError(
189 "internal-server-error",
190 f"Slidge is not privileged for {self._user.jid}",
191 )
192 user_roster = await self.__get_user_roster()
193 if user_roster is None:
194 raise XMPPError("internal-server-error", "Could not get your roster.")
195 with self._orm() as orm:
196 contacts = list(
197 orm.scalars(
198 sa.select(Contact).where(
199 Contact.user_account_id == self._session.user_pk,
200 Contact.is_friend,
201 )
202 )
203 )
204 added = 0
205 removed = 0
206 updated = 0
207 for contact in contacts:
208 item = user_roster.pop(contact.jid.bare, None)
209 if item is None:
210 added += 1
211 await self.__roster_push_contact(contact)
212 else:
213 if await self.__roster_push_contact(contact, item):
214 updated += 1
215 for jid_str in user_roster:
216 jid = JID(jid_str)
217 if jid.local and jid.domain == self.xmpp.boundjid.bare:
218 await self.__roster_push(jid_str, {"subscription": "remove"})
219 removed += 1
220 with self._orm() as orm:
221 orm.execute(
222 sa.update(Contact)
223 .where(Contact.id.in_([c.id for c in contacts]))
224 .values({Contact.added_to_roster: True})
225 )
226 orm.commit()
227 return added, removed, updated
229 async def __roster_push(self, jid: str, item: _RosterItemDict) -> None:
230 try:
231 await self._session.xmpp.plugin["xep_0356"].set_roster(
232 self._session.user_jid.bare, {jid: item}
233 )
234 except (IqError, IqTimeout) as e:
235 self._session.log.error("Error while pushing to roster: %r", e)
237 async def __roster_push_contact(
238 self, contact: Contact | LegacyContact, item: _RosterItemDict | None = None
239 ) -> bool:
240 if isinstance(contact, LegacyContact):
241 contact = contact._stored
242 if item is None:
243 item = self.__get_roster_item(contact)
244 elif not self.__update_roster_item(item):
245 return False
246 await self.__roster_push(contact.jid.bare, item)
247 return True
249 async def __get_user_roster(self) -> dict[str, _RosterItemDict] | None:
250 try:
251 iq = await self._session.xmpp.plugin["xep_0356"].get_roster(
252 self._session.user_jid.bare
253 )
254 except (IqError, IqTimeout) as e:
255 self._session.log.error("Error while getting roster: %r", e)
256 return None
257 else:
258 return iq["roster"]["items"] # type:ignore[no-any-return]
260 def __update_roster_item(self, item: _RosterItemDict) -> bool:
261 groups = {self.xmpp.ROSTER_GROUP, *item["groups"]}
262 change = False
263 if set(item["groups"]) != groups:
264 item["groups"] = list(groups)
265 change = True
266 if item["subscription"] != "both":
267 item["subscription"] = "both"
268 change = True
269 return change
271 def _batch_load(self) -> sa.Select[Contact]:
272 return (
273 super()
274 ._batch_load()
275 .options(
276 sa.orm.undefer_group("presence"),
277 sa.orm.undefer_group("vcard"),
278 sa.orm.undefer_group("last_msg"),
279 )
280 )
283_VCARD_FIELDS = {
284 "full_name",
285 "birthday",
286 "bio",
287 "url",
288 "email",
289 "country",
290 "locality",
291 "pronouns",
292}
295log = logging.getLogger(__name__)