由session改为token

redis/test_user_redis.py

@@ -3,11 +3,14 @@ import random
 import string
 import uuid
 import time
-from flask import Flask, request, jsonify, session
+import jwt
+import datetime
+from flask import Flask, request, jsonify, Request
 from redis import Redis
 
+SECERT_KEY = "8U2LL1"
+
 app = Flask(__name__)
-app.secret_key = '333888'
 redis = Redis(host='192.168.3.229', port=6379, password='lizhen-redis')
 # redis = Redis(host='10.254.13.87', port=6379)
 # redis = Redis(host='localhost', port=6379)
@@ -15,7 +18,6 @@ redis = Redis(host='192.168.3.229', port=6379, password='lizhen-redis')
 
 # 生成验证码
 def generate_verification_code():
-    # code = ''.join(random.choices(string.ascii_uppercase + string.digits, k=6))
     code = ''.join(random.choices(string.digits, k=6))
     return code
 
@@ -60,7 +62,7 @@ def register():
     # 检查用户名是否已被注册
     if redis.hexists('users', username):
         return jsonify({'message': 'Username already exists'}), 400
-    
+
     # 生成唯一的用户ID
     user_id = str(uuid.uuid4())
 
@@ -73,7 +75,10 @@ def register():
     }
     redis.hset('users', username, json.dumps(user_data))
 
-    return jsonify({'code': 0, 'message': 'Registration successful'})
+    return jsonify({
+        'code': 0,
+        'message': 'Registration successful'
+    })
 
 
 # 用户登录
@@ -92,20 +97,24 @@ def login():
     if password != eval(user_data)['password']:
         return jsonify({'code': 400, 'message': 'Invalid password'})
 
-    # 登录验证通过，将用户信息存储到会话中
-    session['username'] = request.json.get('username')
-
-    return jsonify({'code': 0, 'message': 'Login successful'})
+    # 生成令牌
+    token = generate_token(eval(user_data)['user_id'], username)
+    return jsonify({
+        'code': 0,
+        'message': 'Login successful',
+        'data': {
+            'token': token
+        }
+    })
 
 
 # 需要验证登录状态的接口
 @app.route('/protected', methods=['GET'])
 def protected():
-    # 检查会话中的用户信息
-    if 'username' in session:
-        username = session['username']
-        # 其他业务逻辑...
-        return jsonify({'code': 0, 'message': f'Hello, {username}! This is a protected endpoint.'})
+    token = parse_token(request)
+    # 验证令牌
+    if not validate_token(token):
+        return jsonify({'code': 401, 'message': 'Invalid token'}), 200
 
     # 如果用户未登录，则返回未授权的响应
     return jsonify({'code': 401, 'message': 'Unauthorized'})
@@ -114,8 +123,11 @@ def protected():
 # 用户注销
 @app.route('/logout', methods=['POST'])
 def logout():
-    # 清除会话中的用户信息
-    session.pop('username', None)
+    token = parse_token(request)
+    # 验证令牌
+    if not validate_token(token):
+        # 将令牌添加到 Redis 黑名单
+        redis.set(token, 'revoked')
     return jsonify({'code': 0, 'message': 'Logout successful'})
 
 
@@ -123,18 +135,18 @@ def logout():
 @app.route('/purchase', methods=['POST'])
 def purchase():
     package_id = request.json.get('package_id')
+    token = parse_token(request)
+
+    # 验证令牌
+    if not validate_token(token):
+        return jsonify({'code': 401, 'message': 'Invalid token'})
 
     # 根据套餐ID获取套餐信息
     package = get_package_by_id(package_id)
     if not package:
         return jsonify({'code': 400, 'message': 'Invalid package ID'})
-    
-    # 根据用户名查询用户ID
-    username = session.get('username')
-    if not username:
-        return jsonify({'code': 400, 'message': 'User not logged in'})
 
-    user_id = get_user_id_by_username(username)
+    user_id = get_user_id_from_token(token)
     if not user_id:
         return jsonify({'code': 400, 'message': 'User not found'})
 
@@ -145,7 +157,7 @@ def purchase():
     # 检查如果用户已经支付了高级套餐，则不能支付比高级套餐更低级的基础套餐
     if not is_package_expired(user_id) and has_purchased_advanced_package(user_id) and package_id == '1':
         return jsonify({'code': 400, 'message': 'Cannot purchase lower level package'})
-    
+
     # 存储用户套餐信息到Redis
     store_user_package(user_id, package)
 
@@ -155,13 +167,17 @@ def purchase():
 # 验证用户聊天次数
 @app.route('/validate', methods=['POST'])
 def validate():
-    # 根据用户名查询用户ID
-    username = session.get('username')
-    user_id = get_user_id_by_username(username)
-    
+    token = parse_token(request)
+
+    # 验证令牌
+    if not validate_token(token):
+        return jsonify({'code': 401, 'message': 'Invalid token'})
+
+    user_id = get_user_id_from_token(token)
+
     if not user_id:
         return jsonify({'code': 400, 'message': 'User not found'})
-    
+
     # 获取用户套餐信息
     package = get_user_package(user_id)
     if not package:
@@ -174,6 +190,64 @@ def validate():
     return jsonify({'code': 0, 'message': 'Chat limit not exceeded'})
 
 
+def parse_token(request: Request):
+    token_with_bearer = request.headers.get('Authorization')
+
+    if token_with_bearer is not None and token_with_bearer.startswith('Bearer '):
+        token = token_with_bearer.split(' ')[1]
+    else:
+        # 处理未包含 "Bearer" 前缀的情况
+        token = token_with_bearer
+    return token
+
+
+# 生成令牌
+def generate_token(user_id, username):
+    # 构造包含用户信息的负载
+    payload = {
+        'user_id': user_id,
+        'username': username,
+        'exp': datetime.datetime.utcnow() + datetime.timedelta(hours=1)
+        }
+
+    # 在这里，您可以使用您的密钥（secret key）来签署令牌
+    # 选择适当的签名算法，并设置适当的过期时间等参数
+    # 仅使用 HS256 算法和过期时间为1小时
+    token = jwt.encode(payload, SECERT_KEY, algorithm='HS256')
+    return token
+
+
+# 验证令牌
+def validate_token(token):
+    try:
+        print(token)
+        # 使用密钥进行解码
+        payload = jwt.decode(token, SECERT_KEY, algorithms=['HS256'])
+        print(payload)
+
+        # 检查令牌的过期时间
+        if 'exp' in payload and datetime.datetime.utcnow() > datetime.datetime.fromtimestamp(payload['exp']):
+            return False
+
+        return True
+    except (jwt.DecodeError, jwt.InvalidTokenError):
+        return False
+
+
+def get_user_id_from_token(token):
+    try:
+        decoded_token = jwt.decode(
+            token, SECERT_KEY, algorithms=['HS256'])
+        user_id = decoded_token.get('user_id')
+        return user_id
+    except jwt.ExpiredSignatureError:
+        # 处理过期的令牌
+        return None
+    except (jwt.DecodeError, jwt.InvalidTokenError):
+        # 处理解码或无效的令牌
+        return None
+
+
 # 获取用户ID通过用户名
 def get_user_id_by_username(username):
     user_data = redis.hget('users', username)
@@ -257,12 +331,15 @@ def get_package_expiration(user_id):
     return expiration
 
 # 检查用户聊天次数是否超过限制
+
+
 def exceeded_chat_limit(user_id, package):
     user_basic_chat_key = f'user:{user_id}:basic_chat'
     user_advanced_chat_key = f'user:{user_id}:advanced_chat'
 
     basic_chat_limit = int(package.get(b'basic_chat_limit', 0).decode('utf-8'))
-    advanced_chat_limit = int(package.get(b'advanced_chat_limit', 0).decode('utf-8'))
+    advanced_chat_limit = int(package.get(
+        b'advanced_chat_limit', 0).decode('utf-8'))
 
     if basic_chat_limit >= 0 and int(redis.get(user_basic_chat_key) or 0) >= basic_chat_limit:
         return True