update:用户登录全流程

4 files changed, 108 insertions, 48 deletions

diff --git a/server/app.py b/server/app.py
index e63f551..b47fceb 100644
--- a/server/app.py
+++ b/server/app.py
@@ -1,15 +1,14 @@
 from apiflask import APIFlask
-from flask_cors import CORS
-from flask_migrate import Migrate
+
 
 import settings
 from apps.agentcomm import bp as agentbp
 from apps.login import bp as loginbp
 from apps.sysinfo import bp as sysbp
-from apps.sysmange import bp as mangbp
+# from apps.sysmange import bp as mangbp
 from apps.target import bp as targetbp
 from apps.task import bp as taskbp
-from exts import db
+from exts import db, login_manager, CORS, Migrate
 
 # 注册蓝图
 app = APIFlask(__name__, template_folder='./static/templates')
@@ -21,8 +20,8 @@ app.register_blueprint(targetbp)
 app.register_blueprint(agentbp)
 # 任务管理接口
 app.register_blueprint(taskbp)
-# 系统管理接口
-app.register_blueprint(mangbp)
+# # 系统管理接口
+# app.register_blueprint(mangbp)
 # 系统状态获取接口
 app.register_blueprint(sysbp)
 # 登录退出接口
@@ -31,10 +30,12 @@ app.register_blueprint(loginbp)
 app.config[
     "SQLALCHEMY_DATABASE_URI"] = f"mysql+pymysql://root:{settings.MYSQL_PAWD}@{settings.MYSQL_HOST}:{settings.MYSQL_PORT}/{settings.MYSQL_DATADB}"
 app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///project.db"
+app.config['SECRET_KEY'] = 'your-secret-key'
 db.init_app(app)
-
 migrate = Migrate(app, db)
 
+login_manager.init_app(app)
+
 @app.get('/')
 @app.doc("获取测试页面")
 def hello():
diff --git a/server/apps/login.py b/server/apps/login.py
index 4be3ef0..fc2ebca 100644
--- a/server/apps/login.py
+++ b/server/apps/login.py
@@ -1,57 +1,102 @@
 # 登录
+from flask import jsonify, request
+
 from apiflask import APIBlueprint
 from apiflask.fields import String, Integer, Boolean
-
+from uuid import uuid4
 from model import User
-from exts import db
-
+from exts import db, login_manager
+from flask_login import login_user, logout_user, login_required, current_user
 bp = APIBlueprint("登录退出接口集合", __name__, url_prefix="/auth")
 
-# TODO 使用cookie
+@login_manager.user_loader
+def load_user(user_id):
+    return User.query.get(user_id)
 
-# 用户登录处理函数
[email protected]("/login")
[email protected]("登录接口",
-        description="输入用户名和密码，返回响应码code，消息msg以及是否成功登陆标识status</br>若成功登录则msg返回用户编号")
+
[email protected]('/register', methods=['POST'])
[email protected]("系统用户创建接口", description="接收参数包括 新帐户名、新用户密码以及当前执行该注册操作的账户名")
 @bp.input({
-    "username": String(),
-    "pwd": String()
+    "username": String(required=True),
+    "password": String(required=True),
+    "group": String(default="Normal")
 })
 @bp.output({
-    # 响应码
     "code": Integer(),
-    # 返回消息
-    "msg": String(),
-    # 是否成功登录
-    "status": Boolean(),
+    "message": String()
 })
+@login_required
+def register(json_data):
+    data = json_data
+    username = data.get('username')
+    password = data.get('password')
 
-def login(json_data):
-    users = db.session.execute(db.select(User).filter_by(user_name=json_data["username"])).scalars()
-    for user in users:
-        if user.user_pwd_hash == json_data["pwd"]:
-            return {"code":200, "status":True, "msg":user.user_id}
-    return {"code":200, "status":False, "msg":""}
+    if User.query.filter_by(user_name=username).first():
+        return jsonify({
+            'code': 400,
+            'message': 'User already exists'}),
+
+    user = User(
+        user_id=str(uuid4()),
+        user_name=data["username"],
+        # created_by='Root',
+        created_by=current_user.user_name,
+        user_group='Normal'
+    )
+    user.set_password(password)
+    db.session.add(user)
+    db.session.commit()
+    return jsonify({
+        'code': 201,
+        'message': 'User registered successfully'})
 
 
-# 用户退出处理函数
[email protected]("/logout")
[email protected]("退出登录接口", description="输入用户名和用户编号，返回响应码code，消息msg以及是否成功退出标识status")
[email protected]('/login', methods=['POST'])
[email protected]("登录接口",
+        description="输入用户名和密码，返回响应码code，消息message")
 @bp.input({
-    "id": String(),
-    "username": String()
+    "username": String(required=True),
+    "password": String(required=True)
 })
 @bp.output({
-    # 响应码
-    "code": Integer(),
-    # 返回消息
-    "msg": String(),
-    # 是否成功退出
-    "status": Boolean(),
+    "message": String()
+})
+def login(json_data):
+    data = json_data
+    username = data.get('username')
+    password = data.get('password')
+    user = User.query.filter_by(user_name=username).first()
+    if user and user.check_password(password):
+        login_user(user)
+        return jsonify({
+            'code': 200,
+            'message': 'Logged in successfully'})
+    return jsonify({
+        'code': 401,
+        'message': 'Invalid credentials'})
+
+
[email protected]('/logout', methods=['POST'])
[email protected]("登出接口",
+        description="返回响应码code，消息message")
[email protected]({
+    "message": String()
+})
+@login_required
+def logout():
+    logout_user()
+    return jsonify({
+        'code': 200,
+        'message': 'Logged out successfully'})
+
[email protected]('/protected', methods=['GET'])
[email protected]("登陆状态测试接口",
+        description="返回响应码code，消息message")
[email protected]({
+    "message": String()
 })
-def logout(json_data):
-    return {
-        "code": 200,
-        "msg":"退出成功",
-        "status":True
-    }
\ No newline at end of file
+@login_required
+def protected():
+    return jsonify({
+        'code': 200,
+        'message': f'Logged in as: {current_user.user_name}'})
diff --git a/server/exts.py b/server/exts.py
index 2e1eeb6..84a79cb 100644
--- a/server/exts.py
+++ b/server/exts.py
@@ -1,3 +1,7 @@
 from flask_sqlalchemy import SQLAlchemy
+from flask_login import LoginManager
+from flask_cors import CORS
+from flask_migrate import Migrate
 
+login_manager = LoginManager()
 db = SQLAlchemy()
\ No newline at end of file
diff --git a/server/model.py b/server/model.py
index bf73bd2..01c196b 100644
--- a/server/model.py
+++ b/server/model.py
@@ -2,11 +2,12 @@ from exts import db
 from settings import *
 import json
 import datetime
+from flask_login import UserMixin
+from werkzeug.security import generate_password_hash, check_password_hash
 
-
-class User(db.Model):
+class User(db.Model, UserMixin):
     __tablename__ = MYSQL_TAB_USER
-    user_id = db.Column(db.String(255), primary_key=True, nullable=False)
+    user_id = db.Column(db.String(255), primary_key=True, nullable=False, )
     user_name = db.Column(db.String(255))
     created_by = db.Column(db.String(255))
     created_time = db.Column(db.TIMESTAMP, server_default=db.func.now())
@@ -15,6 +16,15 @@ class User(db.Model):
 
     syslogs = db.relationship('Syslog', backref='user')
     tasks = db.relationship('Task', backref='user')
+    
+    def get_id(self):
+        return (self.user_id)
+
+    def set_password(self, password):
+        self.user_pwd_hash = generate_password_hash(password)
+
+    def check_password(self, password):
+        return check_password_hash(self.user_pwd_hash, password)
 
 
 class Task(db.Model):