pydsdb: Also accept ldb.MessageElement values to dsdb routines

This shows the correct way to accept a value that may be a list of strings
or a proper ldb.MessageElement.

Andrew Bartlett

Signed-off-by: Andrew Bartlett <abartlet@samba.org>
Reviewed-by: Garming Sam <garming@catalyst.net.nz>
(cherry picked from commit b48776d78b446ad4abd4a6bc2ba6b488a29b11d2)
BUG: https://bugzilla.samba.org/show_bug.cgi?id=11789

2 files changed, 63 insertions, 54 deletions

diff --git a/python/samba/dbchecker.py b/python/samba/dbchecker.py
index 74e9678367f..9e3dca922a9 100644
--- a/python/samba/dbchecker.py
+++ b/python/samba/dbchecker.py
@@ -1214,8 +1214,8 @@ newSuperior: %s""" % (str(from_dn), str(to_rdn), str(to_base)))
                 continue
 
             if str(attrname).lower() == 'objectclass':
-                normalised = self.samdb.dsdb_normalise_attributes(self.samdb_schema, attrname, list(obj[attrname]))
-                if list(normalised) != list(obj[attrname]):
+                normalised = self.samdb.dsdb_normalise_attributes(self.samdb_schema, attrname, obj[attrname])
+                if normalised != obj[attrname]:
                     self.err_normalise_mismatch_replace(dn, attrname, list(obj[attrname]))
                     error_count += 1
                 continue
diff --git a/source4/dsdb/pydsdb.c b/source4/dsdb/pydsdb.c
index 0a2b86eed6c..8e7bff556ca 100644
--- a/source4/dsdb/pydsdb.c
+++ b/source4/dsdb/pydsdb.c
@@ -529,11 +529,6 @@ static PyObject *py_dsdb_DsReplicaAttribute(PyObject *self, PyObject *args)
 
 	PyErr_LDB_OR_RAISE(py_ldb, ldb);
 
-	if (!PyList_Check(el_list)) {
-		PyErr_Format(PyExc_TypeError, "ldif_elements must be a list");
-		return NULL;
-	}
-
 	schema = dsdb_get_schema(ldb, NULL);
 	if (!schema) {
 		PyErr_SetString(PyExc_RuntimeError, "Failed to find a schema from ldb");
@@ -555,32 +550,42 @@ static PyObject *py_dsdb_DsReplicaAttribute(PyObject *self, PyObject *args)
 		return NULL;
 	}
 
-	el = talloc_zero(tmp_ctx, struct ldb_message_element);
-	if (el == NULL) {
-		PyErr_NoMemory();
-		talloc_free(tmp_ctx);
-		return NULL;
-	}
-
-	el->name = ldap_display_name;
-	el->num_values = PyList_Size(el_list);
+	/* If we were not given an LdbMessageElement */
+	if (!PyList_Check(el_list)) {
+		if (!py_check_dcerpc_type(el_list, "ldb", "MessageElement")) {
+			PyErr_SetString(py_ldb_get_exception(),
+					"list of strings or ldb MessageElement object required");
+			return NULL;
+		}
+		el = pyldb_MessageElement_AsMessageElement(el_list);
+	} else {
+		el = talloc_zero(tmp_ctx, struct ldb_message_element);
+		if (el == NULL) {
+			PyErr_NoMemory();
+			talloc_free(tmp_ctx);
+			return NULL;
+		}
 
-	el->values = talloc_array(el, struct ldb_val, el->num_values);
-	if (el->values == NULL) {
-		PyErr_NoMemory();
-		talloc_free(tmp_ctx);
-		return NULL;
-	}
+		el->name = ldap_display_name;
+		el->num_values = PyList_Size(el_list);
 
-	for (i = 0; i < el->num_values; i++) {
-		PyObject *item = PyList_GetItem(el_list, i);
-		if (!PyString_Check(item)) {
-			PyErr_Format(PyExc_TypeError, "ldif_elements should be strings");
+		el->values = talloc_array(el, struct ldb_val, el->num_values);
+		if (el->values == NULL) {
+			PyErr_NoMemory();
 			talloc_free(tmp_ctx);
 			return NULL;
 		}
-		el->values[i].data = (uint8_t *)PyString_AsString(item);
-		el->values[i].length = PyString_Size(item);
+
+		for (i = 0; i < el->num_values; i++) {
+			PyObject *item = PyList_GetItem(el_list, i);
+			if (!PyString_Check(item)) {
+				PyErr_Format(PyExc_TypeError, "ldif_elements should be strings");
+				talloc_free(tmp_ctx);
+				return NULL;
+			}
+			el->values[i].data = (uint8_t *)PyString_AsString(item);
+			el->values[i].length = PyString_Size(item);
+		}
 	}
 
 	attr = talloc_zero(tmp_ctx, struct drsuapi_DsReplicaAttribute);
@@ -624,11 +629,6 @@ static PyObject *py_dsdb_normalise_attributes(PyObject *self, PyObject *args)
 
 	PyErr_LDB_OR_RAISE(py_ldb, ldb);
 
-	if (!PyList_Check(el_list)) {
-		PyErr_Format(PyExc_TypeError, "ldif_elements must be a list");
-		return NULL;
-	}
-
 	schema = dsdb_get_schema(ldb, NULL);
 	if (!schema) {
 		PyErr_SetString(PyExc_RuntimeError, "Failed to find a schema from ldb");
@@ -650,32 +650,41 @@ static PyObject *py_dsdb_normalise_attributes(PyObject *self, PyObject *args)
 		return NULL;
 	}
 
-	el = talloc_zero(tmp_ctx, struct ldb_message_element);
-	if (el == NULL) {
-		PyErr_NoMemory();
-		talloc_free(tmp_ctx);
-		return NULL;
-	}
-
-	el->name = ldap_display_name;
-	el->num_values = PyList_Size(el_list);
+	if (!PyList_Check(el_list)) {
+		if (!py_check_dcerpc_type(el_list, "ldb", "MessageElement")) {
+			PyErr_SetString(py_ldb_get_exception(),
+					"list of strings or ldb MessageElement object required");
+			return NULL;
+		}
+		el = pyldb_MessageElement_AsMessageElement(el_list);
+	} else {
+		el = talloc_zero(tmp_ctx, struct ldb_message_element);
+		if (el == NULL) {
+			PyErr_NoMemory();
+			talloc_free(tmp_ctx);
+			return NULL;
+		}
 
-	el->values = talloc_array(el, struct ldb_val, el->num_values);
-	if (el->values == NULL) {
-		PyErr_NoMemory();
-		talloc_free(tmp_ctx);
-		return NULL;
-	}
+		el->name = ldap_display_name;
+		el->num_values = PyList_Size(el_list);
 
-	for (i = 0; i < el->num_values; i++) {
-		PyObject *item = PyList_GetItem(el_list, i);
-		if (!PyString_Check(item)) {
-			PyErr_Format(PyExc_TypeError, "ldif_elements should be strings");
+		el->values = talloc_array(el, struct ldb_val, el->num_values);
+		if (el->values == NULL) {
+			PyErr_NoMemory();
 			talloc_free(tmp_ctx);
 			return NULL;
 		}
-		el->values[i].data = (uint8_t *)PyString_AsString(item);
-		el->values[i].length = PyString_Size(item);
+
+		for (i = 0; i < el->num_values; i++) {
+			PyObject *item = PyList_GetItem(el_list, i);
+			if (!PyString_Check(item)) {
+				PyErr_Format(PyExc_TypeError, "ldif_elements should be strings");
+				talloc_free(tmp_ctx);
+				return NULL;
+			}
+			el->values[i].data = (uint8_t *)PyString_AsString(item);
+			el->values[i].length = PyString_Size(item);
+		}
 	}
 
 	/* Normalise "objectClass" attribute if needed */