DreamWeaver/dreamweaver
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

wip: snapshot full local workspace state
Some checks are pending
Build and Push Docker Images / changes (push) Waiting to run
Details
Build and Push Docker Images / build-backend (push) Blocked by required conditions
Details
Build and Push Docker Images / build-frontend (push) Blocked by required conditions
Details
Build and Push Docker Images / build-admin-frontend (push) Blocked by required conditions
Details

Browse Source
This commit is contained in:
torin
2026-04-17 18:58:11 +08:00
parent fea4ef012f
commit b8d3cb4644
181 changed files with 16964 additions and 17486 deletions
Download Patch File Download Diff File Expand all files Collapse all files

120
backend/app/core/admin_auth.py
View File

@@ -1,60 +1,60 @@
import secrets
from fastapi import Depends, HTTPException, Request, status
from fastapi.security import HTTPBasic, HTTPBasicCredentials
from app.core.config import settings
from app.core.rate_limiter import (
clear_failed_attempts,
is_locked_out,
record_failed_attempt,
)
security = HTTPBasic()
MAX_ATTEMPTS = 5
LOCKOUT_SECONDS = 900 # 15分钟
def _get_client_ip(request: Request) -> str:
forwarded = request.headers.get("x-forwarded-for")
if forwarded:
return forwarded.split(",")[0].strip()
if request.client and request.client.host:
return request.client.host
return "unknown"
async def admin_guard(
request: Request,
credentials: HTTPBasicCredentials = Depends(security),
):
client_ip = _get_client_ip(request)
lockout_key = f"admin_login:{client_ip}"
# 检查是否被锁定
remaining = await is_locked_out(lockout_key, MAX_ATTEMPTS, LOCKOUT_SECONDS)
if remaining > 0:
raise HTTPException(
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
detail=f"登录尝试过多，请 {remaining} 秒后重试",
)
# 使用 secrets.compare_digest 防止时序攻击
username_ok = secrets.compare_digest(
credentials.username.encode(), settings.admin_username.encode()
)
password_ok = secrets.compare_digest(
credentials.password.encode(), settings.admin_password.encode()
)
if not (username_ok and password_ok):
await record_failed_attempt(lockout_key, MAX_ATTEMPTS, LOCKOUT_SECONDS)
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="用户名或密码错误",
)
# 登录成功，清除失败记录
await clear_failed_attempts(lockout_key)
return True
import secrets
from fastapi import Depends, HTTPException, Request, status
from fastapi.security import HTTPBasic, HTTPBasicCredentials
from app.core.config import settings
from app.core.rate_limiter import (
clear_failed_attempts,
is_locked_out,
record_failed_attempt,
)
security = HTTPBasic()
MAX_ATTEMPTS = 5
LOCKOUT_SECONDS = 900 # 15分钟
def _get_client_ip(request: Request) -> str:
forwarded = request.headers.get("x-forwarded-for")
if forwarded:
return forwarded.split(",")[0].strip()
if request.client and request.client.host:
return request.client.host
return "unknown"
async def admin_guard(
request: Request,
credentials: HTTPBasicCredentials = Depends(security),
):
client_ip = _get_client_ip(request)
lockout_key = f"admin_login:{client_ip}"
# 检查是否被锁定
remaining = await is_locked_out(lockout_key, MAX_ATTEMPTS, LOCKOUT_SECONDS)
if remaining > 0:
raise HTTPException(
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
detail=f"登录尝试过多，请 {remaining} 秒后重试",
)
# 使用 secrets.compare_digest 防止时序攻击
username_ok = secrets.compare_digest(
credentials.username.encode(), settings.admin_username.encode()
)
password_ok = secrets.compare_digest(
credentials.password.encode(), settings.admin_password.encode()
)
if not (username_ok and password_ok):
await record_failed_attempt(lockout_key, MAX_ATTEMPTS, LOCKOUT_SECONDS)
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="用户名或密码错误",
)
# 登录成功，清除失败记录
await clear_failed_attempts(lockout_key)
return True

96
backend/app/core/logging.py
View File

@@ -1,48 +1,48 @@
"""结构化日志配置。"""
import logging
import sys
import structlog
from app.core.config import settings
def setup_logging():
"""配置 structlog 结构化日志。"""
shared_processors = [
structlog.contextvars.merge_contextvars,
structlog.stdlib.add_log_level,
structlog.stdlib.add_logger_name,
structlog.processors.TimeStamper(fmt="iso"),
structlog.processors.StackInfoRenderer(),
]
if settings.debug:
processors = shared_processors + [
structlog.dev.ConsoleRenderer(colors=True),
]
else:
processors = shared_processors + [
structlog.processors.format_exc_info,
structlog.processors.JSONRenderer(),
]
structlog.configure(
processors=processors,
wrapper_class=structlog.stdlib.BoundLogger,
context_class=dict,
logger_factory=structlog.stdlib.LoggerFactory(),
cache_logger_on_first_use=True,
)
logging.basicConfig(
format="%(message)s",
stream=sys.stdout,
level=logging.DEBUG if settings.debug else logging.INFO,
)
def get_logger(name: str) -> structlog.stdlib.BoundLogger:
"""获取结构化日志器。"""
return structlog.get_logger(name)
"""结构化日志配置。"""
import logging
import sys
import structlog
from app.core.config import settings
def setup_logging():
"""配置 structlog 结构化日志。"""
shared_processors = [
structlog.contextvars.merge_contextvars,
structlog.stdlib.add_log_level,
structlog.stdlib.add_logger_name,
structlog.processors.TimeStamper(fmt="iso"),
structlog.processors.StackInfoRenderer(),
]
if settings.debug:
processors = shared_processors + [
structlog.dev.ConsoleRenderer(colors=True),
]
else:
processors = shared_processors + [
structlog.processors.format_exc_info,
structlog.processors.JSONRenderer(),
]
structlog.configure(
processors=processors,
wrapper_class=structlog.stdlib.BoundLogger,
context_class=dict,
logger_factory=structlog.stdlib.LoggerFactory(),
cache_logger_on_first_use=True,
)
logging.basicConfig(
format="%(message)s",
stream=sys.stdout,
level=logging.DEBUG if settings.debug else logging.INFO,
)
def get_logger(name: str) -> structlog.stdlib.BoundLogger:
"""获取结构化日志器。"""
return structlog.get_logger(name)

380
backend/app/core/prompts.py
View File

@@ -1,190 +1,190 @@
# ruff: noqa: E501
"""AI 提示词模板 (Modernized)"""
# 随机元素列表：为故事注入不可预测的魔法
RANDOM_ELEMENTS = [
"一个会打喷嚏的云朵",
"一本地图上找不到的神秘图书馆",
"一只能实现小愿望的彩色蜗牛",
"一扇通往颠倒世界的门",
"一顶能听懂动物说话的旧帽子",
"一个装着星星的玻璃罐",
"一棵结满笑声果实的树",
"一只能在水上画画的画笔",
"一个怕黑的影子",
"一只收集回声的瓶子",
"一双会自己跳舞的红鞋子",
"一个只能在月光下看见的邮筒",
"一张会改变模样的全家福",
"一把可以打开梦境的钥匙",
"一个喜欢讲冷笑话的冰箱",
"一条通往星期八的秘密小径"
]
# ==============================================================================
# Model A: 故事生成 (Story Generation)
# ==============================================================================
SYSTEM_INSTRUCTION_STORYTELLER = """
# Role
You are "**Dream Weaver**", a world-class children's storyteller with the imagination of Pixar and the warmth of Miyazaki.
Your mission is to create engaging, safe, and educational stories for children (ages 3-8).
# Core Philosophy
1. **Show, Don't Tell**: Don't preach the lesson. Let the character's actions and the plot demonstrate the theme.
2. **Safety First**: No violence, horror, or scary elements. Conflict should be emotional or situational, not physical.
3. **Vivid Imagery**: Use sensory details (colors, sounds, smells) that appeal to children.
4. **Empowerment**: The child protagonist should solve the problem using wit, kindness, or courage, not just luck.
# Continuity & Memory (CRITICAL)
- **Universal Context**: The story takes place in the child's established "Story Universe". Respect existing world rules.
- **Character Consistency**: If "Child Profile" or "Sidekicks" are provided, you MUST use their specific names and traits. Do NOT invent new main characters unless asked.
- **Callback**: If "Past Memories" are provided, try to make a natural, one-sentence reference to a past adventure to build a sense of continuity (e.g., "Just like when we found the lost star...").
# Output Format
You MUST return a pure JSON object with NO markdown formatting (no ```json code blocks).
The JSON object must have the following schema:
{
"mode": "generated",
"title": "A catchy, imaginative title",
"story_text": "The full story text. Use \\n\\n for paragraph breaks.",
"cover_prompt_suggestion": "A detailed English image generation prompt for the story cover. Style: whimsical, children's book illustration, soft lighting, vibrant colors."
}
"""
USER_PROMPT_GENERATION = """
# Task: Write a Children's Story
## Contextual Memory (Use these if provided)
{memory_context}
## Inputs
- **Keywords/Topic**: {keywords}
- **Educational Theme**: {education_theme}
- **Magic Element (Must Incorporate)**: {random_element}
## Constraints
- Length: 300-600 words.
- Structure: Beginning (Hook) -> Middle (Challenge) -> End (Resolution & Growth).
"""
# ==============================================================================
# Model B: 故事润色 (Story Enhancement)
# ==============================================================================
SYSTEM_INSTRUCTION_ENHANCER = """
# Role
You are "**Dream Weaver Editor**", an expert children's book editor who turns rough drafts into polished gems.
# Mission
Analyze the user's input story and rewrite it to be:
1. **More Engaging**: Enhance the plot with a "Magic Element" to add surprise.
2. **More Educational**: Weave the "Educational Theme" deeper into the narrative arc.
3. **Better Written**: Polish the sentences for rhythm and flow (suitable for reading aloud).
4. **Safe**: Remove any inappropriate content (violence, scary interaction) and replace it with constructive solutions.
# Output Format
You MUST return a pure JSON object with NO markdown formatting (no ```json code blocks).
The JSON object must have the following schema:
{
"mode": "enhanced",
"title": "An improved title (or the original if perfect)",
"story_text": "The rewritten story text. Use \\n\\n for paragraph breaks.",
"cover_prompt_suggestion": "A detailed English image generation prompt for the cover."
}
"""
USER_PROMPT_ENHANCEMENT = """
# Task: Enhance This Story
## Contextual Memory
{memory_context}
## Inputs
- **Original Story**: {full_story}
- **Target Theme**: {education_theme}
- **Magic Element to Add**: {random_element}
## Constraints
- Length: 300-600 words.
- Keep the original character names if possible, but feel free to upgrade the plot.
"""
# ==============================================================================
# Model C: 成就提取 (Achievement Extraction)
# ==============================================================================
# 保持简单，暂不使用 System Instruction，沿用单次提示
ACHIEVEMENT_EXTRACTION_PROMPT = """
Analyze the story and extract key growth moments or achievements for the child protagonist.
# Story
{story_text}
# Target Categories (Examples)
- **Courage**: Overcoming fear, trying something new.
- **Kindness**: Helping others, sharing, empathy.
- **Curiosity**: Asking questions, exploring, learning.
- **Resilience**: Not giving up, handling failure.
- **Wisdom**: Problem-solving, honesty, patience.
# Output Format
Return a pure JSON object (no markdown):
{{
"achievements": [
{{
"type": "Category Name",
"description": "Brief reason (max 10 words)",
"score": 8 // 1-10 intensity
}}
]
}}
"""
# ==============================================================================
# Model D: 绘本生成 (Storybook Generation)
# ==============================================================================
SYSTEM_INSTRUCTION_STORYBOOK = """
# Role
You are "**Dream Weaver Illustrator**", a creative children's book author and visual director.
Your mission is to create a paginated picture book for children (ages 3-8), where each page has text and a matching illustration prompt.
# Core Philosophy
1. **Pacing**: The story must flow logically across the specified number of pages.
2. **Visual Consistency**: Define the "Main Character" and "Art Style" once, and ensure all image prompts adhere to them.
3. **Language**: The story text MUST be in **Chinese (Simplified)**. The image prompts MUST be in **English**.
4. **Memory**: If a memory context is provided, incorporate known characters or references naturally.
# Output Format
You MUST return a pure JSON object using the following schema (no markdown):
{
"title": "Story Title (Chinese)",
"main_character": "Description of the protagonist (e.g., 'A small blue robot with rusty gears')",
"art_style": "Visual style description (e.g., 'Watercolor, soft pastel colors, whimsical')",
"pages": [
{
"page_number": 1,
"text": "Page text in Chinese (30-60 chars).",
"image_prompt": "Detailed English image prompt describing the scene. Include 'main_character' reference."
}
],
"cover_prompt": "English image prompt for the book cover."
}
"""
USER_PROMPT_STORYBOOK = """
# Task: Create a {page_count}-Page Storybook
## Contextual Memory
{memory_context}
## Inputs
- **Keywords/Topic**: {keywords}
- **Educational Theme**: {education_theme}
- **Magic Element**: {random_element}
## Constraints
- Pages: Exactly {page_count} pages.
- Structure: Intro -> Development -> Climax -> Resolution.
"""
# ruff: noqa: E501
"""AI 提示词模板 (Modernized)"""
# 随机元素列表：为故事注入不可预测的魔法
RANDOM_ELEMENTS = [
"一个会打喷嚏的云朵",
"一本地图上找不到的神秘图书馆",
"一只能实现小愿望的彩色蜗牛",
"一扇通往颠倒世界的门",
"一顶能听懂动物说话的旧帽子",
"一个装着星星的玻璃罐",
"一棵结满笑声果实的树",
"一只能在水上画画的画笔",
"一个怕黑的影子",
"一只收集回声的瓶子",
"一双会自己跳舞的红鞋子",
"一个只能在月光下看见的邮筒",
"一张会改变模样的全家福",
"一把可以打开梦境的钥匙",
"一个喜欢讲冷笑话的冰箱",
"一条通往星期八的秘密小径"
]
# ==============================================================================
# Model A: 故事生成 (Story Generation)
# ==============================================================================
SYSTEM_INSTRUCTION_STORYTELLER = """
# Role
You are "**Dream Weaver**", a world-class children's storyteller with the imagination of Pixar and the warmth of Miyazaki.
Your mission is to create engaging, safe, and educational stories for children (ages 3-8).
# Core Philosophy
1. **Show, Don't Tell**: Don't preach the lesson. Let the character's actions and the plot demonstrate the theme.
2. **Safety First**: No violence, horror, or scary elements. Conflict should be emotional or situational, not physical.
3. **Vivid Imagery**: Use sensory details (colors, sounds, smells) that appeal to children.
4. **Empowerment**: The child protagonist should solve the problem using wit, kindness, or courage, not just luck.
# Continuity & Memory (CRITICAL)
- **Universal Context**: The story takes place in the child's established "Story Universe". Respect existing world rules.
- **Character Consistency**: If "Child Profile" or "Sidekicks" are provided, you MUST use their specific names and traits. Do NOT invent new main characters unless asked.
- **Callback**: If "Past Memories" are provided, try to make a natural, one-sentence reference to a past adventure to build a sense of continuity (e.g., "Just like when we found the lost star...").
# Output Format
You MUST return a pure JSON object with NO markdown formatting (no ```json code blocks).
The JSON object must have the following schema:
{
"mode": "generated",
"title": "A catchy, imaginative title",
"story_text": "The full story text. Use \\n\\n for paragraph breaks.",
"cover_prompt_suggestion": "A detailed English image generation prompt for the story cover. Style: whimsical, children's book illustration, soft lighting, vibrant colors."
}
"""
USER_PROMPT_GENERATION = """
# Task: Write a Children's Story
## Contextual Memory (Use these if provided)
{memory_context}
## Inputs
- **Keywords/Topic**: {keywords}
- **Educational Theme**: {education_theme}
- **Magic Element (Must Incorporate)**: {random_element}
## Constraints
- Length: 300-600 words.
- Structure: Beginning (Hook) -> Middle (Challenge) -> End (Resolution & Growth).
"""
# ==============================================================================
# Model B: 故事润色 (Story Enhancement)
# ==============================================================================
SYSTEM_INSTRUCTION_ENHANCER = """
# Role
You are "**Dream Weaver Editor**", an expert children's book editor who turns rough drafts into polished gems.
# Mission
Analyze the user's input story and rewrite it to be:
1. **More Engaging**: Enhance the plot with a "Magic Element" to add surprise.
2. **More Educational**: Weave the "Educational Theme" deeper into the narrative arc.
3. **Better Written**: Polish the sentences for rhythm and flow (suitable for reading aloud).
4. **Safe**: Remove any inappropriate content (violence, scary interaction) and replace it with constructive solutions.
# Output Format
You MUST return a pure JSON object with NO markdown formatting (no ```json code blocks).
The JSON object must have the following schema:
{
"mode": "enhanced",
"title": "An improved title (or the original if perfect)",
"story_text": "The rewritten story text. Use \\n\\n for paragraph breaks.",
"cover_prompt_suggestion": "A detailed English image generation prompt for the cover."
}
"""
USER_PROMPT_ENHANCEMENT = """
# Task: Enhance This Story
## Contextual Memory
{memory_context}
## Inputs
- **Original Story**: {full_story}
- **Target Theme**: {education_theme}
- **Magic Element to Add**: {random_element}
## Constraints
- Length: 300-600 words.
- Keep the original character names if possible, but feel free to upgrade the plot.
"""
# ==============================================================================
# Model C: 成就提取 (Achievement Extraction)
# ==============================================================================
# 保持简单，暂不使用 System Instruction，沿用单次提示
ACHIEVEMENT_EXTRACTION_PROMPT = """
Analyze the story and extract key growth moments or achievements for the child protagonist.
# Story
{story_text}
# Target Categories (Examples)
- **Courage**: Overcoming fear, trying something new.
- **Kindness**: Helping others, sharing, empathy.
- **Curiosity**: Asking questions, exploring, learning.
- **Resilience**: Not giving up, handling failure.
- **Wisdom**: Problem-solving, honesty, patience.
# Output Format
Return a pure JSON object (no markdown):
{{
"achievements": [
{{
"type": "Category Name",
"description": "Brief reason (max 10 words)",
"score": 8 // 1-10 intensity
}}
]
}}
"""
# ==============================================================================
# Model D: 绘本生成 (Storybook Generation)
# ==============================================================================
SYSTEM_INSTRUCTION_STORYBOOK = """
# Role
You are "**Dream Weaver Illustrator**", a creative children's book author and visual director.
Your mission is to create a paginated picture book for children (ages 3-8), where each page has text and a matching illustration prompt.
# Core Philosophy
1. **Pacing**: The story must flow logically across the specified number of pages.
2. **Visual Consistency**: Define the "Main Character" and "Art Style" once, and ensure all image prompts adhere to them.
3. **Language**: The story text MUST be in **Chinese (Simplified)**. The image prompts MUST be in **English**.
4. **Memory**: If a memory context is provided, incorporate known characters or references naturally.
# Output Format
You MUST return a pure JSON object using the following schema (no markdown):
{
"title": "Story Title (Chinese)",
"main_character": "Description of the protagonist (e.g., 'A small blue robot with rusty gears')",
"art_style": "Visual style description (e.g., 'Watercolor, soft pastel colors, whimsical')",
"pages": [
{
"page_number": 1,
"text": "Page text in Chinese (30-60 chars).",
"image_prompt": "Detailed English image prompt describing the scene. Include 'main_character' reference."
}
],
"cover_prompt": "English image prompt for the book cover."
}
"""
USER_PROMPT_STORYBOOK = """
# Task: Create a {page_count}-Page Storybook
## Contextual Memory
{memory_context}
## Inputs
- **Keywords/Topic**: {keywords}
- **Educational Theme**: {education_theme}
- **Magic Element**: {random_element}
## Constraints
- Pages: Exactly {page_count} pages.
- Structure: Intro -> Development -> Climax -> Resolution.
"""

282
backend/app/core/rate_limiter.py
View File

@@ -1,141 +1,141 @@
"""Redis-backed rate limiter with in-memory fallback.
Uses a fixed-window counter pattern via Redis INCR + EXPIRE.
Falls back to an in-memory TTLCache when Redis is unavailable,
preserving identical behavior for dev/test environments.
"""
import time
from cachetools import TTLCache
from fastapi import HTTPException
from app.core.logging import get_logger
from app.core.redis import get_redis
logger = get_logger(__name__)
# ── In-memory fallback caches ──────────────────────────────────────────────
_local_rate_cache: TTLCache[str, int] = TTLCache(maxsize=10000, ttl=120)
_local_lockout_cache: TTLCache[str, tuple[int, float]] = TTLCache(maxsize=1000, ttl=900)
async def check_rate_limit(key: str, limit: int, window_seconds: int) -> None:
"""Check and increment a sliding-window rate counter.
Args:
key: Unique identifier (e.g. ``"story:<user_id>"``).
limit: Maximum requests allowed within the window.
window_seconds: Window duration in seconds.
Raises:
HTTPException: 429 when the limit is exceeded.
"""
try:
redis = await get_redis()
# Fixed-window bucket: key + minute boundary
bucket = int(time.time() // window_seconds)
redis_key = f"ratelimit:{key}:{bucket}"
count = await redis.incr(redis_key)
if count == 1:
await redis.expire(redis_key, window_seconds)
if count > limit:
raise HTTPException(
status_code=429,
detail="Too many requests, please slow down.",
)
return
except HTTPException:
raise
except Exception as exc:
logger.warning("rate_limit_redis_fallback", error=str(exc))
# ── Fallback: in-memory counter ────────────────────────────────────────
count = _local_rate_cache.get(key, 0) + 1
_local_rate_cache[key] = count
if count > limit:
raise HTTPException(
status_code=429,
detail="Too many requests, please slow down.",
)
async def record_failed_attempt(
key: str,
max_attempts: int,
lockout_seconds: int,
) -> bool:
"""Record a failed login attempt and return whether the key is locked out.
Args:
key: Unique identifier (e.g. ``"admin_login:<ip>"``).
max_attempts: Number of failures before lockout.
lockout_seconds: Duration of lockout in seconds.
Returns:
``True`` if the key is now locked out, ``False`` otherwise.
"""
try:
redis = await get_redis()
redis_key = f"lockout:{key}"
count = await redis.incr(redis_key)
if count == 1:
await redis.expire(redis_key, lockout_seconds)
return count >= max_attempts
except Exception as exc:
logger.warning("lockout_redis_fallback", error=str(exc))
# ── Fallback ───────────────────────────────────────────────────────────
if key in _local_lockout_cache:
attempts, first_fail = _local_lockout_cache[key]
_local_lockout_cache[key] = (attempts + 1, first_fail)
return (attempts + 1) >= max_attempts
else:
_local_lockout_cache[key] = (1, time.time())
return 1 >= max_attempts
async def is_locked_out(key: str, max_attempts: int, lockout_seconds: int) -> int:
"""Check if a key is currently locked out.
Returns:
Remaining lockout seconds (> 0 means locked), 0 means not locked.
"""
try:
redis = await get_redis()
redis_key = f"lockout:{key}"
count = await redis.get(redis_key)
if count is not None and int(count) >= max_attempts:
ttl = await redis.ttl(redis_key)
return max(ttl, 0)
return 0
except Exception as exc:
logger.warning("lockout_check_redis_fallback", error=str(exc))
# ── Fallback ───────────────────────────────────────────────────────────
if key in _local_lockout_cache:
attempts, first_fail = _local_lockout_cache[key]
if attempts >= max_attempts:
remaining = int(lockout_seconds - (time.time() - first_fail))
if remaining > 0:
return remaining
else:
del _local_lockout_cache[key]
return 0
async def clear_failed_attempts(key: str) -> None:
"""Clear lockout state on successful login."""
try:
redis = await get_redis()
await redis.delete(f"lockout:{key}")
except Exception as exc:
logger.warning("lockout_clear_redis_fallback", error=str(exc))
# Always clear local cache too
_local_lockout_cache.pop(key, None)
"""Redis-backed rate limiter with in-memory fallback.
Uses a fixed-window counter pattern via Redis INCR + EXPIRE.
Falls back to an in-memory TTLCache when Redis is unavailable,
preserving identical behavior for dev/test environments.
"""
import time
from cachetools import TTLCache
from fastapi import HTTPException
from app.core.logging import get_logger
from app.core.redis import get_redis
logger = get_logger(__name__)
# ── In-memory fallback caches ──────────────────────────────────────────────
_local_rate_cache: TTLCache[str, int] = TTLCache(maxsize=10000, ttl=120)
_local_lockout_cache: TTLCache[str, tuple[int, float]] = TTLCache(maxsize=1000, ttl=900)
async def check_rate_limit(key: str, limit: int, window_seconds: int) -> None:
"""Check and increment a sliding-window rate counter.
Args:
key: Unique identifier (e.g. ``"story:<user_id>"``).
limit: Maximum requests allowed within the window.
window_seconds: Window duration in seconds.
Raises:
HTTPException: 429 when the limit is exceeded.
"""
try:
redis = await get_redis()
# Fixed-window bucket: key + minute boundary
bucket = int(time.time() // window_seconds)
redis_key = f"ratelimit:{key}:{bucket}"
count = await redis.incr(redis_key)
if count == 1:
await redis.expire(redis_key, window_seconds)
if count > limit:
raise HTTPException(
status_code=429,
detail="Too many requests, please slow down.",
)
return
except HTTPException:
raise
except Exception as exc:
logger.warning("rate_limit_redis_fallback", error=str(exc))
# ── Fallback: in-memory counter ────────────────────────────────────────
count = _local_rate_cache.get(key, 0) + 1
_local_rate_cache[key] = count
if count > limit:
raise HTTPException(
status_code=429,
detail="Too many requests, please slow down.",
)
async def record_failed_attempt(
key: str,
max_attempts: int,
lockout_seconds: int,
) -> bool:
"""Record a failed login attempt and return whether the key is locked out.
Args:
key: Unique identifier (e.g. ``"admin_login:<ip>"``).
max_attempts: Number of failures before lockout.
lockout_seconds: Duration of lockout in seconds.
Returns:
``True`` if the key is now locked out, ``False`` otherwise.
"""
try:
redis = await get_redis()
redis_key = f"lockout:{key}"
count = await redis.incr(redis_key)
if count == 1:
await redis.expire(redis_key, lockout_seconds)
return count >= max_attempts
except Exception as exc:
logger.warning("lockout_redis_fallback", error=str(exc))
# ── Fallback ───────────────────────────────────────────────────────────
if key in _local_lockout_cache:
attempts, first_fail = _local_lockout_cache[key]
_local_lockout_cache[key] = (attempts + 1, first_fail)
return (attempts + 1) >= max_attempts
else:
_local_lockout_cache[key] = (1, time.time())
return 1 >= max_attempts
async def is_locked_out(key: str, max_attempts: int, lockout_seconds: int) -> int:
"""Check if a key is currently locked out.
Returns:
Remaining lockout seconds (> 0 means locked), 0 means not locked.
"""
try:
redis = await get_redis()
redis_key = f"lockout:{key}"
count = await redis.get(redis_key)
if count is not None and int(count) >= max_attempts:
ttl = await redis.ttl(redis_key)
return max(ttl, 0)
return 0
except Exception as exc:
logger.warning("lockout_check_redis_fallback", error=str(exc))
# ── Fallback ───────────────────────────────────────────────────────────
if key in _local_lockout_cache:
attempts, first_fail = _local_lockout_cache[key]
if attempts >= max_attempts:
remaining = int(lockout_seconds - (time.time() - first_fail))
if remaining > 0:
return remaining
else:
del _local_lockout_cache[key]
return 0
async def clear_failed_attempts(key: str) -> None:
"""Clear lockout state on successful login."""
try:
redis = await get_redis()
await redis.delete(f"lockout:{key}")
except Exception as exc:
logger.warning("lockout_clear_redis_fallback", error=str(exc))
# Always clear local cache too
_local_lockout_cache.pop(key, None)

50
backend/app/core/security.py
View File

@@ -1,25 +1,25 @@
from datetime import datetime, timedelta, timezone
from jose import JWTError, jwt
from app.core.config import settings
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_DAYS = 7
def create_access_token(data: dict) -> str:
"""创建 JWT token"""
to_encode = data.copy()
expire = datetime.now(timezone.utc) + timedelta(days=ACCESS_TOKEN_EXPIRE_DAYS)
to_encode.update({"exp": expire})
return jwt.encode(to_encode, settings.secret_key, algorithm=ALGORITHM)
def decode_access_token(token: str) -> dict | None:
"""解码 JWT token"""
try:
payload = jwt.decode(token, settings.secret_key, algorithms=[ALGORITHM])
return payload
except JWTError:
return None
from datetime import datetime, timedelta, timezone
from jose import JWTError, jwt
from app.core.config import settings
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_DAYS = 7
def create_access_token(data: dict) -> str:
"""创建 JWT token"""
to_encode = data.copy()
expire = datetime.now(timezone.utc) + timedelta(days=ACCESS_TOKEN_EXPIRE_DAYS)
to_encode.update({"exp": expire})
return jwt.encode(to_encode, settings.secret_key, algorithm=ALGORITHM)
def decode_access_token(token: str) -> dict | None:
"""解码 JWT token"""
try:
payload = jwt.decode(token, settings.secret_key, algorithms=[ALGORITHM])
return payload
except JWTError:
return None