Coverage for slidge/core/dispatcher/registration.py: 75%

1from __future__ import annotations 

2 

3import logging 

4from typing import TYPE_CHECKING, Optional 

5 

6from slixmpp import JID, Iq 

7from slixmpp.exceptions import XMPPError 

8 

9from ...db import GatewayUser 

10from .. import config 

11from .util import DispatcherMixin 

12 

13if TYPE_CHECKING: 

14 from slidge.core.gateway import BaseGateway 

15 

16 

17class RegistrationMixin(DispatcherMixin): 

18 def __init__(self, xmpp: "BaseGateway"): 

19 super().__init__(xmpp) 

20 self.xmpp = xmpp 

21 xmpp["xep_0077"].api.register( 

22 self.xmpp.make_registration_form, "make_registration_form" 

23 ) 

24 xmpp["xep_0077"].api.register(self._user_get, "user_get") 

25 xmpp["xep_0077"].api.register(self._user_validate, "user_validate") 

26 xmpp["xep_0077"].api.register(self._user_modify, "user_modify") 

27 # kept for slixmpp internal API compat 

28 # TODO: either fully use slixmpp internal API or rewrite registration without it at all 

29 xmpp["xep_0077"].api.register(lambda *a: None, "user_remove") 

30 

31 xmpp.add_event_handler("user_register", self._on_user_register) 

32 xmpp.add_event_handler("user_unregister", self._on_user_unregister) 

33 

34 def get_user(self, jid: JID) -> GatewayUser | None: 

35 return self.xmpp.store.users.get(jid) 

36 

37 async def _user_get( 

38 self, _gateway_jid, _node, ifrom: JID, iq: Iq 

39 ) -> GatewayUser | None: 

40 if ifrom is None: 

41 ifrom = iq.get_from() 

42 return self.get_user(ifrom) 

43 

44 async def _user_validate(self, _gateway_jid, _node, ifrom: JID, iq: Iq): 

45 xmpp = self.xmpp 

46 log.debug("User validate: %s", ifrom.bare) 

47 form_dict = {f.var: iq.get(f.var) for f in xmpp.REGISTRATION_FIELDS} 

48 xmpp.raise_if_not_allowed_jid(ifrom) 

49 legacy_module_data = await xmpp.user_prevalidate(ifrom, form_dict) 

50 if legacy_module_data is None: 

51 legacy_module_data = form_dict 

52 user = self.xmpp.store.users.new( 

53 jid=ifrom, 

54 legacy_module_data=legacy_module_data, # type:ignore 

55 ) 

56 log.info("New user: %s", user) 

57 

58 async def _user_modify( 

59 self, _gateway_jid, _node, ifrom: JID, form_dict: dict[str, Optional[str]] 

60 ): 

61 await self.xmpp.user_prevalidate(ifrom, form_dict) 

62 log.debug("Modify user: %s", ifrom) 

63 user = self.xmpp.store.users.get(ifrom) 

64 if user is None: 

65 raise XMPPError("internal-server-error", "User not found") 

66 user.legacy_module_data.update(form_dict) 

67 self.xmpp.store.users.update(user) 

68 

69 async def _on_user_register(self, iq: Iq): 

70 session = await self._get_session(iq, wait_for_ready=False) 

71 for jid in config.ADMINS: 

72 self.xmpp.send_message( 

73 mto=jid, 

74 mbody=f"{iq.get_from()} has registered", 

75 mtype="chat", 

76 mfrom=self.xmpp.boundjid.bare, 

77 ) 

78 session.send_gateway_message(self.xmpp.WELCOME_MESSAGE) 

79 await self.xmpp.login_wrap(session) 

80 

81 async def _on_user_unregister(self, iq: Iq): 

82 await self.xmpp.session_cls.kill_by_jid(iq.get_from()) 

83 

84 

85log = logging.getLogger(__name__)