Flask-Sockets和Flask-Login联合实现websocket的登录认证功能

flask_login 提供了一个方便的方式来管理用户会话。当你在 Flask 的 HTTP 视图中使用它时，你可以简单地使用 @login_required 装饰器来确保用户已登录。

但是，flask_sockets 并没有直接与 flask_login 集成。如果你想在建立 WebSocket 连接时检查用户是否已登录，你需要采取一些额外的步骤。

以下是一个示例，说明如何在 flask_sockets 路由中使用 flask_login 进行身份验证：

初始化 Flask、Flask-Login 和 Flask-Sockets：

from flask import Flask, request, session
from flask_login import LoginManager, current_user, UserMixin, login_required
from flask_sockets import Socketsapp = Flask(__name__)
app.config['SECRET_KEY'] = 'your-secret-key'
sockets = Sockets(app)
login_manager = LoginManager()
login_manager.init_app(app)

定义用户模型和加载用户的回调：

class User(UserMixin):# For simplicity, this example does not use a real database.# Instead, it uses this dictionary to simulate user data.users = {"1": {"id": "1", "username": "user1", "password": "pass1"}}def __init__(self, id_, username):self.id = id_self.username = username@classmethoddef get(cls, id_):user_data = cls.users.get(id_)if not user_data:return Nonereturn User(id_=user_data["id"], username=user_data["username"])@login_manager.user_loader
def load_user(user_id):return User.get(user_id)

定义 WebSocket 路由并检查登录状态：

为了检查用户是否已登录，我们会读取 Flask 的会话数据，因为 flask_login 在其中存储了有关用户身份的信息。

@sockets.route('/socket')
def chat_socket(ws):# Flask's context might not be available in the WebSocket route,# so we manually load the user using Flask-Login's method.user = load_user(session.get("user_id"))# If user is not loaded (i.e., not logged in), we close the connection.if user is None:ws.close(reason="User not logged in.")returnwhile not ws.closed:message = ws.receive()if message:ws.send(f"Hello {user.username}, you said: {message}")

创建一些基本的登录和注销路由：

from flask import render_template, redirect, url_for, request@app.route('/login', methods=["GET", "POST"])
def login():if request.method == "POST":username = request.form.get("username")password = request.form.get("password")user = [u for u in User.users.values() if u["username"] == username and u["password"] == password]if user:user = User.get(user[0]["id"])login_user(user)return redirect(url_for("index"))return render_template("login.html")@app.route('/logout')
@login_required
def logout():logout_user()return redirect(url_for("login"))@app.route('/')
@login_required
def index():return "Logged in as " + current_user.username

运行应用程序：

if __name__ == "__main__":from gevent.pywsgi import WSGIServerfrom geventwebsocket.handler import WebSocketHandlerserver = WSGIServer(('127.0.0.1', 5000), app, handler_class=WebSocketHandler)server.serve_forever()