Coverage for slidge/contact/roster.py: 72%
137 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-10 04:45 +0000
1import asyncio
2import logging
3import warnings
4from collections.abc import AsyncIterator, Iterator
6from slixmpp import JID
7from slixmpp.exceptions import IqError, IqTimeout, XMPPError
8from sqlalchemy.orm import Session
9from sqlalchemy.orm import Session as OrmSession
11from ..db.models import Contact, GatewayUser
12from ..util.jid_escaping import EscapeMixin
13from ..util.lock import NamedLockMixin
14from ..util.types import AnySession, HoleBound
15from ..util.util import derive_wired_class, timeit
16from .contact import LegacyContact
19class ContactIsUser(Exception):
20 pass
23class LegacyRoster[LegacyContactType: LegacyContact](
24 NamedLockMixin,
25 EscapeMixin,
26):
27 """
28 Virtual roster of a gateway user that allows to represent all
29 of their contacts as singleton instances (if used properly and not too bugged).
31 Every :class:`.BaseSession` instance will have its own :class:`.LegacyRoster` instance
32 accessible via the :attr:`.BaseSession.contacts` attribute.
34 Typically, you will mostly use the :meth:`.LegacyRoster.by_legacy_id` function to
35 retrieve a contact instance.
37 You might need to override :meth:`.LegacyRoster.legacy_id_to_jid_username` and/or
38 :meth:`.LegacyRoster.jid_username_to_legacy_id` to incorporate some custom logic
39 if you need some characters when translation JID user parts and legacy IDs.
40 """
42 contact_cls: type[LegacyContactType]
43 """
44 The concrete :class:`.LegacyContact` subclass this roster produces.
46 Derived automatically from the generic parameter, e.g.,
47 ``class Roster(LegacyRoster[Contact])`` produces ``Contact`` instances.
48 """
50 def __init_subclass__(cls, **kwargs: object) -> None:
51 super().__init_subclass__(**kwargs)
52 derive_wired_class(cls, LegacyRoster, "contact_cls")
54 def __init__(self, session: AnySession) -> None:
55 super().__init__()
57 self.log = logging.getLogger(f"{session.user_jid.bare}:roster")
58 self.user_legacy_id: str | None = None
59 self.ready: asyncio.Future[bool] = session.xmpp.loop.create_future()
61 self.session = session
62 self.__filling = False
64 @property
65 def user(self) -> GatewayUser:
66 return self.session.user
68 def orm(self) -> Session:
69 return self.session.xmpp.store.session()
71 def from_store(self, stored: Contact) -> LegacyContactType:
72 return self.contact_cls(self.session, stored=stored)
74 def __repr__(self) -> str:
75 return f"<Roster of {self.session.user_jid}>"
77 def __iter__(self) -> Iterator[LegacyContactType]:
78 with self.orm() as orm:
79 contacts = orm.query(Contact).filter_by(user=self.user, updated=True).all()
80 for stored in contacts:
81 yield self.from_store(stored)
83 def known_contacts(self, only_friends: bool = True) -> dict[str, LegacyContactType]:
84 if only_friends:
85 return {c.jid.bare: c for c in self if c.is_friend}
86 return {c.jid.bare: c for c in self}
88 async def by_jid(
89 self, contact_jid: JID, *update_info_args: object
90 ) -> LegacyContactType:
91 # """
92 # Retrieve a contact by their JID
93 #
94 # If the contact was not instantiated before, it will be created
95 # using :meth:`slidge.LegacyRoster.jid_username_to_legacy_id` to infer their
96 # legacy user ID.
97 #
98 # :param contact_jid:
99 # :return:
100 # """
101 username = contact_jid.node
102 if not username:
103 raise XMPPError(
104 "bad-request", "Contacts must have a local part in their JID"
105 )
106 contact_jid = JID(contact_jid.bare)
107 async with self.lock(("username", username)):
108 legacy_id = await self.jid_username_to_legacy_id(username)
109 if legacy_id == self.user_legacy_id:
110 raise ContactIsUser
111 if self.get_lock(("legacy_id", legacy_id)):
112 self.log.debug("Already updating %s via by_legacy_id()", contact_jid)
113 return await self.by_legacy_id(legacy_id)
115 with self.orm() as orm:
116 stored = (
117 orm.query(Contact)
118 .filter_by(user=self.user, jid_localpart=username)
119 .one_or_none()
120 )
121 if stored is None:
122 stored = Contact(
123 user_account_id=self.session.user_pk,
124 legacy_id=legacy_id,
125 jid_localpart=username,
126 )
127 return await self.__update_if_needed(stored, *update_info_args)
129 async def __update_if_needed(
130 self, stored: Contact, *update_info_args: object
131 ) -> LegacyContactType:
132 contact = self.from_store(stored)
133 if contact.stored.updated and not update_info_args:
134 return contact
136 with contact.updating_info():
137 await contact.update_info(*update_info_args)
138 if contact.is_friend and not self.__filling:
139 await contact.add_to_roster()
141 if contact.cached_presence is not None:
142 contact._store_last_presence(contact.cached_presence)
143 return contact
145 def by_jid_only_if_exists(self, contact_jid: JID) -> LegacyContactType | None:
146 with self.orm() as orm:
147 stored = (
148 orm.query(Contact)
149 .filter_by(user=self.user, jid_localpart=contact_jid.local)
150 .one_or_none()
151 )
152 if stored is not None and stored.updated:
153 return self.from_store(stored)
154 return None
156 @timeit
157 async def by_legacy_id(
158 self, /, legacy_id: str, *update_info_args: object
159 ) -> LegacyContactType:
160 """
161 Retrieve a contact by their legacy_id
163 If the contact was not instantiated before, it will be created
164 using :meth:`slidge.LegacyRoster.legacy_id_to_jid_username` to infer their
165 legacy user ID.
167 :param legacy_id:
168 :return:
169 """
170 if legacy_id == self.user_legacy_id:
171 raise ContactIsUser
172 async with self.lock(("legacy_id", legacy_id)):
173 username = await self.legacy_id_to_jid_username(legacy_id)
174 if self.get_lock(("username", username)):
175 self.log.debug("Already updating %s via by_jid()", username)
177 return await self.by_jid(
178 JID(username + "@" + self.session.xmpp.boundjid.bare),
179 *update_info_args,
180 )
182 with self.orm() as orm:
183 stored = (
184 orm.query(Contact)
185 .filter_by(user=self.user, legacy_id=str(legacy_id))
186 .one_or_none()
187 )
188 if stored is None:
189 stored = Contact(
190 user_account_id=self.session.user_pk,
191 legacy_id=str(legacy_id),
192 jid_localpart=username,
193 )
194 return await self.__update_if_needed(stored, *update_info_args)
196 @timeit
197 async def _fill(self, orm: OrmSession) -> None:
198 try:
199 if hasattr(self.session.xmpp, "TEST_MODE"):
200 # dirty hack to avoid mocking xmpp server replies to this
201 # during tests
202 raise PermissionError
203 iq = await self.session.xmpp["xep_0356"].get_roster(
204 self.session.user_jid.bare
205 )
206 user_roster = iq["roster"]["items"]
207 except (PermissionError, IqError, IqTimeout):
208 user_roster = None
210 self.__filling = True
211 async for contact in self.fill():
212 since = (
213 HoleBound(
214 id=contact.stored.last_sent_msg_legacy_id,
215 timestamp=contact.stored.last_sent_msg_date,
216 )
217 if contact.stored.last_sent_msg_legacy_id is not None
218 and contact.stored.last_sent_msg_date is not None
219 else None
220 )
221 try:
222 await contact.backfill(since)
223 except NotImplementedError:
224 break
225 except Exception:
226 contact.log.exception("Error during backfill")
227 if user_roster is None:
228 continue
229 item = contact.get_roster_item()
230 old = user_roster.get(contact.jid.bare)
231 if old is not None and all(
232 old[k] == item[contact.jid.bare].get(k)
233 for k in ("subscription", "groups", "name")
234 ):
235 self.log.debug("No need to update roster")
236 continue
237 self.log.debug("Updating roster")
238 if not contact.is_friend:
239 continue
240 if not self.session.user.preferences.get("roster_push", True):
241 continue
242 try:
243 await self.session.xmpp["xep_0356"].set_roster(
244 self.session.user_jid.bare,
245 item,
246 )
247 except (PermissionError, IqError, IqTimeout) as e:
248 warnings.warn(f"Could not add to roster: {e}")
249 else:
250 contact.added_to_roster = True
251 contact.send_last_presence(force=True)
252 orm.commit()
253 self.__filling = False
255 async def fill(self) -> AsyncIterator[LegacyContactType]:
256 """
257 Populate slidge's "virtual roster".
259 This should yield contacts that are meant to be added to the user's
260 roster, typically by using ``await self.by_legacy_id(contact_id)``.
261 Setting the contact nicknames, avatar, etc. should be in
262 :meth:`LegacyContact.update_info()`
264 It's not mandatory to override this method, but it is recommended way
265 to populate "friends" of the user. Calling
266 ``await (await self.by_legacy_id(contact_id)).add_to_roster()``
267 accomplishes the same thing, but doing it in here allows to batch
268 DB queries and is better performance-wise.
270 """
271 return
272 yield
274 async def on_create_group(
275 self, name: str, contacts: list[LegacyContactType]
276 ) -> str:
277 """
278 Called when the user request the creation of a group via the dedicated
279 :term:`Command`.
281 :param name: Name of the group.
282 :param contacts: list of contacts that should be members of the group.
283 :return: Legacy ID of the newly created group.
284 """
285 raise NotImplementedError
288# References to `LegacyRoster` need for this to be defined.
289LegacyRoster.contact_cls = LegacyContact # type:ignore[misc]
291log = logging.getLogger(__name__)