From cfb0937605c20b6c2004625542031ebaba0f1101 Mon Sep 17 00:00:00 2001
From: Ashesh Vashi <ashesh.vashi@enterprisedb.com>
Date: Mon, 17 Jul 2017 10:32:56 +0530
Subject: [PATCH 01/11] Added support for preparing the encrypted password of a
 PostgreSQL password using the libpq functions - 'PQencryptPasswordConn', and
 'PQencryptPassword'.

---
 doc/src/connection.rst    | 27 +++++++++++++++
 psycopg/connection_type.c | 69 +++++++++++++++++++++++++++++++++++++++
 tests/test_connection.py  | 56 ++++++++++++++++++++++++++++++-
 3 files changed, 151 insertions(+), 1 deletion(-)

diff --git a/doc/src/connection.rst b/doc/src/connection.rst
index fbbc53e1..0e6b0b13 100644
--- a/doc/src/connection.rst
+++ b/doc/src/connection.rst
@@ -681,6 +681,33 @@ The ``connection`` class
             .. __: http://www.postgresql.org/docs/current/static/libpq-status.html#LIBPQ-PQTRANSACTIONSTATUS
 
 
+    .. method:: encrypt_password(password, user, [algorithm])
+
+        Returns the encrypted form of a PostgreSQL password based on the
+        current password encryption algorithm.
+
+        Raises `~psycopg2.NotSupportedError` if the ``psycopg2`` module was
+        compiled with a ``libpq`` version lesser than 10 (which can be detected
+        by the `~psycopg2.__libpq_version__` constant), when encryption
+        algorithm other than 'md5' is specified for the server version greater
+        than, or equal to 10.
+
+        Ignores the encrytion algorithm for servers version less than 10, and
+        always uses 'md5' as encryption algorithm.
+
+        .. seealso:: libpq docs for `PQencryptPasswordConn()`__ for details.
+
+            .. __: https://www.postgresql.org/docs/devel/static/libpq-misc.html#libpq-pqencryptpasswordconn
+
+        .. seealso:: libpq docs for `PQencryptPassword()`__ for details.
+
+            .. __: https://www.postgresql.org/docs/devel/static/libpq-misc.html#libpq-pqencryptpassword
+
+        .. seealso:: libpq docs for `password_encryption`__ for details.
+
+            .. __: https://www.postgresql.org/docs/devel/static/runtime-config-connection.html#guc-password-encryption
+
+
     .. index::
         pair: Protocol; Version
 
diff --git a/psycopg/connection_type.c b/psycopg/connection_type.c
index 8c5085b5..7d365f3f 100644
--- a/psycopg/connection_type.c
+++ b/psycopg/connection_type.c
@@ -547,6 +547,73 @@ do { \
     EXC_IF_TPC_PREPARED(self, what); \
 } while(0)
 
+/* encrypt_password - Prepare the encrypted password form */
+#define psyco_encrypt_password_doc \
+"encrypt_password('password', 'user', ...) -- Prepares the encrypted form of a PostgreSQL password.\n\n" \
+"Accepted arguments are 'algorithm'."
+
+static PyObject *
+psyco_encrypt_password(connectionObject *self, PyObject *args, PyObject *kwargs)
+{
+    const char *password = NULL,
+               *user = NULL,
+               *algorithm = NULL;
+    char *encrypted = NULL;
+
+    PyObject *res = Py_None;
+
+    static char *kwlist[] = {"password", "user", "algorithm", NULL};
+
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "ss|s", kwlist, &password, &user, &algorithm)) {
+        return NULL;
+    }
+
+    if (self->server_version < 100000 ||
+        (algorithm && strcmp(algorithm, "md5") == 0)
+    ) {
+        encrypted = PQencryptPassword(password, user);
+
+        if (encrypted != NULL)
+        {
+            res = PyString_FromString(encrypted);
+            free(encrypted);
+        }
+        return res;
+    }
+    else
+    {
+#if PG_VERSION_NUM >= 100000
+        encrypted = PQencryptPasswordConn(self->pgconn, password, user, algorithm);
+
+        if (!encrypted)
+        {
+            const char *msg;
+            msg = PQerrorMessage(self->pgconn);
+            if (msg && *msg) {
+                PyErr_Format(
+                    ProgrammingError,
+                    "Error encrypting the password!\n%s",
+                    msg
+                );
+                return NULL;
+            }
+        }
+        else
+        {
+            res = PyString_FromString(encrypted);
+        }
+
+        return res;
+#else
+        PyErr_SetString(
+            NotSupportedError,
+            "Password encryption (other than 'md5' algorithm) is not supported for the server version >= 10 in libpq < 10"
+        );
+#endif
+    }
+    return NULL;
+}
+
 /* set_session - set default transaction characteristics */
 
 #define psyco_conn_set_session_doc \
@@ -1176,6 +1243,8 @@ static struct PyMethodDef connectionObject_methods[] = {
      METH_NOARGS, psyco_conn_isexecuting_doc},
     {"cancel", (PyCFunction)psyco_conn_cancel,
      METH_NOARGS, psyco_conn_cancel_doc},
+    {"encrypt_password", (PyCFunction)psyco_encrypt_password,
+     METH_VARARGS|METH_KEYWORDS, psyco_encrypt_password_doc},
     {NULL}
 };
 
diff --git a/tests/test_connection.py b/tests/test_connection.py
index 42a406ce..b9944ae8 100755
--- a/tests/test_connection.py
+++ b/tests/test_connection.py
@@ -37,7 +37,9 @@ from psycopg2 import extensions as ext
 from testutils import (
     script_to_py3, 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
 
@@ -1382,6 +1384,58 @@ class TransactionControlTests(ConnectingTestCase):
         cur.execute("SHOW default_transaction_read_only;")
         self.assertEqual(cur.fetchone()[0], 'off')
 
+    @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(
+            self.conn.encrypt_password('psycopg2', 'ashesh', 'md5'),
+            'md594839d658c28a357126f105b9cb14cfc'
+        )
+
+        if libpq_version() < 100000:
+            if server_encryption_algorithm == 'md5':
+                self.assertEqual(
+                    self.conn.encrypt_password('psycopg2', 'ashesh'),
+                    'md594839d658c28a357126f105b9cb14cfc'
+                )
+            else:
+                self.assertRaises(
+                    psycopg2.ProgrammingError,
+                    self.conn.encrypt_password, 'psycopg2', 'ashesh'
+                )
+        else:
+            enc_password = self.conn.encrypt_password('psycopg2', 'ashesh')
+            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(
+                self.conn.encrypt_password(
+                    'psycopg2', 'ashesh', 'scram-sha-256'
+                )[:14], 'SCRAM-SHA-256$'
+            )
+
+    @skip_after_postgres(10)
+    def test_encrypt_password_pre_10(self):
+        self.assertEqual(
+            self.conn.encrypt_password('psycopg2', 'ashesh'),
+            'md594839d658c28a357126f105b9cb14cfc'
+        )
+
+        # Encryption algorithm will be ignored for postgres version < 10, it
+        # will always use MD5.
+        self.assertEqual(
+            self.conn.encrypt_password('psycopg2', 'ashesh', 'abc'),
+            'md594839d658c28a357126f105b9cb14cfc'
+        )
+
 
 class AutocommitTests(ConnectingTestCase):
     def test_closed(self):

From 78eb80d0cfb305d7eb467f7f25fdea50d379dd0b Mon Sep 17 00:00:00 2001
From: Ashesh Vashi <ashesh.vashi@enterprisedb.com>
Date: Mon, 17 Jul 2017 10:54:50 +0530
Subject: [PATCH 02/11] Using 'Text_FromUTF8' macro for transforming the
 encrypted C string to Python string to make it Python 3 compatible.

---
 psycopg/connection_type.c | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/psycopg/connection_type.c b/psycopg/connection_type.c
index 7d365f3f..52c1a9f8 100644
--- a/psycopg/connection_type.c
+++ b/psycopg/connection_type.c
@@ -575,8 +575,8 @@ psyco_encrypt_password(connectionObject *self, PyObject *args, PyObject *kwargs)
 
         if (encrypted != NULL)
         {
-            res = PyString_FromString(encrypted);
-            free(encrypted);
+            res = Text_FromUTF8(encrypted);
+            PQfreemem(encrypted);
         }
         return res;
     }
@@ -600,10 +600,11 @@ psyco_encrypt_password(connectionObject *self, PyObject *args, PyObject *kwargs)
         }
         else
         {
-            res = PyString_FromString(encrypted);
+            res = Text_FromUTF8(encrypted);
+            PQfreemem(encrypted);
         }
-
         return res;
+
 #else
         PyErr_SetString(
             NotSupportedError,

From 2c1966a7f61e6623798395c02a8bce5027e258f7 Mon Sep 17 00:00:00 2001
From: Ashesh Vashi <ashesh.vashi@enterprisedb.com>
Date: Mon, 17 Jul 2017 11:06:55 +0530
Subject: [PATCH 03/11] When compiled with libpq version < 10, it raises
 'psycopg2.NotSupportedError' (not, psycopg2.ProgrammingError).

---
 tests/test_connection.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_connection.py b/tests/test_connection.py
index b9944ae8..5894f1d2 100755
--- a/tests/test_connection.py
+++ b/tests/test_connection.py
@@ -1404,7 +1404,7 @@ class TransactionControlTests(ConnectingTestCase):
                 )
             else:
                 self.assertRaises(
-                    psycopg2.ProgrammingError,
+                    psycopg2.NotSupportedError,
                     self.conn.encrypt_password, 'psycopg2', 'ashesh'
                 )
         else:

From e089d94c8858d75db523d48738efa409d0ac4167 Mon Sep 17 00:00:00 2001
From: Ashesh Vashi <ashesh.vashi@enterprisedb.com>
Date: Mon, 17 Jul 2017 11:46:44 +0530
Subject: [PATCH 04/11] 'encrypt_password' raises 'psycopg2.NotSupportedErorr'
 exception for server version >= 10, when compiled using libpq version < 10,
 when no algorithm is specified.

---
 tests/test_connection.py | 14 ++++----------
 1 file changed, 4 insertions(+), 10 deletions(-)

diff --git a/tests/test_connection.py b/tests/test_connection.py
index 5894f1d2..4e1d5ccc 100755
--- a/tests/test_connection.py
+++ b/tests/test_connection.py
@@ -1397,16 +1397,10 @@ class TransactionControlTests(ConnectingTestCase):
         )
 
         if libpq_version() < 100000:
-            if server_encryption_algorithm == 'md5':
-                self.assertEqual(
-                    self.conn.encrypt_password('psycopg2', 'ashesh'),
-                    'md594839d658c28a357126f105b9cb14cfc'
-                )
-            else:
-                self.assertRaises(
-                    psycopg2.NotSupportedError,
-                    self.conn.encrypt_password, 'psycopg2', 'ashesh'
-                )
+            self.assertRaises(
+                psycopg2.NotSupportedError,
+                self.conn.encrypt_password, 'psycopg2', 'ashesh'
+            )
         else:
             enc_password = self.conn.encrypt_password('psycopg2', 'ashesh')
             if server_encryption_algorithm == 'md5':

From 84d405894c4236c1d2bb54c673a71e4a8e78a2ec Mon Sep 17 00:00:00 2001
From: Ashesh Vashi <ashesh.vashi@enterprisedb.com>
Date: Thu, 14 Sep 2017 23:42:37 +0530
Subject: [PATCH 05/11] Moving the encrypt_password method from the connection
 class to the psycopgmodule, and exported it from psycopg2.extensions as per
 review comments.

---
 lib/extensions.py         |   2 +-
 psycopg/connection_type.c |  70 -------------------------
 psycopg/psycopgmodule.c   | 105 ++++++++++++++++++++++++++++++++++++++
 tests/test_connection.py  |  17 +++---
 4 files changed, 116 insertions(+), 78 deletions(-)

diff --git a/lib/extensions.py b/lib/extensions.py
index 91b81331..47155bb7 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/connection_type.c b/psycopg/connection_type.c
index 52c1a9f8..8c5085b5 100644
--- a/psycopg/connection_type.c
+++ b/psycopg/connection_type.c
@@ -547,74 +547,6 @@ do { \
     EXC_IF_TPC_PREPARED(self, what); \
 } while(0)
 
-/* encrypt_password - Prepare the encrypted password form */
-#define psyco_encrypt_password_doc \
-"encrypt_password('password', 'user', ...) -- Prepares the encrypted form of a PostgreSQL password.\n\n" \
-"Accepted arguments are 'algorithm'."
-
-static PyObject *
-psyco_encrypt_password(connectionObject *self, PyObject *args, PyObject *kwargs)
-{
-    const char *password = NULL,
-               *user = NULL,
-               *algorithm = NULL;
-    char *encrypted = NULL;
-
-    PyObject *res = Py_None;
-
-    static char *kwlist[] = {"password", "user", "algorithm", NULL};
-
-    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "ss|s", kwlist, &password, &user, &algorithm)) {
-        return NULL;
-    }
-
-    if (self->server_version < 100000 ||
-        (algorithm && strcmp(algorithm, "md5") == 0)
-    ) {
-        encrypted = PQencryptPassword(password, user);
-
-        if (encrypted != NULL)
-        {
-            res = Text_FromUTF8(encrypted);
-            PQfreemem(encrypted);
-        }
-        return res;
-    }
-    else
-    {
-#if PG_VERSION_NUM >= 100000
-        encrypted = PQencryptPasswordConn(self->pgconn, password, user, algorithm);
-
-        if (!encrypted)
-        {
-            const char *msg;
-            msg = PQerrorMessage(self->pgconn);
-            if (msg && *msg) {
-                PyErr_Format(
-                    ProgrammingError,
-                    "Error encrypting the password!\n%s",
-                    msg
-                );
-                return NULL;
-            }
-        }
-        else
-        {
-            res = Text_FromUTF8(encrypted);
-            PQfreemem(encrypted);
-        }
-        return res;
-
-#else
-        PyErr_SetString(
-            NotSupportedError,
-            "Password encryption (other than 'md5' algorithm) is not supported for the server version >= 10 in libpq < 10"
-        );
-#endif
-    }
-    return NULL;
-}
-
 /* set_session - set default transaction characteristics */
 
 #define psyco_conn_set_session_doc \
@@ -1244,8 +1176,6 @@ static struct PyMethodDef connectionObject_methods[] = {
      METH_NOARGS, psyco_conn_isexecuting_doc},
     {"cancel", (PyCFunction)psyco_conn_cancel,
      METH_NOARGS, psyco_conn_cancel_doc},
-    {"encrypt_password", (PyCFunction)psyco_encrypt_password,
-     METH_VARARGS|METH_KEYWORDS, psyco_encrypt_password_doc},
     {NULL}
 };
 
diff --git a/psycopg/psycopgmodule.c b/psycopg/psycopgmodule.c
index 6c95bd69..a6f7240e 100644
--- a/psycopg/psycopgmodule.c
+++ b/psycopg/psycopgmodule.c
@@ -403,6 +403,109 @@ psyco_libpq_version(PyObject *self)
 #endif
 }
 
+/* encrypt_password - Prepare the encrypted password form */
+#define psyco_encrypt_password_doc \
+"encrypt_password(password, user, [conn_or_curs], [algorithm]) -- Prepares the encrypted form of a PostgreSQL password.\n\n"
+
+static PyObject *
+psyco_encrypt_password(PyObject *self, PyObject *args)
+{
+    char *encrypted = NULL;
+
+    PyObject *obj = NULL,
+             *res = Py_None,
+             *password = NULL,
+             *user = NULL,
+             *algorithm = NULL;
+
+    connectionObject *conn = NULL;
+
+    static char *kwlist[] = {"password", "user", "scope", "algorithm", NULL};
+
+    if (!PyArg_ParseTuple(args, "OO|OO",
+            &password, &user, &obj, &algorithm)) {
+        return NULL;
+    }
+
+    if (obj != NULL && obj != Py_None) {
+        if (PyObject_TypeCheck(obj, &cursorType)) {
+            conn = ((cursorObject*)obj)->conn;
+        }
+        else if (PyObject_TypeCheck(obj, &connectionType)) {
+            conn = (connectionObject*)obj;
+        }
+        else {
+            PyErr_SetString(PyExc_TypeError,
+                    "argument 3 must be a connection or a cursor");
+            return NULL;
+        }
+    }
+
+    /* for ensure_bytes */
+    Py_INCREF(user);
+    Py_INCREF(password);
+    if (algorithm) {
+        Py_INCREF(algorithm);
+    }
+
+    if (!(user = psycopg_ensure_bytes(user))) { goto exit; }
+    if (!(password = psycopg_ensure_bytes(password))) { goto exit; }
+    if (algorithm && !(algorithm = psycopg_ensure_bytes(algorithm))) {
+        goto exit;
+    }
+
+    /* Use the libpq API 'PQencryptPassword', when no connection object is
+       available, or the algorithm is set to as 'md5', or the database server
+       version < 10 */
+    if (conn == NULL || conn->server_version < 100000 ||
+            (algorithm != NULL && algorithm != Py_None &&
+             strcmp(Bytes_AS_STRING(algorithm), "md5") == 0)) {
+        encrypted = PQencryptPassword(Bytes_AS_STRING(password),
+            Bytes_AS_STRING(user));
+
+        if (encrypted != NULL) {
+            res = Text_FromUTF8(encrypted);
+            PQfreemem(encrypted);
+        }
+        goto exit;
+    }
+
+#if PG_VERSION_NUM >= 100000
+    encrypted = PQencryptPasswordConn(conn->pgconn, Bytes_AS_STRING(password),
+            Bytes_AS_STRING(user),
+            algorithm ?  Bytes_AS_STRING(algorithm) : NULL);
+
+    if (!encrypted) {
+        const char *msg = PQerrorMessage(conn->pgconn);
+        if (msg && *msg) {
+            PyErr_Format(ProgrammingError, msg);
+            res = NULL;
+            goto exit;
+        }
+    }
+    else {
+        res = Text_FromUTF8(encrypted);
+        PQfreemem(encrypted);
+    }
+#else
+    PyErr_SetString(
+            NotSupportedError,
+            "Password encryption (other than 'md5' algorithm) is not supported for the server version >= 10 in libpq < 10"
+            );
+    res = NULL;
+#endif
+
+exit:
+    Py_XDECREF(user);
+    Py_XDECREF(password);
+    if (algorithm) {
+        Py_XDECREF(algorithm);
+    }
+
+    return res;
+}
+
+
 /* psyco_encodings_fill
 
    Fill the module's postgresql<->python encoding table */
@@ -852,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, psyco_encrypt_password_doc},
 
     {NULL, NULL, 0, NULL}        /* Sentinel */
 };
diff --git a/tests/test_connection.py b/tests/test_connection.py
index cc9330b2..8cc2db5f 100755
--- a/tests/test_connection.py
+++ b/tests/test_connection.py
@@ -1392,17 +1392,20 @@ class TransactionControlTests(ConnectingTestCase):
 
         # MD5 algorithm
         self.assertEqual(
-            self.conn.encrypt_password('psycopg2', 'ashesh', 'md5'),
+            ext.encrypt_password('psycopg2', 'ashesh', self.conn, 'md5'),
             'md594839d658c28a357126f105b9cb14cfc'
         )
 
         if libpq_version() < 100000:
             self.assertRaises(
                 psycopg2.NotSupportedError,
-                self.conn.encrypt_password, 'psycopg2', 'ashesh'
+                ext.encrypt_password, 'psycopg2', 'ashesh', self.conn,
+                'scram-sha-256'
             )
         else:
-            enc_password = self.conn.encrypt_password('psycopg2', 'ashesh')
+            enc_password = ext.encrypt_password(
+                'psycopg2', 'ashesh', self.conn
+            )
             if server_encryption_algorithm == 'md5':
                 self.assertEqual(
                     enc_password, 'md594839d658c28a357126f105b9cb14cfc'
@@ -1411,22 +1414,22 @@ class TransactionControlTests(ConnectingTestCase):
                 self.assertEqual(enc_password[:14], 'SCRAM-SHA-256$')
 
             self.assertEqual(
-                self.conn.encrypt_password(
-                    'psycopg2', 'ashesh', 'scram-sha-256'
+                ext.encrypt_password(
+                    'psycopg2', 'ashesh', self.conn, 'scram-sha-256'
                 )[:14], 'SCRAM-SHA-256$'
             )
 
     @skip_after_postgres(10)
     def test_encrypt_password_pre_10(self):
         self.assertEqual(
-            self.conn.encrypt_password('psycopg2', 'ashesh'),
+            ext.encrypt_password('psycopg2', 'ashesh', self.conn),
             'md594839d658c28a357126f105b9cb14cfc'
         )
 
         # Encryption algorithm will be ignored for postgres version < 10, it
         # will always use MD5.
         self.assertEqual(
-            self.conn.encrypt_password('psycopg2', 'ashesh', 'abc'),
+            ext.encrypt_password('psycopg2', 'ashesh', self.conn, 'abc'),
             'md594839d658c28a357126f105b9cb14cfc'
         )
 

From 36f0db81d228b612a66abf6c880d1d6629da77fc Mon Sep 17 00:00:00 2001
From: Ashesh Vashi <ashesh.vashi@enterprisedb.com>
Date: Tue, 8 May 2018 15:29:16 +0530
Subject: [PATCH 06/11] Fixed the string format error reported by Travis-CI.

Reference: https://travis-ci.org/psycopg/psycopg2/jobs/376288585
---
 psycopg/psycopgmodule.c | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/psycopg/psycopgmodule.c b/psycopg/psycopgmodule.c
index 8b6623af..5ed1c1c6 100644
--- a/psycopg/psycopgmodule.c
+++ b/psycopg/psycopgmodule.c
@@ -424,8 +424,6 @@ psyco_encrypt_password(PyObject *self, PyObject *args)
 
     connectionObject *conn = NULL;
 
-    static char *kwlist[] = {"password", "user", "scope", "algorithm", NULL};
-
     if (!PyArg_ParseTuple(args, "OO|OO",
             &password, &user, &obj, &algorithm)) {
         return NULL;
@@ -482,7 +480,7 @@ psyco_encrypt_password(PyObject *self, PyObject *args)
     if (!encrypted) {
         const char *msg = PQerrorMessage(conn->pgconn);
         if (msg && *msg) {
-            PyErr_Format(ProgrammingError, msg);
+            PyErr_Format(ProgrammingError, "%s", msg);
             res = NULL;
             goto exit;
         }

From a3063900eeb1442b2c58ebb1b89a924548c8db6d Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Sun, 20 May 2018 19:18:42 +0100
Subject: [PATCH 07/11] Fixed code flow in encrypt_password()

Fixed several shortcomings highlighted in #576 and not fixed as
requested.

Also fixed broken behaviour of ignoring the algorithm if the connection
is missing.
---
 psycopg/psycopgmodule.c  | 119 +++++++++++++++++++--------------------
 tests/test_connection.py |  28 ++++-----
 2 files changed, 73 insertions(+), 74 deletions(-)

diff --git a/psycopg/psycopgmodule.c b/psycopg/psycopgmodule.c
index 5ed1c1c6..121bf133 100644
--- a/psycopg/psycopgmodule.c
+++ b/psycopg/psycopgmodule.c
@@ -409,100 +409,99 @@ psyco_libpq_version(PyObject *self)
 
 /* encrypt_password - Prepare the encrypted password form */
 #define psyco_encrypt_password_doc \
-"encrypt_password(password, user, [conn_or_curs], [algorithm]) -- Prepares the encrypted form of a PostgreSQL password.\n\n"
+"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)
+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;
 
-    PyObject *obj = NULL,
-             *res = Py_None,
-             *password = NULL,
-             *user = NULL,
-             *algorithm = NULL;
+    static char *kwlist[] = {"password", "user", "scope", "algorithm", NULL};
 
     connectionObject *conn = NULL;
 
-    if (!PyArg_ParseTuple(args, "OO|OO",
-            &password, &user, &obj, &algorithm)) {
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "OO|OO", kwlist,
+            &password, &user, &scope, &algorithm)) {
         return NULL;
     }
 
-    if (obj != NULL && obj != Py_None) {
-        if (PyObject_TypeCheck(obj, &cursorType)) {
-            conn = ((cursorObject*)obj)->conn;
+    if (scope != Py_None) {
+        if (PyObject_TypeCheck(scope, &cursorType)) {
+            conn = ((cursorObject*)scope)->conn;
         }
-        else if (PyObject_TypeCheck(obj, &connectionType)) {
-            conn = (connectionObject*)obj;
+        else if (PyObject_TypeCheck(scope, &connectionType)) {
+            conn = (connectionObject*)scope;
         }
         else {
             PyErr_SetString(PyExc_TypeError,
-                    "argument 3 must be a connection or a cursor");
-            return NULL;
+                    "the scope must be a connection or a cursor");
+            goto exit;
         }
     }
 
     /* for ensure_bytes */
     Py_INCREF(user);
     Py_INCREF(password);
-    if (algorithm) {
-        Py_INCREF(algorithm);
-    }
+    Py_INCREF(algorithm);
 
     if (!(user = psycopg_ensure_bytes(user))) { goto exit; }
     if (!(password = psycopg_ensure_bytes(password))) { goto exit; }
-    if (algorithm && !(algorithm = psycopg_ensure_bytes(algorithm))) {
-        goto exit;
-    }
-
-    /* Use the libpq API 'PQencryptPassword', when no connection object is
-       available, or the algorithm is set to as 'md5', or the database server
-       version < 10 */
-    if (conn == NULL || conn->server_version < 100000 ||
-            (algorithm != NULL && algorithm != Py_None &&
-             strcmp(Bytes_AS_STRING(algorithm), "md5") == 0)) {
-        encrypted = PQencryptPassword(Bytes_AS_STRING(password),
-            Bytes_AS_STRING(user));
-
-        if (encrypted != NULL) {
-            res = Text_FromUTF8(encrypted);
-            PQfreemem(encrypted);
-        }
-        goto exit;
-    }
-
-#if PG_VERSION_NUM >= 100000
-    encrypted = PQencryptPasswordConn(conn->pgconn, Bytes_AS_STRING(password),
-            Bytes_AS_STRING(user),
-            algorithm ?  Bytes_AS_STRING(algorithm) : NULL);
-
-    if (!encrypted) {
-        const char *msg = PQerrorMessage(conn->pgconn);
-        if (msg && *msg) {
-            PyErr_Format(ProgrammingError, "%s", msg);
-            res = NULL;
+    if (algorithm != Py_None) {
+        if (!(algorithm = psycopg_ensure_bytes(algorithm))) {
             goto exit;
         }
     }
-    else {
-        res = Text_FromUTF8(encrypted);
-        PQfreemem(encrypted);
+
+    /* 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 = 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) is not supported for the server version >= 10 in libpq < 10"
-            );
-    res = NULL;
+        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);
-    if (algorithm) {
-        Py_XDECREF(algorithm);
-    }
+    Py_XDECREF(algorithm);
 
     return res;
 }
diff --git a/tests/test_connection.py b/tests/test_connection.py
index 6d696a41..f8d2d1d2 100755
--- a/tests/test_connection.py
+++ b/tests/test_connection.py
@@ -1397,6 +1397,18 @@ class TransactionControlTests(ConnectingTestCase):
         cur.execute("SHOW default_transaction_read_only;")
         self.assertEqual(cur.fetchone()[0], 'off')
 
+    def test_idempotence_check(self):
+        self.conn.autocommit = False
+        self.conn.readonly = True
+        self.conn.autocommit = True
+        self.conn.readonly = True
+
+        cur = self.conn.cursor()
+        cur.execute("SHOW transaction_read_only")
+        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()
@@ -1441,20 +1453,8 @@ class TransactionControlTests(ConnectingTestCase):
 
         # Encryption algorithm will be ignored for postgres version < 10, it
         # will always use MD5.
-        self.assertEqual(
-            ext.encrypt_password('psycopg2', 'ashesh', self.conn, 'abc'),
-            'md594839d658c28a357126f105b9cb14cfc'
-        )
-
-    def test_idempotence_check(self):
-        self.conn.autocommit = False
-        self.conn.readonly = True
-        self.conn.autocommit = True
-        self.conn.readonly = True
-
-        cur = self.conn.cursor()
-        cur.execute("SHOW transaction_read_only")
-        self.assertEqual(cur.fetchone()[0], 'on')
+        self.assertRaises(psycopg2.ProgrammingError,
+            ext.encrypt_password, 'psycopg2', 'ashesh', self.conn, 'abc')
 
 
 class AutocommitTests(ConnectingTestCase):

From 9e4f89a2a1e154d05f48b996a20fef1bdbbe9d51 Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Sun, 20 May 2018 20:13:04 +0100
Subject: [PATCH 08/11] encrypt_password docs moved to extension module and
 updated

---
 doc/src/connection.rst | 27 ---------------------------
 doc/src/extensions.rst | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 32 insertions(+), 27 deletions(-)

diff --git a/doc/src/connection.rst b/doc/src/connection.rst
index d656fa1b..2910f301 100644
--- a/doc/src/connection.rst
+++ b/doc/src/connection.rst
@@ -676,33 +676,6 @@ The ``connection`` class
             .. __: http://www.postgresql.org/docs/current/static/libpq-status.html#LIBPQ-PQTRANSACTIONSTATUS
 
 
-    .. method:: encrypt_password(password, user, [algorithm])
-
-        Returns the encrypted form of a PostgreSQL password based on the
-        current password encryption algorithm.
-
-        Raises `~psycopg2.NotSupportedError` if the ``psycopg2`` module was
-        compiled with a ``libpq`` version lesser than 10 (which can be detected
-        by the `~psycopg2.__libpq_version__` constant), when encryption
-        algorithm other than 'md5' is specified for the server version greater
-        than, or equal to 10.
-
-        Ignores the encrytion algorithm for servers version less than 10, and
-        always uses 'md5' as encryption algorithm.
-
-        .. seealso:: libpq docs for `PQencryptPasswordConn()`__ for details.
-
-            .. __: https://www.postgresql.org/docs/devel/static/libpq-misc.html#libpq-pqencryptpasswordconn
-
-        .. seealso:: libpq docs for `PQencryptPassword()`__ for details.
-
-            .. __: https://www.postgresql.org/docs/devel/static/libpq-misc.html#libpq-pqencryptpassword
-
-        .. seealso:: libpq docs for `password_encryption`__ for details.
-
-            .. __: https://www.postgresql.org/docs/devel/static/runtime-config-connection.html#guc-password-encryption
-
-
     .. index::
         pair: Protocol; Version
 
diff --git a/doc/src/extensions.rst b/doc/src/extensions.rst
index 8545fcfa..34d53a7e 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

From abca14d6015dc88a8e66c4baa374ab7102798ee4 Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Sun, 20 May 2018 20:50:04 +0100
Subject: [PATCH 09/11] Fixed keywords support for encrypt_password and tests
 completed

---
 psycopg/psycopgmodule.c  |  2 +-
 tests/test_connection.py | 30 ++++++++++++++++++++++++++++--
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/psycopg/psycopgmodule.c b/psycopg/psycopgmodule.c
index 121bf133..1b0567b3 100644
--- a/psycopg/psycopgmodule.c
+++ b/psycopg/psycopgmodule.c
@@ -957,7 +957,7 @@ static PyMethodDef psycopgMethods[] = {
     {"get_wait_callback",  (PyCFunction)psyco_get_wait_callback,
      METH_NOARGS, psyco_get_wait_callback_doc},
     {"encrypt_password", (PyCFunction)psyco_encrypt_password,
-     METH_VARARGS, psyco_encrypt_password_doc},
+     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 f8d2d1d2..7861ab71 100755
--- a/tests/test_connection.py
+++ b/tests/test_connection.py
@@ -1421,6 +1421,13 @@ class TestEncryptPassword(ConnectingTestCase):
             '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,
@@ -1444,6 +1451,9 @@ class TestEncryptPassword(ConnectingTestCase):
                 )[: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(
@@ -1451,11 +1461,27 @@ class TestEncryptPassword(ConnectingTestCase):
             'md594839d658c28a357126f105b9cb14cfc'
         )
 
-        # Encryption algorithm will be ignored for postgres version < 10, it
-        # will always use MD5.
         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')
+
 
 class AutocommitTests(ConnectingTestCase):
     def test_closed(self):

From 9cf658ec6e2844c50fb7824a42f71cf965e607c8 Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Sun, 20 May 2018 21:18:36 +0100
Subject: [PATCH 10/11] Fixed refcount handling in encrypt_password

Added tests to check bad types, which discovered the above problem: on
type error we would have decref'd on exit something that was only
borrowed (because we wouldn't have performed matching increfs).
---
 psycopg/psycopgmodule.c  | 17 ++++++++---------
 tests/test_connection.py | 11 +++++++++++
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/psycopg/psycopgmodule.c b/psycopg/psycopgmodule.c
index 1b0567b3..23e648d2 100644
--- a/psycopg/psycopgmodule.c
+++ b/psycopg/psycopgmodule.c
@@ -418,16 +418,20 @@ psyco_encrypt_password(PyObject *self, PyObject *args, PyObject *kwargs)
     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};
 
-    connectionObject *conn = 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;
@@ -437,16 +441,11 @@ psyco_encrypt_password(PyObject *self, PyObject *args, PyObject *kwargs)
         }
         else {
             PyErr_SetString(PyExc_TypeError,
-                    "the scope must be a connection or a cursor");
+                "the scope must be a connection or a cursor");
             goto exit;
         }
     }
 
-    /* for ensure_bytes */
-    Py_INCREF(user);
-    Py_INCREF(password);
-    Py_INCREF(algorithm);
-
     if (!(user = psycopg_ensure_bytes(user))) { goto exit; }
     if (!(password = psycopg_ensure_bytes(password))) { goto exit; }
     if (algorithm != Py_None) {
@@ -473,7 +472,7 @@ psyco_encrypt_password(PyObject *self, PyObject *args, PyObject *kwargs)
             goto exit;
         }
 
-        /* TODO: algo = will block: forbid on async/green conn? */
+        /* 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);
diff --git a/tests/test_connection.py b/tests/test_connection.py
index 7861ab71..13635f1f 100755
--- a/tests/test_connection.py
+++ b/tests/test_connection.py
@@ -1482,6 +1482,17 @@ class TestEncryptPassword(ConnectingTestCase):
                 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):

From 9eb3e0cb79aa58eaa566ff8d7ca080038f6c670a Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Sun, 20 May 2018 22:31:22 +0100
Subject: [PATCH 11/11] encrypt_password() reported in the news file

---
 NEWS | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/NEWS b/NEWS
index c44639da..77d547c5 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.