letterm commited on
Commit
c8f5c91
·
verified ·
1 Parent(s): 6928332

Delete token_manager.py

Browse files
Files changed (1) hide show
  1. token_manager.py +0 -344
token_manager.py DELETED
@@ -1,344 +0,0 @@
1
- """
2
- 多Token管理器模块
3
- 支持多个refresh token的管理、负载均衡和自动刷新
4
- """
5
- import time
6
- import threading
7
- import requests
8
- from typing import List, Optional, Dict, Any
9
- from dataclasses import dataclass
10
- from collections import defaultdict
11
- from concurrent.futures import ThreadPoolExecutor, as_completed
12
- from loguru import logger
13
-
14
- from config import Config
15
- from utils import Utils
16
-
17
-
18
- @dataclass
19
- class TokenInfo:
20
- """Token信息数据类"""
21
- refresh_token: str
22
- token_id: str # token标识符
23
- access_token: Optional[str] = None
24
- last_refresh_time: Optional[float] = None
25
- refresh_count: int = 0
26
- error_count: int = 0
27
- is_active: bool = True
28
-
29
-
30
- class MultiTokenManager:
31
- """多Token管理器,支持负载均衡和自动刷新"""
32
-
33
- def __init__(self):
34
- # Token存储
35
- self.tokens: Dict[str, TokenInfo] = {}
36
- self.token_lock = threading.RLock()
37
-
38
- # 负载均衡相关
39
- self.current_index = 0
40
- self.usage_count = defaultdict(int)
41
-
42
- # 自动刷新相关
43
- self.refresh_timer = None
44
- self.refresh_executor = ThreadPoolExecutor(max_workers=5, thread_name_prefix="token-refresh")
45
- self.session = requests.Session()
46
-
47
- # 初始化
48
- self._initialize_tokens()
49
-
50
- def _initialize_tokens(self):
51
- """初始化token列表"""
52
- # 从环境变量获取refresh token列表(支持分号分割)
53
- env_tokens = Config.get_refresh_tokens()
54
-
55
- if env_tokens:
56
- # 过滤有效的token
57
- valid_tokens = [token for token in env_tokens if Utils.validate_refresh_token_format(token)]
58
-
59
- if valid_tokens:
60
- self.add_refresh_tokens(valid_tokens)
61
- logger.info(f"✅ 从环境变量初始化了 {len(valid_tokens)} 个有效token")
62
-
63
- if len(env_tokens) > len(valid_tokens):
64
- invalid_count = len(env_tokens) - len(valid_tokens)
65
- logger.warning(f"⚠️ 跳过了 {invalid_count} 个格式无效的token")
66
- else:
67
- logger.warning("⚠️ 环境变量中的所有token格式都无效")
68
- else:
69
- logger.warning("⚠️ 未找到有效的refresh token,请通过环境变量 WARP_REFRESH_TOKEN 设置或在管理界面添加")
70
- logger.info("💡 环境变量支持多个token,使用分号(;)分割:token1;token2;token3")
71
-
72
- def add_refresh_token(self, refresh_token: str) -> bool:
73
- """添加单个refresh token"""
74
- return self.add_refresh_tokens([refresh_token])
75
-
76
- def add_refresh_tokens(self, refresh_tokens: List[str]) -> bool:
77
- """添加多个refresh token"""
78
- with self.token_lock:
79
- added_count = 0
80
-
81
- for token in refresh_tokens:
82
- if not Utils.validate_refresh_token_format(token):
83
- logger.warning(f"⚠️ Token格式无效: {token[:20]}...")
84
- continue
85
-
86
- if token in self.tokens:
87
- logger.info(f"⚠️ Token已存在,跳过: {self.tokens[token].token_id}")
88
- continue
89
-
90
- token_id = Utils.generate_token_id(token)
91
- self.tokens[token] = TokenInfo(
92
- refresh_token=token,
93
- token_id=token_id
94
- )
95
- added_count += 1
96
- logger.info(f"✅ 添加新token: {token_id}")
97
-
98
- logger.info(f"📊 成功添加 {added_count} 个新token,当前总数: {len(self.tokens)}")
99
- return added_count > 0
100
-
101
- def remove_refresh_token(self, refresh_token: str) -> bool:
102
- """移除refresh token"""
103
- with self.token_lock:
104
- if refresh_token in self.tokens:
105
- token_id = self.tokens[refresh_token].token_id
106
- del self.tokens[refresh_token]
107
- logger.info(f"🗑️ 移除token: {token_id}")
108
- return True
109
- return False
110
-
111
- def get_refresh_token_from_env(self, refresh_token: str) -> bool:
112
- """检查token是否来自环境变量"""
113
- env_tokens = Config.get_refresh_tokens()
114
- return refresh_token in env_tokens
115
-
116
- def remove_duplicate_tokens(self, new_tokens: List[str]) -> List[str]:
117
- """移除重复的token,优先保留环境变量中的token"""
118
- with self.token_lock:
119
- unique_tokens = []
120
-
121
- for token in new_tokens:
122
- # 检查是否与现有token重复
123
- if token in self.tokens:
124
- # 如果现有token来自环境变量,跳过新token
125
- if self.get_refresh_token_from_env(token):
126
- logger.info(f"⚠️ 跳过重复token(环境变量优先): {self.tokens[token].token_id}")
127
- continue
128
- else:
129
- # 如果新token优先级更高,移除旧token
130
- self.remove_refresh_token(token)
131
-
132
- unique_tokens.append(token)
133
-
134
- return unique_tokens
135
-
136
- def get_access_token(self, refresh_token: str) -> Optional[str]:
137
- """通过refresh token获取新的access token"""
138
- token_id = self.tokens.get(refresh_token, TokenInfo("", "")).token_id
139
-
140
- try:
141
- headers = {
142
- 'Accept-Encoding': 'gzip, br',
143
- 'Content-Type': 'application/x-www-form-urlencoded',
144
- 'x-warp-client-version': Config.WARP_CLIENT_VERSION,
145
- 'x-warp-os-category': Config.WARP_OS_CATEGORY,
146
- 'x-warp-os-name': Config.WARP_OS_NAME,
147
- 'x-warp-os-version': Config.WARP_OS_VERSION
148
- }
149
-
150
- data = {
151
- 'grant_type': 'refresh_token',
152
- 'refresh_token': refresh_token
153
- }
154
-
155
- logger.debug(f"🔄 正在刷新access token: {token_id}")
156
-
157
- response = self.session.post(
158
- Config.GOOGLE_TOKEN_URL,
159
- headers=headers,
160
- data=data,
161
- timeout=30
162
- )
163
-
164
- if response.status_code == 200:
165
- token_data = response.json()
166
- access_token = token_data.get('access_token')
167
- if access_token:
168
- logger.success(f"✅ Access token获取成功: {token_id}")
169
- return access_token
170
- else:
171
- logger.error(f"❌ 响应中未找到access_token字段: {token_id}")
172
- else:
173
- logger.error(f"❌ 获取token失败,状态码 {response.status_code}: {token_id}")
174
-
175
- except Exception as e:
176
- logger.error(f"❌ 获取access token时出错: {token_id} - {e}")
177
-
178
- return None
179
-
180
- def refresh_single_token(self, refresh_token: str) -> bool:
181
- """刷新单个token"""
182
- with self.token_lock:
183
- if refresh_token not in self.tokens:
184
- return False
185
-
186
- token_info = self.tokens[refresh_token]
187
-
188
- try:
189
- access_token = self.get_access_token(refresh_token)
190
-
191
- with self.token_lock:
192
- if access_token:
193
- token_info.access_token = access_token
194
- token_info.last_refresh_time = time.time()
195
- token_info.refresh_count += 1
196
- token_info.error_count = 0 # 重置错误计数
197
- token_info.is_active = True
198
-
199
- logger.success(f"🔄 Token刷新成功: {token_info.token_id} (第{token_info.refresh_count}次)")
200
- return True
201
- else:
202
- token_info.error_count += 1
203
- if token_info.error_count >= 3:
204
- token_info.is_active = False
205
- logger.warning(f"⚠️ Token连续失败3次,标记为不可用: {token_info.token_id}")
206
-
207
- return False
208
-
209
- except Exception as e:
210
- with self.token_lock:
211
- token_info.error_count += 1
212
-
213
- logger.error(f"❌ 刷新token时出错: {token_info.token_id} - {e}")
214
- return False
215
-
216
- def refresh_all_tokens(self):
217
- """并发刷新所有token"""
218
- with self.token_lock:
219
- refresh_tokens = list(self.tokens.keys())
220
-
221
- if not refresh_tokens:
222
- logger.warning("⚠️ 没有可刷新的token")
223
- return
224
-
225
- logger.info(f"🔄 开始并发刷新 {len(refresh_tokens)} 个token")
226
-
227
- # 使用线程池并发刷新
228
- futures = {
229
- self.refresh_executor.submit(self.refresh_single_token, token): token
230
- for token in refresh_tokens
231
- }
232
-
233
- success_count = 0
234
- for future in as_completed(futures):
235
- token = futures[future]
236
- try:
237
- success = future.result(timeout=30)
238
- if success:
239
- success_count += 1
240
- except Exception as e:
241
- token_id = self.tokens.get(token, TokenInfo("", "")).token_id
242
- logger.error(f"❌ 刷新token时出错: {token_id} - {e}")
243
-
244
- logger.info(f"📊 Token刷新完成: {success_count}/{len(refresh_tokens)} 成功")
245
-
246
- def get_current_access_token(self) -> Optional[str]:
247
- """获取当前可用的access token(负��均衡)"""
248
- with self.token_lock:
249
- # 获取所有活跃token
250
- active_tokens = [
251
- (token, info) for token, info in self.tokens.items()
252
- if info.is_active and info.access_token
253
- ]
254
-
255
- if not active_tokens:
256
- logger.warning("⚠️ 没有可用的access token!")
257
- logger.warning("💡 请确保:")
258
- logger.warning(" 1. 设置了有效的 WARP_REFRESH_TOKEN 环境变量")
259
- logger.warning(" 2. 或通过管理界面添加有效的 refresh token")
260
- logger.warning(" 3. refresh token 已成功刷新获得 access token")
261
- return None
262
-
263
- # 简单的轮询负载均衡
264
- if self.current_index >= len(active_tokens):
265
- self.current_index = 0
266
-
267
- selected_token, token_info = active_tokens[self.current_index]
268
- self.current_index = (self.current_index + 1) % len(active_tokens)
269
-
270
- # 更新使用计数
271
- self.usage_count[selected_token] += 1
272
-
273
- logger.info(f"🎯 轮询选择token: {token_info.token_id} (使用次数: {self.usage_count[selected_token]})")
274
- return token_info.access_token
275
-
276
- def get_token_status(self) -> Dict[str, Any]:
277
- """获取所有token的状态信息"""
278
- with self.token_lock:
279
- status = {
280
- 'total_tokens': len(self.tokens),
281
- 'active_tokens': sum(1 for info in self.tokens.values() if info.is_active),
282
- 'tokens_with_access': sum(1 for info in self.tokens.values() if info.access_token),
283
- 'tokens': []
284
- }
285
-
286
- for token, info in self.tokens.items():
287
- token_status = {
288
- 'refresh_token': info.token_id, # 使用token_id而不是实际token
289
- 'has_access_token': bool(info.access_token),
290
- 'is_active': info.is_active,
291
- 'refresh_count': info.refresh_count,
292
- 'error_count': info.error_count,
293
- 'last_refresh_time': info.last_refresh_time,
294
- 'usage_count': self.usage_count.get(token, 0)
295
- }
296
- status['tokens'].append(token_status)
297
-
298
- return status
299
-
300
- def start_auto_refresh(self):
301
- """启动自动刷新服务"""
302
- logger.info(f"🚀 启动Token自动刷新服务,间隔: {Config.TOKEN_REFRESH_INTERVAL // 60} 分钟")
303
-
304
- # 立即进行一次刷新
305
- refresh_thread = threading.Thread(target=self._initial_refresh, daemon=True)
306
- refresh_thread.start()
307
-
308
- def _initial_refresh(self):
309
- """初始刷新(在后台线程中执行)"""
310
- logger.info("🔄 正在后台执行初始token刷新...")
311
- self.refresh_all_tokens()
312
- self._schedule_next_refresh()
313
-
314
- def _schedule_next_refresh(self):
315
- """安排下次刷新"""
316
- if self.refresh_timer:
317
- self.refresh_timer.cancel()
318
-
319
- self.refresh_timer = threading.Timer(
320
- Config.TOKEN_REFRESH_INTERVAL,
321
- self._auto_refresh_callback
322
- )
323
- self.refresh_timer.daemon = True
324
- self.refresh_timer.start()
325
-
326
- # 计算下次刷新时间
327
- next_refresh_time = time.time() + Config.TOKEN_REFRESH_INTERVAL
328
- next_refresh_str = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(next_refresh_time))
329
- logger.info(f"⏰ 下次token刷新时间: {next_refresh_str}")
330
-
331
- def _auto_refresh_callback(self):
332
- """自动刷新回调"""
333
- logger.info("🔄 开始定时自动刷新所有token...")
334
- self.refresh_all_tokens()
335
- self._schedule_next_refresh()
336
-
337
- def stop_auto_refresh(self):
338
- """停止自动刷新"""
339
- if self.refresh_timer:
340
- self.refresh_timer.cancel()
341
- self.refresh_timer = None
342
-
343
- self.refresh_executor.shutdown(wait=False)
344
- logger.info("⏹️ Token自动刷新服务已停止")