Coverage for aipyapp/aipy/trustoken.py: 16%
104 statements
« prev ^ index » next coverage.py v7.10.3, created at 2025-08-11 13:03 +0200
1import os
2import time
3import webbrowser
5import requests
6import qrcode
8from .. import T
10POLL_INTERVAL = 5 # 轮询间隔（秒）
12class TrustTokenAPI:
13 """Handles all HTTP operations for TrustToken binding and authentication."""
15 def __init__(self, coordinator_url=None):
16 """
17 Initialize the TrustToken API handler.
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")
24 def request_binding(self):
25 """Request binding from the coordinator server.
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
43 def check_status(self, request_id):
44 """Check the binding status from the coordinator server.
46 Args:
47 request_id (str): The request ID to check status for.
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
66class TrustToken:
67 """A class to handle TrustToken binding and authentication processes."""
69 def __init__(self, coordinator_url=None, poll_interval=5):
70 """
71 Initialize the TrustToken handler.
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
80 def request_binding(self, qrcode=False):
81 """Request binding from the coordinator server.
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
90 approval_url = data['approval_url']
91 request_id = data['request_id']
92 expires_in = data['expires_in']
94 print(T("""Binding request sent successfully.
95Request ID: {}
97>>> Please open this URL in your browser on an authenticated device to approve:
98>>> {}
100(This link expires in {} seconds)""").format(request_id, approval_url, expires_in))
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
119 def poll_status(self, request_id, save_func=None):
120 """Poll the binding status from the coordinator server.
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.
126 Returns:
127 bool: True if binding was successful, False otherwise.
128 """
129 start_time = time.time()
130 polling_timeout = 310
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
140 status = data.get('status')
141 if status == 'pending':
142 print('.', end='', flush=True)
143 time.sleep(self.poll_interval)
144 continue
146 print()
147 print(T("Current status: {}...").format(status))
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
160 time.sleep(self.poll_interval)
161 except KeyboardInterrupt:
162 print(T("Polling cancelled by user."))
163 return False
165 print(T("Polling timed out."))
166 return False
168 def fetch_token(self, save_func):
169 """Fetch a token from the coordinator server.
171 Args:
172 save_func (callable): Function to save the token when approved.
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
190if __name__ == "__main__":
191 tt = TrustToken()
192 tt.fetch_token(lambda token: print(f"Token: {token}"))