Coverage for slidge/contact/roster.py: 76%
130 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-18 04:30 +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
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(self, contact_jid: JID) -> LegacyContactType:
89 # """
90 # Retrieve a contact by their JID
91 #
92 # If the contact was not instantiated before, it will be created
93 # using :meth:`slidge.LegacyRoster.jid_username_to_legacy_id` to infer their
94 # legacy user ID.
95 #
96 # :param contact_jid:
97 # :return:
98 # """
99 username = contact_jid.node
100 if not username:
101 raise XMPPError(
102 "bad-request", "Contacts must have a local part in their JID"
103 )
104 contact_jid = JID(contact_jid.bare)
105 async with self.lock(("username", username)):
106 legacy_id = await self.jid_username_to_legacy_id(username)
107 if legacy_id == self.user_legacy_id:
108 raise ContactIsUser
109 if self.get_lock(("legacy_id", legacy_id)):
110 self.log.debug("Already updating %s via by_legacy_id()", contact_jid)
111 return await self.by_legacy_id(legacy_id)
113 with self.orm() as orm:
114 stored = (
115 orm.query(Contact)
116 .filter_by(user=self.user, jid_localpart=username)
117 .one_or_none()
118 )
119 if stored is None:
120 stored = Contact(
121 user_account_id=self.session.user_pk,
122 legacy_id=legacy_id,
123 jid_localpart=username,
124 )
125 return await self.__update_if_needed(stored)
127 async def __update_if_needed(self, stored: Contact) -> LegacyContactType:
128 contact = self.from_store(stored)
129 if contact.stored.updated:
130 return contact
132 with contact.updating_info():
133 await contact.update_info()
134 if contact.is_friend and not self.__filling:
135 await contact.add_to_roster()
137 if contact.cached_presence is not None:
138 contact._store_last_presence(contact.cached_presence)
139 return contact
141 def by_jid_only_if_exists(self, contact_jid: JID) -> LegacyContactType | None:
142 with self.orm() as orm:
143 stored = (
144 orm.query(Contact)
145 .filter_by(user=self.user, jid_localpart=contact_jid.local)
146 .one_or_none()
147 )
148 if stored is not None and stored.updated:
149 return self.from_store(stored)
150 return None
152 @timeit
153 async def by_legacy_id(self, /, legacy_id: str) -> LegacyContactType:
154 """
155 Retrieve a contact by their legacy_id
157 If the contact was not instantiated before, it will be created
158 using :meth:`slidge.LegacyRoster.legacy_id_to_jid_username` to infer their
159 legacy user ID.
161 :param legacy_id:
162 :return:
163 """
164 if legacy_id == self.user_legacy_id:
165 raise ContactIsUser
166 async with self.lock(("legacy_id", legacy_id)):
167 username = await self.legacy_id_to_jid_username(legacy_id)
168 if self.get_lock(("username", username)):
169 self.log.debug("Already updating %s via by_jid()", username)
171 return await self.by_jid(
172 JID(username + "@" + self.session.xmpp.boundjid.bare)
173 )
175 with self.orm() as orm:
176 stored = (
177 orm.query(Contact)
178 .filter_by(user=self.user, legacy_id=str(legacy_id))
179 .one_or_none()
180 )
181 if stored is None:
182 stored = Contact(
183 user_account_id=self.session.user_pk,
184 legacy_id=str(legacy_id),
185 jid_localpart=username,
186 )
187 return await self.__update_if_needed(stored)
189 @timeit
190 async def _fill(self, orm: OrmSession) -> None:
191 try:
192 if hasattr(self.session.xmpp, "TEST_MODE"):
193 # dirty hack to avoid mocking xmpp server replies to this
194 # during tests
195 raise PermissionError
196 iq = await self.session.xmpp["xep_0356"].get_roster(
197 self.session.user_jid.bare
198 )
199 user_roster = iq["roster"]["items"]
200 except (PermissionError, IqError, IqTimeout):
201 user_roster = None
203 self.__filling = True
204 async for contact in self.fill():
205 if user_roster is None:
206 continue
207 item = contact.get_roster_item()
208 old = user_roster.get(contact.jid.bare)
209 if old is not None and all(
210 old[k] == item[contact.jid.bare].get(k)
211 for k in ("subscription", "groups", "name")
212 ):
213 self.log.debug("No need to update roster")
214 continue
215 self.log.debug("Updating roster")
216 if not contact.is_friend:
217 continue
218 if not self.session.user.preferences.get("roster_push", True):
219 continue
220 try:
221 await self.session.xmpp["xep_0356"].set_roster(
222 self.session.user_jid.bare,
223 item,
224 )
225 except (PermissionError, IqError, IqTimeout) as e:
226 warnings.warn(f"Could not add to roster: {e}")
227 else:
228 contact.added_to_roster = True
229 contact.send_last_presence(force=True)
230 orm.commit()
231 self.__filling = False
233 async def fill(self) -> AsyncIterator[LegacyContactType]:
234 """
235 Populate slidge's "virtual roster".
237 This should yield contacts that are meant to be added to the user's
238 roster, typically by using ``await self.by_legacy_id(contact_id)``.
239 Setting the contact nicknames, avatar, etc. should be in
240 :meth:`LegacyContact.update_info()`
242 It's not mandatory to override this method, but it is recommended way
243 to populate "friends" of the user. Calling
244 ``await (await self.by_legacy_id(contact_id)).add_to_roster()``
245 accomplishes the same thing, but doing it in here allows to batch
246 DB queries and is better performance-wise.
248 """
249 return
250 yield
252 async def on_create_group(
253 self, name: str, contacts: list[LegacyContactType]
254 ) -> str:
255 """
256 Called when the user request the creation of a group via the dedicated
257 :term:`Command`.
259 :param name: Name of the group.
260 :param contacts: list of contacts that should be members of the group.
261 :return: Legacy ID of the newly created group.
262 """
263 raise NotImplementedError
266# References to `LegacyRoster` need for this to be defined.
267LegacyRoster.contact_cls = LegacyContact # type:ignore[misc]
269log = logging.getLogger(__name__)