141 lines
5.0 KiB
Python
141 lines
5.0 KiB
Python
# -*- coding:utf-8 -*-
|
||
"""
|
||
@Author : xuxingchen
|
||
@Contact : xuxingchen@sinochem.com
|
||
@Desc : 登陆界面控制逻辑
|
||
"""
|
||
import hashlib
|
||
import os.path
|
||
import random
|
||
import io
|
||
import time
|
||
from typing import Optional
|
||
from starlette.responses import StreamingResponse
|
||
|
||
from pydantic import BaseModel
|
||
from fastapi import APIRouter, Query, Request
|
||
from PIL import Image, ImageDraw, ImageFont
|
||
|
||
from utils.database import get_table_handler
|
||
from utils import logger
|
||
from utils.misc import generate_captcha_text
|
||
from models.sessions import SessionsTable
|
||
|
||
router = APIRouter()
|
||
|
||
|
||
class LoginItem(BaseModel):
|
||
username: str
|
||
password: str
|
||
session_id: str
|
||
captcha: str
|
||
|
||
|
||
class LoginResp(BaseModel):
|
||
status: bool
|
||
message: str
|
||
token: Optional[str] = None
|
||
|
||
|
||
def generate_captcha_image(text, size=(120, 40), font_path='./data/mvboli.ttf', font_size=24):
|
||
"""生成验证码图片"""
|
||
image = Image.new('RGB', size, (73, 109, 137)) # 设置背景色
|
||
d = ImageDraw.Draw(image)
|
||
assert os.path.exists(font_path), "字体文件不存在"
|
||
font = ImageFont.truetype(font_path, font_size)
|
||
# d.text((5, 5), text, font=font, fill=(255, 255, 0)) # 设置字体颜色和位置
|
||
font_box = font.getbbox(text)
|
||
text_width, text_height = font_box[2] - font_box[0], font_box[3] - font_box[1]
|
||
text_x = (size[0] - text_width) // 2
|
||
text_y = (size[1] - text_height) // 2
|
||
d.text((text_x, text_y - 8), text, font=font, fill=(255, 255, 0))
|
||
|
||
# 添加噪点
|
||
for _ in range(100): # 可以根据需要调整噪点数量
|
||
x = random.randint(0, size[0] - 1)
|
||
y = random.randint(0, size[1] - 1)
|
||
image.putpixel((x, y), (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)))
|
||
|
||
# 添加线条
|
||
for _ in range(3): # 可以根据需要调整线条数量
|
||
x1, y1 = random.randint(0, size[0] - 1), random.randint(0, size[1] - 1)
|
||
x2, y2 = random.randint(0, size[0] - 1), random.randint(0, size[1] - 1)
|
||
d.line([(x1, y1), (x2, y2)], fill=(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)),
|
||
width=2)
|
||
|
||
# 将图片保存为字节流
|
||
bytes_io = io.BytesIO()
|
||
image.save(bytes_io, format='PNG')
|
||
# image.save("./data/tmp.png", format='PNG')
|
||
bytes_io.seek(0)
|
||
return bytes_io
|
||
|
||
|
||
def hash_password(password):
|
||
"""对密码进行 SHA-256 哈希加密"""
|
||
password_bytes = password.encode('utf-8')
|
||
sha256_hash = hashlib.sha256()
|
||
sha256_hash.update(password_bytes)
|
||
return sha256_hash.hexdigest()
|
||
|
||
|
||
def verify_password(password, hashed_password):
|
||
"""比对密码和哈希值"""
|
||
password_hash = hash_password(password)
|
||
return password_hash.lower() == hashed_password.lower()
|
||
|
||
|
||
def authenticate_token(token: str):
|
||
th = get_table_handler()
|
||
timeout_timestamp = str(int(time.time() * 1000 - 4 * 60 * 60 * 1000))
|
||
return SessionsTable.check_token(th, timeout_timestamp, token)
|
||
|
||
|
||
@router.get("/generateCaptcha", summary="生成验证码")
|
||
async def generate_captcha_endpoint(request: Request,
|
||
session_id: str = Query(None, description="Session ID")):
|
||
"""生成验证码图片并返回"""
|
||
if session_id:
|
||
captcha_text = generate_captcha_text()
|
||
logger.Logger.debug(f"{request.url.path} 验证码生成：{captcha_text}")
|
||
# session_id 入库
|
||
th = get_table_handler()
|
||
SessionsTable.insert(th, session_id, captcha_text)
|
||
captcha_image = generate_captcha_image(captcha_text)
|
||
# 设置HTTP响应的头部，指定内容类型为PNG图片
|
||
headers = {"Content-Type": "image/png"}
|
||
return StreamingResponse(captcha_image, headers=headers)
|
||
else:
|
||
logger.Logger.error(f"{request.url.path} 请求参数缺少 session_id - {request.client.host}")
|
||
return {"status": False, "message": "请求参数缺少 session_id"}
|
||
|
||
|
||
@router.post("/login", response_model=LoginResp, response_model_exclude_none=True, summary="登录请求")
|
||
async def login(item: LoginItem):
|
||
"""简单的账户登录逻辑"""
|
||
status = False
|
||
token = None
|
||
th = get_table_handler()
|
||
# 验证码校验
|
||
right_captcha = SessionsTable.get_captcha(th, item.session_id)
|
||
if right_captcha is not None:
|
||
if item.captcha.lower() == right_captcha.lower():
|
||
# 账户密码校验
|
||
if item.username == "admin":
|
||
if verify_password(item.password, "B12AD23C7E230E2CA365508DD16635DD3D7214BCD9BEA27457A356FD5C15F8BF"):
|
||
status = True
|
||
message = "校验通过"
|
||
token = generate_captcha_text(16)
|
||
SessionsTable.update(th, item.session_id, item.username, token)
|
||
else:
|
||
message = "密码错误"
|
||
else:
|
||
message = "账户不存在"
|
||
else:
|
||
message = "验证码错误"
|
||
else:
|
||
message = "无效 session_id"
|
||
resp = LoginResp(status=status, message=message, token=token)
|
||
logger.Logger.debug("/login " + str(resp.__dict__))
|
||
return resp
|