From: Daniele Varrazzo Date: Mon, 2 Aug 2021 00:47:52 +0000 (+0200) Subject: Add pgconn.encrypt_password function X-Git-Tag: 3.0.dev2~12 X-Git-Url: http://git.ipfire.org/gitweb.cgi?a=commitdiff_plain;h=7dc1f1697058db5a84b5a4efb6243eb985a9656d;p=thirdparty%2Fpsycopg.git Add pgconn.encrypt_password function --- diff --git a/psycopg/psycopg/pq/_pq_ctypes.py b/psycopg/psycopg/pq/_pq_ctypes.py index 606423745..e34f5cc32 100644 --- a/psycopg/psycopg/pq/_pq_ctypes.py +++ b/psycopg/psycopg/pq/_pq_ctypes.py @@ -557,6 +557,10 @@ PQfreemem = pq.PQfreemem PQfreemem.argtypes = [c_void_p] PQfreemem.restype = None +PQencryptPasswordConn = pq.PQencryptPasswordConn +PQencryptPasswordConn.argtypes = [PGconn_ptr, c_char_p, c_char_p, c_char_p] +PQencryptPasswordConn.restype = POINTER(c_char) + PQmakeEmptyPGresult = pq.PQmakeEmptyPGresult PQmakeEmptyPGresult.argtypes = [PGconn_ptr, c_int] PQmakeEmptyPGresult.restype = PGresult_ptr diff --git a/psycopg/psycopg/pq/_pq_ctypes.pyi b/psycopg/psycopg/pq/_pq_ctypes.pyi index dffdd89e7..e55bfc2e9 100644 --- a/psycopg/psycopg/pq/_pq_ctypes.pyi +++ b/psycopg/psycopg/pq/_pq_ctypes.pyi @@ -110,6 +110,12 @@ def PQsetResultAttrs( arg2: int, arg3: Array[PGresAttDesc_struct], # type: ignore ) -> int: ... +def PQencryptPasswordConn( + arg1: Optional[PGconn_struct], + arg2: bytes, + arg3: bytes, + arg4: Optional[bytes], +) -> bytes: ... # fmt: off # autogenerated: start diff --git a/psycopg/psycopg/pq/abc.py b/psycopg/psycopg/pq/abc.py index 9b4b6ce3d..385e37323 100644 --- a/psycopg/psycopg/pq/abc.py +++ b/psycopg/psycopg/pq/abc.py @@ -238,6 +238,11 @@ class PGconn(Protocol): def get_copy_data(self, async_: int) -> Tuple[int, memoryview]: ... + def encrypt_password( + self, passwd: bytes, user: bytes, algorithm: Optional[bytes] = None + ) -> bytes: + ... + def make_empty_result(self, exec_status: int) -> "PGresult": ... diff --git a/psycopg/psycopg/pq/pq_ctypes.py b/psycopg/psycopg/pq/pq_ctypes.py index c98138078..7a6f55547 100644 --- a/psycopg/psycopg/pq/pq_ctypes.py +++ b/psycopg/psycopg/pq/pq_ctypes.py @@ -603,6 +603,21 @@ class PGconn: else: return nbytes, memoryview(b"") + def encrypt_password( + self, passwd: bytes, user: bytes, algorithm: Optional[bytes] = None + ) -> bytes: + out = impl.PQencryptPasswordConn( + self._pgconn_ptr, passwd, user, algorithm + ) + if not out: + raise e.OperationalError( + f"password encryption failed: {error_message(self)}" + ) + + rv = string_at(out) + impl.PQfreemem(out) + return rv + def make_empty_result(self, exec_status: int) -> "PGresult": rv = impl.PQmakeEmptyPGresult(self._pgconn_ptr, exec_status) if not rv: diff --git a/psycopg_c/psycopg_c/pq/libpq.pxd b/psycopg_c/psycopg_c/pq/libpq.pxd index d266a8d5d..4aece5e2e 100644 --- a/psycopg_c/psycopg_c/pq/libpq.pxd +++ b/psycopg_c/psycopg_c/pq/libpq.pxd @@ -248,6 +248,8 @@ cdef extern from "libpq-fe.h": # 33.11. Miscellaneous Functions void PQfreemem(void *ptr) nogil void PQconninfoFree(PQconninfoOption *connOptions) + char *PQencryptPasswordConn( + PGconn *conn, const char *passwd, const char *user, const char *algorithm); PGresult *PQmakeEmptyPGresult(PGconn *conn, ExecStatusType status) int PQsetResultAttrs(PGresult *res, int numAttributes, PGresAttDesc *attDescs) int PQlibVersion() diff --git a/psycopg_c/psycopg_c/pq/pgconn.pyx b/psycopg_c/psycopg_c/pq/pgconn.pyx index 570b9ab04..71ffbb004 100644 --- a/psycopg_c/psycopg_c/pq/pgconn.pyx +++ b/psycopg_c/psycopg_c/pq/pgconn.pyx @@ -501,6 +501,23 @@ cdef class PGconn: else: return nbytes, b"" # won't parse it, doesn't really be memoryview + def encrypt_password( + self, const char *passwd, const char *user, algorithm = None + ) -> bytes: + cdef char *out + cdef const char *calgo = NULL + if algorithm: + calgo = algorithm + out = libpq.PQencryptPasswordConn(self._pgconn_ptr, passwd, user, calgo) + if not out: + raise e.OperationalError( + f"password encryption failed: {error_message(self)}" + ) + + rv = bytes(out) + libpq.PQfreemem(out) + return rv + def make_empty_result(self, int exec_status) -> PGresult: cdef libpq.PGresult *rv = libpq.PQmakeEmptyPGresult( self._pgconn_ptr, exec_status) diff --git a/tests/pq/test_pgconn.py b/tests/pq/test_pgconn.py index 2fcfaee2e..2dd348751 100644 --- a/tests/pq/test_pgconn.py +++ b/tests/pq/test_pgconn.py @@ -462,6 +462,39 @@ def test_notice_error(pgconn, caplog): assert "hello error" in rec.message +def test_encrypt_password(pgconn): + enc = pgconn.encrypt_password(b"psycopg2", b"ashesh", b"md5") + assert enc == b"md594839d658c28a357126f105b9cb14cfc" + + +def test_encrypt_password_scram(pgconn): + enc = pgconn.encrypt_password(b"psycopg2", b"ashesh", b"scram-sha-256") + assert enc.startswith(b"SCRAM-SHA-256$") + + +def test_encrypt_password_badalgo(pgconn): + with pytest.raises(psycopg.OperationalError): + assert pgconn.encrypt_password(b"psycopg2", b"ashesh", b"wat") + + +def test_encrypt_password_query(pgconn): + res = pgconn.exec_(b"set password_encryption to 'md5'") + assert res.status == pq.ExecStatus.COMMAND_OK + enc = pgconn.encrypt_password(b"psycopg2", b"ashesh") + assert enc == b"md594839d658c28a357126f105b9cb14cfc" + + res = pgconn.exec_(b"set password_encryption to 'scram-sha-256'") + assert res.status == pq.ExecStatus.COMMAND_OK + enc = pgconn.encrypt_password(b"psycopg2", b"ashesh") + assert enc.startswith(b"SCRAM-SHA-256$") + + +def test_encrypt_password_closed(pgconn): + pgconn.finish() + with pytest.raises(psycopg.OperationalError): + assert pgconn.encrypt_password(b"psycopg2", b"ashesh") + + def test_str(pgconn, dsn): assert "[IDLE]" in str(pgconn) pgconn.finish()