Coverage for aipyapp/gui/trustoken.py: 0%

1 import time

2

3 import wx

4 import wx.adv

5 import threading

6 import qrcode

7 import io

8 from PIL import Image

9

10 from .. import T

11 from ..aipy.trustoken import TrustTokenAPI

12

13 class TrustTokenAuthDialog(wx.Dialog):

14 """A dialog for TrustToken authentication with QR code display."""

15

16 def __init__(self, parent=None, coordinator_url=None, poll_interval=5):

17 """

18 Initialize the authentication dialog.

19

20 Args:

21 parent: Parent window

22 coordinator_url (str, optional): The coordinator server URL

23 poll_interval (int, optional): Polling interval in seconds

24 """

25 super().__init__(parent, title=T("TrustToken Authentication"),

26 size=(400, 500),

27 style=wx.DEFAULT_DIALOG_STYLE )

28

29 self.api = TrustTokenAPI(coordinator_url)

30 self.poll_interval = poll_interval

31 self.request_id = None

32 self.polling_thread = None

33 self.stop_polling = False

34 self.start_time = None

35 self.polling_timeout = 310 # 5 minutes and 10 seconds

36

37 self._init_ui()

38 self.Centre()

39

40 def _init_ui(self):

41 """Initialize the user interface."""

42 main_sizer = wx.BoxSizer(wx.VERTICAL)

43

44 # Status text

45 self.status_text = wx.StaticText(self, label='')

46 main_sizer.Add(self.status_text, 0, wx.TOP | wx.ALIGN_CENTER, 5)

47

48 # QR code display

49 self.qr_bitmap = wx.StaticBitmap(self, size=(300, 300))

50 main_sizer.Add(self.qr_bitmap, 0, wx.ALL | wx.ALIGN_CENTER, 5)

51

52 self.other_text = wx.adv.HyperlinkCtrl(self, label='扫码绑定AiPy大脑，您也可以配置其它大模型大脑', url='https://d.aipy.app/d/77')

53 main_sizer.Add(self.other_text, 0, wx.TOP | wx.ALIGN_CENTER, 5)

54

55 # Progress bar

56 self.progress_bar = wx.Gauge(self, range=100)

57 main_sizer.Add(self.progress_bar, 0, wx.ALL | wx.EXPAND, 5)

58

59 # Time remaining text

60 self.time_text = wx.StaticText(self, label='')

61 self.time_text.Hide()

62 main_sizer.Add(self.time_text, 0, wx.ALL | wx.ALIGN_CENTER, 5)

63

64 # Buttons

65 self.cancel_button = wx.Button(self, wx.ID_CANCEL, T("Cancel"))

66 main_sizer.Add(self.cancel_button, 0, wx.ALIGN_CENTER | wx.ALL, 5)

67

68 self.SetSizer(main_sizer)

69

70 # Bind events

71 self.Bind(wx.EVT_BUTTON, self._on_cancel, self.cancel_button)

72 self.Bind(wx.EVT_CLOSE, self._on_close)

73

74 def _update_progress(self):

75 """Update the progress bar and time remaining."""

76 if not self.start_time:

77 return

78

79 elapsed = time.time() - self.start_time

80 if elapsed >= self.polling_timeout:

81 progress = 100

82 time_remaining = 0

83 else:

84 progress = int((elapsed / self.polling_timeout) * 100)

85 time_remaining = int(self.polling_timeout - elapsed)

86

87 wx.CallAfter(self.progress_bar.SetValue, progress)

88 wx.CallAfter(self.time_text.SetLabel, T("Time remaining: {} seconds", time_remaining))

89

90 def _poll_status(self, save_func):

91 """Poll the binding status in a separate thread."""

92 self.start_time = time.time()

93 self.time_text.Show()

94 while not self.stop_polling and time.time() - self.start_time < self.polling_timeout:

95 self._update_progress()

96

97 data = self.api.check_status(self.request_id)

98 if not data:

99 time.sleep(self.poll_interval)

100 continue

101

102 status = data.get('status')

103 wx.CallAfter(self._update_status, T("Current status: {}...", T(status)))

104

105 if status == 'approved':

106 if save_func:

107 save_func(data['secret_token'])

108 wx.CallAfter(self.EndModal, wx.ID_OK)

109 return True

110 elif status == 'expired':

111 wx.CallAfter(self._update_status, T("Binding request expired."))

112 wx.CallAfter(self.EndModal, wx.ID_CANCEL)

113 return False

114 elif status == 'pending':

115 pass

116 else:

117 wx.CallAfter(self._update_status, T("Received unknown status: {}", status))

118 wx.CallAfter(self.EndModal, wx.ID_CANCEL)

119 return False

120

121 time.sleep(self.poll_interval)

122

123 if not self.stop_polling:

124 wx.CallAfter(self._update_status, T("Polling timed out."))

125 wx.CallAfter(self.EndModal, wx.ID_CANCEL)

126 return False

127

128 def _on_cancel(self, event):

129 """Handle cancel button click."""

130 self.stop_polling = True

131 if self.polling_thread:

132 self.polling_thread.join()

133 self.EndModal(wx.ID_CANCEL)

134

135 def _on_close(self, event):

136 """Handle dialog close."""

137 self.stop_polling = True

138 if self.polling_thread:

139 self.polling_thread.join()

140 event.Skip()

141

142 def _update_qr_code(self, url):

143 """Update the QR code display with the given URL."""

144 try:

145 qr = qrcode.QRCode(

146 error_correction=qrcode.constants.ERROR_CORRECT_L,

147 border=1

148 )

149 qr.add_data(url)

150 qr.make(fit=True)

151

152 # Convert QR code to wx.Bitmap

153 img = qr.make_image(fill_color="black", back_color="white")

154 img_byte_arr = io.BytesIO()

155 img.save(img_byte_arr, format='PNG')

156 img_byte_arr = img_byte_arr.getvalue()

157

158 wx_img = wx.Image(io.BytesIO(img_byte_arr))

159 wx_img = wx_img.Scale(300, 300, wx.IMAGE_QUALITY_HIGH)

160 self.qr_bitmap.SetBitmap(wx.Bitmap(wx_img))

161 self.Layout()

162 except Exception as e:

163 wx.MessageBox(T("(Could not display QR code: {})", e), T("Error"), wx.OK | wx.ICON_ERROR)

164

165 def _update_status(self, status):

166 """Update the status text."""

167 self.status_text.SetLabel(status)

168 self.Layout()

169

170 def fetch_token(self, save_func):

171 """Start the token fetching process.

172

173 Args:

174 save_func (callable): Function to save the token when approved.

175

176 Returns:

177 bool: True if token was successfully fetched and saved, False otherwise.

178 """

179 self._update_status(T('requesting_binding'))

180 data = self.api.request_binding()

181 if not data:

182 wx.MessageBox(T("Failed to initiate binding request.", None), T("Error"), wx.OK | wx.ICON_ERROR)

183 return False

184

185 approval_url = data['approval_url']

186 self.request_id = data['request_id']

187 expires_in = data['expires_in']

188 self.polling_timeout = expires_in

189 self._update_status(T("Current status: {}...", T("Browser has opened the Trustoken website, please register or login to authorize")))

190 self._update_qr_code(approval_url)

191

192 # Start polling in a separate thread

193 self.polling_thread = threading.Thread(

194 target=self._poll_status, args=(save_func,))

195 self.polling_thread.daemon = True

196 self.polling_thread.start()

197

198 # Show the dialog

199 result = self.ShowModal()

200 return result == wx.ID_OK

201

202 if __name__ == "__main__":

203 # Test the TrustTokenAuthDialog

204 app = wx.App()

205

206 def save_token(token):

207 print(f"Token received: {token}")

208 # Here you would typically save the token to your configuration

209 # For example:

210 # config_manager.save_tt_config(token)

211

212 # Create and show the dialog

213 dialog = TrustTokenAuthDialog(None)

214 if dialog.fetch_token(save_token):

215 print("Authentication successful!")

216 else:

217 print("Authentication failed or was cancelled.")

218

219 app.MainLoop()