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

1""" 

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

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

4""" 

5 

6import asyncio 

7import functools 

8import tempfile 

9from enum import IntEnum 

10from typing import Any 

11 

12import qrcode 

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 qr = qrcode.make(qr_text) 

140 with tempfile.NamedTemporaryFile(suffix=".png") as f: 

141 qr.save(f.name) 

142 img_url, _ = await self.xmpp.send_file(f.name, mto=ifrom) 

143 if img_url is None: 

144 raise XMPPError( 

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

146 ) 

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

148 return Form( 

149 title="Flash this", 

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

151 fields=[ 

152 FormField( 

153 "qr_img", 

154 type="fixed", 

155 value=qr_text, 

156 image_url=img_url, 

157 ), 

158 FormField( 

159 "qr_text", 

160 type="fixed", 

161 value=qr_text, 

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

163 ), 

164 FormField( 

165 "qr_img_url", 

166 type="fixed", 

167 value=img_url, 

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

169 ), 

170 ], 

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

172 ) 

173 

174 async def two_fa( 

175 self, 

176 form_values: FormValues, 

177 _session: None, 

178 _ifrom: JID, 

179 user: GatewayUser, 

180 ) -> Form: 

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

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

183 if data is not None: 

184 user.legacy_module_data.update(data) 

185 return await self.preferences(user) 

186 

187 async def qr( 

188 self, 

189 _form_values: FormValues, 

190 _session: None, 

191 _ifrom: JID, 

192 user: GatewayUser, 

193 ) -> Form: 

194 try: 

195 data = await asyncio.wait_for( 

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

197 config.QR_TIMEOUT, 

198 ) 

199 except TimeoutError: 

200 raise XMPPError( 

201 "remote-server-timeout", 

202 ( 

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

204 "or you took too much time" 

205 ), 

206 ) 

207 if data is not None: 

208 user.legacy_module_data.update(data) 

209 return await self.preferences(user) 

210 

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

212 return Form( 

213 title="Preferences", 

214 instructions=Preferences.HELP, 

215 fields=self.xmpp.PREFERENCES, 

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

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

218 ) 

219 

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

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

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

223 self.xmpp.send_message( 

224 mfrom=self.xmpp.boundjid.bare, 

225 mto=user.jid, 

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

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

228 + self.SUCCESS_MESSAGE, 

229 )