Coverage for slidge/contact/roster.py: 75%
126 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-28 18:29 +0000
1import asyncio
2import logging
3import warnings
4from collections.abc import AsyncIterator, Iterator
5from typing import Generic
7from slixmpp import JID
8from slixmpp.exceptions import IqError, IqTimeout, XMPPError
9from sqlalchemy.orm import Session
10from sqlalchemy.orm import Session as OrmSession
12from ..db.models import Contact, GatewayUser
13from ..util import SubclassableOnce
14from ..util.jid_escaping import EscapeMixin
15from ..util.lock import NamedLockMixin
16from ..util.types import AnySession, LegacyContactType
17from ..util.util import timeit
20class ContactIsUser(Exception):
21 pass
24class LegacyRoster(
25 NamedLockMixin,
26 EscapeMixin,
27 SubclassableOnce,
28 Generic[LegacyContactType],
29):
30 """
31 Virtual roster of a gateway user that allows to represent all
32 of their contacts as singleton instances (if used properly and not too bugged).
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.by_legacy_id` function to
38 retrieve a contact instance.
40 You might need to override :meth:`.LegacyRoster.legacy_id_to_jid_username` and/or
41 :meth:`.LegacyRoster.jid_username_to_legacy_id` to incorporate some custom logic
42 if you need some characters when translation JID user parts and legacy IDs.
43 """
45 _contact_cls: type[LegacyContactType]
47 def __init__(self, session: AnySession) -> None:
48 super().__init__()
50 self.log = logging.getLogger(f"{session.user_jid.bare}:roster")
51 self.user_legacy_id: str | None = None
52 self.ready: asyncio.Future[bool] = session.xmpp.loop.create_future()
54 self.session = session
55 self.__filling = False
57 @property
58 def user(self) -> GatewayUser:
59 return self.session.user
61 def orm(self) -> Session:
62 return self.session.xmpp.store.session()
64 def from_store(self, stored: Contact) -> LegacyContactType:
65 return self._contact_cls(self.session, stored=stored)
67 def __repr__(self) -> str:
68 return f"<Roster of {self.session.user_jid}>"
70 def __iter__(self) -> Iterator[LegacyContactType]:
71 with self.orm() as orm:
72 contacts = orm.query(Contact).filter_by(user=self.user, updated=True).all()
73 for stored in contacts:
74 yield self.from_store(stored)
76 def known_contacts(self, only_friends: bool = True) -> dict[str, LegacyContactType]:
77 if only_friends:
78 return {c.jid.bare: c for c in self if c.is_friend}
79 return {c.jid.bare: c for c in self}
81 async def by_jid(self, contact_jid: JID) -> LegacyContactType:
82 # """
83 # Retrieve a contact by their JID
84 #
85 # If the contact was not instantiated before, it will be created
86 # using :meth:`slidge.LegacyRoster.jid_username_to_legacy_id` to infer their
87 # legacy user ID.
88 #
89 # :param contact_jid:
90 # :return:
91 # """
92 username = contact_jid.node
93 if not username:
94 raise XMPPError(
95 "bad-request", "Contacts must have a local part in their JID"
96 )
97 contact_jid = JID(contact_jid.bare)
98 async with self.lock(("username", username)):
99 legacy_id = await self.jid_username_to_legacy_id(username)
100 if legacy_id == self.user_legacy_id:
101 raise ContactIsUser
102 if self.get_lock(("legacy_id", legacy_id)):
103 self.log.debug("Already updating %s via by_legacy_id()", contact_jid)
104 return await self.by_legacy_id(legacy_id)
106 with self.orm() as orm:
107 stored = (
108 orm.query(Contact)
109 .filter_by(user=self.user, jid_localpart=username)
110 .one_or_none()
111 )
112 if stored is None:
113 stored = Contact(
114 user_account_id=self.session.user_pk,
115 legacy_id=legacy_id,
116 jid_localpart=username,
117 )
118 return await self.__update_if_needed(stored)
120 async def __update_if_needed(self, stored: Contact) -> LegacyContactType:
121 contact = self.from_store(stored)
122 if contact.stored.updated:
123 return contact
125 with contact.updating_info():
126 await contact.update_info()
127 if contact.is_friend and not self.__filling:
128 await contact.add_to_roster()
130 if contact.cached_presence is not None:
131 contact._store_last_presence(contact.cached_presence)
132 return contact
134 def by_jid_only_if_exists(self, contact_jid: JID) -> LegacyContactType | None:
135 with self.orm() as orm:
136 stored = (
137 orm.query(Contact)
138 .filter_by(user=self.user, jid_localpart=contact_jid.local)
139 .one_or_none()
140 )
141 if stored is not None and stored.updated:
142 return self.from_store(stored)
143 return None
145 @timeit
146 async def by_legacy_id(self, /, legacy_id: str) -> LegacyContactType:
147 """
148 Retrieve a contact by their legacy_id
150 If the contact was not instantiated before, it will be created
151 using :meth:`slidge.LegacyRoster.legacy_id_to_jid_username` to infer their
152 legacy user ID.
154 :param legacy_id:
155 :return:
156 """
157 if legacy_id == self.user_legacy_id:
158 raise ContactIsUser
159 async with self.lock(("legacy_id", legacy_id)):
160 username = await self.legacy_id_to_jid_username(legacy_id)
161 if self.get_lock(("username", username)):
162 self.log.debug("Already updating %s via by_jid()", username)
164 return await self.by_jid(
165 JID(username + "@" + self.session.xmpp.boundjid.bare)
166 )
168 with self.orm() as orm:
169 stored = (
170 orm.query(Contact)
171 .filter_by(user=self.user, legacy_id=str(legacy_id))
172 .one_or_none()
173 )
174 if stored is None:
175 stored = Contact(
176 user_account_id=self.session.user_pk,
177 legacy_id=str(legacy_id),
178 jid_localpart=username,
179 )
180 return await self.__update_if_needed(stored)
182 @timeit
183 async def _fill(self, orm: OrmSession) -> None:
184 try:
185 if hasattr(self.session.xmpp, "TEST_MODE"):
186 # dirty hack to avoid mocking xmpp server replies to this
187 # during tests
188 raise PermissionError
189 iq = await self.session.xmpp["xep_0356"].get_roster(
190 self.session.user_jid.bare
191 )
192 user_roster = iq["roster"]["items"]
193 except (PermissionError, IqError, IqTimeout):
194 user_roster = None
196 self.__filling = True
197 async for contact in self.fill():
198 if user_roster is None:
199 continue
200 item = contact.get_roster_item()
201 old = user_roster.get(contact.jid.bare)
202 if old is not None and all(
203 old[k] == item[contact.jid.bare].get(k)
204 for k in ("subscription", "groups", "name")
205 ):
206 self.log.debug("No need to update roster")
207 continue
208 self.log.debug("Updating roster")
209 if not contact.is_friend:
210 continue
211 if not self.session.user.preferences.get("roster_push", True):
212 continue
213 try:
214 await self.session.xmpp["xep_0356"].set_roster(
215 self.session.user_jid.bare,
216 item,
217 )
218 except (PermissionError, IqError, IqTimeout) as e:
219 warnings.warn(f"Could not add to roster: {e}")
220 else:
221 contact.added_to_roster = True
222 contact.send_last_presence(force=True)
223 orm.commit()
224 self.__filling = False
226 async def fill(self) -> AsyncIterator[LegacyContactType]:
227 """
228 Populate slidge's "virtual roster".
230 This should yield contacts that are meant to be added to the user's
231 roster, typically by using ``await self.by_legacy_id(contact_id)``.
232 Setting the contact nicknames, avatar, etc. should be in
233 :meth:`LegacyContact.update_info()`
235 It's not mandatory to override this method, but it is recommended way
236 to populate "friends" of the user. Calling
237 ``await (await self.by_legacy_id(contact_id)).add_to_roster()``
238 accomplishes the same thing, but doing it in here allows to batch
239 DB queries and is better performance-wise.
241 """
242 return
243 yield
246log = logging.getLogger(__name__)