【Flask】使用 Token 作为身份验证和会话管理

【Flask】使用 Token 作为身份验证和会话管理

在 Flask 中生成 token 通常用于身份验证和会话管理。最常用的方法是使用其扩展，如 Flask-JWT 或 Flask-Security。

假设我们已经写好了简单的登录注册接口：

@app.route('/api/register', methods=['POST'])
def register():
    data = request.get_json()

    # 确认数据有效性
    username = data.get('username')
    password = data.get('password')
    if not username or not password:
        return jsonify({'message': 'Missing username or password'}), 400

    # 检查用户名是否已存在
    if User.query.filter_by(username=username).first():
        return jsonify({'message': 'Username already exists'}), 400

    # 创建新用户
    hashed_password = generate_password_hash(password)
    new_user = User(username=username, password=hashed_password)
    db.session.add(new_user)
    db.session.commit()

    return jsonify({'message': 'User created successfully', 'user': new_user.to_dict()}), 201

@app.route('/api/login', methods=['POST'])
def login():
    data = request.get_json()

    username = data.get('username')
    password = data.get('password')
    if not username or not password:
        return jsonify({'message': 'Missing username or password'}), 400

    # 检查用户是否存在
    user = User.query.filter_by(username=username).first()
    if user and check_password_hash(user.password, password):
        # 在这里，你应该生成并返回一个认证令牌
        return jsonify({'message': 'Login successful', 'user': user.to_dict()}), 200
    else:
        return jsonify({'message': 'Invalid credentials'}), 401

if __name__ == '__main__':
    app.run(debug=True)

为了添加 token，首先，需要一个方式来生成和验证 JWT（JSON Web Tokens）。通过安装 Flask-JWT-Extended 来实现：

pip install Flask-JWT-Extended

修改 Flask 应用以集成 JWT：
1. 配置 Flask-JWT-Extended:

from flask import Flask, jsonify, request
from flask_jwt_extended import JWTManager, create_access_token, jwt_required, get_jwt_identity
from werkzeug.security import generate_password_hash, check_password_hash
# 其他必要的导入

app = Flask(__name__)
app.config['JWT_SECRET_KEY'] = 'your_jwt_secret_key'  # 使用一个安全的密钥
jwt = JWTManager(app)

2. 生成 Token:

@app.route('/api/login', methods=['POST'])
def login():
    # 现有代码...

    if user and check_password_hash(user.password, password):
        access_token = create_access_token(identity=username)
        return jsonify({'message': 'Login successful', 'access_token': access_token, 'user': user.to_dict()}), 200
    else:
        return jsonify({'message': 'Invalid credentials'}), 401

在这里，create_access_token 函数用于生成 JWT。identity 参数通常是用户的唯一标识（如用户名）。

3. 保护路由:
对于需要认证的路由，使用 @jwt_required() 装饰器。例如，如果您有一个需要用户登录才能访问的路由：

@app.route('/protected')
@jwt_required()
def protected_route():
    current_user = get_jwt_identity()
    return jsonify(logged_in_as=current_user), 200

4. 数据库设置:
   确保您的数据库模型 User 已经设置好并且包含了 username 和 password 字段。这通常在 Flask 应用的模型定义部分完成，例如使用 Flask-SQLAlchemy。

from flask_sqlalchemy import SQLAlchemy

db = SQLAlchemy(app)

class User(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    username = db.Column(db.String(50), unique=True, nullable=False)
    password = db.Column(db.String(80), nullable=False)

    # 可选：将用户对象转换为字典形式
    def to_dict(self):
        return {"id": self.id, "username": self.username}

示例代码

from flask import Flask, jsonify, request
from flask_sqlalchemy import SQLAlchemy
from flask_jwt_extended import JWTManager, create_access_token, jwt_required, get_jwt_identity
from werkzeug.security import generate_password_hash, check_password_hash

# 应用和数据库配置
app = Flask(__name__)
app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///users.db'  # 使用 SQLite 数据库
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
app.config['JWT_SECRET_KEY'] = 'your_jwt_secret_key'  # JWT 秘钥

db = SQLAlchemy(app)
jwt = JWTManager(app)

# 用户模型
class User(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    username = db.Column(db.String(50), unique=True, nullable=False)
    password = db.Column(db.String(80), nullable=False)

# 创建数据库
@app.before_first_request
def create_tables():
    db.create_all()

# 注册接口
@app.route('/register', methods=['POST'])
def register():
    data = request.get_json()
    username = data.get('username')
    password = data.get('password')

    if not username or not password:
        return jsonify({'message': 'Missing username or password'}), 400

    if User.query.filter_by(username=username).first():
        return jsonify({'message': 'Username already exists'}), 400

    hashed_password = generate_password_hash(password)
    new_user = User(username=username, password=hashed_password)
    db.session.add(new_user)
    db.session.commit()

    return jsonify({'message': 'User created successfully'}), 201

# 登录接口
@app.route('/login', methods=['POST'])
def login():
    data = request.get_json()
    username = data.get('username')
    password = data.get('password')

    if not username or not password:
        return jsonify({'message': 'Missing username or password'}), 400

    user = User.query.filter_by(username=username).first()
    if user and check_password_hash(user.password, password):
        access_token = create_access_token(identity=username)
        return jsonify({'message': 'Login successful', 'access_token': access_token}), 200
    else:
        return jsonify({'message': 'Invalid credentials'}), 401

# 受保护的路由
@app.route('/protected', methods=['GET'])
@jwt_required()
def protected():
    current_user = get_jwt_identity()
    return jsonify({'logged_in_as': current_user}), 200

if __name__ == '__main__':
    app.run(debug=True)

如何获取secret_key

加密 token 的秘钥（通常称为 “JWT Secret Key”）是自己定义的一个字符串。这个秘钥用于加密和解密 JSON Web Tokens（JWT），因此保持其安全性非常重要。秘钥应该足够复杂，以防止猜测或暴力破解攻击。
秘钥应该是一个长且随机的字符串。可以使用 Python 的 secrets 模块来生成一个安全的秘钥，例如：

import secrets
secret_key = secrets.token_urlsafe(64)
print(secret_key)

通常，最佳实践是将 JWT Secret Key 存储在环境变量中。这样，可以在不同的部署环境中使用不同的秘钥，而无需更改代码。例如，在 Flask 应用中，可以这样设置：

import os
app.config['JWT_SECRET_KEY'] = os.environ.get('JWT_SECRET_KEY', 'your_default_secret_key')

这里，如果环境变量 JWT_SECRET_KEY 存在，则使用它；如果不存在，使用默认值 ‘your_default_secret_key’。