Fix OpenLDAP custom queries to use search_scope parameter

openldap/changelog.d/21385.fixed

@@ -0,0 +1 @@
+Fix custom queries to properly use search_scope parameter

openldap/datadog_checks/openldap/openldap.py

@@ -12,6 +12,12 @@
 class OpenLDAP(AgentCheck):
     METRIC_PREFIX = 'openldap'
     SERVICE_CHECK_CONNECT = '{}.can_connect'.format(METRIC_PREFIX)
+    DEFAULT_SEARCH_SCOPE = ldap3.SUBTREE
+    SEARCH_SCOPE_MAPPING = {
+        "base": ldap3.BASE,
+        "level": ldap3.LEVEL,
+        "subtree": ldap3.SUBTREE,
+    }
 
     SEARCH_BASE = 'cn=Monitor'
     SEARCH_FILTER = '(objectClass=*)'
@@ -155,6 +161,15 @@ def _perform_custom_queries(self, conn, custom_queries, tags, instance):
                 self.log.error("`search_filter` field is required for custom query #%s", name)
                 continue
             attrs = query.get("attributes")
+            search_scope_str = query.get("search_scope", self.DEFAULT_SEARCH_SCOPE).lower()
+            search_scope = self.SEARCH_SCOPE_MAPPING.get(search_scope_str)
+
+            if search_scope is None:
+                self.log.error(
+                    "`search_scope` field needs to be one of 'base', 'level' or 'subtree'. Value provided: '%s'",
+                    search_scope_str,
+                )
+                continue
             if "username" in query:
                 username = query.get("username")
                 password = query.get("password")
@@ -181,7 +196,7 @@ def _perform_custom_queries(self, conn, custom_queries, tags, instance):
 
             try:
                 # Perform the search query
-                conn.search(search_base, search_filter, attributes=attrs)
+                conn.search(search_base, search_filter, search_scope=search_scope, attributes=attrs)
             except ldap3.core.exceptions.LDAPException:
                 self.log.exception("Unable to perform search query for %s", name)
                 continue

openldap/tests/test_unit.py

@@ -225,7 +225,7 @@ def test__perform_custom_queries(check, mocker):
     _, _, _, _, queries, tags = check._get_instance_params(instance)
     check._perform_custom_queries(conn_mock, queries, tags, instance)
     conn_mock.rebind.assert_called_once_with(user="user", password="pass", authentication=ldap3.SIMPLE)
-    conn_mock.search.assert_called_once_with("base", "filter", attributes=None)
+    conn_mock.search.assert_called_once_with("base", "filter", search_scope=ldap3.SUBTREE, attributes=None)
     log_mock.error.assert_not_called()  # No error logged
 
     # Check query rebind different user
@@ -248,10 +248,107 @@ def test__perform_custom_queries(check, mocker):
     _, _, _, _, queries, tags = check._get_instance_params(instance)
     check._perform_custom_queries(conn_mock, queries, tags, instance)
     conn_mock.rebind.assert_called_once_with(user="user2", password="pass2", authentication=ldap3.SIMPLE)
-    conn_mock.search.assert_called_once_with("base", "filter", attributes=["*"])
+    conn_mock.search.assert_called_once_with("base", "filter", search_scope=ldap3.SUBTREE, attributes=["*"])
     log_mock.error.assert_not_called()  # No error logged
 
 
+@pytest.mark.parametrize(
+    "search_scope_str,expected_scope",
+    [
+        ("base", ldap3.BASE),
+        ("level", ldap3.LEVEL),
+        ("subtree", ldap3.SUBTREE),
+        ("BASE", ldap3.BASE),
+        ("LEVEL", ldap3.LEVEL),
+        ("SUBTREE", ldap3.SUBTREE),
+    ],
+)
+def test_custom_query_search_scope_valid(check, mocker, search_scope_str, expected_scope):
+    """Test that valid search_scope parameters are properly converted"""
+    instance = {
+        "url": "url",
+        "custom_queries": [
+            {
+                "name": "test_query",
+                "search_base": "ou=users,dc=example,dc=com",
+                "search_filter": "(objectClass=person)",
+                "search_scope": search_scope_str,
+            }
+        ],
+    }
+
+    log_mock = mocker.MagicMock()
+    check.log = log_mock
+    conn_mock = mocker.MagicMock()
+    conn_mock.rebind.return_value = True
+    conn_mock.entries = []
+
+    _, _, _, _, queries, tags = check._get_instance_params(instance)
+    check._perform_custom_queries(conn_mock, queries, tags, instance)
+
+    conn_mock.search.assert_called_once_with(
+        "ou=users,dc=example,dc=com", "(objectClass=person)", search_scope=expected_scope, attributes=None
+    )
+
+
+def test_custom_query_search_scope_default(check, mocker):
+    """Test that missing search_scope defaults to SUBTREE"""
+    instance = {
+        "url": "url",
+        "custom_queries": [
+            {
+                "name": "test_query",
+                "search_base": "ou=users,dc=example,dc=com",
+                "search_filter": "(objectClass=person)",
+                # No search_scope specified
+            }
+        ],
+    }
+
+    log_mock = mocker.MagicMock()
+    check.log = log_mock
+    conn_mock = mocker.MagicMock()
+    conn_mock.rebind.return_value = True
+    conn_mock.entries = []
+
+    _, _, _, _, queries, tags = check._get_instance_params(instance)
+    check._perform_custom_queries(conn_mock, queries, tags, instance)
+
+    conn_mock.search.assert_called_once_with(
+        "ou=users,dc=example,dc=com", "(objectClass=person)", search_scope=ldap3.SUBTREE, attributes=None
+    )
+
+
+def test_custom_query_search_scope_invalid(check, mocker):
+    """Test that invalid search_scope logs error and skips the query"""
+    instance = {
+        "url": "url",
+        "custom_queries": [
+            {
+                "name": "test_query",
+                "search_base": "ou=users,dc=example,dc=com",
+                "search_filter": "(objectClass=person)",
+                "search_scope": "invalid_scope",
+            }
+        ],
+    }
+
+    log_mock = mocker.MagicMock()
+    check.log = log_mock
+    conn_mock = mocker.MagicMock()
+    conn_mock.rebind.return_value = True
+    conn_mock.entries = []
+
+    _, _, _, _, queries, tags = check._get_instance_params(instance)
+    check._perform_custom_queries(conn_mock, queries, tags, instance)
+
+    log_mock.error.assert_called_once_with(
+        "`search_scope` field needs to be one of 'base', 'level' or 'subtree'. Value provided: '%s'", "invalid_scope"
+    )
+    # Verify that search was NOT called due to invalid scope
+    conn_mock.search.assert_not_called()
+
+
 def test__extract_common_name(check):
     # Check lower case and spaces converted
     dn = "cn=Max File Descriptors,cn=Connections,cn=Monitor"