Coverage for slidge/core/dispatcher/util.py: 96%

1from __future__ import annotations 

2 

3import logging 

4from collections.abc import Awaitable, Callable 

5from functools import wraps 

6from typing import TypeVar 

7 

8from slixmpp import JID, Iq, Message, Presence 

9from slixmpp.exceptions import XMPPError 

10from slixmpp.xmlstream import StanzaBase 

11 

12from ...contact import LegacyContact 

13from ...core.session import BaseSession 

14from ...util.types import AnyGateway, AnyRecipient, AnyRoom, AnySession 

15from ...util.util import recover_replaced_legacy_msg_id 

16 

17 

18class Ignore(BaseException): 

19 pass 

20 

21 

22class DispatcherMixin: 

23 __slots__: list[str] = [] 

24 xmpp: AnyGateway 

25 

26 def __init__(self, xmpp: AnyGateway) -> None: 

27 self.xmpp = xmpp # type:ignore[misc] 

28 

29 async def _get_session( 

30 self, 

31 stanza: Message | Presence | Iq, 

32 timeout: int | None = 10, 

33 wait_for_ready: bool = True, 

34 logged: bool = False, 

35 ) -> BaseSession: 

36 xmpp = self.xmpp 

37 if stanza.get_from().server == xmpp.boundjid.bare: 

38 log.debug("Ignoring echo") 

39 raise Ignore 

40 if ( 

41 isinstance(stanza, Message) 

42 and stanza.get_type() == "chat" 

43 and stanza.get_to() == xmpp.boundjid.bare 

44 ): 

45 log.debug("Ignoring message to component") 

46 raise Ignore 

47 session = await self._get_session_from_jid( 

48 stanza.get_from(), timeout, wait_for_ready, logged 

49 ) 

50 if isinstance(stanza, Message) and _ignore(session, stanza): 

51 raise Ignore 

52 return session 

53 

54 async def _get_session_from_jid( 

55 self, 

56 jid: JID, 

57 timeout: int | None = 10, 

58 wait_for_ready: bool = True, 

59 logged: bool = False, 

60 ) -> AnySession: 

61 session = self.xmpp.get_session_from_jid(jid) 

62 if session is None: 

63 raise XMPPError("registration-required") 

64 if logged: 

65 session.raise_if_not_logged() 

66 if wait_for_ready: 

67 await session.wait_for_ready(timeout) 

68 return session 

69 

70 async def get_muc_from_stanza(self, iq: Iq | Message | Presence) -> AnyRoom: 

71 ito = iq.get_to() 

72 if ito == self.xmpp.boundjid.bare: 

73 raise XMPPError("bad-request", text="This is only handled for Rooms") 

74 

75 session = await self._get_session(iq, logged=True) 

76 muc = await session.rooms.get_or_create_by_jid(ito) 

77 return muc 

78 

79 def _xmpp_msg_id_to_legacy( 

80 self, 

81 xmpp_id: str, 

82 recipient: AnyRecipient, 

83 origin: bool = False, 

84 ) -> str: 

85 with self.xmpp.store.session() as orm: 

86 legacy_id = self.xmpp.store.id_map.get_legacy( 

87 orm, recipient._stored.id, xmpp_id, recipient.is_group, origin 

88 ) 

89 if legacy_id is None: 

90 legacy_id = xmpp_id 

91 

92 return legacy_id if origin else recover_replaced_legacy_msg_id(legacy_id) 

93 

94 async def _get_recipient_and_thread( 

95 self, msg: Message 

96 ) -> tuple[AnyRecipient, str | None]: 

97 session = await self._get_session(msg) 

98 e: AnyRecipient = await get_recipient(session, msg) 

99 legacy_thread = await self._xmpp_to_legacy_thread(session, msg, e) 

100 return e, legacy_thread 

101 

102 async def _xmpp_to_legacy_thread[T: LegacyContact | AnyRoom]( 

103 self, session: AnySession, msg: Message, recipient: T 

104 ) -> str | None: 

105 if not self.xmpp.THREADS: 

106 return None 

107 

108 xmpp_thread = msg["thread"] 

109 if not xmpp_thread: 

110 return None 

111 

112 with session.xmpp.store.session() as orm: 

113 legacy_thread_str = session.xmpp.store.id_map.get_thread( 

114 orm, recipient._stored.id, xmpp_thread, recipient.is_group 

115 ) 

116 if legacy_thread_str is not None: 

117 return legacy_thread_str 

118 async with session.thread_creation_lock: 

119 if "reply" in msg and "id" in msg["reply"]: 

120 root_legacy_id = self._xmpp_msg_id_to_legacy( 

121 msg["reply"]["id"], recipient 

122 ) 

123 else: 

124 root_legacy_id = self._xmpp_msg_id_to_legacy( 

125 xmpp_thread, 

126 recipient, 

127 ) 

128 legacy_thread = await recipient.create_thread( 

129 xmpp_thread, 

130 root_legacy_id, 

131 ) 

132 with session.xmpp.store.session() as orm: 

133 session.xmpp.store.id_map.set_thread( 

134 orm, 

135 recipient._stored.id, 

136 str(legacy_thread), 

137 xmpp_thread, 

138 recipient.is_group, 

139 ) 

140 orm.commit() 

141 return legacy_thread 

142 

143 

144def _ignore(session: AnySession, msg: Message) -> bool: 

145 i = msg.get_id() 

146 if i.startswith("slidge-carbon-"): 

147 return True 

148 if i not in session.ignore_messages: 

149 return False 

150 session.log.debug("Ignored sent carbon: %s", i) 

151 session.ignore_messages.remove(i) 

152 return True 

153 

154 

155async def get_recipient(session: BaseSession, m: Message) -> AnyRecipient: 

156 session.raise_if_not_logged() 

157 if m.get_type() == "groupchat": 

158 muc = await session.rooms.get_or_create_by_jid(m.get_to()) 

159 r = m.get_from().resource 

160 if r not in muc.get_user_resources(): 

161 session.create_task(muc.kick_resource(r), name=f"kick of {r} from {muc}") 

162 raise XMPPError("not-acceptable", "You are not connected to this chat") 

163 return muc 

164 else: 

165 return await session.contacts.get_or_create_by_jid(m.get_to()) 

166 

167 

168SelfType = TypeVar("SelfType") 

169StanzaType = TypeVar("StanzaType", bound=StanzaBase) 

170HandlerType = Callable[[SelfType, StanzaType], Awaitable[None]] 

171 

172 

173def exceptions_to_xmpp_errors[SelfType, StanzaType: StanzaBase]( 

174 cb: HandlerType[SelfType, StanzaType], 

175) -> HandlerType[SelfType, StanzaType]: 

176 @wraps(cb) 

177 async def wrapped(self: SelfType, stanza: StanzaType) -> None: 

178 try: 

179 await cb(self, stanza) 

180 except Ignore: 

181 pass 

182 except XMPPError: 

183 raise 

184 except NotImplementedError as e: 

185 log.debug("NotImplementedError raised in %s", cb) 

186 tb = e.__traceback__ 

187 assert tb is not None 

188 while tb.tb_next is not None: 

189 tb = tb.tb_next 

190 frame = tb.tb_frame 

191 method_name = frame.f_code.co_name 

192 self_obj = frame.f_locals.get("self") 

193 if self_obj is not None: 

194 method_name = f"{self_obj.__class__.__name__}.{method_name}" 

195 raise XMPPError( 

196 "feature-not-implemented", 

197 f"{method_name} is not implemented by the legacy module", 

198 clear=False, 

199 ) 

200 except Exception as e: 

201 log.error( 

202 "Failed to handle incoming stanza: %s - %s", self, stanza, exc_info=e 

203 ) 

204 raise XMPPError("internal-server-error", str(e)) 

205 

206 return wrapped 

207 

208 

209log = logging.getLogger(__name__)