添加支付

redis/test_user_redis.py

@@ -1,6 +1,7 @@
 import json
 import random
 import string
+import uuid
 from flask import Flask, request, jsonify, session
 from redis import Redis
 
@@ -58,9 +59,14 @@ def register():
     # 检查用户名是否已被注册
     if redis.hexists('users', username):
         return jsonify({'message': 'Username already exists'}), 400
+    
+    # 生成唯一的用户ID
+    user_id = str(uuid.uuid4())
 
     # 保存用户信息到Redis
     user_data = {
+        'user_id': user_id,
+        'username': username,
         'email': email,
         'password': password
     }
@@ -99,7 +105,7 @@ def protected():
         username = session['username']
         # 其他业务逻辑...
         return jsonify({'message': f'Hello, {username}! This is a protected endpoint.'})
-    
+
     # 如果用户未登录，则返回未授权的响应
     return jsonify({'message': 'Unauthorized'}), 401
 
@@ -112,5 +118,115 @@ def logout():
     return jsonify({'message': 'Logout successful'})
 
 
+# 购买支付套餐
+@app.route('/purchase', methods=['POST'])
+def purchase():
+    package_id = request.json.get('package_id')
+
+    # 根据套餐ID获取套餐信息
+    package = get_package_by_id(package_id)
+    if not package:
+        return jsonify({'message': 'Invalid package ID'}), 400
+    
+    # 根据用户名查询用户ID
+    username = session.get('username')
+    if not username:
+        return jsonify({'message': 'User not logged in'}), 400
+
+    user_id = get_user_id_by_username(username)
+    if not user_id:
+        return jsonify({'message': 'User not found'}), 400
+
+    # 存储用户套餐信息到Redis
+    store_user_package(user_id, package)
+
+    return jsonify({'message': 'Purchase successful'})
+
+
+# 验证用户聊天次数
+@app.route('/validate', methods=['POST'])
+def validate():
+    # 根据用户名查询用户ID
+    username = session.get('username')
+    user_id = get_user_id_by_username(username)
+    
+    if not user_id:
+        return jsonify({'message': 'User not found'}), 400
+    # 获取用户套餐信息
+    package = get_user_package(user_id)
+    if not package:
+        return jsonify({'message': 'User has not purchased any package'}), 400
+
+    # 检查用户聊天次数是否超过限制
+    if exceeded_chat_limit(user_id, package):
+        return jsonify({'message': 'Chat limit exceeded'}), 400
+
+    return jsonify({'message': 'Chat limit not exceeded'})
+
+
+# 获取用户ID通过用户名
+def get_user_id_by_username(username):
+    user_data = redis.hget('users', username)
+    if user_data:
+        user_data = json.loads(user_data.decode('utf-8'))
+        user_id = user_data.get('user_id')
+        return user_id
+    return None
+
+
+# 根据套餐ID获取套餐信息
+def get_package_by_id(package_id):
+    packages = {
+        '1': {
+            'name': 'Package 1',
+            'basic_chat_limit': 10,
+            'advanced_chat_limit': 10,
+            'price': 10
+        },
+        '2': {
+            'name': 'Package 2',
+            'basic_chat_limit': -1,  # -1 表示无限次
+            'advanced_chat_limit': -1,
+            'price': 100
+        }
+    }
+
+    return packages.get(package_id)
+
+
+# 存储用户套餐信息到Redis
+def store_user_package(user_id, package):
+    user_package_key = f'user:{user_id}:package'
+    redis.hset(user_package_key, 'name', package['name'])
+    redis.hset(user_package_key, 'basic_chat_limit',
+               package['basic_chat_limit'])
+    redis.hset(user_package_key, 'advanced_chat_limit',
+               package['advanced_chat_limit'])
+
+
+# 获取用户套餐信息
+def get_user_package(user_id):
+    user_package_key = f'user:{user_id}:package'
+    package = redis.hgetall(user_package_key)
+    return package
+
+
+# 检查用户聊天次数是否超过限制
+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'))
+
+    if basic_chat_limit >= 0 and int(redis.get(user_basic_chat_key) or 0) >= basic_chat_limit:
+        return True
+
+    if advanced_chat_limit >= 0 and int(redis.get(user_advanced_chat_key) or 0) >= advanced_chat_limit:
+        return True
+
+    return False
+
+
 if __name__ == '__main__':
     app.run(debug=True)