Merge branch 'encrypt-pass'

5 files changed, 227 insertions, 2 deletions

diff --git a/NEWS b/NEWS
index c44639d..77d547c 100644
--- a/NEWS
+++ b/NEWS
@@ -4,6 +4,10 @@ Current release
 What's new in psycopg 2.8
 -------------------------
 
+New features:
+
+- Added `~psycopg2.extensions.encrypt_password()` function (:ticket:`#576`).
+
 Other changes:
 
 - Dropped support for Python 2.6, 3.2, 3.3.
diff --git a/doc/src/extensions.rst b/doc/src/extensions.rst
index 8545fcf..34d53a7 100644
--- a/doc/src/extensions.rst
+++ b/doc/src/extensions.rst
@@ -555,6 +555,38 @@ Other functions
         .. __: http://www.postgresql.org/docs/current/static/libpq-exec.html#LIBPQ-PQESCAPEIDENTIFIER
 
 
+.. method:: encrypt_password(password, user, scope=None, algorithm=None)
+
+    Return the encrypted form of a PostgreSQL password.
+
+    :param password: the cleartext password to encrypt
+    :param user: the name of the user to use the password for
+    :param scope: the scope to encrypt the password into; if *algorithm* is
+        ``md5`` it can be `!None`
+    :type scope: `connection` or `cursor`
+    :param algorithm: the password encryption algorithm to use
+
+    The *algorithm* ``md5`` is always supported. Other algorithms are only
+    supported if the client libpq version is at least 10 and may require a
+    compatible server version: check the `PostgreSQL encryption
+    documentation`__ to know the algorithms supported by your server.
+
+    .. __: https://www.postgresql.org/docs/current/static/encryption-options.html
+
+    Using `!None` as *algorithm* will result in querying the server to know the
+    current server password encryption setting, which is a blocking operation:
+    query the server separately and specify a value for *algorithm* if you
+    want to maintain a non-blocking behaviour.
+
+    .. versionadded:: 2.8
+
+    .. seealso:: PostgreSQL docs for the `password_encryption`__ setting, libpq `PQencryptPasswordConn()`__, `PQencryptPassword()`__ functions.
+
+        .. __: https://www.postgresql.org/docs/current/static/runtime-config-connection.html#GUC-PASSWORD-ENCRYPTION
+        .. __: https://www.postgresql.org/docs/current/static/libpq-misc.html#LIBPQ-PQENCRYPTPASSWORDCONN
+        .. __: https://www.postgresql.org/docs/current/static/libpq-misc.html#LIBPQ-PQENCRYPTPASSWORD
+
+
 
 .. index::
     pair: Isolation level; Constants
diff --git a/lib/extensions.py b/lib/extensions.py
index d15f76c..3c0e225 100644
--- a/lib/extensions.py
+++ b/lib/extensions.py
@@ -63,7 +63,7 @@ from psycopg2._psycopg import (                             # noqa
     string_types, binary_types, new_type, new_array_type, register_type,
     ISQLQuote, Notify, Diagnostics, Column,
     QueryCanceledError, TransactionRollbackError,
-    set_wait_callback, get_wait_callback, )
+    set_wait_callback, get_wait_callback, encrypt_password, )
 
 
 """Isolation level values."""
diff --git a/psycopg/psycopgmodule.c b/psycopg/psycopgmodule.c
index 5deaa16..23e648d 100644
--- a/psycopg/psycopgmodule.c
+++ b/psycopg/psycopgmodule.c
@@ -407,6 +407,105 @@ psyco_libpq_version(PyObject *self)
 #endif
 }
 
+/* encrypt_password - Prepare the encrypted password form */
+#define psyco_encrypt_password_doc \
+"encrypt_password(password, user, [scope], [algorithm]) -- Prepares the encrypted form of a PostgreSQL password.\n\n"
+
+static PyObject *
+psyco_encrypt_password(PyObject *self, PyObject *args, PyObject *kwargs)
+{
+    char *encrypted = NULL;
+    PyObject *password = NULL, *user = NULL;
+    PyObject *scope = Py_None, *algorithm = Py_None;
+    PyObject *res = NULL;
+    connectionObject *conn = NULL;
+
+    static char *kwlist[] = {"password", "user", "scope", "algorithm", NULL};
+
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "OO|OO", kwlist,
+            &password, &user, &scope, &algorithm)) {
+        return NULL;
+    }
+
+    /* for ensure_bytes */
+    Py_INCREF(user);
+    Py_INCREF(password);
+    Py_INCREF(algorithm);
+
+    if (scope != Py_None) {
+        if (PyObject_TypeCheck(scope, &cursorType)) {
+            conn = ((cursorObject*)scope)->conn;
+        }
+        else if (PyObject_TypeCheck(scope, &connectionType)) {
+            conn = (connectionObject*)scope;
+        }
+        else {
+            PyErr_SetString(PyExc_TypeError,
+                "the scope must be a connection or a cursor");
+            goto exit;
+        }
+    }
+
+    if (!(user = psycopg_ensure_bytes(user))) { goto exit; }
+    if (!(password = psycopg_ensure_bytes(password))) { goto exit; }
+    if (algorithm != Py_None) {
+        if (!(algorithm = psycopg_ensure_bytes(algorithm))) {
+            goto exit;
+        }
+    }
+
+    /* If we have to encrypt md5 we can use the libpq < 10 API */
+    if (algorithm != Py_None &&
+            strcmp(Bytes_AS_STRING(algorithm), "md5") == 0) {
+        encrypted = PQencryptPassword(
+            Bytes_AS_STRING(password), Bytes_AS_STRING(user));
+    }
+
+    /* If the algorithm is not md5 we have to use the API available from
+     * libpq 10. */
+    else {
+#if PG_VERSION_NUM >= 100000
+        if (!conn) {
+            PyErr_SetString(ProgrammingError,
+                "password encryption (other than 'md5' algorithm)"
+                " requires a connection or cursor");
+            goto exit;
+        }
+
+        /* TODO: algo = None will block: forbid on async/green conn? */
+        encrypted = PQencryptPasswordConn(conn->pgconn,
+            Bytes_AS_STRING(password), Bytes_AS_STRING(user),
+            algorithm != Py_None ? Bytes_AS_STRING(algorithm) : NULL);
+#else
+        PyErr_SetString(NotSupportedError,
+            "password encryption (other than 'md5' algorithm)"
+            " requires libpq 10");
+        goto exit;
+#endif
+    }
+
+    if (encrypted) {
+        res = Text_FromUTF8(encrypted);
+    }
+    else {
+        const char *msg = PQerrorMessage(conn->pgconn);
+        PyErr_Format(ProgrammingError,
+            "password encryption failed: %s", msg ? msg : "no reason given");
+        goto exit;
+    }
+
+exit:
+    if (encrypted) {
+        PQfreemem(encrypted);
+    }
+    Py_XDECREF(user);
+    Py_XDECREF(password);
+    Py_XDECREF(algorithm);
+
+    return res;
+}
+
+
 /* psyco_encodings_fill
 
    Fill the module's postgresql<->python encoding table */
@@ -856,6 +955,8 @@ static PyMethodDef psycopgMethods[] = {
      METH_O, psyco_set_wait_callback_doc},
     {"get_wait_callback",  (PyCFunction)psyco_get_wait_callback,
      METH_NOARGS, psyco_get_wait_callback_doc},
+    {"encrypt_password", (PyCFunction)psyco_encrypt_password,
+     METH_VARARGS|METH_KEYWORDS, psyco_encrypt_password_doc},
 
     {NULL, NULL, 0, NULL}        /* Sentinel */
 };
diff --git a/tests/test_connection.py b/tests/test_connection.py
index 4625e7e..13635f1 100755
--- a/tests/test_connection.py
+++ b/tests/test_connection.py
@@ -37,7 +37,9 @@ from psycopg2 import extensions as ext
 from .testutils import (
     unittest, decorate_all_tests, skip_if_no_superuser,
     skip_before_postgres, skip_after_postgres, skip_before_libpq,
-    ConnectingTestCase, skip_if_tpc_disabled, skip_if_windows, slow)
+    ConnectingTestCase, skip_if_tpc_disabled, skip_if_windows, slow,
+    libpq_version
+)
 
 from .testconfig import dsn, dbname
 
@@ -1406,6 +1408,92 @@ class TransactionControlTests(ConnectingTestCase):
         self.assertEqual(cur.fetchone()[0], 'on')
 
 
+class TestEncryptPassword(ConnectingTestCase):
+    @skip_before_postgres(10)
+    def test_encrypt_password_post_9_6(self):
+        cur = self.conn.cursor()
+        cur.execute("SHOW password_encryption;")
+        server_encryption_algorithm = cur.fetchone()[0]
+
+        # MD5 algorithm
+        self.assertEqual(
+            ext.encrypt_password('psycopg2', 'ashesh', self.conn, 'md5'),
+            'md594839d658c28a357126f105b9cb14cfc'
+        )
+
+        # keywords
+        self.assertEqual(
+            ext.encrypt_password(
+                password='psycopg2', user='ashesh',
+                scope=self.conn, algorithm='md5'),
+            'md594839d658c28a357126f105b9cb14cfc'
+        )
+        if libpq_version() < 100000:
+            self.assertRaises(
+                psycopg2.NotSupportedError,
+                ext.encrypt_password, 'psycopg2', 'ashesh', self.conn,
+                'scram-sha-256'
+            )
+        else:
+            enc_password = ext.encrypt_password(
+                'psycopg2', 'ashesh', self.conn
+            )
+            if server_encryption_algorithm == 'md5':
+                self.assertEqual(
+                    enc_password, 'md594839d658c28a357126f105b9cb14cfc'
+                )
+            elif server_encryption_algorithm == 'scram-sha-256':
+                self.assertEqual(enc_password[:14], 'SCRAM-SHA-256$')
+
+            self.assertEqual(
+                ext.encrypt_password(
+                    'psycopg2', 'ashesh', self.conn, 'scram-sha-256'
+                )[:14], 'SCRAM-SHA-256$'
+            )
+
+        self.assertRaises(psycopg2.ProgrammingError,
+            ext.encrypt_password, 'psycopg2', 'ashesh', self.conn, 'abc')
+
+    @skip_after_postgres(10)
+    def test_encrypt_password_pre_10(self):
+        self.assertEqual(
+            ext.encrypt_password('psycopg2', 'ashesh', self.conn),
+            'md594839d658c28a357126f105b9cb14cfc'
+        )
+
+        self.assertRaises(psycopg2.ProgrammingError,
+            ext.encrypt_password, 'psycopg2', 'ashesh', self.conn, 'abc')
+
+    def test_encrypt_md5(self):
+        self.assertEqual(
+            ext.encrypt_password('psycopg2', 'ashesh', algorithm='md5'),
+            'md594839d658c28a357126f105b9cb14cfc'
+        )
+
+    def test_encrypt_scram(self):
+        if libpq_version() >= 100000:
+            self.assert_(
+                ext.encrypt_password(
+                    'psycopg2', 'ashesh', self.conn, 'scram-sha-256')
+                .startswith('SCRAM-SHA-256$'))
+        else:
+            self.assertRaises(psycopg2.NotSupportedError,
+                ext.encrypt_password,
+                password='psycopg2', user='ashesh',
+                scope=self.conn, algorithm='scram-sha-256')
+
+    def test_bad_types(self):
+        self.assertRaises(TypeError, ext.encrypt_password)
+        self.assertRaises(TypeError, ext.encrypt_password,
+            'password', 42, self.conn, 'md5')
+        self.assertRaises(TypeError, ext.encrypt_password,
+            42, 'user', self.conn, 'md5')
+        self.assertRaises(TypeError, ext.encrypt_password,
+            42, 'user', 'wat', 'abc')
+        self.assertRaises(TypeError, ext.encrypt_password,
+            'password', 'user', 'wat', 42)
+
+
 class AutocommitTests(ConnectingTestCase):
     def test_closed(self):
         self.conn.close()