6 files changed, 44 insertions, 33 deletions
diff --git a/source4/auth/pyauth.c b/source4/auth/pyauth.c
index 865609e8e3..f81b449540 100644
--- a/source4/auth/pyauth.c
+++ b/source4/auth/pyauth.c
@@ -107,7 +107,8 @@ void initauth(void)
 	if (PyType_Ready(&PyAuthSession) < 0)
 		return;
 
-	m = Py_InitModule3("auth", py_auth_methods, "Authentication and authorization support.");
+	m = Py_InitModule3("auth", py_auth_methods,
+					   "Authentication and authorization support.");
 	if (m == NULL)
 		return;
 
diff --git a/source4/auth/pyauth.h b/source4/auth/pyauth.h
index f3c9cea73d..2d0379e6e8 100644
--- a/source4/auth/pyauth.h
+++ b/source4/auth/pyauth.h
@@ -23,7 +23,6 @@
 #include "lib/talloc/pytalloc.h"
 #include "auth/session.h"
 
-PyAPI_DATA(PyTypeObject) PyAuthSession;
 #define PyAuthSession_AsSession(obj) py_talloc_get_type(obj, struct auth_session_info)
 #define PyAuthSession_Check(obj) PyObject_TypeCheck(obj, &PyAuthSession)
 struct auth_session_info *PyObject_AsSession(PyObject *obj);
diff --git a/source4/lib/ldb-samba/pyldb.c b/source4/lib/ldb-samba/pyldb.c
index b5ce7f3191..084afb7603 100644
--- a/source4/lib/ldb-samba/pyldb.c
+++ b/source4/lib/ldb-samba/pyldb.c
@@ -27,6 +27,7 @@
 #include "auth/credentials/pycredentials.h"
 #include "ldb_wrap.h"
 #include "lib/ldb-samba/ldif_handlers.h"
+#include "auth/pyauth.h"
 
 static PyObject *pyldb_module;
 static PyObject *py_ldb_error;
@@ -164,6 +165,40 @@ static PyObject *py_ldb_set_utf8_casefold(PyObject *self)
 	Py_RETURN_NONE;
 }
 
+static PyObject *py_ldb_set_session_info(PyObject *self, PyObject *args)
+{
+	PyObject *py_session_info;
+	struct auth_session_info *info;
+	struct ldb_context *ldb;
+	PyObject *mod_samba_auth;
+	PyObject *PyAuthSession_Type;
+	bool ret;
+
+	mod_samba_auth = PyImport_ImportModule("samba.auth");
+	if (mod_samba_auth == NULL)
+		return NULL;
+
+	PyAuthSession_Type = PyObject_GetAttrString(mod_samba_auth, "AuthSession");
+	if (PyAuthSession_Type == NULL)
+		return NULL;
+
+	ret = PyArg_ParseTuple(args, "O!", PyAuthSession_Type, &py_session_info);
+
+	Py_DECREF(PyAuthSession_Type);
+	Py_DECREF(mod_samba_auth);
+
+	if (!ret)
+		return NULL;
+
+	ldb = PyLdb_AsLdbContext(self);
+
+	info = PyAuthSession_AsSession(py_session_info);
+
+	ldb_set_opaque(ldb, "sessionInfo", info);
+
+	Py_RETURN_NONE;
+}
+
 static PyObject *py_ldb_register_samba_handlers(PyObject *self)
 {
 	struct ldb_context *ldb;
@@ -196,6 +231,9 @@ static PyMethodDef py_samba_ldb_methods[] = {
 		METH_NOARGS,
 		"register_samba_handlers()\n"
 		"Register Samba-specific LDB modules and schemas." },
+	{ "set_session_info", (PyCFunction)py_ldb_set_session_info, METH_VARARGS,
+		"set_session_info(session_info)\n"
+		"Set session info to use when connecting." },
 	{ NULL },
 };
 
diff --git a/source4/scripting/python/pyglue.c b/source4/scripting/python/pyglue.c
index b044796b52..8639958561 100644
--- a/source4/scripting/python/pyglue.c
+++ b/source4/scripting/python/pyglue.c
@@ -117,27 +117,6 @@ static PyObject *py_set_debug_level(PyObject *self, PyObject *args)
 	Py_RETURN_NONE;
 }
 
-static PyObject *py_ldb_set_session_info(PyObject *self, PyObject *args)
-{
-	PyObject *py_session_info, *py_ldb;
-	struct auth_session_info *info;
-	struct ldb_context *ldb;
-	if (!PyArg_ParseTuple(args, "OO", &py_ldb, &py_session_info))
-		return NULL;
-
-	PyErr_LDB_OR_RAISE(py_ldb, ldb);
-	/*if (!PyAuthSession_Check(py_session_info)) {
-		PyErr_SetString(PyExc_TypeError, "Expected session info object");
-		return NULL;
-	}*/
-
-	info = PyAuthSession_AsSession(py_session_info);
-
-	ldb_set_opaque(ldb, "sessionInfo", info);
-
-	Py_RETURN_NONE;
-}
-
 static PyObject *py_samdb_set_domain_sid(PyLdbObject *self, PyObject *args)
 { 
 	PyObject *py_ldb, *py_sid;
@@ -463,9 +442,6 @@ static PyMethodDef py_misc_methods[] = {
 		"Generate random password with a length >= min and <= max." },
 	{ "unix2nttime", (PyCFunction)py_unix2nttime, METH_VARARGS,
 		"unix2nttime(timestamp) -> nttime" },
-	{ "ldb_set_session_info", (PyCFunction)py_ldb_set_session_info, METH_VARARGS,
-		"ldb_set_session_info(ldb, session_info)\n"
-		"Set session info to use when connecting." },
 	{ "samdb_set_domain_sid", (PyCFunction)py_samdb_set_domain_sid, METH_VARARGS,
 		"samdb_set_domain_sid(samdb, sid)\n"
 		"Set SID of domain to use." },
diff --git a/source4/scripting/python/samba/__init__.py b/source4/scripting/python/samba/__init__.py
index 6fe1b929e3..9d0d71101b 100644
--- a/source4/scripting/python/samba/__init__.py
+++ b/source4/scripting/python/samba/__init__.py
@@ -110,9 +110,6 @@ class Ldb(_Ldb):
         if url is not None:
             self.connect(url, flags, options)
 
-    def set_session_info(self, session_info):
-        glue.ldb_set_session_info(self, session_info)
-
     def set_create_perms(self, perms=0600):
         # we usually want Samba databases to be private. If we later find we
         # need one public, we will have to change this here
diff --git a/source4/scripting/python/samba/provision.py b/source4/scripting/python/samba/provision.py
index 6bae97bd97..5ac635d7f5 100644
--- a/source4/scripting/python/samba/provision.py
+++ b/source4/scripting/python/samba/provision.py
@@ -857,14 +857,14 @@ def setup_samdb(path, setup_path, session_info, provision_backend, lp,
 
     # Also wipes the database
     setup_samdb_partitions(path, setup_path, message=message, lp=lp,
-                           provision_backend=provision_backend, session_info=session_info,
-                           names=names, 
-                           serverrole=serverrole, schema=schema)
+        provision_backend=provision_backend, session_info=session_info,
+        names=names, serverrole=serverrole, schema=schema)
 
     if (schema == None):
         schema = Schema(setup_path, domainsid, schemadn=names.schemadn, serverdn=names.serverdn)
 
-    # Load the database, but importantly, use Ldb not SamDB as we don't want to load the global schema
+    # Load the database, but importantly, use Ldb not SamDB as we don't want to
+    # load the global schema
     samdb = Ldb(session_info=session_info, 
                 credentials=provision_backend.credentials, lp=lp)