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

1import time 

2 

3import wx 

4import wx.adv 

5import threading 

6import qrcode 

7import io 

8from PIL import Image 

9 

10from .. import T 

11from ..aipy.trustoken import TrustTokenAPI 

12 

13class 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 

202if __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()