bakery-ia/services/auth/app/core/security.py

# services/auth/app/core/security.py - FIXED VERSION
"""
Security utilities for authentication service
FIXED VERSION - Consistent password hashing using passlib
"""

import re
import hashlib
from datetime import datetime, timedelta, timezone
from typing import Optional, Dict, Any, List
from shared.redis_utils import get_redis_client
from fastapi import HTTPException, status
import structlog
from passlib.context import CryptContext

from app.core.config import settings
from shared.auth.jwt_handler import JWTHandler

logger = structlog.get_logger()

# ✅ FIX: Use passlib for consistent password hashing
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")

# Initialize JWT handler with SAME configuration as gateway
jwt_handler = JWTHandler(settings.JWT_SECRET_KEY, settings.JWT_ALGORITHM)

# Note: Redis client is now accessed via get_redis_client() from shared.redis_utils

class SecurityManager:
    """Security utilities for authentication - FIXED VERSION"""
    
    @staticmethod
    def validate_password(password: str) -> bool:
        """Validate password strength"""
        if len(password) < settings.PASSWORD_MIN_LENGTH:
            return False
        
        if len(password) > 128:  # Max length from Pydantic schema
            return False
        
        if settings.PASSWORD_REQUIRE_UPPERCASE and not re.search(r'[A-Z]', password):
            return False
        
        if settings.PASSWORD_REQUIRE_LOWERCASE and not re.search(r'[a-z]', password):
            return False
        
        if settings.PASSWORD_REQUIRE_NUMBERS and not re.search(r'\d', password):
            return False
        
        if settings.PASSWORD_REQUIRE_SYMBOLS and not re.search(r'[!@#$%^&*(),.?":{}|<>]', password):
            return False
        
        return True
    
    @staticmethod
    def get_password_validation_errors(password: str) -> List[str]:
        """Get detailed password validation errors for better UX"""
        errors = []
        
        if len(password) < settings.PASSWORD_MIN_LENGTH:
            errors.append(f"Password must be at least {settings.PASSWORD_MIN_LENGTH} characters long")
        
        if len(password) > 128:
            errors.append("Password cannot exceed 128 characters")
        
        if settings.PASSWORD_REQUIRE_UPPERCASE and not re.search(r'[A-Z]', password):
            errors.append("Password must contain at least one uppercase letter")
        
        if settings.PASSWORD_REQUIRE_LOWERCASE and not re.search(r'[a-z]', password):
            errors.append("Password must contain at least one lowercase letter")
        
        if settings.PASSWORD_REQUIRE_NUMBERS and not re.search(r'\d', password):
            errors.append("Password must contain at least one number")
        
        if settings.PASSWORD_REQUIRE_SYMBOLS and not re.search(r'[!@#$%^&*(),.?":{}|<>]', password):
            errors.append("Password must contain at least one symbol (!@#$%^&*(),.?\":{}|<>)")
        
        return errors
    
    @staticmethod
    def hash_password(password: str) -> str:
        """Hash password using passlib bcrypt - FIXED"""
        return pwd_context.hash(password)
    
    @staticmethod
    def verify_password(password: str, hashed_password: str) -> bool:
        """Verify password against hash using passlib - FIXED"""
        try:
            return pwd_context.verify(password, hashed_password)
        except Exception as e:
            logger.error(f"Password verification error: {e}")
            return False
    
    @staticmethod
    def create_access_token(user_data: Dict[str, Any]) -> str:
        """
        Create JWT ACCESS token with proper payload structure
        ✅ FIXED: Only creates access tokens
        """
        
        # Validate required fields for access token
        if "user_id" not in user_data:
            raise ValueError("user_id required for access token creation")
        
        if "email" not in user_data:
            raise ValueError("email required for access token creation")
        
        try:
            expire = datetime.now(timezone.utc) + timedelta(minutes=settings.JWT_ACCESS_TOKEN_EXPIRE_MINUTES)
            
            # ✅ FIX 1: ACCESS TOKEN payload structure
            payload = {
                "sub": user_data["user_id"],
                "user_id": user_data["user_id"],
                "email": user_data["email"],
                "type": "access",  # ✅ EXPLICITLY set as access token
                "exp": expire,
                "iat": datetime.now(timezone.utc),
                "iss": "bakery-auth"
            }
            
            # Add optional fields for access tokens
            if "full_name" in user_data:
                payload["full_name"] = user_data["full_name"]
            if "is_verified" in user_data:
                payload["is_verified"] = user_data["is_verified"]
            if "is_active" in user_data:
                payload["is_active"] = user_data["is_active"]
                
            # ✅ CRITICAL FIX: Include role in access token!
            if "role" in user_data:
                payload["role"] = user_data["role"]
            else:
                payload["role"] = "admin"  # Default role if not specified
            
            # NEW: Add subscription data to JWT payload
            if "tenant_id" in user_data:
                payload["tenant_id"] = user_data["tenant_id"]
            
            if "tenant_role" in user_data:
                payload["tenant_role"] = user_data["tenant_role"]
            
            if "subscription" in user_data:
                payload["subscription"] = user_data["subscription"]
            
            if "tenant_access" in user_data:
                # Limit tenant_access to 10 entries to prevent JWT size explosion
                tenant_access = user_data["tenant_access"]
                if tenant_access and len(tenant_access) > 10:
                    tenant_access = tenant_access[:10]
                    logger.warning(f"Truncated tenant_access to 10 entries for user {user_data['user_id']}")
                payload["tenant_access"] = tenant_access
            
            logger.debug(f"Creating access token with payload keys: {list(payload.keys())}")
            
            # ✅ FIX 2: Use JWT handler to create access token
            token = jwt_handler.create_access_token_from_payload(payload)
            logger.debug(f"Access token created successfully for user {user_data['email']}")
            return token
            
        except Exception as e:
            logger.error(f"Access token creation failed for {user_data.get('email', 'unknown')}: {e}")
            raise ValueError(f"Failed to create access token: {str(e)}")
    
    @staticmethod
    def create_refresh_token(user_data: Dict[str, Any]) -> str:
        """
        Create JWT REFRESH token with minimal payload structure
        ✅ FIXED: Only creates refresh tokens, different from access tokens
        """
        
        # Validate required fields for refresh token
        if "user_id" not in user_data:
            raise ValueError("user_id required for refresh token creation")
        
        if not user_data.get("user_id"):
            raise ValueError("user_id cannot be empty")
            
        try:
            expire = datetime.now(timezone.utc) + timedelta(days=settings.JWT_REFRESH_TOKEN_EXPIRE_DAYS)
            
            # ✅ FIX 3: REFRESH TOKEN payload structure (minimal, different from access)
            payload = {
                "sub": user_data["user_id"],
                "user_id": user_data["user_id"],
                "type": "refresh",  # ✅ EXPLICITLY set as refresh token
                "exp": expire,
                "iat": datetime.now(timezone.utc),
                "iss": "bakery-auth"
            }
            
            # Add unique JTI for refresh tokens to prevent duplicates
            if "jti" in user_data:
                payload["jti"] = user_data["jti"]
            else:
                import uuid
                payload["jti"] = str(uuid.uuid4())
            
            # Include email only if available (optional for refresh tokens)
            if "email" in user_data and user_data["email"]:
                payload["email"] = user_data["email"]
            
            logger.debug(f"Creating refresh token with payload keys: {list(payload.keys())}")
            
            # ✅ FIX 4: Use JWT handler to create REFRESH token (not access token!)
            token = jwt_handler.create_refresh_token_from_payload(payload)
            logger.debug(f"Refresh token created successfully for user {user_data['user_id']}")
            return token
            
        except Exception as e:
            logger.error(f"Refresh token creation failed for {user_data.get('user_id', 'unknown')}: {e}")
            raise ValueError(f"Failed to create refresh token: {str(e)}")
    
    @staticmethod
    def verify_token(token: str) -> Optional[Dict[str, Any]]:
        """Verify JWT token with enhanced error handling"""
        try:
            payload = jwt_handler.verify_token(token)
            if payload:
                logger.debug(f"Token verified successfully for user: {payload.get('email', 'unknown')}")
            return payload
        except Exception as e:
            logger.warning(f"Token verification failed: {e}")
            return None
    
    @staticmethod
    def decode_token(token: str) -> Dict[str, Any]:
        """Decode JWT token without verification (for refresh token handling)"""
        try:
            payload = jwt_handler.decode_token_no_verify(token)
            return payload
        except Exception as e:
            logger.error(f"Token decoding failed: {e}")
            raise ValueError("Invalid token format")
    
    @staticmethod
    def generate_secure_hash(data: str) -> str:
        """Generate secure hash for token storage"""
        return hashlib.sha256(data.encode()).hexdigest()

    @staticmethod
    def create_service_token(service_name: str, tenant_id: Optional[str] = None) -> str:
        """
        Create JWT service token for inter-service communication
        ✅ UNIFIED: Uses shared JWT handler for consistent token creation
        ✅ ENHANCED: Supports tenant context for tenant-scoped operations
        
        Args:
            service_name: Name of the service (e.g., 'auth-service', 'tenant-service')
            tenant_id: Optional tenant ID for tenant-scoped service operations
        
        Returns:
            Encoded JWT service token
        """
        try:
            # Use unified JWT handler to create service token
            token = jwt_handler.create_service_token(
                service_name=service_name,
                tenant_id=tenant_id
            )
            logger.debug(f"Created service token for {service_name}", tenant_id=tenant_id)
            return token
            
        except Exception as e:
            logger.error(f"Failed to create service token for {service_name}: {e}")
            raise ValueError(f"Failed to create service token: {str(e)}")
    
    @staticmethod
    async def track_login_attempt(email: str, ip_address: str, success: bool) -> None:
        """Track login attempts for security monitoring"""
        try:
            # This would use Redis for production
            # For now, just log the attempt
            logger.info(f"Login attempt tracked: email={email}, ip={ip_address}, success={success}")
        except Exception as e:
            logger.warning(f"Failed to track login attempt: {e}")
    
    @staticmethod
    def is_token_expired(token: str) -> bool:
        """Check if token is expired"""
        try:
            payload = SecurityManager.decode_token(token)
            exp_timestamp = payload.get("exp")
            if exp_timestamp:
                exp_datetime = datetime.fromtimestamp(exp_timestamp, tz=timezone.utc)
                return datetime.now(timezone.utc) > exp_datetime
            return True
        except Exception:
            return True
    @staticmethod
    def verify_token(token: str) -> Optional[Dict[str, Any]]:
        """Verify JWT token with enhanced error handling"""
        try:
            payload = jwt_handler.verify_token(token)
            if payload:
                logger.debug(f"Token verified successfully for user: {payload.get('email', 'unknown')}")
            return payload
        except Exception as e:
            logger.warning(f"Token verification failed: {e}")
            return None
    
    @staticmethod
    async def track_login_attempt(email: str, ip_address: str, success: bool) -> None:
        """Track login attempts for security monitoring"""
        try:
            redis_client = await get_redis_client()
            key = f"login_attempts:{email}:{ip_address}"

            if success:
                # Clear failed attempts on successful login
                await redis_client.delete(key)
            else:
                # Increment failed attempts
                attempts = await redis_client.incr(key)
                if attempts == 1:
                    # Set expiration on first failed attempt
                    await redis_client.expire(key, settings.LOCKOUT_DURATION_MINUTES * 60)

                if attempts >= settings.MAX_LOGIN_ATTEMPTS:
                    logger.warning(f"Account locked for {email} from {ip_address}")
                    raise HTTPException(
                        status_code=status.HTTP_429_TOO_MANY_REQUESTS,
                        detail=f"Too many failed login attempts. Try again in {settings.LOCKOUT_DURATION_MINUTES} minutes."
                    )

        except HTTPException:
            raise  # Re-raise HTTPException
        except Exception as e:
            logger.error(f"Failed to track login attempt: {e}")
    
    @staticmethod
    async def is_account_locked(email: str, ip_address: str) -> bool:
        """Check if account is locked due to failed login attempts"""
        try:
            redis_client = await get_redis_client()
            key = f"login_attempts:{email}:{ip_address}"
            attempts = await redis_client.get(key)

            if attempts:
                attempts = int(attempts)
                return attempts >= settings.MAX_LOGIN_ATTEMPTS

        except Exception as e:
            logger.error(f"Failed to check account lock status: {e}")

        return False
    
    @staticmethod
    def hash_api_key(api_key: str) -> str:
        """Hash API key for storage"""
        return hashlib.sha256(api_key.encode()).hexdigest()
    
    @staticmethod
    def generate_secure_token(length: int = 32) -> str:
        """Generate secure random token"""
        import secrets
        return secrets.token_urlsafe(length)

    @staticmethod
    def generate_reset_token() -> str:
        """Generate a secure password reset token"""
        import secrets
        return secrets.token_urlsafe(32)
    
    @staticmethod
    def mask_sensitive_data(data: str, visible_chars: int = 4) -> str:
        """Mask sensitive data for logging"""
        if not data or len(data) <= visible_chars:
            return "*" * len(data) if data else ""
        
        return data[:visible_chars] + "*" * (len(data) - visible_chars)
    
    @staticmethod
    async def check_login_attempts(email: str) -> bool:
        """Check if user has exceeded login attempts"""
        try:
            redis_client = await get_redis_client()
            key = f"login_attempts:{email}"
            attempts = await redis_client.get(key)

            if attempts is None:
                return True

            return int(attempts) < settings.MAX_LOGIN_ATTEMPTS
        except Exception as e:
            logger.error(f"Error checking login attempts: {e}")
            return True  # Allow on error

    @staticmethod
    async def increment_login_attempts(email: str) -> None:
        """Increment login attempts for email"""
        try:
            redis_client = await get_redis_client()
            key = f"login_attempts:{email}"
            await redis_client.incr(key)
            await redis_client.expire(key, settings.LOCKOUT_DURATION_MINUTES * 60)
        except Exception as e:
            logger.error(f"Error incrementing login attempts: {e}")

    @staticmethod
    async def clear_login_attempts(email: str) -> None:
        """Clear login attempts for email after successful login"""
        try:
            redis_client = await get_redis_client()
            key = f"login_attempts:{email}"
            await redis_client.delete(key)
            logger.debug(f"Cleared login attempts for {email}")
        except Exception as e:
            logger.error(f"Error clearing login attempts: {e}")
    
    @staticmethod
    async def store_refresh_token(user_id: str, token: str) -> None:
        """Store refresh token in Redis"""
        try:
            redis_client = await get_redis_client()
            token_hash = SecurityManager.hash_api_key(token)  # Reuse hash method
            key = f"refresh_token:{user_id}:{token_hash}"

            # Store with expiration matching JWT refresh token expiry
            expire_seconds = settings.JWT_REFRESH_TOKEN_EXPIRE_DAYS * 24 * 60 * 60
            await redis_client.setex(key, expire_seconds, "valid")

        except Exception as e:
            logger.error(f"Error storing refresh token: {e}")

    @staticmethod
    async def is_refresh_token_valid(user_id: str, token: str) -> bool:
        """Check if refresh token is still valid in Redis"""
        try:
            redis_client = await get_redis_client()
            token_hash = SecurityManager.hash_api_key(token)
            key = f"refresh_token:{user_id}:{token_hash}"

            exists = await redis_client.exists(key)
            return bool(exists)

        except Exception as e:
            logger.error(f"Error checking refresh token validity: {e}")
            return False

    @staticmethod
    async def revoke_refresh_token(user_id: str, token: str) -> None:
        """Revoke refresh token by removing from Redis"""
        try:
            redis_client = await get_redis_client()
            token_hash = SecurityManager.hash_api_key(token)
            key = f"refresh_token:{user_id}:{token_hash}"

            await redis_client.delete(key)
            logger.debug(f"Revoked refresh token for user {user_id}")

        except Exception as e:
            logger.error(f"Error revoking refresh token: {e}")