yangtb24 commited on
Commit
09735d0
1 Parent(s): a776e9e

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +24 -1
app.py CHANGED
@@ -200,6 +200,26 @@ def select_key(request_type):
200
  key = available_keys[int(time.time() * 1000) % len(available_keys)]
201
  return key
202
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
203
  # 创建一个后台调度器
204
  scheduler = BackgroundScheduler()
205
 
@@ -243,8 +263,11 @@ def check_tokens():
243
  @app.route('/handsome/v1/chat/completions', methods=['POST'])
244
  def handsome_chat_completions():
245
  """
246
- 处理 /handsome/v1/chat/completions 路由的请求。
247
  """
 
 
 
248
  data = request.get_json()
249
  if not data or 'model' not in data:
250
  return jsonify({"error": "Invalid request data"}), 400
 
200
  key = available_keys[int(time.time() * 1000) % len(available_keys)]
201
  return key
202
 
203
+ def check_authorization(request):
204
+ """
205
+ 检查请求头中的 Authorization 字段是否匹配环境变量 AUTHORIZATION_KEY。
206
+ """
207
+ authorization_key = os.environ.get("AUTHORIZATION_KEY")
208
+ if not authorization_key:
209
+ logging.warning("环境变量 AUTHORIZATION_KEY 未设置,请设置后重试。")
210
+ return False
211
+
212
+ auth_header = request.headers.get('Authorization')
213
+ if not auth_header:
214
+ logging.warning("请求头中缺少 Authorization 字段。")
215
+ return False
216
+
217
+ if auth_header != f"Bearer {authorization_key}":
218
+ logging.warning(f"无效的 Authorization 密钥:{auth_header}")
219
+ return False
220
+
221
+ return True
222
+
223
  # 创建一个后台调度器
224
  scheduler = BackgroundScheduler()
225
 
 
263
  @app.route('/handsome/v1/chat/completions', methods=['POST'])
264
  def handsome_chat_completions():
265
  """
266
+ 处理 /handsome/v1/chat/completions 路由的请求,添加鉴权。
267
  """
268
+ if not check_authorization(request):
269
+ return jsonify({"error": "Unauthorized"}), 401
270
+
271
  data = request.get_json()
272
  if not data or 'model' not in data:
273
  return jsonify({"error": "Invalid request data"}), 400