Coverage for slidge/contact/roster.py: 72%
138 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-29 05:05 +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 contact.stored.ensure_loaded("last_msg", self.orm)
213 since = (
214 HoleBound(
215 id=contact.stored.last_sent_msg_legacy_id,
216 timestamp=contact.stored.last_sent_msg_date,
217 )
218 if contact.stored.last_sent_msg_legacy_id is not None
219 and contact.stored.last_sent_msg_date is not None
220 else None
221 )
222 try:
223 await contact.backfill(since)
224 except NotImplementedError:
225 break
226 except Exception:
227 contact.log.exception("Error during backfill")
228 if user_roster is None:
229 continue
230 item = contact.get_roster_item()
231 old = user_roster.get(contact.jid.bare)
232 if old is not None and all(
233 old[k] == item[contact.jid.bare].get(k)
234 for k in ("subscription", "groups", "name")
235 ):
236 self.log.debug("No need to update roster")
237 continue
238 self.log.debug("Updating roster")
239 if not contact.is_friend:
240 continue
241 if not self.session.user.preferences.get("roster_push", True):
242 continue
243 try:
244 await self.session.xmpp["xep_0356"].set_roster(
245 self.session.user_jid.bare,
246 item,
247 )
248 except (PermissionError, IqError, IqTimeout) as e:
249 warnings.warn(f"Could not add to roster: {e}")
250 else:
251 contact.added_to_roster = True
252 contact.send_last_presence(force=True)
253 orm.commit()
254 self.__filling = False
256 async def fill(self) -> AsyncIterator[LegacyContactType]:
257 """
258 Populate slidge's "virtual roster".
260 This should yield contacts that are meant to be added to the user's
261 roster, typically by using ``await self.by_legacy_id(contact_id)``.
262 Setting the contact nicknames, avatar, etc. should be in
263 :meth:`LegacyContact.update_info()`
265 It's not mandatory to override this method, but it is recommended way
266 to populate "friends" of the user. Calling
267 ``await (await self.by_legacy_id(contact_id)).add_to_roster()``
268 accomplishes the same thing, but doing it in here allows to batch
269 DB queries and is better performance-wise.
271 """
272 return
273 yield
275 async def on_create_group(
276 self, name: str, contacts: list[LegacyContactType]
277 ) -> str:
278 """
279 Called when the user request the creation of a group via the dedicated
280 :term:`Command`.
282 :param name: Name of the group.
283 :param contacts: list of contacts that should be members of the group.
284 :return: Legacy ID of the newly created group.
285 """
286 raise NotImplementedError
289# References to `LegacyRoster` need for this to be defined.
290LegacyRoster.contact_cls = LegacyContact # type:ignore[misc]
292log = logging.getLogger(__name__)