import os
from datetime import timedelta
import uvicorn
from fastapi import FastAPI, Depends, Request, status
from fastapi.responses import HTMLResponse, Response, RedirectResponse
from fastapi.security import OAuth2PasswordRequestForm
from fastapi_login import LoginManager
from fastapi_login.exceptions import InvalidCredentialsException
from loguru import logger as log
app = FastAPI()
# 处理登录错误的异常
class NotAuthenticatedException(Exception):
pass
# noinspection PyUnusedLocal
def exc_handler(request, exc):
"""未授权时重定向到登录页面,上面的两个参数为必需项"""
return RedirectResponse(url='/login')
# You also have to add an exception handler to your app instance
app.add_exception_handler(NotAuthenticatedException, exc_handler)
SECRET = os.urandom(24).hex()
manager = LoginManager(
secret=SECRET,
token_url='/auth/token',
use_cookie=True,
custom_exception=NotAuthenticatedException
)
fake_db = {'lzwang': {'password': '123456'}}
@manager.user_loader()
async def load_user(user_id: str):
user = fake_db.get(user_id)
return user
# noinspection PyUnusedLocal
@app.get("/login", response_class=HTMLResponse)
async def login_page(request: Request):
"""简易登录页面"""
return """
<html>
<head>
<title>登录页面</title>
<link href="/static/css/style.css" rel="stylesheet">
</head>
<body>
<h1>用户登录</h1>
<form method="post" action="/auth/token">
<label for="username">用户名:</label><br>
<input type="text" id="username" name="username"><br>
<label for="password">密码:</label><br>
<input type="password" id="password" name="password"><br><br>
<input type="submit" value="登录">
</form>
</body>
</html>
"""
# noinspection PyUnusedLocal
@app.post('/auth/token')
async def login(response: Response, input_data: OAuth2PasswordRequestForm = Depends()):
# 获取用户输入凭据
log.info(f'Login: User Input = {input_data.username}, {input_data.password}')
# 获取数据库正确的凭据
db_data = await load_user(input_data.username)
log.info(f'Login: Database user_data = {db_data}')
# 凭据校验
if not db_data:
raise InvalidCredentialsException
elif input_data.password != db_data['password']:
raise InvalidCredentialsException
# 创建 token ,配置自动过期时间
access_token = manager.create_access_token(
data={"sub": input_data.username}, expires=timedelta(seconds=15)
)
log.success(f"Login SUCCESS, token: {access_token}")
# 将此 token 设置到 response 上,重定向到首页
response = RedirectResponse(url="/", status_code=status.HTTP_303_SEE_OTHER)
manager.set_cookie(response, access_token)
return response
# noinspection PyUnusedLocal
@app.get('/')
async def must_auth_page(user=Depends(manager)):
"""测试必须登录的页面"""
return f"Visit Page with AUTH OK, Login SUCCESS"
@app.get('/{item}')
async def no_need_auth_page(item: int):
"""测试无需登录的页面"""
return f"Visit Page without AUTH OK, item={item}"
if __name__ == "__main__":
uvicorn.run(
app=app,
host='127.0.0.1',
port=7777,
log_level='debug'
)