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' )