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：

python 复制代码

from flask import Flask, request, session
from flask_login import LoginManager, current_user, UserMixin, login_required
from flask_sockets import Sockets

app = Flask(__name__)
app.config['SECRET_KEY'] = 'your-secret-key'
sockets = Sockets(app)
login_manager = LoginManager()
login_manager.init_app(app)

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

python 复制代码

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

    @classmethod
    def get(cls, id_):
        user_data = cls.users.get(id_)
        if not user_data:
            return None
        return 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 在其中存储了有关用户身份的信息。

python 复制代码

@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.")
        return

    while not ws.closed:
        message = ws.receive()
        if message:
            ws.send(f"Hello {user.username}, you said: {message}")

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

python 复制代码

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

运行应用程序：

python 复制代码

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

注意：在上述代码中，我们手动加载用户并检查他们是否已登录，而不是使用 @login_required 装饰器。在 WebSocket 路由中，由于 Flask 的请求上下文可能不可用，所以这是必要的。