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