Add login graphic captcha (#124)

* add the get verification code interface

* add the login interface captcha

* clean up the debugging code

Author: wu-clan

backend/app/api/v1/auth/__init__.py

@@ -2,7 +2,9 @@
 # -*- coding: utf-8 -*-
 from fastapi import APIRouter
 from backend.app.api.v1.auth.auth import router as auth_router
+from backend.app.api.v1.auth.captcha import router as captcha_router
 
 router = APIRouter(prefix='/auth', tags=['认证'])
 
 router.include_router(auth_router)
+router.include_router(captcha_router)

backend/app/api/v1/auth/auth.py

@@ -10,7 +10,7 @@
 from backend.app.common.jwt import DependsJwtAuth
 from backend.app.common.response.response_schema import response_base
 from backend.app.schemas.token import GetLoginToken, GetSwaggerToken, GetNewToken
-from backend.app.schemas.user import Auth
+from backend.app.schemas.user import AuthLogin
 from backend.app.services.auth_service import AuthService
 
 router = APIRouter()
@@ -28,7 +28,7 @@ async def swagger_user_login(form_data: OAuth2PasswordRequestForm = Depends()) -
     description='json 格式登录, 仅支持在第三方api工具调试接口, 例如: postman',
     dependencies=[Depends(RateLimiter(times=5, minutes=15))],
 )
-async def user_login(request: Request, obj: Auth, background_tasks: BackgroundTasks):
+async def user_login(request: Request, obj: AuthLogin, background_tasks: BackgroundTasks):
     access_token, refresh_token, access_expire, refresh_expire, user = await AuthService().login(
         request=request, obj=obj, background_tasks=background_tasks
     )

backend/app/api/v1/auth/captcha.py

@@ -0,0 +1,30 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+from fast_captcha import img_captcha
+from fastapi import APIRouter, Depends, Request
+from fastapi_limiter.depends import RateLimiter
+from starlette.concurrency import run_in_threadpool
+
+from backend.app.common.redis import redis_client
+from backend.app.common.response.response_schema import response_base
+from backend.app.core.conf import settings
+
+router = APIRouter()
+
+
+@router.get(
+    '/captcha',
+    summary='获取登录验证码',
+    dependencies=[Depends(RateLimiter(times=5, seconds=10))],
+)
+async def get_captcha(request: Request):
+    """
+    此接口可能存在性能损耗，尽管是异步接口，但是验证码生成是同步IO事件，使用线程池处理尽量减少性能损耗
+    """
+    img_type: str = 'base64'
+    img, code = await run_in_threadpool(img_captcha, img_byte=img_type)
+    ip = request.state.ip
+    await redis_client.set(
+        f'{settings.CAPTCHA_LOGIN_REDIS_PREFIX}:{ip}', code, ex=settings.CAPTCHA_LOGIN_EXPIRE_SECONDS
+    )
+    return await response_base.success(data={'image_type': img_type, 'image': img})

backend/app/api/v1/menu.py

@@ -13,6 +13,7 @@
 
 router = APIRouter()
 
+
 @router.get('/{pk}', summary='获取目录详情', dependencies=[DependsJwtAuth])
 async def get_menu(pk: int):
     menu = await MenuService.get(pk=pk)
@@ -22,8 +23,8 @@ async def get_menu(pk: int):
 
 @router.get('', summary='获取所有目录展示树', dependencies=[DependsJwtAuth])
 async def get_all_menus(
-        name: Annotated[str | None, Query()] = None,
-        status: Annotated[bool | None, Query()] = None,
+    name: Annotated[str | None, Query()] = None,
+    status: Annotated[bool | None, Query()] = None,
 ):
     menu = await MenuService.get_select(name=name, status=status)
     return await response_base.success(data=menu)
@@ -49,6 +50,3 @@ async def delete_menu(pk: int):
     if count > 0:
         return await response_base.success()
     return await response_base.fail()
-
-
-

backend/app/core/conf.py

@@ -87,6 +87,10 @@ def validator_api_url(cls, values):
     TOKEN_REDIS_PREFIX: str = 'fba_token'
     TOKEN_REFRESH_REDIS_PREFIX: str = 'fba_refresh_token'
 
+    # captcha
+    CAPTCHA_LOGIN_REDIS_PREFIX: str = 'fba_login_captcha'
+    CAPTCHA_LOGIN_EXPIRE_SECONDS: int = 60 * 5  # 过期时间，单位：秒
+
     # Log
     LOG_STDOUT_FILENAME: str = 'fba_access.log'
     LOG_STDERR_FILENAME: str = 'fba_error.log'

backend/app/crud/crud_dept.py

@@ -2,7 +2,7 @@
 # -*- coding: utf-8 -*-
 from typing import Any
 
-from sqlalchemy import select, desc, and_, asc
+from sqlalchemy import select, and_, asc
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import selectinload
 

backend/app/schemas/role.py

@@ -10,7 +10,7 @@
 
 class RoleBase(BaseModel):
     name: str
-    data_scope: int | None = Field(default=RoleDataScope.custom, description='数据范围（1：全部数据权限 2：自定数据权限）')
+    data_scope: int | None = Field(default=RoleDataScope.custom, description='数据范围（1：全部数据权限 2：自定数据权限）')  # noqa: E501
     status: bool
     remark: str | None = None
 

backend/app/schemas/user.py

@@ -14,6 +14,10 @@ class Auth(BaseModel):
     password: str
 
 
+class AuthLogin(Auth):
+    captcha: str
+
+
 class CreateUser(Auth):
     dept_id: int
     roles: list[int]

backend/app/services/auth_service.py

@@ -13,10 +13,11 @@
 from backend.app.common.exception import errors
 from backend.app.common.jwt import get_token
 from backend.app.common.redis import redis_client
+from backend.app.common.response.response_code import CustomCode
 from backend.app.core.conf import settings
 from backend.app.crud.crud_user import UserDao
 from backend.app.database.db_mysql import async_db_session
-from backend.app.schemas.user import Auth
+from backend.app.schemas.user import AuthLogin
 from backend.app.services.login_log_service import LoginLogService
 
 
@@ -40,7 +41,7 @@ async def swagger_login(self, *, form_data: OAuth2PasswordRequestForm):
             access_token, _ = await jwt.create_access_token(str(user.id), multi_login=user.is_multi_login)
             return access_token, user
 
-    async def login(self, *, request: Request, obj: Auth, background_tasks: BackgroundTasks):
+    async def login(self, *, request: Request, obj: AuthLogin, background_tasks: BackgroundTasks):
         async with async_db_session() as db:
             try:
                 current_user = await UserDao.get_by_username(db, obj.username)
@@ -50,6 +51,11 @@ async def login(self, *, request: Request, obj: Auth, background_tasks: Backgrou
                     raise errors.AuthorizationError(msg='密码错误')
                 elif not current_user.is_active:
                     raise errors.AuthorizationError(msg='用户已锁定, 登陆失败')
+                captcha_code = await redis_client.get(f'{settings.CAPTCHA_LOGIN_REDIS_PREFIX}:{request.state.ip}')
+                if not captcha_code:
+                    raise errors.AuthorizationError(msg='验证码失效，请重新获取')
+                if captcha_code.lower() != obj.captcha.lower():
+                    raise errors.CustomError(error=CustomCode.CAPTCHA_ERROR)
                 await UserDao.update_login_time(db, obj.username, self.login_time)
                 user = await UserDao.get(db, current_user.id)
                 access_token, access_token_expire_time = await jwt.create_access_token(
@@ -60,7 +66,7 @@ async def login(self, *, request: Request, obj: Auth, background_tasks: Backgrou
                 )
             except errors.NotFoundError as e:
                 raise errors.NotFoundError(msg=e.msg)
-            except errors.AuthorizationError as e:
+            except (errors.AuthorizationError, errors.CustomError) as e:
                 err_log_info = dict(
                     db=db,
                     request=request,

requirements.txt

@@ -10,7 +10,7 @@ casbin_async_sqlalchemy_adapter==1.1.0
 cryptography==41.0.0
 email-validator==1.1.3
 Faker==9.7.1
-fast-captcha==0.1.3
+fast-captcha==0.2.1
 fastapi==0.95.2
 fastapi-limiter==0.1.5
 fastapi-pagination==0.12.1