jude-cehua/utils.py

import re
import time
import hashlib
from typing import Dict, List, Optional, Tuple
from functools import wraps
from flask import request, jsonify
from http import HTTPStatus
from datetime import datetime, timezone, timedelta

def get_beijing_time() -> datetime:
    """
    获取北京时间（UTC+8）
    
    Returns:
        北京时间的datetime对象
    """
    beijing_tz = timezone(timedelta(hours=8))
    return datetime.now(beijing_tz)

def sanitize_input(text: str, max_length: int = 20000) -> Tuple[str, Optional[str]]:
    """
    清理和验证用户输入
    
    Args:
        text: 输入文本
        max_length: 最大长度限制
        
    Returns:
        (清理后的文本, 错误信息)
    """
    if not text or not text.strip():
        return "", "消息不能为空"
    
    # 移除首尾空白
    text = text.strip()
    
    # 检查长度
    if len(text) > max_length:
        return "", f"消息长度不能超过{max_length}字符"
    
    # 移除潜在的恶意脚本标签
    text = re.sub(r'<script[^>]*>.*?</script>', '', text, flags=re.IGNORECASE | re.DOTALL)
    text = re.sub(r'<iframe[^>]*>.*?</iframe>', '', text, flags=re.IGNORECASE | re.DOTALL)
    
    # 限制连续换行
    text = re.sub(r'\n{3,}', '\n\n', text)
    
    return text, None

def validate_user_id(user_id: str) -> Tuple[str, Optional[str]]:
    """
    验证用户ID
    
    Args:
        user_id: 用户ID
        
    Returns:
        (验证后的用户ID, 错误信息)
    """
    if not user_id:
        return "", "用户ID不能为空"
    
    # 检查用户ID格式（UUID格式）
    uuid_pattern = re.compile(
        r'^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$',
        re.IGNORECASE
    )
    
    if not uuid_pattern.match(user_id):
        return "", "用户ID格式无效"
    
    return user_id, None

class RateLimiter:
    """简单的内存速率限制器"""
    
    def __init__(self, max_requests: int, window_seconds: int):
        self.max_requests = max_requests
        self.window_seconds = window_seconds
        self.requests: Dict[str, List[float]] = {}
    
    def is_allowed(self, identifier: str) -> bool:
        """
        检查请求是否被允许
        
        Args:
            identifier: 请求标识符（如IP地址或用户ID）
            
        Returns:
            是否允许请求
        """
        now = time.time()
        
        # 清理过期的请求记录
        if identifier in self.requests:
            self.requests[identifier] = [
                req_time for req_time in self.requests[identifier]
                if now - req_time < self.window_seconds
            ]
        else:
            self.requests[identifier] = []
        
        # 检查是否超过限制
        if len(self.requests[identifier]) >= self.max_requests:
            return False
        
        # 记录新请求
        self.requests[identifier].append(now)
        return True
    
    def get_remaining(self, identifier: str) -> int:
        """获取剩余请求次数"""
        now = time.time()
        if identifier in self.requests:
            valid_requests = [
                req_time for req_time in self.requests[identifier]
                if now - req_time < self.window_seconds
            ]
            return max(0, self.max_requests - len(valid_requests))
        return self.max_requests

def rate_limit(max_requests: int, window_seconds: int, identifier_func=None):
    """
    速率限制装饰器
    
    Args:
        max_requests: 最大请求数
        window_seconds: 时间窗口（秒）
        identifier_func: 获取标识符的函数，默认为IP地址
    """
    limiter = RateLimiter(max_requests, window_seconds)
    
    def decorator(f):
        @wraps(f)
        def decorated_function(*args, **kwargs):
            # 获取请求标识符
            if identifier_func:
                identifier = identifier_func()
            else:
                identifier = request.remote_addr or 'unknown'
            
            if not limiter.is_allowed(identifier):
                remaining_time = window_seconds
                return jsonify({
                    'success': False,
                    'error': f'请求过于频繁，请等待{remaining_time}秒后重试',
                    'rate_limit': {
                        'max_requests': max_requests,
                        'window_seconds': window_seconds,
                        'remaining': 0
                    }
                }), HTTPStatus.TOO_MANY_REQUESTS
            
            # 添加速率限制信息到响应头
            response = f(*args, **kwargs)
            if hasattr(response, 'headers'):
                response.headers['X-RateLimit-Limit'] = str(max_requests)
                response.headers['X-RateLimit-Remaining'] = str(limiter.get_remaining(identifier))
                response.headers['X-RateLimit-Reset'] = str(int(time.time() + window_seconds))
            
            return response
        return decorated_function
    return decorator

def truncate_conversation_history(history: List[Dict[str, str]], max_length: int) -> List[Dict[str, str]]:
    """
    截断对话历史，保持最新的消息
    
    Args:
        history: 对话历史列表
        max_length: 最大历史长度
        
    Returns:
        截断后的对话历史
    """
    if len(history) <= max_length:
        return history
    
    # 保留最新的消息，但确保至少保留用户和助手的各一条消息
    if max_length >= 2:
        return history[-max_length:]
    else:
        return history[-1:] if history else []

def generate_session_id(user_id: int, conversation_id: int = None) -> str:
    """
    生成基于时间和用户ID的session_id
    
    Args:
        user_id: 用户ID
        conversation_id: 对话ID（可选）
        
    Returns:
        生成的session_id字符串（16位）
    """
    # 获取当前时间戳（毫秒）
    timestamp = int(time.time() * 1000)
    
    # 构建基础字符串
    base_string = f"{user_id}_{timestamp}"
    if conversation_id:
        base_string += f"_{conversation_id}"
    
    # 使用MD5生成哈希值，然后取前16位
    hash_object = hashlib.md5(base_string.encode())
    session_id = hash_object.hexdigest()[:16]  # 只取前16位
    
    return session_id

def log_request_info(message: str, user_id: str, request_method: str = None):
    """记录请求信息的辅助函数"""
    if not request_method:
        request_method = request.method
    
    from app import logger
    logger.info(
        f"收到{request_method}请求: "
        f"消息长度={len(message)}字符, "
        f"user_id={user_id}, "
        f"IP={request.remote_addr or 'unknown'}"
    )