Coverage for slidge/core/dispatcher/muc/mam.py: 92%

1import asyncio 

2from typing import TYPE_CHECKING 

3 

4from slixmpp import CoroutineCallback, Iq, StanzaPath 

5from slixmpp.exceptions import XMPPError 

6from slixmpp.xmlstream import StanzaBase 

7 

8from ... import config 

9from ..util import DispatcherMixin, exceptions_to_xmpp_errors 

10 

11if TYPE_CHECKING: 

12 from slidge import BaseGateway 

13 

14 

15class MamMixin(DispatcherMixin): 

16 __slots__: list[str] = [] 

17 __mam_cleanup_task: asyncio.Task[None] 

18 

19 def __init__(self, xmpp: "BaseGateway") -> None: 

20 super().__init__(xmpp) 

21 self.__mam_cleanup_task = xmpp.loop.create_task( # type:ignore[misc] 

22 self.__mam_cleanup(), 

23 name="MAM cleanup", 

24 ) 

25 xmpp.register_handler( 

26 CoroutineCallback( 

27 "MAM_query", 

28 StanzaPath("iq@type=set/mam"), 

29 self.__handle_mam, # type:ignore[arg-type] 

30 ) 

31 ) 

32 xmpp.register_handler( 

33 CoroutineCallback( 

34 "MAM_get_from", 

35 StanzaPath("iq@type=get/mam"), 

36 self.__handle_mam_get_form, 

37 ) 

38 ) 

39 xmpp.register_handler( 

40 CoroutineCallback( 

41 "MAM_get_meta", 

42 StanzaPath("iq@type=get/mam_metadata"), 

43 self.__handle_mam_metadata, # type:ignore[arg-type] 

44 ) 

45 ) 

46 

47 async def __mam_cleanup(self) -> None: 

48 if not config.MAM_MAX_DAYS: 

49 return 

50 while True: 

51 await asyncio.sleep(3600 * 6) 

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

53 self.xmpp.store.mam.nuke_older_than(orm, config.MAM_MAX_DAYS) 

54 orm.commit() 

55 

56 @exceptions_to_xmpp_errors 

57 async def __handle_mam(self, iq: Iq) -> None: 

58 muc = await self.get_muc_from_stanza(iq) 

59 await muc.send_mam(iq) 

60 

61 async def __handle_mam_get_form(self, iq: StanzaBase) -> None: 

62 assert isinstance(iq, Iq) 

63 ito = iq.get_to() 

64 

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

66 raise XMPPError( 

67 text="No MAM on the component itself, use a JID with a resource" 

68 ) 

69 

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

71 await session.bookmarks.by_jid(ito) 

72 

73 reply = iq.reply() 

74 form = self.xmpp.plugin["xep_0004"].make_form() 

75 form.add_field(ftype="hidden", var="FORM_TYPE", value="urn:xmpp:mam:2") 

76 form.add_field(ftype="jid-single", var="with") 

77 form.add_field(ftype="text-single", var="start") 

78 form.add_field(ftype="text-single", var="end") 

79 form.add_field(ftype="text-single", var="before-id") 

80 form.add_field(ftype="text-single", var="after-id") 

81 form.add_field(ftype="boolean", var="include-groupchat") 

82 field = form.add_field(ftype="list-multi", var="ids") 

83 field["validate"]["datatype"] = "xs:string" 

84 field["validate"]["open"] = True 

85 reply["mam"].append(form) 

86 reply.send() 

87 

88 @exceptions_to_xmpp_errors 

89 async def __handle_mam_metadata(self, iq: Iq) -> None: 

90 muc = await self.get_muc_from_stanza(iq) 

91 await muc.send_mam_metadata(iq)