auth.py

"""
Authentication module supporting both local and LDAP (Azure Entra ID) authentication
"""

from flask import Blueprint, request, jsonify, current_app
from flask_login import login_user, logout_user, login_required, current_user
import jwt
import bcrypt
import ldap3
from datetime import datetime, timedelta
from models import db, User, UserRole, AuthType
from functools import wraps

auth_bp = Blueprint('auth', __name__)

class AuthService:
    """Authentication service handling both local and LDAP authentication"""
    
    @staticmethod
    def hash_password(password):
        """Hash password for local users"""
        return bcrypt.hashpw(password.encode('utf-8'), bcrypt.gensalt()).decode('utf-8')
    
    @staticmethod
    def verify_password(password, password_hash):
        """Verify password for local users"""
        return bcrypt.checkpw(password.encode('utf-8'), password_hash.encode('utf-8'))
    
    @staticmethod
    def authenticate_ldap(username, password):
        """Authenticate user against Azure Entra ID LDAP"""
        try:
            # LDAP server configuration
            server = ldap3.Server(
                current_app.config['LDAP_SERVER'],
                use_ssl=True,
                get_info=ldap3.ALL
            )
            
            # First, bind with service account to search for user
            bind_conn = ldap3.Connection(
                server,
                user=current_app.config['LDAP_BIND_DN'],
                password=current_app.config['LDAP_BIND_PASSWORD'],
                auto_bind=True
            )
            
            # Search for user
            user_search_filter = current_app.config['LDAP_USER_SEARCH'].format(username=username)
            bind_conn.search(
                current_app.config['LDAP_BASE_DN'],
                user_search_filter,
                attributes=['cn', 'mail', 'sAMAccountName', 'memberOf', 'distinguishedName']
            )
            
            if not bind_conn.entries:
                return None, "User not found in LDAP"
            
            user_entry = bind_conn.entries[0]
            user_dn = str(user_entry.distinguishedName)
            
            # Try to authenticate with user credentials
            user_conn = ldap3.Connection(
                server,
                user=user_dn,
                password=password,
                auto_bind=True
            )
            
            # If we get here, authentication was successful
            user_info = {
                'username': str(user_entry.sAMAccountName),
                'email': str(user_entry.mail) if user_entry.mail else f"{username}@company.com",
                'dn': user_dn,
                'groups': [str(group) for group in user_entry.memberOf] if user_entry.memberOf else []
            }
            
            user_conn.unbind()
            bind_conn.unbind()
            
            return user_info, None
            
        except Exception as e:
            current_app.logger.error(f"LDAP authentication error: {str(e)}")
            return None, f"LDAP authentication failed: {str(e)}"
    
    @staticmethod
    def determine_role_from_groups(groups):
        """Determine user role based on LDAP groups"""
        # Define group mappings (customize based on your Azure AD groups)
        group_role_mapping = {
            'CN=Multicaster-Admins': UserRole.ADMIN,
            'CN=Multicaster-Managers': UserRole.MANAGER,
            'CN=Multicaster-Operators': UserRole.OPERATOR,
            'CN=Network-Admins': UserRole.ADMIN,
            'CN=IT-Staff': UserRole.OPERATOR
        }
        
        # Check groups for highest privilege level
        for group in groups:
            for group_pattern, role in group_role_mapping.items():
                if group_pattern in group:
                    if role == UserRole.ADMIN:
                        return UserRole.ADMIN
                    elif role == UserRole.MANAGER:
                        return UserRole.MANAGER
                    elif role == UserRole.OPERATOR:
                        return UserRole.OPERATOR
        
        # Default role
        return UserRole.VIEWER
    
    @staticmethod
    def create_or_update_ldap_user(user_info):
        """Create or update LDAP user in local database"""
        user = User.query.filter_by(username=user_info['username']).first()
        
        if not user:
            # Create new LDAP user
            role = AuthService.determine_role_from_groups(user_info['groups'])
            user = User(
                username=user_info['username'],
                email=user_info['email'],
                role=role,
                auth_type=AuthType.LDAP,
                ldap_dn=user_info['dn']
            )
            db.session.add(user)
        else:
            # Update existing user
            user.email = user_info['email']
            user.ldap_dn = user_info['dn']
            user.role = AuthService.determine_role_from_groups(user_info['groups'])
            user.last_login = datetime.utcnow()
        
        db.session.commit()
        return user
    
    @staticmethod
    def generate_token(user):
        """Generate JWT token for authenticated user"""
        payload = {
            'user_id': user.id,
            'username': user.username,
            'role': user.role.value,
            'exp': datetime.utcnow() + timedelta(seconds=current_app.config['JWT_ACCESS_TOKEN_EXPIRES']),
            'iat': datetime.utcnow()
        }
        
        token = jwt.encode(
            payload,
            current_app.config['JWT_SECRET_KEY'],
            algorithm='HS256'
        )
        
        return token
    
    @staticmethod
    def verify_token(token):
        """Verify JWT token and return user"""
        try:
            payload = jwt.decode(
                token,
                current_app.config['JWT_SECRET_KEY'],
                algorithms=['HS256']
            )
            
            user = User.query.get(payload['user_id'])
            if not user or not user.is_active:
                return None
            
            return user
            
        except jwt.ExpiredSignatureError:
            return None
        except jwt.InvalidTokenError:
            return None

def require_role(min_role):
    """Decorator to require minimum user role"""
    def decorator(f):
        @wraps(f)
        def decorated_function(*args, **kwargs):
            token = request.headers.get('Authorization')
            if not token:
                return jsonify({'error': 'No authorization token provided'}), 401
            
            if token.startswith('Bearer '):
                token = token[7:]
            
            user = AuthService.verify_token(token)
            if not user:
                return jsonify({'error': 'Invalid or expired token'}), 401
            
            # Role hierarchy check
            role_hierarchy = {
                UserRole.VIEWER: 1,
                UserRole.OPERATOR: 2,
                UserRole.MANAGER: 3,
                UserRole.ADMIN: 4
            }
            
            if role_hierarchy.get(user.role, 0) < role_hierarchy.get(min_role, 999):
                return jsonify({'error': 'Insufficient permissions'}), 403
            
            # Add user to request context
            request.current_user = user
            return f(*args, **kwargs)
        
        return decorated_function
    return decorator

@auth_bp.route('/login', methods=['POST'])
def login():
    """Login endpoint supporting both local and LDAP authentication"""
    data = request.get_json()
    
    if not data or not data.get('username') or not data.get('password'):
        return jsonify({'error': 'Username and password required'}), 400
    
    username = data['username']
    password = data['password']
    auth_type = data.get('auth_type', 'auto')  # 'local', 'ldap', or 'auto'
    
    user = None
    error = None
    
    # Try local authentication first (if requested or auto)
    if auth_type in ['local', 'auto']:
        local_user = User.query.filter_by(
            username=username,
            auth_type=AuthType.LOCAL,
            is_active=True
        ).first()
        
        if local_user and AuthService.verify_password(password, local_user.password_hash):
            user = local_user
            user.last_login = datetime.utcnow()
            db.session.commit()
    
    # Try LDAP authentication if local failed or LDAP requested
    if not user and auth_type in ['ldap', 'auto'] and current_app.config.get('LDAP_SERVER'):
        user_info, ldap_error = AuthService.authenticate_ldap(username, password)
        
        if user_info:
            user = AuthService.create_or_update_ldap_user(user_info)
        else:
            error = ldap_error
    
    if not user:
        return jsonify({
            'error': error or 'Invalid username or password'
        }), 401
    
    # Generate token
    token = AuthService.generate_token(user)
    
    return jsonify({
        'token': token,
        'user': user.to_dict(),
        'expires_in': current_app.config['JWT_ACCESS_TOKEN_EXPIRES']
    })

@auth_bp.route('/logout', methods=['POST'])
@require_role(UserRole.VIEWER)
def logout():
    """Logout endpoint"""
    # In JWT, logout is handled client-side by discarding the token
    # For enhanced security, you might want to maintain a token blacklist
    return jsonify({'message': 'Logged out successfully'})

@auth_bp.route('/me', methods=['GET'])
@require_role(UserRole.VIEWER)
def get_current_user():
    """Get current user information"""
    return jsonify({'user': request.current_user.to_dict()})

@auth_bp.route('/users', methods=['GET'])
@require_role(UserRole.ADMIN)
def list_users():
    """List all users (admin only)"""
    users = User.query.all()
    return jsonify({
        'users': [user.to_dict() for user in users]
    })

@auth_bp.route('/users', methods=['POST'])
@require_role(UserRole.ADMIN)
def create_user():
    """Create new local user (admin only)"""
    data = request.get_json()
    
    if not data or not all(k in data for k in ['username', 'email', 'password']):
        return jsonify({'error': 'Username, email, and password required'}), 400
    
    # Check if user already exists
    if User.query.filter_by(username=data['username']).first():
        return jsonify({'error': 'Username already exists'}), 409
    
    if User.query.filter_by(email=data['email']).first():
        return jsonify({'error': 'Email already exists'}), 409
    
    # Create new user
    user = User(
        username=data['username'],
        email=data['email'],
        password_hash=AuthService.hash_password(data['password']),
        role=UserRole(data.get('role', 'viewer')),
        auth_type=AuthType.LOCAL
    )
    
    db.session.add(user)
    db.session.commit()
    
    return jsonify({
        'message': 'User created successfully',
        'user': user.to_dict()
    }), 201

@auth_bp.route('/users/<int:user_id>', methods=['PUT'])
@require_role(UserRole.ADMIN)
def update_user(user_id):
    """Update user (admin only)"""
    user = User.query.get_or_404(user_id)
    data = request.get_json()
    
    if 'email' in data:
        user.email = data['email']
    
    if 'role' in data:
        user.role = UserRole(data['role'])
    
    if 'is_active' in data:
        user.is_active = data['is_active']
    
    if 'password' in data and user.auth_type == AuthType.LOCAL:
        user.password_hash = AuthService.hash_password(data['password'])
    
    db.session.commit()
    
    return jsonify({
        'message': 'User updated successfully',
        'user': user.to_dict()
    })

@auth_bp.route('/users/<int:user_id>', methods=['DELETE'])
@require_role(UserRole.ADMIN)
def delete_user(user_id):
    """Delete user (admin only)"""
    user = User.query.get_or_404(user_id)
    
    # Prevent deletion of last admin
    if user.role == UserRole.ADMIN:
        admin_count = User.query.filter_by(role=UserRole.ADMIN, is_active=True).count()
        if admin_count <= 1:
            return jsonify({'error': 'Cannot delete last admin user'}), 400
    
    db.session.delete(user)
    db.session.commit()
    
    return jsonify({'message': 'User deleted successfully'})