Add pgconn.encrypt_password function

diff --git a/psycopg/psycopg/pq/_pq_ctypes.py b/psycopg/psycopg/pq/_pq_ctypes.py
index 60642374553575691b9ba7b6ddb8612d07bcc98e..e34f5cc32d286d0e82420866171937bec794d573 100644 (file)
--- 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 dffdd89e7afb19b83077d9e3f92b41405e78a374..e55bfc2e96c783d548588f05e00f87ae448bf1f5 100644 (file)
--- 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 9b4b6ce3d095b3064bb4cb28d4c57553656e41b3..385e3732347524c515634f849de7f7bee0ef6d8f 100644 (file)
--- 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 c981380781c05dd831744d837a8fd2743a6c6a0b..7a6f55547ed56019d32d51b156db1cee2b2f2d95 100644 (file)
--- 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 d266a8d5d64335e08e78ccf3b21a101aa6faaa39..4aece5e2e4b155a9aa21e26d52fbad758150fd8e 100644 (file)
--- 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 570b9ab0406caab79f8dce9a0fec54703a8ec810..71ffbb004deb3967913906ebaa1fa18fc1fd9197 100644 (file)
--- 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, <libpq.ExecStatusType>exec_status)

diff --git a/tests/pq/test_pgconn.py b/tests/pq/test_pgconn.py
index 2fcfaee2e75a35a9ea0feb5ad698d69a1aa36d60..2dd34875199c9d8aa24f3092fcb5a994aad3213d 100644 (file)
--- 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()