fastapi-login/main.py
2024-01-24 10:23:42 +08:00

69 lines
2.3 KiB
Python

import uvicorn
from passlib.context import CryptContext
from fastapi import FastAPI, Request
from models.User import *
from models.Token import *
from controllers.JWT import secret_key
from starlette.responses import Response
app = FastAPI()
pwd_context = CryptContext(schemes=["bcrypt"])
# 模拟数据库中的用户
users_db = {
"admin": {
"username": "admin",
"password": pwd_context.hash("admin123")
}
}
@app.get("/")
async def index(request: Request, response: Response,):
cookie = request.headers.get("Cookie")
if cookie is None:
token = await generate_token("guest")
response.set_cookie("token", token)
return {"code": 200, "status": "success"}
@app.post("/login")
async def login(request: Request, response: Response, user: UserLogin):
if user.username not in users_db:
response.status_code = 401
return {"code": 401, "status": "error", "detail": "Invalid username"}
stored_user = users_db[user.username]
if not pwd_context.verify(user.password, stored_user["password"]):
response.status_code = 401
return {"code": 401, "status": "error", "detail": "Invalid password"}
token = await generate_token(user.username)
response.set_cookie("token", token)
return {"code": 200, "status": "success"}
@app.get("/users/me")
async def get_user_profile(request: Request, response: Response):
cookie = request.headers.get("Cookie")
token = await get_token_from_cookie(cookie)
if token is None:
response.status_code = 400
return {"code": 400, "status": "error", "detail": "Token is required"}
try:
payload = jwt.decode(token, secret_key, algorithms=["HS256"])
username = payload["username"]
if username == "guest":
return {"code": 200, "status": "success", "username": "guest"}
if username not in users_db:
response.status_code = 401
return {"code": 401, "status": "error", "detail": "Invalid username"}
return {"code": 200, "status": "success", "username": username}
except jwt.DecodeError:
response.status_code = 401
return {"code": 401, "status": "error", "detail": "Invalid token"}
if __name__ == '__main__':
uvicorn.run(app='main:app', host="0.0.0.0", port=80, reload=True)