From 66c9c68badff8e5957960c489b0139359ab6d550 Mon Sep 17 00:00:00 2001
From: Andreas Schneider <asn@samba.org>
Date: Wed, 3 Jun 2020 12:38:30 +0200
Subject: [PATCH] auth:creds: Add python bindings for (get|set)_smb_encryption

Signed-off-by: Andreas Schneider <asn@samba.org>
Reviewed-by: Stefan Metzmacher <metze@samba.org>
---
 auth/credentials/pycredentials.c  | 62 +++++++++++++++++++++++++++++++
 python/samba/tests/credentials.py |  6 +++
 2 files changed, 68 insertions(+)

diff --git a/auth/credentials/pycredentials.c b/auth/credentials/pycredentials.c
index 1a83c506088..628aae6500b 100644
--- a/auth/credentials/pycredentials.c
+++ b/auth/credentials/pycredentials.c
@@ -1022,6 +1022,52 @@ static PyObject *py_creds_set_smb_ipc_signing(PyObject *self, PyObject *args)
 	Py_RETURN_NONE;
 }
 
+static PyObject *py_creds_get_smb_encryption(PyObject *self, PyObject *unused)
+{
+	enum smb_encryption_setting encryption_state;
+	struct cli_credentials *creds = NULL;
+
+	creds = PyCredentials_AsCliCredentials(self);
+	if (creds == NULL) {
+		PyErr_Format(PyExc_TypeError, "Credentials expected");
+		return NULL;
+	}
+
+	encryption_state = cli_credentials_get_smb_encryption(creds);
+	return PyLong_FromLong(encryption_state);
+}
+
+static PyObject *py_creds_set_smb_encryption(PyObject *self, PyObject *args)
+{
+	enum smb_encryption_setting encryption_state;
+	struct cli_credentials *creds = NULL;
+	enum credentials_obtained obt = CRED_SPECIFIED;
+
+	creds = PyCredentials_AsCliCredentials(self);
+	if (creds == NULL) {
+		PyErr_Format(PyExc_TypeError, "Credentials expected");
+		return NULL;
+	}
+	if (!PyArg_ParseTuple(args, "i|i", &encryption_state, &obt)) {
+		return NULL;
+	}
+
+	switch (encryption_state) {
+	case SMB_ENCRYPTION_DEFAULT:
+	case SMB_ENCRYPTION_OFF:
+	case SMB_ENCRYPTION_IF_REQUIRED:
+	case SMB_ENCRYPTION_DESIRED:
+	case SMB_ENCRYPTION_REQUIRED:
+		break;
+	default:
+		PyErr_Format(PyExc_TypeError, "Invalid encryption state value");
+		return NULL;
+	}
+
+	cli_credentials_set_smb_encryption(creds, encryption_state, obt);
+	Py_RETURN_NONE;
+}
+
 static PyMethodDef py_creds_methods[] = {
 	{
 		.ml_name  = "get_username",
@@ -1322,6 +1368,16 @@ static PyMethodDef py_creds_methods[] = {
 		.ml_meth  = py_creds_set_smb_ipc_signing,
 		.ml_flags = METH_VARARGS,
 	},
+	{
+		.ml_name  = "get_smb_encryption",
+		.ml_meth  = py_creds_get_smb_encryption,
+		.ml_flags = METH_NOARGS,
+	},
+	{
+		.ml_name  = "set_smb_encryption",
+		.ml_meth  = py_creds_set_smb_encryption,
+		.ml_flags = METH_VARARGS,
+	},
 	{ .ml_name = NULL }
 };
 
@@ -1414,6 +1470,12 @@ MODULE_INIT_FUNC(credentials)
 	PyModule_AddObject(m, "SMB_SIGNING_DESIRED", PyLong_FromLong(SMB_SIGNING_DESIRED));
 	PyModule_AddObject(m, "SMB_SIGNING_REQUIRED", PyLong_FromLong(SMB_SIGNING_REQUIRED));
 
+	PyModule_AddObject(m, "SMB_ENCRYPTION_DEFAULT", PyLong_FromLong(SMB_ENCRYPTION_DEFAULT));
+	PyModule_AddObject(m, "SMB_ENCRYPTION_OFF", PyLong_FromLong(SMB_ENCRYPTION_OFF));
+	PyModule_AddObject(m, "SMB_ENCRYPTION_IF_REQUIRED", PyLong_FromLong(SMB_ENCRYPTION_IF_REQUIRED));
+	PyModule_AddObject(m, "SMB_ENCRYPTION_DESIRED", PyLong_FromLong(SMB_ENCRYPTION_DESIRED));
+	PyModule_AddObject(m, "SMB_ENCRYPTION_REQUIRED", PyLong_FromLong(SMB_ENCRYPTION_REQUIRED));
+
 	Py_INCREF(&PyCredentials);
 	PyModule_AddObject(m, "Credentials", (PyObject *)&PyCredentials);
 	Py_INCREF(&PyCredentialCacheContainer);
diff --git a/python/samba/tests/credentials.py b/python/samba/tests/credentials.py
index 48b263295f7..fcd430f5ab4 100644
--- a/python/samba/tests/credentials.py
+++ b/python/samba/tests/credentials.py
@@ -460,3 +460,9 @@ class CredentialsTests(samba.tests.TestCaseInTempDir):
         self.assertEqual(creds.get_smb_ipc_signing(), credentials.SMB_SIGNING_REQUIRED)
         creds.set_smb_ipc_signing(credentials.SMB_SIGNING_OFF)
         self.assertEqual(creds.get_smb_ipc_signing(), credentials.SMB_SIGNING_OFF)
+
+    def test_smb_encryption(self):
+        creds = credentials.Credentials()
+        self.assertEqual(creds.get_smb_encryption(), credentials.SMB_ENCRYPTION_DEFAULT)
+        creds.set_smb_encryption(credentials.SMB_ENCRYPTION_REQUIRED)
+        self.assertEqual(creds.get_smb_encryption(), credentials.SMB_ENCRYPTION_REQUIRED)