Add graphic captcha before user SMS login

ai_knowledge_base_v2/apps/backend/app/services/auth_service.py

@@ -8,6 +8,7 @@ from sqlalchemy.orm import Session
 
 from app.core.security import create_access_token
 from app.models.user import User
+from app.services.captcha_service import CaptchaService
 from app.services.sms_code_service import SmsCodeService
 
 
@@ -15,7 +16,8 @@ class AuthService:
     _revoked_token_jti: set[str] = set()
 
     @classmethod
-    def send_sms_code(cls, db: Session, phone: str) -> None:
+    def send_sms_code(cls, db: Session, phone: str, captcha_id: str, captcha_code: str) -> None:
+        CaptchaService.verify(captcha_id, captcha_code)
         user = cls._get_existing_user(db, phone)
         cls._ensure_user_can_login(user)
         SmsCodeService.send_code(db, phone)

ai_knowledge_base_v2/apps/backend/app/services/captcha_service.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+import base64
+from dataclasses import dataclass
+from datetime import UTC, datetime, timedelta
+from io import BytesIO
+import logging
+import random
+import string
+from uuid import uuid4
+
+from fastapi import HTTPException, status
+from PIL import Image, ImageDraw, ImageFilter, ImageFont
+
+logger = logging.getLogger(__name__)
+
+CAPTCHA_EXPIRES_SECONDS = 300
+CAPTCHA_WIDTH = 168
+CAPTCHA_HEIGHT = 64
+CAPTCHA_LENGTH = 4
+
+
+@dataclass
+class CaptchaRecord:
+    code: str
+    expires_at: datetime
+
+
+class CaptchaService:
+    _records: dict[str, CaptchaRecord] = {}
+
+    @classmethod
+    def create(cls) -> dict:
+        cls._cleanup_expired()
+        captcha_id = str(uuid4())
+        code = _random_code()
+        cls._records[captcha_id] = CaptchaRecord(
+            code=code,
+            expires_at=datetime.now(UTC) + timedelta(seconds=CAPTCHA_EXPIRES_SECONDS),
+        )
+        return {
+            "captchaId": captcha_id,
+            "imageBase64": _render_captcha_image(code),
+            "expiresInSeconds": CAPTCHA_EXPIRES_SECONDS,
+        }
+
+    @classmethod
+    def verify(cls, captcha_id: str, captcha_code: str) -> None:
+        record = cls._records.get(captcha_id)
+        if record is None:
+            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="图形验证码不存在或已过期")
+        if record.expires_at < datetime.now(UTC):
+            cls._records.pop(captcha_id, None)
+            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="图形验证码已过期")
+        if record.code.lower() != captcha_code.strip().lower():
+            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="图形验证码错误")
+        cls._records.pop(captcha_id, None)
+
+    @classmethod
+    def _cleanup_expired(cls) -> None:
+        now = datetime.now(UTC)
+        expired_ids = [captcha_id for captcha_id, record in cls._records.items() if record.expires_at < now]
+        for captcha_id in expired_ids:
+            cls._records.pop(captcha_id, None)
+
+
+def _random_code() -> str:
+    alphabet = "".join(ch for ch in string.ascii_uppercase + string.digits if ch not in "0O1IL")
+    return "".join(random.SystemRandom().choice(alphabet) for _ in range(CAPTCHA_LENGTH))
+
+
+def _render_captcha_image(code: str) -> str:
+    image = Image.new("RGB", (CAPTCHA_WIDTH, CAPTCHA_HEIGHT), "#f7fbfa")
+    draw = ImageDraw.Draw(image)
+    rng = random.SystemRandom()
+
+    for _ in range(70):
+        x = rng.randint(0, CAPTCHA_WIDTH - 1)
+        y = rng.randint(0, CAPTCHA_HEIGHT - 1)
+        color = rng.choice(["#d7e5df", "#c7ddd5", "#e4ece9"])
+        draw.point((x, y), fill=color)
+
+    for _ in range(4):
+        start = (rng.randint(0, CAPTCHA_WIDTH // 2), rng.randint(0, CAPTCHA_HEIGHT))
+        end = (rng.randint(CAPTCHA_WIDTH // 2, CAPTCHA_WIDTH), rng.randint(0, CAPTCHA_HEIGHT))
+        draw.line([start, end], fill=rng.choice(["#b7cec5", "#cbdad5"]), width=1)
+
+    font = _load_font()
+    char_width = CAPTCHA_WIDTH // CAPTCHA_LENGTH
+    for index, char in enumerate(code):
+        layer = Image.new("RGBA", (char_width, CAPTCHA_HEIGHT), (255, 255, 255, 0))
+        layer_draw = ImageDraw.Draw(layer)
+        bbox = layer_draw.textbbox((0, 0), char, font=font)
+        text_width = bbox[2] - bbox[0]
+        text_height = bbox[3] - bbox[1]
+        x = (char_width - text_width) // 2
+        y = (CAPTCHA_HEIGHT - text_height) // 2 - 3
+        layer_draw.text((x, y), char, font=font, fill=rng.choice(["#0f735d", "#173f35", "#31564d"]))
+        layer = layer.rotate(rng.randint(-14, 14), resample=Image.Resampling.BICUBIC, expand=False)
+        image.paste(layer, (index * char_width, 0), layer)
+
+    image = image.filter(ImageFilter.SMOOTH)
+    output = BytesIO()
+    image.save(output, format="PNG")
+    encoded = base64.b64encode(output.getvalue()).decode("ascii")
+    return f"data:image/png;base64,{encoded}"
+
+
+def _load_font():
+    for path in (
+        "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
+        "/usr/local/share/fonts/DejaVuSans-Bold.ttf",
+        "/System/Library/Fonts/Supplemental/Arial Bold.ttf",
+        "/Library/Fonts/Arial Bold.ttf",
+    ):
+        try:
+            return ImageFont.truetype(path, 38)
+        except OSError:
+            continue
+    logger.warning("Captcha font fallback to Pillow default font")
+    return ImageFont.load_default()