style: 使用ruff格式化登录模块，优化导入

ruoyi-fastapi-backend/module_admin/service/login_service.py

@@ -1,24 +1,28 @@
-from fastapi import Request, Form
-from fastapi import Depends
-from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
-from jose import JWTError, jwt
 import random
 import uuid
-from datetime import timedelta
-from module_admin.service.user_service import *
-from module_admin.entity.vo.login_vo import *
-from module_admin.entity.vo.common_vo import CrudResponseModel
-from module_admin.dao.login_dao import *
-from exceptions.exception import LoginException, AuthException, ServiceException
+from datetime import datetime, timedelta
+from fastapi import Depends, Form, Request
+from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
+from jose import JWTError, jwt
+from sqlalchemy.ext.asyncio import AsyncSession
+from typing import Dict, List, Optional, Union
 from config.constant import CommonConstant, MenuConstant
 from config.env import AppConfig, JwtConfig, RedisInitKeyConfig
 from config.get_db import get_db
+from exceptions.exception import LoginException, AuthException, ServiceException
+from module_admin.dao.login_dao import login_by_account
+from module_admin.dao.user_dao import UserDao
+from module_admin.entity.do.menu_do import SysMenu
+from module_admin.entity.vo.common_vo import CrudResponseModel
+from module_admin.entity.vo.login_vo import MenuTreeModel, MetaModel, RouterModel, SmsCode, UserLogin, UserRegister
+from module_admin.entity.vo.user_vo import AddUserModel, CurrentUserModel, ResetUserModel, TokenData, UserInfoModel
+from module_admin.service.user_service import UserService
 from utils.common_util import CamelCaseUtil
-from utils.pwd_util import *
-from utils.response_util import *
-from utils.message_util import *
+from utils.log_util import logger
+from utils.message_util import message_service
+from utils.pwd_util import PwdUtil
 
-oauth2_scheme = OAuth2PasswordBearer(tokenUrl="login")
+oauth2_scheme = OAuth2PasswordBearer(tokenUrl='login')
 
 
 class CustomOAuth2PasswordRequestForm(OAuth2PasswordRequestForm):
@@ -27,19 +31,25 @@ class CustomOAuth2PasswordRequestForm(OAuth2PasswordRequestForm):
     """
 
     def __init__(
-            self,
-            grant_type: str = Form(default=None, regex="password"),
-            username: str = Form(),
-            password: str = Form(),
-            scope: str = Form(default=""),
-            client_id: Optional[str] = Form(default=None),
-            client_secret: Optional[str] = Form(default=None),
-            code: Optional[str] = Form(default=""),
-            uuid: Optional[str] = Form(default=""),
-            login_info: Optional[Dict[str, str]] = Form(default=None)
+        self,
+        grant_type: str = Form(default=None, regex='password'),
+        username: str = Form(),
+        password: str = Form(),
+        scope: str = Form(default=''),
+        client_id: Optional[str] = Form(default=None),
+        client_secret: Optional[str] = Form(default=None),
+        code: Optional[str] = Form(default=''),
+        uuid: Optional[str] = Form(default=''),
+        login_info: Optional[Dict[str, str]] = Form(default=None),
     ):
-        super().__init__(grant_type=grant_type, username=username, password=password,
-                         scope=scope, client_id=client_id, client_secret=client_secret)
+        super().__init__(
+            grant_type=grant_type,
+            username=username,
+            password=password,
+            scope=scope,
+            client_id=client_id,
+            client_secret=client_secret,
+        )
         self.code = code
         self.uuid = uuid
         self.login_info = login_info
@@ -61,47 +71,61 @@ class LoginService:
         """
         await cls.__check_login_ip(request)
         account_lock = await request.app.state.redis.get(
-            f"{RedisInitKeyConfig.ACCOUNT_LOCK.get('key')}:{login_user.user_name}")
+            f"{RedisInitKeyConfig.ACCOUNT_LOCK.get('key')}:{login_user.user_name}"
+        )
         if login_user.user_name == account_lock:
-            logger.warning("账号已锁定，请稍后再试")
-            raise LoginException(data="", message="账号已锁定，请稍后再试")
+            logger.warning('账号已锁定，请稍后再试')
+            raise LoginException(data='', message='账号已锁定，请稍后再试')
         # 判断请求是否来自于api文档，如果是返回指定格式的结果，用于修复api文档认证成功后token显示undefined的bug
-        request_from_swagger = request.headers.get('referer').endswith('docs') if request.headers.get('referer') else False
-        request_from_redoc = request.headers.get('referer').endswith('redoc') if request.headers.get('referer') else False
+        request_from_swagger = (
+            request.headers.get('referer').endswith('docs') if request.headers.get('referer') else False
+        )
+        request_from_redoc = (
+            request.headers.get('referer').endswith('redoc') if request.headers.get('referer') else False
+        )
         # 判断是否开启验证码，开启则验证，否则不验证（dev模式下来自API文档的登录请求不检验）
-        if not login_user.captcha_enabled or ((request_from_swagger or request_from_redoc) and AppConfig.app_env == 'dev'):
+        if not login_user.captcha_enabled or (
+            (request_from_swagger or request_from_redoc) and AppConfig.app_env == 'dev'
+        ):
             pass
         else:
             await cls.__check_login_captcha(request, login_user)
         user = await login_by_account(query_db, login_user.user_name)
         if not user:
-            logger.warning("用户不存在")
-            raise LoginException(data="", message="用户不存在")
+            logger.warning('用户不存在')
+            raise LoginException(data='', message='用户不存在')
         if not PwdUtil.verify_password(login_user.password, user[0].password):
             cache_password_error_count = await request.app.state.redis.get(
-                f"{RedisInitKeyConfig.PASSWORD_ERROR_COUNT.get('key')}:{login_user.user_name}")
+                f"{RedisInitKeyConfig.PASSWORD_ERROR_COUNT.get('key')}:{login_user.user_name}"
+            )
             password_error_counted = 0
             if cache_password_error_count:
                 password_error_counted = cache_password_error_count
             password_error_count = int(password_error_counted) + 1
             await request.app.state.redis.set(
-                f"{RedisInitKeyConfig.PASSWORD_ERROR_COUNT.get('key')}:{login_user.user_name}", password_error_count,
-                ex=timedelta(minutes=10))
+                f"{RedisInitKeyConfig.PASSWORD_ERROR_COUNT.get('key')}:{login_user.user_name}",
+                password_error_count,
+                ex=timedelta(minutes=10),
+            )
             if password_error_count > 5:
                 await request.app.state.redis.delete(
-                    f"{RedisInitKeyConfig.PASSWORD_ERROR_COUNT.get('key')}:{login_user.user_name}")
+                    f"{RedisInitKeyConfig.PASSWORD_ERROR_COUNT.get('key')}:{login_user.user_name}"
+                )
                 await request.app.state.redis.set(
-                    f"{RedisInitKeyConfig.ACCOUNT_LOCK.get('key')}:{login_user.user_name}", login_user.user_name,
-                    ex=timedelta(minutes=10))
-                logger.warning("10分钟内密码已输错超过5次，账号已锁定，请10分钟后再试")
-                raise LoginException(data="", message="10分钟内密码已输错超过5次，账号已锁定，请10分钟后再试")
-            logger.warning("密码错误")
-            raise LoginException(data="", message="密码错误")
+                    f"{RedisInitKeyConfig.ACCOUNT_LOCK.get('key')}:{login_user.user_name}",
+                    login_user.user_name,
+                    ex=timedelta(minutes=10),
+                )
+                logger.warning('10分钟内密码已输错超过5次，账号已锁定，请10分钟后再试')
+                raise LoginException(data='', message='10分钟内密码已输错超过5次，账号已锁定，请10分钟后再试')
+            logger.warning('密码错误')
+            raise LoginException(data='', message='密码错误')
         if user[0].status == '1':
-            logger.warning("用户已停用")
-            raise LoginException(data="", message="用户已停用")
+            logger.warning('用户已停用')
+            raise LoginException(data='', message='用户已停用')
         await request.app.state.redis.delete(
-            f"{RedisInitKeyConfig.PASSWORD_ERROR_COUNT.get('key')}:{login_user.user_name}")
+            f"{RedisInitKeyConfig.PASSWORD_ERROR_COUNT.get('key')}:{login_user.user_name}"
+        )
         return user
 
     @classmethod
@@ -112,11 +136,12 @@ class LoginService:
         :return: 校验结果
         """
         black_ip_value = await request.app.state.redis.get(
-            f"{RedisInitKeyConfig.SYS_CONFIG.get('key')}:sys.login.blackIPList")
+            f"{RedisInitKeyConfig.SYS_CONFIG.get('key')}:sys.login.blackIPList"
+        )
         black_ip_list = black_ip_value.split(',') if black_ip_value else []
         if request.headers.get('X-Forwarded-For') in black_ip_list:
-            logger.warning("当前IP禁止登录")
-            raise LoginException(data="", message="当前IP禁止登录")
+            logger.warning('当前IP禁止登录')
+            raise LoginException(data='', message='当前IP禁止登录')
         return True
 
     @classmethod
@@ -128,13 +153,14 @@ class LoginService:
         :return: 校验结果
         """
         captcha_value = await request.app.state.redis.get(
-            f"{RedisInitKeyConfig.CAPTCHA_CODES.get('key')}:{login_user.uuid}")
+            f"{RedisInitKeyConfig.CAPTCHA_CODES.get('key')}:{login_user.uuid}"
+        )
         if not captcha_value:
-            logger.warning("验证码已失效")
-            raise LoginException(data="", message="验证码已失效")
+            logger.warning('验证码已失效')
+            raise LoginException(data='', message='验证码已失效')
         if login_user.code != str(captcha_value):
-            logger.warning("验证码错误")
-            raise LoginException(data="", message="验证码错误")
+            logger.warning('验证码错误')
+            raise LoginException(data='', message='验证码错误')
         return True
 
     @classmethod
@@ -150,13 +176,14 @@ class LoginService:
             expire = datetime.utcnow() + expires_delta
         else:
             expire = datetime.utcnow() + timedelta(minutes=30)
-        to_encode.update({"exp": expire})
+        to_encode.update({'exp': expire})
         encoded_jwt = jwt.encode(to_encode, JwtConfig.jwt_secret_key, algorithm=JwtConfig.jwt_algorithm)
         return encoded_jwt
 
     @classmethod
-    async def get_current_user(cls, request: Request = Request, token: str = Depends(oauth2_scheme),
-                               query_db: AsyncSession = Depends(get_db)):
+    async def get_current_user(
+        cls, request: Request = Request, token: str = Depends(oauth2_scheme), query_db: AsyncSession = Depends(get_db)
+    ):
         """
         根据token获取当前用户信息
         :param request: Request对象
@@ -172,31 +199,41 @@ class LoginService:
             if token.startswith('Bearer'):
                 token = token.split(' ')[1]
             payload = jwt.decode(token, JwtConfig.jwt_secret_key, algorithms=[JwtConfig.jwt_algorithm])
-            user_id: str = payload.get("user_id")
-            session_id: str = payload.get("session_id")
+            user_id: str = payload.get('user_id')
+            session_id: str = payload.get('session_id')
             if user_id is None:
-                logger.warning("用户token不合法")
-                raise AuthException(data="", message="用户token不合法")
+                logger.warning('用户token不合法')
+                raise AuthException(data='', message='用户token不合法')
             token_data = TokenData(user_id=int(user_id))
         except JWTError:
-            logger.warning("用户token已失效，请重新登录")
-            raise AuthException(data="", message="用户token已失效，请重新登录")
+            logger.warning('用户token已失效，请重新登录')
+            raise AuthException(data='', message='用户token已失效，请重新登录')
         query_user = await UserDao.get_user_by_id(query_db, user_id=token_data.user_id)
         if query_user.get('user_basic_info') is None:
-            logger.warning("用户token不合法")
-            raise AuthException(data="", message="用户token不合法")
+            logger.warning('用户token不合法')
+            raise AuthException(data='', message='用户token不合法')
         if AppConfig.app_same_time_login:
-            redis_token = await request.app.state.redis.get(f"{RedisInitKeyConfig.ACCESS_TOKEN.get('key')}:{session_id}")
+            redis_token = await request.app.state.redis.get(
+                f"{RedisInitKeyConfig.ACCESS_TOKEN.get('key')}:{session_id}"
+            )
         else:
             # 此方法可实现同一账号同一时间只能登录一次
-            redis_token = await request.app.state.redis.get(f"{RedisInitKeyConfig.ACCESS_TOKEN.get('key')}:{query_user.get('user_basic_info').user_id}")
+            redis_token = await request.app.state.redis.get(
+                f"{RedisInitKeyConfig.ACCESS_TOKEN.get('key')}:{query_user.get('user_basic_info').user_id}"
+            )
         if token == redis_token:
             if AppConfig.app_same_time_login:
-                await request.app.state.redis.set(f"{RedisInitKeyConfig.ACCESS_TOKEN.get('key')}:{session_id}", redis_token,
-                                                  ex=timedelta(minutes=JwtConfig.jwt_redis_expire_minutes))
+                await request.app.state.redis.set(
+                    f"{RedisInitKeyConfig.ACCESS_TOKEN.get('key')}:{session_id}",
+                    redis_token,
+                    ex=timedelta(minutes=JwtConfig.jwt_redis_expire_minutes),
+                )
             else:
-                await request.app.state.redis.set(f"{RedisInitKeyConfig.ACCESS_TOKEN.get('key')}:{query_user.get('user_basic_info').user_id}", redis_token,
-                                                  ex=timedelta(minutes=JwtConfig.jwt_redis_expire_minutes))
+                await request.app.state.redis.set(
+                    f"{RedisInitKeyConfig.ACCESS_TOKEN.get('key')}:{query_user.get('user_basic_info').user_id}",
+                    redis_token,
+                    ex=timedelta(minutes=JwtConfig.jwt_redis_expire_minutes),
+                )
 
             role_id_list = [item.role_id for item in query_user.get('user_role_info')]
             if 1 in role_id_list:
@@ -215,13 +252,13 @@ class LoginService:
                     postIds=post_ids,
                     roleIds=role_ids,
                     dept=CamelCaseUtil.transform_result(query_user.get('user_dept_info')),
-                    role=CamelCaseUtil.transform_result(query_user.get('user_role_info'))
-                )
+                    role=CamelCaseUtil.transform_result(query_user.get('user_role_info')),
+                ),
             )
             return current_user
         else:
-            logger.warning("用户token已失效，请重新登录")
-            raise AuthException(data="", message="用户token已失效，请重新登录")
+            logger.warning('用户token已失效，请重新登录')
+            raise AuthException(data='', message='用户token已失效，请重新登录')
 
     @classmethod
     async def get_current_user_routers(cls, user_id: int, query_db: AsyncSession):
@@ -232,7 +269,14 @@ class LoginService:
         :return: 当前用户路由信息对象
         """
         query_user = await UserDao.get_user_by_id(query_db, user_id=user_id)
-        user_router_menu = sorted([row for row in query_user.get('user_menu_info') if row.menu_type in [MenuConstant.TYPE_DIR, MenuConstant.TYPE_MENU]], key=lambda x: x.order_num)
+        user_router_menu = sorted(
+            [
+                row
+                for row in query_user.get('user_menu_info')
+                if row.menu_type in [MenuConstant.TYPE_DIR, MenuConstant.TYPE_MENU]
+            ],
+            key=lambda x: x.order_num,
+        )
         menus = cls.__generate_menus(0, user_router_menu)
         user_router = cls.__generate_user_router_menu(menus)
         return [router.model_dump(exclude_unset=True, by_alias=True) for router in user_router]
@@ -275,8 +319,8 @@ class LoginService:
                     title=permission.menu_name,
                     icon=permission.icon,
                     noCache=True if permission.is_cache == 1 else False,
-                    link=permission.path if RouterUtil.is_http(permission.path) else None
-                )
+                    link=permission.path if RouterUtil.is_http(permission.path) else None,
+                ),
             )
             c_menus = permission.children
             if c_menus and permission.menu_type == MenuConstant.TYPE_DIR:
@@ -294,17 +338,14 @@ class LoginService:
                         title=permission.menu_name,
                         icon=permission.icon,
                         noCache=True if permission.is_cache == 1 else False,
-                        link=permission.path if RouterUtil.is_http(permission.path) else None
+                        link=permission.path if RouterUtil.is_http(permission.path) else None,
                     ),
-                    query=permission.query
+                    query=permission.query,
                 )
                 children_list.append(children)
                 router.children = children_list
             elif permission.parent_id == 0 and RouterUtil.is_inner_link(permission):
-                router.meta = MetaModel(
-                    title=permission.menu_name,
-                    icon=permission.icon
-                )
+                router.meta = MetaModel(title=permission.menu_name, icon=permission.icon)
                 router.path = '/'
                 children_list: List[RouterModel] = []
                 router_path = RouterUtil.inner_link_replace_each(permission.path)
@@ -315,8 +356,8 @@ class LoginService:
                     meta=MetaModel(
                         title=permission.menu_name,
                         icon=permission.icon,
-                        link=permission.path if RouterUtil.is_http(permission.path) else None
-                    )
+                        link=permission.path if RouterUtil.is_http(permission.path) else None,
+                    ),
                 )
                 children_list.append(children)
                 router.children = children_list
@@ -334,15 +375,26 @@ class LoginService:
         :param user_register: 注册用户对象
         :return: 注册结果
         """
-        register_enabled = True if await request.app.state.redis.get(
-            f"{RedisInitKeyConfig.SYS_CONFIG.get('key')}:sys.account.registerUser") == 'true' else False
-        captcha_enabled = True if await request.app.state.redis.get(
-            f"{RedisInitKeyConfig.SYS_CONFIG.get('key')}:sys.account.captchaEnabled") == 'true' else False
+        register_enabled = (
+            True
+            if await request.app.state.redis.get(f"{RedisInitKeyConfig.SYS_CONFIG.get('key')}:sys.account.registerUser")
+            == 'true'
+            else False
+        )
+        captcha_enabled = (
+            True
+            if await request.app.state.redis.get(
+                f"{RedisInitKeyConfig.SYS_CONFIG.get('key')}:sys.account.captchaEnabled"
+            )
+            == 'true'
+            else False
+        )
         if user_register.password == user_register.confirm_password:
             if register_enabled:
                 if captcha_enabled:
                     captcha_value = await request.app.state.redis.get(
-                        f"{RedisInitKeyConfig.CAPTCHA_CODES.get('key')}:{user_register.uuid}")
+                        f"{RedisInitKeyConfig.CAPTCHA_CODES.get('key')}:{user_register.uuid}"
+                    )
                     if not captcha_value:
                         raise ServiceException(message='验证码已失效')
                     elif user_register.code != str(captcha_value):
@@ -350,7 +402,7 @@ class LoginService:
                 add_user = AddUserModel(
                     userName=user_register.username,
                     nickName=user_register.username,
-                    password=PwdUtil.get_password_hash(user_register.password)
+                    password=PwdUtil.get_password_hash(user_register.password),
                 )
                 result = await UserService.add_user_services(query_db, add_user)
                 return result
@@ -369,15 +421,17 @@ class LoginService:
         :return: 短信验证码对象
         """
         redis_sms_result = await request.app.state.redis.get(
-            f"{RedisInitKeyConfig.SMS_CODE.get('key')}:{user.session_id}")
+            f"{RedisInitKeyConfig.SMS_CODE.get('key')}:{user.session_id}"
+        )
         if redis_sms_result:
             return SmsCode(**dict(is_success=False, sms_code='', session_id='', message='短信验证码仍在有效期内'))
         is_user = await UserDao.get_user_by_name(query_db, user.user_name)
         if is_user:
             sms_code = str(random.randint(100000, 999999))
             session_id = str(uuid.uuid4())
-            await request.app.state.redis.set(f"{RedisInitKeyConfig.SMS_CODE.get('key')}:{session_id}", sms_code,
-                                              ex=timedelta(minutes=2))
+            await request.app.state.redis.set(
+                f"{RedisInitKeyConfig.SMS_CODE.get('key')}:{session_id}", sms_code, ex=timedelta(minutes=2)
+            )
             # 此处模拟调用短信服务
             message_service(sms_code)
 
@@ -395,7 +449,8 @@ class LoginService:
         :return: 重置结果
         """
         redis_sms_result = await request.app.state.redis.get(
-            f"{RedisInitKeyConfig.SMS_CODE.get('key')}:{forget_user.session_id}")
+            f"{RedisInitKeyConfig.SMS_CODE.get('key')}:{forget_user.session_id}"
+        )
         if forget_user.sms_code == redis_sms_result:
             forget_user.password = PwdUtil.get_password_hash(forget_user.password)
             forget_user.user_id = (await UserDao.get_user_by_name(query_db, forget_user.user_name)).user_id
@@ -484,7 +539,9 @@ class RouterUtil:
         :param menu: 菜单数对象
         :return: 是否为菜单内部跳转
         """
-        return menu.parent_id == 0 and menu.menu_type == MenuConstant.TYPE_MENU and menu.is_frame == MenuConstant.NO_FRAME
+        return (
+            menu.parent_id == 0 and menu.menu_type == MenuConstant.TYPE_MENU and menu.is_frame == MenuConstant.NO_FRAME
+        )
 
     @classmethod
     def is_inner_link(cls, menu: MenuTreeModel):
@@ -520,8 +577,8 @@ class RouterUtil:
         :param path: 内链域名
         :return: 替换后的内链域名
         """
-        old_values = [CommonConstant.HTTP, CommonConstant.HTTPS, CommonConstant.WWW, ".", ":"]
-        new_values = ["", "", "", "/", "/"]
+        old_values = [CommonConstant.HTTP, CommonConstant.HTTPS, CommonConstant.WWW, '.', ':']
+        new_values = ['', '', '', '/', '/']
         for old, new in zip(old_values, new_values):
             path = path.replace(old, new)
         return path