aiteamutils 0.2.52__tar.gz → 0.2.53__tar.gz

{aiteamutils-0.2.52 → aiteamutils-0.2.53}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aiteamutils
-Version: 0.2.52
+Version: 0.2.53
 Summary: AI Team Utilities
 Project-URL: Homepage, https://github.com/yourusername/aiteamutils
 Project-URL: Issues, https://github.com/yourusername/aiteamutils/issues

{aiteamutils-0.2.52 → aiteamutils-0.2.53}/aiteamutils/security.py

@@ -1,12 +1,10 @@
 """보안 관련 유틸리티."""
 from datetime import datetime, timedelta, UTC
-from typing import Dict, Any, Optional, Literal, Callable, Type
-from sqlalchemy.orm import DeclarativeBase as Base
+from typing import Dict, Any, Optional, Literal, Callable, TYPE_CHECKING
 from fastapi import Request, HTTPException, status
 from functools import wraps
 from jose import jwt, JWTError
 from passlib.context import CryptContext
-from sqlalchemy.ext.asyncio import AsyncSession
 import logging
 
 from .exceptions import CustomException, ErrorCode
@@ -20,29 +18,12 @@ _rate_limits: Dict[str, Dict[str, Any]] = {}
 
 class RateLimitExceeded(CustomException):
     """Rate limit 초과 예외."""
-    
-    def __init__(
-        self,
-        detail: str,
-        source_function: str,
-        remaining_seconds: float,
-        max_requests: int,
-        window_seconds: int
-    ):
-        """Rate limit 초과 예외를 초기화합니다."""
+    def __init__(self, detail: str, source_function: str):
         super().__init__(
             ErrorCode.RATE_LIMIT_EXCEEDED,
             detail=detail,
-            source_function=source_function,
-            metadata={
-                "remaining_seconds": remaining_seconds,
-                "max_requests": max_requests,
-                "window_seconds": window_seconds
-            }
+            source_function=source_function
         )
-        self.remaining_seconds = remaining_seconds
-        self.max_requests = max_requests
-        self.window_seconds = window_seconds
 
 def verify_password(plain_password: str, hashed_password: str) -> bool:
     """비밀번호를 검증합니다."""
@@ -68,28 +49,97 @@ def hash_password(password: str) -> str:
             original_error=e
         )
 
+def rate_limit(
+    max_requests: int,
+    window_seconds: int,
+    key_func: Optional[Callable] = None
+):
+    """Rate limiting 데코레이터."""
+    def decorator(func: Callable) -> Callable:
+        @wraps(func)
+        async def wrapper(*args, **kwargs):
+            # Request 객체 찾기
+            request = None
+            for arg in args:
+                if isinstance(arg, Request):
+                    request = arg
+                    break
+            if not request:
+                for arg in kwargs.values():
+                    if isinstance(arg, Request):
+                        request = arg
+                        break
+            if not request:
+                raise CustomException(
+                    ErrorCode.INTERNAL_ERROR,
+                    detail="Request object not found",
+                    source_function="rate_limit"
+                )
+            
+            # 레이트 리밋 키 생성
+            if key_func:
+                rate_limit_key = f"rate_limit:{key_func(request)}"
+            else:
+                client_ip = request.client.host
+                rate_limit_key = f"rate_limit:{client_ip}:{func.__name__}"
+            
+            try:
+                now = datetime.now(UTC)
+                
+                # 현재 rate limit 정보 가져오기
+                rate_info = _rate_limits.get(rate_limit_key)
+                
+                if rate_info is None or (now - rate_info["start_time"]).total_seconds() >= window_seconds:
+                    # 새로운 rate limit 설정
+                    _rate_limits[rate_limit_key] = {
+                        "count": 1,
+                        "start_time": now
+                    }
+                else:
+                    # 기존 rate limit 업데이트
+                    if rate_info["count"] >= max_requests:
+                        # rate limit 초과
+                        remaining_seconds = window_seconds - (now - rate_info["start_time"]).total_seconds()
+                        raise CustomException(
+                            ErrorCode.RATE_LIMIT_EXCEEDED,
+                            detail=f"{int(remaining_seconds)}",
+                            source_function=func.__name__
+                        )
+                    rate_info["count"] += 1
+                
+                try:
+                    # 원래 함수 실행
+                    return await func(*args, **kwargs)
+                except CustomException as e:
+                    raise e
+                except Exception as e:
+                    raise CustomException(
+                        ErrorCode.INTERNAL_ERROR,
+                        detail=str(e),
+                        source_function=func.__name__,
+                        original_error=e
+                    )
+                    
+            except CustomException as e:
+                raise e
+            except Exception as e:
+                raise CustomException(
+                    ErrorCode.INTERNAL_ERROR,
+                    detail=str(e),
+                    source_function="rate_limit",
+                    original_error=e
+                )
+                
+        return wrapper
+    return decorator
+
 async def create_jwt_token(
     user_data: Dict[str, Any],
     token_type: Literal["access", "refresh"],
-    session: AsyncSession,
-    log_model: Type[Base],
+    db_service: Any,
     request: Optional[Request] = None
 ) -> str:
-    """JWT 토큰을 생성합니다.
-    
-    Args:
-        user_data (Dict[str, Any]): 사용자 데이터
-        token_type (Literal["access", "refresh"]): 토큰 타입
-        session (AsyncSession): 데이터베이스 세션
-        log_model (Type[Base]): 로그 모델
-        request (Optional[Request], optional): FastAPI 요청 객체. Defaults to None.
-        
-    Returns:
-        str: 생성된 JWT 토큰
-        
-    Raises:
-        CustomException: 토큰 생성 실패 시
-    """
+    """JWT 토큰을 생성하고 로그를 기록합니다."""
     try:
         settings = get_settings()
         
@@ -98,7 +148,7 @@ async def create_jwt_token(
             token_data = {
                 # 등록 클레임
                 "iss": settings.token_issuer,
-                "sub": user_data["ulid"],
+                "sub": user_data["username"],
                 "aud": settings.token_audience,
                 "exp": expires_at,
                 
@@ -123,7 +173,7 @@ async def create_jwt_token(
             expires_at = datetime.now(UTC) + timedelta(days=14)
             token_data = {
                 "iss": settings.token_issuer,
-                "sub": user_data["ulid"],
+                "sub": user_data["username"],
                 "exp": expires_at,
                 "token_type": token_type,
                 "user_ulid": user_data["ulid"]
@@ -146,13 +196,14 @@ async def create_jwt_token(
         # 로그 생성
         try:
             activity_type = ActivityType.ACCESS_TOKEN_ISSUED if token_type == "access" else ActivityType.REFRESH_TOKEN_ISSUED
-            log_entry = log_model(
-                type=activity_type,
-                user_ulid=user_data["ulid"],
-                token=token
+            await db_service.create_log(
+                {
+                    "type": activity_type,
+                    "user_ulid": user_data["ulid"],
+                    "token": token
+                },
+                request
             )
-            session.add(log_entry)
-            await session.flush()
         except Exception as e:
             # 로그 생성 실패는 토큰 생성에 영향을 주지 않음
             logging.error(f"Failed to create token log: {str(e)}")
@@ -173,18 +224,7 @@ async def verify_jwt_token(
     token: str,
     expected_type: Optional[Literal["access", "refresh"]] = None
 ) -> Dict[str, Any]:
-    """JWT 토큰을 검증합니다.
-    
-    Args:
-        token: 검증할 JWT 토큰
-        expected_type: 예상되는 토큰 타입
-        
-    Returns:
-        Dict[str, Any]: 토큰 페이로드
-        
-    Raises:
-        CustomException: 토큰 검증 실패 시
-    """
+    """JWT 토큰을 검증합니다."""
     try:
         settings = get_settings()
         
@@ -202,23 +242,14 @@ async def verify_jwt_token(
         if not token_type:
             raise CustomException(
                 ErrorCode.INVALID_TOKEN,
-                detail="Token type is missing",
+                detail=token,
                 source_function="security.verify_jwt_token"
             )
             
         if expected_type and token_type != expected_type:
             raise CustomException(
                 ErrorCode.INVALID_TOKEN,
-                detail=f"Expected {expected_type} token but got {token_type}",
-                source_function="security.verify_jwt_token"
-            )
-            
-        # 사용자 식별자 검증
-        user_ulid = payload.get("user_ulid")
-        if not user_ulid:
-            raise CustomException(
-                ErrorCode.INVALID_TOKEN,
-                detail="User identifier is missing",
+                detail=token,
                 source_function="security.verify_jwt_token"
             )
             
@@ -227,7 +258,7 @@ async def verify_jwt_token(
     except JWTError as e:
         raise CustomException(
             ErrorCode.INVALID_TOKEN,
-            detail=str(e),
+            detail=token,
             source_function="security.verify_jwt_token",
             original_error=e
         )
@@ -239,4 +270,22 @@ async def verify_jwt_token(
             detail=str(e),
             source_function="security.verify_jwt_token",
             original_error=e
+        )
+
+def validate_token(token: str) -> Dict[str, Any]:
+    """JWT 토큰을 검증하고 페이로드를 반환합니다."""
+    try:
+        settings = get_settings()
+        payload = jwt.decode(
+            token, 
+            settings.jwt_secret, 
+            algorithms=[settings.jwt_algorithm]
+        )
+        return payload
+    except JWTError as e:
+        raise CustomException(
+            ErrorCode.INVALID_TOKEN,
+            detail=token,
+            source_function="security.validate_token",
+            original_error=e
         ) 

aiteamutils-0.2.53/aiteamutils/version.py

@@ -0,0 +1,2 @@
+"""버전 정보"""
+__version__ = "0.2.53"

aiteamutils-0.2.52/aiteamutils/version.py

@@ -1,2 +0,0 @@
-"""버전 정보"""
-__version__ = "0.2.52"