Coverage for aipyapp/aipy/trustoken.py: 16%

1 import os

2 import time

3 import webbrowser

4

5 import requests

6 import qrcode

7

8 from .. import T

9

10 POLL_INTERVAL = 5 # 轮询间隔（秒）

11

12 class TrustTokenAPI:

13 """Handles all HTTP operations for TrustToken binding and authentication."""

14

15 def __init__(self, coordinator_url=None):

16 """

17 Initialize the TrustToken API handler.

18

19 Args:

20 coordinator_url (str, optional): The coordinator server URL. Defaults to None.

21 """

22 self.coordinator_url = coordinator_url or T("https://www.trustoken.ai/api")

23

24 def request_binding(self):

25 """Request binding from the coordinator server.

26

27 Returns:

28 dict: Response data containing approval_url, request_id, and expires_in if successful

29 None: If the request failed

30 """

31 url = f"{self.coordinator_url}/request_bind"

32 try:

33 response = requests.post(url, timeout=10)

34 response.raise_for_status()

35 return response.json()

36 except requests.exceptions.RequestException as e:

37 print(T("Error connecting to coordinator or during request: {}").format(e))

38 return None

39 except Exception as e:

40 print(T("An unexpected error occurred during request: {}").format(e))

41 return None

42

43 def check_status(self, request_id):

44 """Check the binding status from the coordinator server.

45

46 Args:

47 request_id (str): The request ID to check status for.

48

49 Returns:

50 dict: Response data containing status and secret_token if approved

51 None: If the request failed

52 """

53 url = f"{self.coordinator_url}/check_status"

54 params = {'request_id': request_id}

55 try:

56 response = requests.get(url, params=params, timeout=10)

57 response.raise_for_status()

58 return response.json()

59 except requests.exceptions.RequestException as e:

60 print(T("Error connecting to coordinator during polling: {}").format(e))

61 return None

62 except Exception as e:

63 print(T("An unexpected error occurred during polling: {}").format(e))

64 return None

65

66 class TrustToken:

67 """A class to handle TrustToken binding and authentication processes."""

68

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

70 """

71 Initialize the TrustToken handler.

72

73 Args:

74 coordinator_url (str, optional): The coordinator server URL. Defaults to None.

75 poll_interval (int, optional): Polling interval in seconds. Defaults to 5.

76 """

77 self.api = TrustTokenAPI(coordinator_url)

78 self.poll_interval = poll_interval or POLL_INTERVAL

79

80 def request_binding(self, qrcode=False):

81 """Request binding from the coordinator server.

82

83 Returns:

84 str: The request ID if successful, None otherwise.

85 """

86 data = self.api.request_binding()

87 if not data:

88 return None

89

90 approval_url = data['approval_url']

91 request_id = data['request_id']

92 expires_in = data['expires_in']

93

94 print(T("""Binding request sent successfully.

95 Request ID: {}

96

97>>> Please open this URL in your browser on an authenticated device to approve:

98>>> {}

99

100(This link expires in {} seconds)""").format(request_id, approval_url, expires_in))

101

102 if qrcode:

103 print(T("Or scan the QR code below:"))

104 try:

105 qr = qrcode.QRCode(

106 error_correction=qrcode.constants.ERROR_CORRECT_L,

107 border=1

108 )

109 qr.add_data(approval_url)

110 qr.make(fit=True)

111 qr.print_ascii(tty=True)

112 print(T("We recommend you scan the QR code to bind the AiPy brain, you can also configure a third-party large model brain, details refer to: https://d.aipy.app/d/77"))

113 except Exception as e:

114 print(T("(Could not display QR code: {})").format(e))

115 else:

116 webbrowser.open(approval_url)

117 return request_id

118

119 def poll_status(self, request_id, save_func=None):

120 """Poll the binding status from the coordinator server.

121

122 Args:

123 request_id (str): The request ID to check status for.

124 save_func (callable, optional): Function to save the token when approved.

125

126 Returns:

127 bool: True if binding was successful, False otherwise.

128 """

129 start_time = time.time()

130 polling_timeout = 310

131

132 print(T("Browser has opened the Trustoken website, please register or login to authorize"), end='', flush=True)

133 try:

134 while time.time() - start_time < polling_timeout:

135 data = self.api.check_status(request_id)

136 if not data:

137 time.sleep(self.poll_interval)

138 continue

139

140 status = data.get('status')

141 if status == 'pending':

142 print('.', end='', flush=True)

143 time.sleep(self.poll_interval)

144 continue

145

146 print()

147 print(T("Current status: {}...").format(status))

148

149 if status == 'approved':

150 if save_func:

151 save_func(data['secret_token'])

152 return True

153 elif status == 'expired':

154 print(T("Binding request expired."))

155 return False

156 else:

157 print(T("Received unknown status: {}").format(status))

158 return False

159

160 time.sleep(self.poll_interval)

161 except KeyboardInterrupt:

162 print(T("Polling cancelled by user."))

163 return False

164

165 print(T("Polling timed out."))

166 return False

167

168 def fetch_token(self, save_func):

169 """Fetch a token from the coordinator server.

170

171 Args:

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

173

174 Returns:

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

176 """

177 print(T("The current environment lacks the required configuration file. Starting the configuration initialization process to bind with the Trustoken account..."))

178 req_id = self.request_binding()

179 if req_id:

180 if self.poll_status(req_id, save_func):

181 print(T("Binding process completed successfully."))

182 return True

183 else:

184 print(T("Binding process failed or was not completed."))

185 return False

186 else:

187 print(T("Failed to initiate binding request."))

188 return False

189

190 if __name__ == "__main__":

191 tt = TrustToken()

192 tt.fetch_token(lambda token: print(f"Token: {token}"))