psycopg/psycopg2
1
1
Fork 0
mirror of https://github.com/psycopg/psycopg2.git synced 2024-11-25 10:23:43 +03:00
Code Issues Packages Projects Releases Wiki Activity

Moving the encrypt_password method from the connection class to the

Browse Source 
psycopgmodule, and exported it from psycopg2.extensions as per review
comments.
This commit is contained in:
Ashesh Vashi 2017-09-14 23:42:37 +05:30
parent 6e0edf7779
commit 84d405894c
4 changed files with 116 additions and 78 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
lib/extensions.py
View File

@ -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."""

70
psycopg/connection_type.c
View File

@ -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}
};

105
psycopg/psycopgmodule.c
View File

@ -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 */
};

17
tests/test_connection.py
View File

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