Coverage for slidge/command/register.py: 42%

1""" 

2This module handles the registration :term:`Command`, which is a necessary 

3step for a JID to become a slidge :term:`User`. 

4""" 

5 

6from __future__ import annotations 

7 

8import asyncio 

9import functools 

10from enum import IntEnum 

11from typing import Any 

12 

13from slixmpp import JID, Iq 

14from slixmpp.exceptions import XMPPError 

15 

16from ..core import config 

17from ..db import GatewayUser 

18from ..util.types import AnySession, UserPreferences 

19from .base import Command, CommandAccess, Form, FormField, FormSession, FormValues 

20from .user import Preferences 

21 

22 

23class RegistrationType(IntEnum): 

24 """ 

25 An :class:`Enum` to define the registration flow. 

26 """ 

27 

28 SINGLE_STEP_FORM = 0 

29 """ 

30 1 step, 1 form, the only flow compatible with :xep:`0077`. 

31 Using this, the whole flow is defined 

32 by :attr:`slidge.BaseGateway.REGISTRATION_FIELDS` and 

33 :attr:`.REGISTRATION_INSTRUCTIONS`. 

34 """ 

35 

36 QRCODE = 10 

37 """ 

38 The registration requires flashing a QR code in an official client. 

39 See :meth:`slidge.BaseGateway.send_qr`, :meth:`.get_qr_text` 

40 and :meth:`.confirm_qr`. 

41 """ 

42 

43 TWO_FACTOR_CODE = 20 

44 """ 

45 The registration requires confirming login with a 2FA code, 

46 eg something received by email or SMS to finalize the authentication. 

47 See :meth:`.validate_two_factor_code`. 

48 """ 

49 

50 

51class TwoFactorNotRequired(Exception): 

52 """ 

53 Should be raised in :meth:`slidge.BaseGateway.validate` if the code is not 

54 required after all. This can happen for a :term:`Legacy Network` where 2FA 

55 is optional. 

56 """ 

57 

58 

59class Register(Command[AnySession]): 

60 NAME = "📝 Register to the gateway" 

61 HELP = "Link your JID to this gateway" 

62 NODE = "jabber:iq:register" 

63 CHAT_COMMAND = "register" 

64 ACCESS = CommandAccess.NON_USER 

65 

66 SUCCESS_MESSAGE = "Success, welcome!" 

67 

68 def _finalize( 

69 self, 

70 form_values: UserPreferences, 

71 _session: None, 

72 ifrom: JID, 

73 user: GatewayUser, 

74 *_: Any, # noqa:ANN401 

75 ) -> str: 

76 user.preferences = form_values # type: ignore 

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

78 self.xmpp.event("user_register", Iq(sfrom=ifrom.bare)) 

79 return self.SUCCESS_MESSAGE 

80 

81 async def run( 

82 self, 

83 _session: AnySession | None, 

84 ifrom: JID, 

85 *_: str, 

86 ) -> FormSession[AnySession]: 

87 self.xmpp.raise_if_not_allowed_jid(ifrom) 

88 return FormSession( 

89 title=f"Registration to '{self.xmpp.COMPONENT_NAME}'", 

90 instructions=self.xmpp.REGISTRATION_INSTRUCTIONS, 

91 fields=self.xmpp.REGISTRATION_FIELDS, 

92 handler=self.register, 

93 ) 

94 

95 async def register( 

96 self, 

97 form_values: dict[str, Any], 

98 _session: None, 

99 ifrom: JID, 

100 ) -> Form | None: 

101 two_fa_needed = True 

102 try: 

103 data = await self.xmpp.user_prevalidate(ifrom, form_values) 

104 except ValueError as e: 

105 raise XMPPError("bad-request", str(e)) 

106 except TwoFactorNotRequired: 

107 data = None 

108 if self.xmpp.REGISTRATION_TYPE == RegistrationType.TWO_FACTOR_CODE: 

109 two_fa_needed = False 

110 else: 

111 raise 

112 except Exception as e: # noqa: BLE001 

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

114 

115 user = GatewayUser( 

116 jid=JID(ifrom.bare), 

117 legacy_module_data=form_values if data is None else data, 

118 ) 

119 

120 if self.xmpp.REGISTRATION_TYPE == RegistrationType.SINGLE_STEP_FORM or ( 

121 self.xmpp.REGISTRATION_TYPE == RegistrationType.TWO_FACTOR_CODE 

122 and not two_fa_needed 

123 ): 

124 return await self.preferences(user) 

125 

126 if self.xmpp.REGISTRATION_TYPE == RegistrationType.TWO_FACTOR_CODE: 

127 return Form( 

128 title=self.xmpp.REGISTRATION_2FA_TITLE, 

129 instructions=self.xmpp.REGISTRATION_2FA_INSTRUCTIONS, 

130 fields=[FormField("code", label="Code", required=True)], 

131 handler=functools.partial(self.two_fa, user=user), 

132 ) 

133 

134 elif self.xmpp.REGISTRATION_TYPE == RegistrationType.QRCODE: 

135 self.xmpp.qr_pending_registrations[user.jid.bare] = ( 

136 self.xmpp.loop.create_future() 

137 ) 

138 qr_text = await self.xmpp.get_qr_text(user) 

139 img_url = await self.xmpp.send_qr(qr_text, mto=ifrom) 

140 if img_url is None: 

141 raise XMPPError( 

142 "internal-server-error", "Slidge cannot send attachments" 

143 ) 

144 self.xmpp.send_text(qr_text, mto=ifrom) 

145 return Form( 

146 title="Flash this", 

147 instructions="Flash this QR in the appropriate place", 

148 fields=[ 

149 FormField( 

150 "qr_img", 

151 type="fixed", 

152 value=qr_text, 

153 image_url=img_url, 

154 ), 

155 FormField( 

156 "qr_text", 

157 type="fixed", 

158 value=qr_text, 

159 label="Text encoded in the QR code", 

160 ), 

161 FormField( 

162 "qr_img_url", 

163 type="fixed", 

164 value=img_url, 

165 label="URL of the QR code image", 

166 ), 

167 ], 

168 handler=functools.partial(self.qr, user=user), 

169 ) 

170 

171 async def two_fa( 

172 self, 

173 form_values: FormValues, 

174 _session: None, 

175 _ifrom: JID, 

176 user: GatewayUser, 

177 ) -> Form: 

178 assert isinstance(form_values["code"], str) 

179 data = await self.xmpp.validate_two_factor_code(user, form_values["code"]) 

180 if data is not None: 

181 user.legacy_module_data.update(data) 

182 return await self.preferences(user) 

183 

184 async def qr( 

185 self, 

186 _form_values: FormValues, 

187 _session: None, 

188 _ifrom: JID, 

189 user: GatewayUser, 

190 ) -> Form: 

191 try: 

192 data = await asyncio.wait_for( 

193 self.xmpp.qr_pending_registrations[user.jid.bare], 

194 config.QR_TIMEOUT, 

195 ) 

196 except TimeoutError: 

197 raise XMPPError( 

198 "remote-server-timeout", 

199 ( 

200 "It does not seem that the QR code was correctly used, " 

201 "or you took too much time" 

202 ), 

203 ) 

204 if data is not None: 

205 user.legacy_module_data.update(data) 

206 return await self.preferences(user) 

207 

208 async def preferences(self, user: GatewayUser) -> Form: 

209 return Form( 

210 title="Preferences", 

211 instructions=Preferences.HELP, 

212 fields=self.xmpp.PREFERENCES, 

213 handler=functools.partial(self._finalize, user=user), 

214 timeout_handler=functools.partial(self._preferences_timeout, user=user), 

215 ) 

216 

217 def _preferences_timeout(self, user: GatewayUser) -> None: 

218 self.xmpp.event("user_register", Iq(sfrom=user.jid.bare)) 

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

220 self.xmpp.send_message( 

221 mfrom=self.xmpp.boundjid.bare, 

222 mto=user.jid, 

223 mbody="You did not choose your preferences in time, falling back to defaults. " 

224 "You can change preferences later with the 'preferences' command.\n" 

225 + self.SUCCESS_MESSAGE, 

226 )