SCLP

SCLP/routers/login.py

@@ -0,0 +1,140 @@
+# -*- 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