fastapi-rtk 0.2.27__py3-none-any.whl → 1.0.13__py3-none-any.whl

fastapi_rtk/manager.py

@@ -1,7 +1,9 @@
+import ssl
 import typing
 from datetime import UTC, datetime
 from typing import Any, Callable, Coroutine, Dict, Optional, overload
 
+import fastapi
 from fastapi import HTTPException, Request, Response
 from fastapi.security import OAuth2PasswordRequestForm
 from fastapi_users import BaseUserManager, exceptions
@@ -11,6 +13,7 @@ from fastapi_users.password import PasswordHelperProtocol
 from pydantic import BaseModel
 from sqlalchemy import select
 
+from .const import ErrorCode, logger
 from .db import UserDatabase
 from .exceptions import RolesMismatchException
 from .schemas import (
@@ -111,17 +114,205 @@ class UserManager(IDParser, BaseUserManager[User, int]):
             credentials.password, user.hashed_password
         )
         if not verified:
-            await self.user_db.update(
-                user, {"fail_login_count": (user.fail_login_count or 0) + 1}
-            )
+            await self.on_after_authenticate(user, False)
             return None
         # Update password hash to a more robust one if needed
         if updated_password_hash is not None:
             await self.user_db.update(user, {"hashed_password": updated_password_hash})
 
-        await self.user_db.update(user, {"fail_login_count": 0})
+        await self.on_after_authenticate(user)
         return user
 
+    async def authenticate_ldap(self, credentials: OAuth2PasswordRequestForm):
+        """
+        Authenticate a user using LDAP.
+
+        Args:
+            credentials (OAuth2PasswordRequestForm): The credentials to authenticate the user.
+
+        Returns:
+            User | None: The user if the credentials are valid, None otherwise.
+        """
+        try:
+            user = await self.get_by_username(credentials.username)
+        except exceptions.UserNotExists:
+            # Run the hasher to mitigate timing attack
+            # Inspired from Django: https://code.djangoproject.com/ticket/20760
+            self.password_helper.hash(credentials.password)
+            user = None
+
+        # If no username is provided, go away
+        if not credentials.username:
+            return None
+
+        # If user is not active, go away
+        if user and not user.is_active:
+            return None
+
+        # If user is not registered, and not self-registration, go away
+        if not user and not Setting.AUTH_USER_REGISTRATION:
+            logger.warning(
+                "Tried to login using ldap with username %s, but self-registration is disabled.",
+                credentials.username,
+            )
+            return None
+
+        # Ensure ldap3 is installed
+        try:
+            import ldap3
+            import ldap3.core.exceptions
+        except ImportError:
+            logger.error(
+                "You tried to use LDAP authentication, but 'ldap3' is not installed."
+            )
+            return None
+
+        # --- 1. Configure TLS ---
+        # ldap3 centralizes TLS settings into a Tls object.
+        assert Setting.AUTH_LDAP_SERVER
+        tls_config = None
+        use_ssl = Setting.AUTH_LDAP_USE_TLS or "ldaps://" in Setting.AUTH_LDAP_SERVER
+        if use_ssl:
+            validate_cert = ssl.CERT_REQUIRED
+            if Setting.AUTH_LDAP_ALLOW_SELF_SIGNED:
+                validate_cert = ssl.CERT_OPTIONAL
+            elif not Setting.AUTH_LDAP_TLS_DEMAND:
+                # Default behavior is to require certs if TLS is on
+                pass
+
+            tls_config = ldap3.Tls(
+                local_private_key_file=Setting.AUTH_LDAP_TLS_KEYFILE,
+                local_certificate_file=Setting.AUTH_LDAP_TLS_CERTFILE,
+                validate=validate_cert,
+                ca_certs_path=Setting.AUTH_LDAP_TLS_CACERTDIR,
+                ca_certs_file=Setting.AUTH_LDAP_TLS_CACERTFILE,
+                version=ssl.PROTOCOL_TLS,
+            )
+
+        # --- 2. Define the Server ---
+        server = ldap3.Server(
+            Setting.AUTH_LDAP_SERVER,
+            use_ssl=use_ssl,
+            tls=tls_config,
+            get_info=None,  # No need for get_info for authentication
+        )
+
+        try:
+            user_dn = None
+            user_attributes = {}
+
+            # --- Flow 1 (Indirect Search Bind) ---
+            if Setting.AUTH_LDAP_BIND_USER:
+                # Bind with service account, search, then rebind with user credentials.
+                # The 'with' statement handles connection and unbinding.
+                with ldap3.Connection(
+                    server,
+                    Setting.AUTH_LDAP_BIND_USER,
+                    Setting.AUTH_LDAP_BIND_PASSWORD,
+                    auto_bind=True,
+                    auto_referrals=False,  # Equivalent to OPT_REFERRALS = 0
+                ) as conn:
+                    assert Setting.AUTH_LDAP_SEARCH
+
+                    # Search for the user
+                    user_dn, user_attributes = self._ldap3_search(
+                        conn, credentials.username
+                    )
+
+                    if not user_dn:
+                        logger.info(
+                            "No LDAP object found for user '%s'.", credentials.username
+                        )
+                        return None
+
+                    # Validate user credentials by attempting a rebind
+                    if not conn.rebind(user_dn, credentials.password):
+                        logger.info("Login Failed for user: %s", credentials.username)
+                        if user:
+                            await self.on_after_authenticate(user, False)
+                        return None
+
+            # --- Flow 2 (Direct Search Bind) ---
+            else:
+                # Bind directly with the end-user's credentials.
+                bind_username = credentials.username
+                if Setting.AUTH_LDAP_APPEND_DOMAIN:
+                    bind_username = f"{bind_username}@{Setting.AUTH_LDAP_APPEND_DOMAIN}"
+                if Setting.AUTH_LDAP_USERNAME_FORMAT:
+                    bind_username = (
+                        Setting.AUTH_LDAP_USERNAME_FORMAT % credentials.username
+                    )
+
+                with ldap3.Connection(
+                    server,
+                    bind_username,
+                    credentials.password,
+                    auto_bind=True,
+                    auto_referrals=False,
+                ) as conn:
+                    # If we are here, the bind was successful.
+
+                    if Setting.AUTH_LDAP_SEARCH:
+                        # Search for user attributes
+                        user_dn, user_attributes = self._ldap3_search(
+                            conn, credentials.username
+                        )
+                        if not user_dn:
+                            logger.info(
+                                "No LDAP object found for: %s", credentials.username
+                            )
+                            # Bind succeeded but search failed. Decide if this is an error.
+                            # For now, we exit as the original code did.
+                            return None
+
+            # --- 3. Post-Authentication Processing ---
+            # (This logic remains largely the same)
+            if user and user_attributes and Setting.AUTH_ROLES_SYNC_AT_LOGIN:
+                user = await self._ldap3_calculate_user(user, user_attributes)
+                logger.debug(
+                    "Calculated new roles for user='%s' as: %s", user_dn, user.roles
+                )
+
+            if not user and user_attributes and Setting.AUTH_USER_REGISTRATION:
+                user = await self.create(
+                    {
+                        "username": credentials.username,
+                        "password": self.password_helper.generate(),
+                        "first_name": (
+                            user_attributes.get(Setting.AUTH_LDAP_FIRSTNAME_FIELD)
+                            or [""]
+                        )[0],
+                        "last_name": (
+                            user_attributes.get(Setting.AUTH_LDAP_LASTNAME_FIELD)
+                            or [""]
+                        )[0],
+                        "email": (
+                            user_attributes.get(
+                                Setting.AUTH_LDAP_EMAIL_FIELD,
+                            )
+                            or [f"{credentials.username}@email.notfound"]
+                        )[0],
+                    }
+                )
+                user = await self._ldap3_calculate_user(user, user_attributes)
+
+            if user:
+                await self.on_after_authenticate(user)
+                return user
+            else:
+                return None
+
+        except ldap3.core.exceptions.LDAPBindError:
+            # More specific error for failed logins
+            logger.info("Login Failed for user: %s", credentials.username)
+            if user:
+                await self.on_after_authenticate(user, False)
+            return None
+        except ldap3.core.exceptions.LDAPException as e:
+            # Catch all other ldap3 specific errors
+            logger.error("LDAP Error %s", e)
+            return None
+
     async def create(
         self,
         user_create: BaseModel | dict[str, Any],
@@ -309,6 +500,7 @@ class UserManager(IDParser, BaseUserManager[User, int]):
                 # Associate account
                 user = await self.get_by_email(account_email)
                 if not associate_by_email:
+                    await self.on_after_authenticate(user, False)
                     raise exceptions.UserAlreadyExists()
                 user = await self.user_db.add_oauth_account(user, oauth_account_dict)
             except exceptions.UserNotExists:
@@ -366,6 +558,7 @@ class UserManager(IDParser, BaseUserManager[User, int]):
             user = await self._handle_role_keys(
                 user, oauth_callback_params["user_dict"].pop("role_keys", None)
             )
+        await self.on_after_authenticate(user)
         return user
 
     @overload
@@ -386,6 +579,7 @@ class UserManager(IDParser, BaseUserManager[User, int]):
         :param request: Optional FastAPI request that
         triggered the operation, defaults to None.
         :raises UserInactive: The user is inactive.
+        :return: None if in a request context, otherwise returns the reset token.
         """
         if not user.is_active:
             raise exceptions.UserInactive()
@@ -405,30 +599,6 @@ class UserManager(IDParser, BaseUserManager[User, int]):
         if not request:
             return token
 
-    async def on_after_login(
-        self,
-        user: User,
-        request: Request | None = None,
-        response: Response | None = None,
-    ) -> None:
-        """
-        Perform logic after user login.
-
-        Please call await super().on_after_login(user, request, response) to keep the default behavior.
-
-        *You should overload this method to add your own logic.*
-
-        :param user: The user that is logging in
-        :param request: Optional FastAPI request
-        :param response: Optional response built by the transport.
-        Defaults to None
-        """
-        update_fields = {
-            "last_login": datetime.now(UTC).replace(tzinfo=None),
-            "login_count": (user.login_count or 0) + 1,
-        }
-        await self.user_db.update(user, update_fields)
-
     async def on_after_logout(
         self,
         user: User,
@@ -436,16 +606,14 @@ class UserManager(IDParser, BaseUserManager[User, int]):
         response: Response | None = None,
     ) -> None:
         """
-        Custom `on_after_logout` method to handle user logout.
-
         Perform logic after user logout.
 
         *You should overload this method to add your own logic.*
 
-        :param user: The user that is logging out
-        :param request: Optional FastAPI request
-        :param response: Optional response built by the transport.
-        Defaults to None
+        Args:
+            user (User): The user that is logging out.
+            request (Request | None, optional): Optional FastAPI request.
+            response (Response | None, optional): Optional response built by the transport.
         """
         return
 
@@ -453,19 +621,48 @@ class UserManager(IDParser, BaseUserManager[User, int]):
         self, user: User, token: str, request: Request | None = None
     ) -> None:
         if request:
-            raise HTTPException(status_code=501, detail="Not implemented")
+            raise HTTPException(
+                fastapi.status.HTTP_501_NOT_IMPLEMENTED,
+                ErrorCode.FEATURE_NOT_IMPLEMENTED,
+            )
 
     async def on_after_reset_password(
         self, user: User, request: Request | None = None
     ) -> None:
         if request:
-            raise HTTPException(status_code=501, detail="Not implemented")
+            raise HTTPException(
+                fastapi.status.HTTP_501_NOT_IMPLEMENTED,
+                ErrorCode.FEATURE_NOT_IMPLEMENTED,
+            )
 
     async def on_after_request_verify(
         self, user: User, token: str, request: Request | None = None
     ) -> None:
         if request:
-            raise HTTPException(status_code=501, detail="Not implemented")
+            raise HTTPException(
+                fastapi.status.HTTP_501_NOT_IMPLEMENTED,
+                ErrorCode.FEATURE_NOT_IMPLEMENTED,
+            )
+
+    async def on_after_authenticate(self, user: User, success=True):
+        """
+        Perform logic after authentication attempt.
+
+        Please call `await super().on_after_authenticate(user, success)` to keep the default behavior.
+
+        Args:
+            user (User): The user that attempted authentication.
+            success (bool, optional): Whether the authentication was successful. Defaults to True.
+        """
+        update_fields = {}
+        if success:
+            update_fields["fail_login_count"] = 0
+            update_fields["last_login"] = datetime.now(UTC).replace(tzinfo=None)
+            update_fields["login_count"] = (user.login_count or 0) + 1
+        else:
+            update_fields["fail_login_count"] = (user.fail_login_count or 0) + 1
+        if update_fields:
+            await self.user_db.update(user, update_fields)
 
     async def _update(self, user: User, update_dict: Dict[str, Any]) -> User:
         """
@@ -533,3 +730,124 @@ class UserManager(IDParser, BaseUserManager[User, int]):
         else:
             user_or_user_dict["roles"] = await self.get_roles_by_names(role_names)
         return user_or_user_dict
+
+    """
+    --------------------------------------------------------------------------------------------------------
+        LDAP3 HELPER METHODS
+    --------------------------------------------------------------------------------------------------------
+    """
+
+    def _ldap3_search(self, conn, username: str):
+        """
+        Helper function to perform an LDAP search using ldap3.
+
+        Args:
+            conn (ldap3.Connection): The LDAP connection object.
+            username (str): The username to search for.
+
+        Returns:
+            Tuple[str | None, dict | None]: A tuple containing the user's DN and attributes if found, otherwise (None, None).
+        """
+        import ldap3.utils.conv
+
+        # Configuration check remains good practice
+        assert Setting.AUTH_LDAP_SEARCH
+
+        # 1. Filter Construction (with security improvement)
+        # ldap3 provides a utility to escape special characters to prevent LDAP injection
+        safe_username = ldap3.utils.conv.escape_filter_chars(username)
+
+        base_filter = f"({Setting.AUTH_LDAP_UID_FIELD}={safe_username})"
+        if Setting.AUTH_LDAP_SEARCH_FILTER:
+            # Combine the custom filter with the UID search
+            filter_str = f"(&{Setting.AUTH_LDAP_SEARCH_FILTER}{base_filter})"
+        else:
+            filter_str = base_filter
+
+        # 2. Attribute List Construction
+        # Using a set to prevent duplicates if config fields overlap
+        attributes_to_fetch = {
+            Setting.AUTH_LDAP_FIRSTNAME_FIELD,
+            Setting.AUTH_LDAP_LASTNAME_FIELD,
+            Setting.AUTH_LDAP_EMAIL_FIELD,
+        }
+        if Setting.AUTH_ROLES_MAPPING:
+            attributes_to_fetch.add(Setting.AUTH_LDAP_GROUP_FIELD)
+
+        logger.debug(
+            "LDAP search for '%s' with filter '%s' in base '%s'",
+            username,
+            filter_str,
+            Setting.AUTH_LDAP_SEARCH,
+        )
+
+        # 3. The Search Method
+        # The search operation is a method of the connection object
+        conn.search(
+            Setting.AUTH_LDAP_SEARCH,
+            filter_str,
+            search_scope=ldap3.SUBTREE,
+            attributes=list(attributes_to_fetch),
+        )
+        logger.debug("LDAP search returned: %s", conn.entries)
+
+        # 4. Result Processing
+        # Results are stored in the conn.entries property
+        if len(conn.entries) > 1:
+            logger.error("LDAP search for '%s' returned multiple results.", username)
+            return None, None
+
+        if conn.entries:
+            # Get the first (and only) entry
+            user_entry = conn.entries[0]
+
+            # Extract DN and attributes
+            user_dn = user_entry.entry_dn
+            # .entry_attributes_as_dict provides a simple dict format
+            user_attributes = user_entry.entry_attributes_as_dict
+
+            return user_dn, user_attributes
+        else:
+            # No results found
+            return None, None
+
+    async def _ldap3_calculate_user(
+        self, user: User, user_attributes: dict[str, typing.Any]
+    ):
+        """
+        Calculate the roles for the user.
+
+        Args:
+            user (User): User object to update.
+            user_attributes (dict[str, typing.Any]): User attributes from LDAP.
+
+        Returns:
+            User: Updated user object.
+        """
+        # Apply AUTH_ROLES_MAPPING
+        user_role_keys = user_attributes.get(Setting.AUTH_LDAP_GROUP_FIELD, [])
+        user = await self._handle_role_keys(user, user_role_keys)
+
+        # Apply AUTH_USER_REGISTRATION
+        if Setting.AUTH_USER_REGISTRATION:
+            if not Setting.AUTH_USER_REGISTRATION_ROLE:
+                logger.warning(
+                    "AUTH_USER_REGISTRATION is set but AUTH_USER_REGISTRATION_ROLE is not set"
+                )
+                return user
+
+            # lookup registration role in flask db
+            db_roles = await self.get_roles_by_names(
+                Setting.AUTH_USER_REGISTRATION_ROLE
+            )
+            if not db_roles:
+                logger.warning(
+                    "Can't find AUTH_USER_REGISTRATION_ROLE in database: %s",
+                    Setting.AUTH_USER_REGISTRATION_ROLE,
+                )
+                return user
+            await self.update(
+                {}, user, list(set(r.name for r in db_roles + user.roles))
+            )
+
+        return user

fastapi_rtk/mixins.py

@@ -0,0 +1,12 @@
+from .const import logger
+from .utils import lazy
+
+__all__ = ["LoggerMixin"]
+
+
+class LoggerMixin:
+    """
+    A mixin that provides a logger for the class.
+    """
+
+    logger = lazy(lambda self: logger.getChild(self.__class__.__name__))