mcli-framework 7.1.0__py3-none-any.whl → 7.1.2__py3-none-any.whl

mcli/ml/auth/auth_manager.py

@@ -2,21 +2,21 @@
 
 import secrets
 from datetime import datetime, timedelta
-from typing import Optional, Union, Dict, Any
+from typing import Any, Dict, Optional, Union
 from uuid import UUID
 
 import bcrypt
 import jwt
-from fastapi import Depends, HTTPException, status, Request
-from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
-from sqlalchemy.orm import Session
+from fastapi import Depends, HTTPException, Request, status
+from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
 from sqlalchemy import select
+from sqlalchemy.orm import Session
 
 from mcli.ml.config import settings
 from mcli.ml.database.models import User, UserRole
 from mcli.ml.database.session import get_db
-from .models import UserCreate, UserLogin, TokenResponse, TokenData, UserResponse
 
+from .models import TokenData, TokenResponse, UserCreate, UserLogin, UserResponse
 
 # Security scheme
 security = HTTPBearer()
@@ -34,22 +34,15 @@ class AuthManager:
     def hash_password(self, password: str) -> str:
         """Hash a password using bcrypt"""
         salt = bcrypt.gensalt()
-        hashed = bcrypt.hashpw(password.encode('utf-8'), salt)
-        return hashed.decode('utf-8')
+        hashed = bcrypt.hashpw(password.encode("utf-8"), salt)
+        return hashed.decode("utf-8")
 
     def verify_password(self, plain_password: str, hashed_password: str) -> bool:
         """Verify a password against a hash"""
-        return bcrypt.checkpw(
-            plain_password.encode('utf-8'),
-            hashed_password.encode('utf-8')
-        )
+        return bcrypt.checkpw(plain_password.encode("utf-8"), hashed_password.encode("utf-8"))
 
     def create_access_token(
-        self,
-        user_id: str,
-        username: str,
-        role: str,
-        expires_delta: Optional[timedelta] = None
+        self, user_id: str, username: str, role: str, expires_delta: Optional[timedelta] = None
     ) -> str:
         """Create a JWT access token"""
         if expires_delta:
@@ -69,11 +62,7 @@ class AuthManager:
         encoded_jwt = jwt.encode(payload, self.secret_key, algorithm=self.algorithm)
         return encoded_jwt
 
-    def create_refresh_token(
-        self,
-        user_id: str,
-        expires_delta: Optional[timedelta] = None
-    ) -> str:
+    def create_refresh_token(self, user_id: str, expires_delta: Optional[timedelta] = None) -> str:
         """Create a refresh token"""
         if expires_delta:
             expire = datetime.utcnow() + expires_delta
@@ -94,11 +83,7 @@ class AuthManager:
     def verify_token(self, token: str) -> Optional[TokenData]:
         """Verify and decode a JWT token"""
         try:
-            payload = jwt.decode(
-                token,
-                self.secret_key,
-                algorithms=[self.algorithm]
-            )
+            payload = jwt.decode(token, self.secret_key, algorithms=[self.algorithm])
 
             token_data = TokenData(
                 sub=payload.get("sub"),
@@ -124,28 +109,23 @@ class AuthManager:
                 headers={"WWW-Authenticate": "Bearer"},
             )
 
-    async def register_user(
-        self,
-        user_data: UserCreate,
-        db: Session
-    ) -> User:
+    async def register_user(self, user_data: UserCreate, db: Session) -> User:
         """Register a new user"""
         # Check if user already exists
-        existing_user = db.query(User).filter(
-            (User.username == user_data.username) |
-            (User.email == user_data.email)
-        ).first()
+        existing_user = (
+            db.query(User)
+            .filter((User.username == user_data.username) | (User.email == user_data.email))
+            .first()
+        )
 
         if existing_user:
             if existing_user.username == user_data.username:
                 raise HTTPException(
-                    status_code=status.HTTP_400_BAD_REQUEST,
-                    detail="Username already registered"
+                    status_code=status.HTTP_400_BAD_REQUEST, detail="Username already registered"
                 )
             else:
                 raise HTTPException(
-                    status_code=status.HTTP_400_BAD_REQUEST,
-                    detail="Email already registered"
+                    status_code=status.HTTP_400_BAD_REQUEST, detail="Email already registered"
                 )
 
         # Create new user
@@ -168,15 +148,9 @@ class AuthManager:
 
         return new_user
 
-    async def authenticate_user(
-        self,
-        login_data: UserLogin,
-        db: Session
-    ) -> Optional[User]:
+    async def authenticate_user(self, login_data: UserLogin, db: Session) -> Optional[User]:
         """Authenticate a user"""
-        user = db.query(User).filter(
-            User.username == login_data.username
-        ).first()
+        user = db.query(User).filter(User.username == login_data.username).first()
 
         if not user:
             return None
@@ -190,11 +164,7 @@ class AuthManager:
 
         return user
 
-    async def login(
-        self,
-        login_data: UserLogin,
-        db: Session
-    ) -> TokenResponse:
+    async def login(self, login_data: UserLogin, db: Session) -> TokenResponse:
         """Login user and return tokens"""
         user = await self.authenticate_user(login_data, db)
 
@@ -207,15 +177,12 @@ class AuthManager:
 
         if not user.is_active:
             raise HTTPException(
-                status_code=status.HTTP_403_FORBIDDEN,
-                detail="User account is disabled"
+                status_code=status.HTTP_403_FORBIDDEN, detail="User account is disabled"
             )
 
         # Create tokens
         access_token = self.create_access_token(
-            user_id=str(user.id),
-            username=user.username,
-            role=user.role.value
+            user_id=str(user.id), username=user.username, role=user.role.value
         )
 
         refresh_token = self.create_refresh_token(user_id=str(user.id))
@@ -224,26 +191,17 @@ class AuthManager:
             access_token=access_token,
             refresh_token=refresh_token,
             expires_in=self.access_token_expire_minutes * 60,
-            user=UserResponse.from_orm(user)
+            user=UserResponse.from_orm(user),
         )
 
-    async def refresh_access_token(
-        self,
-        refresh_token: str,
-        db: Session
-    ) -> TokenResponse:
+    async def refresh_access_token(self, refresh_token: str, db: Session) -> TokenResponse:
         """Refresh access token using refresh token"""
         try:
-            payload = jwt.decode(
-                refresh_token,
-                self.secret_key,
-                algorithms=[self.algorithm]
-            )
+            payload = jwt.decode(refresh_token, self.secret_key, algorithms=[self.algorithm])
 
             if payload.get("type") != "refresh":
                 raise HTTPException(
-                    status_code=status.HTTP_401_UNAUTHORIZED,
-                    detail="Invalid refresh token"
+                    status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid refresh token"
                 )
 
             user_id = payload.get("sub")
@@ -251,39 +209,34 @@ class AuthManager:
 
             if not user or not user.is_active:
                 raise HTTPException(
-                    status_code=status.HTTP_401_UNAUTHORIZED,
-                    detail="User not found or disabled"
+                    status_code=status.HTTP_401_UNAUTHORIZED, detail="User not found or disabled"
                 )
 
             # Create new access token
             access_token = self.create_access_token(
-                user_id=str(user.id),
-                username=user.username,
-                role=user.role.value
+                user_id=str(user.id), username=user.username, role=user.role.value
             )
 
             return TokenResponse(
                 access_token=access_token,
                 refresh_token=refresh_token,  # Return same refresh token
                 expires_in=self.access_token_expire_minutes * 60,
-                user=UserResponse.from_orm(user)
+                user=UserResponse.from_orm(user),
             )
 
         except jwt.ExpiredSignatureError:
             raise HTTPException(
-                status_code=status.HTTP_401_UNAUTHORIZED,
-                detail="Refresh token has expired"
+                status_code=status.HTTP_401_UNAUTHORIZED, detail="Refresh token has expired"
             )
         except jwt.JWTError:
             raise HTTPException(
-                status_code=status.HTTP_401_UNAUTHORIZED,
-                detail="Invalid refresh token"
+                status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid refresh token"
             )
 
     async def get_current_user(
         self,
         credentials: HTTPAuthorizationCredentials = Depends(security),
-        db: Session = Depends(get_db)
+        db: Session = Depends(get_db),
     ) -> User:
         """Get current authenticated user from JWT token"""
         token = credentials.credentials
@@ -293,28 +246,22 @@ class AuthManager:
         user = db.query(User).filter(User.id == token_data.sub).first()
 
         if not user:
-            raise HTTPException(
-                status_code=status.HTTP_404_NOT_FOUND,
-                detail="User not found"
-            )
+            raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="User not found")
 
         if not user.is_active:
             raise HTTPException(
-                status_code=status.HTTP_403_FORBIDDEN,
-                detail="User account is disabled"
+                status_code=status.HTTP_403_FORBIDDEN, detail="User account is disabled"
             )
 
         return user
 
     def require_role(self, *allowed_roles: UserRole):
         """Decorator/dependency to require specific roles"""
-        async def role_checker(
-            current_user: User = Depends(self.get_current_user)
-        ) -> User:
+
+        async def role_checker(current_user: User = Depends(self.get_current_user)) -> User:
             if current_user.role not in allowed_roles:
                 raise HTTPException(
-                    status_code=status.HTTP_403_FORBIDDEN,
-                    detail="Insufficient permissions"
+                    status_code=status.HTTP_403_FORBIDDEN, detail="Insufficient permissions"
                 )
             return current_user
 
@@ -333,34 +280,25 @@ get_current_user = auth_manager.get_current_user
 require_role = auth_manager.require_role
 
 
-async def get_current_active_user(
-    current_user: User = Depends(get_current_user)
-) -> User:
+async def get_current_active_user(current_user: User = Depends(get_current_user)) -> User:
     """Get current active user"""
     if not current_user.is_active:
-        raise HTTPException(
-            status_code=status.HTTP_403_FORBIDDEN,
-            detail="Inactive user"
-        )
+        raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Inactive user")
     return current_user
 
 
-async def get_admin_user(
-    current_user: User = Depends(get_current_user)
-) -> User:
+async def get_admin_user(current_user: User = Depends(get_current_user)) -> User:
     """Get current admin user"""
     if current_user.role != UserRole.ADMIN:
-        raise HTTPException(
-            status_code=status.HTTP_403_FORBIDDEN,
-            detail="Admin access required"
-        )
+        raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Admin access required")
     return current_user
 
 
+import asyncio
+
 # Rate limiting
 from collections import defaultdict
 from datetime import datetime, timedelta
-import asyncio
 
 
 class RateLimiter:
@@ -379,8 +317,7 @@ class RateLimiter:
 
         # Clean old requests
         self.clients[client_id] = [
-            req_time for req_time in self.clients[client_id]
-            if req_time > minute_ago
+            req_time for req_time in self.clients[client_id] if req_time > minute_ago
         ]
 
         # Check limit
@@ -400,8 +337,7 @@ class RateLimiter:
 
             for client_id in list(self.clients.keys()):
                 self.clients[client_id] = [
-                    req_time for req_time in self.clients[client_id]
-                    if req_time > window_start
+                    req_time for req_time in self.clients[client_id] if req_time > window_start
                 ]
 
                 if not self.clients[client_id]:
@@ -418,8 +354,7 @@ async def check_rate_limit(request: Request):
 
     if not await rate_limiter.check_rate_limit(client_ip):
         raise HTTPException(
-            status_code=status.HTTP_429_TOO_MANY_REQUESTS,
-            detail="Rate limit exceeded"
+            status_code=status.HTTP_429_TOO_MANY_REQUESTS, detail="Rate limit exceeded"
         )
 
-    return True
+    return True

mcli/ml/auth/models.py

@@ -1,41 +1,45 @@
 """Authentication data models"""
 
 from datetime import datetime
-from typing import Optional, List
-from pydantic import BaseModel, EmailStr, Field, validator
+from typing import List, Optional
 from uuid import UUID
 
+from pydantic import BaseModel, EmailStr, Field, validator
+
 
 class UserCreate(BaseModel):
     """User registration model"""
+
     username: str = Field(..., min_length=3, max_length=50, regex="^[a-zA-Z0-9_-]+$")
     email: EmailStr
     password: str = Field(..., min_length=8, max_length=100)
     first_name: Optional[str] = Field(None, max_length=50)
     last_name: Optional[str] = Field(None, max_length=50)
 
-    @validator('password')
+    @validator("password")
     def validate_password(cls, v):
         """Ensure password meets security requirements"""
         if len(v) < 8:
-            raise ValueError('Password must be at least 8 characters long')
+            raise ValueError("Password must be at least 8 characters long")
         if not any(char.isdigit() for char in v):
-            raise ValueError('Password must contain at least one digit')
+            raise ValueError("Password must contain at least one digit")
         if not any(char.isupper() for char in v):
-            raise ValueError('Password must contain at least one uppercase letter')
+            raise ValueError("Password must contain at least one uppercase letter")
         if not any(char.islower() for char in v):
-            raise ValueError('Password must contain at least one lowercase letter')
+            raise ValueError("Password must contain at least one lowercase letter")
         return v
 
 
 class UserLogin(BaseModel):
     """User login model"""
+
     username: str
     password: str
 
 
 class UserResponse(BaseModel):
     """User response model"""
+
     id: UUID
     username: str
     email: str
@@ -53,6 +57,7 @@ class UserResponse(BaseModel):
 
 class TokenResponse(BaseModel):
     """JWT token response"""
+
     access_token: str
     token_type: str = "Bearer"
     expires_in: int
@@ -62,6 +67,7 @@ class TokenResponse(BaseModel):
 
 class TokenData(BaseModel):
     """JWT token payload"""
+
     sub: str  # User ID
     username: str
     role: str
@@ -72,45 +78,50 @@ class TokenData(BaseModel):
 
 class PasswordReset(BaseModel):
     """Password reset request"""
+
     email: EmailStr
 
 
 class PasswordResetConfirm(BaseModel):
     """Password reset confirmation"""
+
     token: str
     new_password: str = Field(..., min_length=8, max_length=100)
 
 
 class PasswordChange(BaseModel):
     """Password change request"""
+
     current_password: str
     new_password: str = Field(..., min_length=8, max_length=100)
 
-    @validator('new_password')
+    @validator("new_password")
     def validate_password(cls, v, values):
         """Ensure new password is different and meets requirements"""
-        if 'current_password' in values and v == values['current_password']:
-            raise ValueError('New password must be different from current password')
+        if "current_password" in values and v == values["current_password"]:
+            raise ValueError("New password must be different from current password")
 
         if len(v) < 8:
-            raise ValueError('Password must be at least 8 characters long')
+            raise ValueError("Password must be at least 8 characters long")
         if not any(char.isdigit() for char in v):
-            raise ValueError('Password must contain at least one digit')
+            raise ValueError("Password must contain at least one digit")
         if not any(char.isupper() for char in v):
-            raise ValueError('Password must contain at least one uppercase letter')
+            raise ValueError("Password must contain at least one uppercase letter")
         if not any(char.islower() for char in v):
-            raise ValueError('Password must contain at least one lowercase letter')
+            raise ValueError("Password must contain at least one lowercase letter")
         return v
 
 
 class APIKeyCreate(BaseModel):
     """API key creation model"""
+
     name: str = Field(..., min_length=1, max_length=100)
     expires_at: Optional[datetime] = None
 
 
 class APIKeyResponse(BaseModel):
     """API key response"""
+
     key: str
     name: str
     created_at: datetime
@@ -120,6 +131,7 @@ class APIKeyResponse(BaseModel):
 
 class UserUpdate(BaseModel):
     """User update model"""
+
     email: Optional[EmailStr] = None
     first_name: Optional[str] = Field(None, max_length=50)
     last_name: Optional[str] = Field(None, max_length=50)
@@ -129,6 +141,7 @@ class UserUpdate(BaseModel):
 
 class UserPermissions(BaseModel):
     """User permissions model"""
+
     user_id: UUID
     permissions: List[str]
     roles: List[str]
@@ -136,6 +149,7 @@ class UserPermissions(BaseModel):
 
 class SessionInfo(BaseModel):
     """Session information"""
+
     session_id: str
     user_id: UUID
     ip_address: str
@@ -147,8 +161,9 @@ class SessionInfo(BaseModel):
 
 class LoginAttempt(BaseModel):
     """Login attempt tracking"""
+
     username: str
     ip_address: str
     success: bool
     timestamp: datetime
-    failure_reason: Optional[str] = None
+    failure_reason: Optional[str] = None

mcli/ml/auth/permissions.py

@@ -2,8 +2,8 @@
 
 from datetime import datetime
 from enum import Enum
-from typing import List, Set, Dict, Any
 from functools import wraps
+from typing import Any, Dict, List, Set
 
 from fastapi import HTTPException, status
 from sqlalchemy.orm import Session
@@ -60,7 +60,6 @@ class Permission(Enum):
 # Role-based permission mapping
 ROLE_PERMISSIONS: Dict[UserRole, Set[Permission]] = {
     UserRole.ADMIN: set(Permission),  # Admin has all permissions
-
     UserRole.ANALYST: {
         Permission.MODEL_VIEW,
         Permission.MODEL_CREATE,
@@ -76,7 +75,6 @@ ROLE_PERMISSIONS: Dict[UserRole, Set[Permission]] = {
         Permission.DATA_EDIT,
         Permission.SYSTEM_STATUS,
     },
-
     UserRole.USER: {
         Permission.MODEL_VIEW,
         Permission.PREDICTION_VIEW,
@@ -88,7 +86,6 @@ ROLE_PERMISSIONS: Dict[UserRole, Set[Permission]] = {
         Permission.DATA_VIEW,
         Permission.SYSTEM_STATUS,
     },
-
     UserRole.VIEWER: {
         Permission.MODEL_VIEW,
         Permission.PREDICTION_VIEW,
@@ -109,8 +106,7 @@ def check_permission(user: User, permission: Permission) -> None:
     """Check permission and raise exception if not allowed"""
     if not has_permission(user, permission):
         raise HTTPException(
-            status_code=status.HTTP_403_FORBIDDEN,
-            detail=f"Permission denied: {permission.value}"
+            status_code=status.HTTP_403_FORBIDDEN, detail=f"Permission denied: {permission.value}"
         )
 
 
@@ -121,12 +117,11 @@ def require_permission(permission: Permission):
         @wraps(func)
         async def wrapper(*args, **kwargs):
             # Extract user from kwargs (assumes it's passed as current_user)
-            current_user = kwargs.get('current_user')
+            current_user = kwargs.get("current_user")
 
             if not current_user:
                 raise HTTPException(
-                    status_code=status.HTTP_401_UNAUTHORIZED,
-                    detail="Authentication required"
+                    status_code=status.HTTP_401_UNAUTHORIZED, detail="Authentication required"
                 )
 
             check_permission(current_user, permission)
@@ -143,18 +138,17 @@ def require_any_permission(*permissions: Permission):
     def decorator(func):
         @wraps(func)
         async def wrapper(*args, **kwargs):
-            current_user = kwargs.get('current_user')
+            current_user = kwargs.get("current_user")
 
             if not current_user:
                 raise HTTPException(
-                    status_code=status.HTTP_401_UNAUTHORIZED,
-                    detail="Authentication required"
+                    status_code=status.HTTP_401_UNAUTHORIZED, detail="Authentication required"
                 )
 
             if not any(has_permission(current_user, p) for p in permissions):
                 raise HTTPException(
                     status_code=status.HTTP_403_FORBIDDEN,
-                    detail=f"Permission denied. Required: {[p.value for p in permissions]}"
+                    detail=f"Permission denied. Required: {[p.value for p in permissions]}",
                 )
 
             return await func(*args, **kwargs)
@@ -170,12 +164,11 @@ def require_all_permissions(*permissions: Permission):
     def decorator(func):
         @wraps(func)
         async def wrapper(*args, **kwargs):
-            current_user = kwargs.get('current_user')
+            current_user = kwargs.get("current_user")
 
             if not current_user:
                 raise HTTPException(
-                    status_code=status.HTTP_401_UNAUTHORIZED,
-                    detail="Authentication required"
+                    status_code=status.HTTP_401_UNAUTHORIZED, detail="Authentication required"
                 )
 
             for permission in permissions:
@@ -257,7 +250,7 @@ class AuditLogger:
         action: str,
         success: bool,
         details: Dict[str, Any] = None,
-        db: Session = None
+        db: Session = None,
     ):
         """Log access attempt"""
         log_entry = {
@@ -267,7 +260,7 @@ class AuditLogger:
             "action": action,
             "success": success,
             "timestamp": datetime.utcnow(),
-            "details": details or {}
+            "details": details or {},
         }
 
         # In production, save to database or logging service
@@ -299,4 +292,4 @@ class PermissionGroup:
         Permission.DATA_EDIT,
     }
 
-    ADMIN = set(Permission)  # All permissions
+    ADMIN = set(Permission)  # All permissions