From 513803c230a2821a2369910944a7aa0750fa3d8a Mon Sep 17 00:00:00 2001
From: Charles Leifer <coleifer@gmail.com>
Date: Mon, 12 Feb 2024 16:01:31 -0600
Subject: [PATCH] Add support for set_key() and reset_key() encryption methods.

These use the library sqlite3_key and sqlite3_rekey if available.

Refs #70
---
 src/connection.c | 56 ++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 56 insertions(+)

diff --git a/src/connection.c b/src/connection.c
index 99c8e26..e4badeb 100644
--- a/src/connection.c
+++ b/src/connection.c
@@ -52,6 +52,10 @@
 #define HAVE_WINDOW_FUNCTION
 #endif
 
+#if SQLITE_HAS_CODEC
+#define HAVE_ENCRYPTION
+#endif
+
 _Py_IDENTIFIER(cursor);
 
 static const char * const begin_statements[] = {
@@ -1982,6 +1986,52 @@ pysqlite_connection_exit(pysqlite_Connection* self, PyObject* args)
     Py_RETURN_FALSE;
 }
 
+#ifdef HAVE_ENCRYPTION
+PyObject* pysqlite_connection_key(pysqlite_Connection *self, PyObject *args)
+{
+    Py_buffer key_buffer;
+    int rc;
+
+    if (!pysqlite_check_connection(self)) {
+        return NULL;
+    }
+    if (!PyArg_ParseTuple(args, "s*", &key_buffer)) {
+        return NULL;
+    }
+
+    rc = sqlite3_key(self->db, key_buffer.buf, key_buffer.len);
+    PyBuffer_Release(&key_buffer);
+
+    if (rc != SQLITE_OK) {
+        PyErr_SetString(pysqlite_OperationalError, sqlite3_errstr(rc));
+        return NULL;
+    }
+    Py_RETURN_NONE;
+}
+
+PyObject* pysqlite_connection_rekey(pysqlite_Connection *self, PyObject *args)
+{
+    Py_buffer key_buffer;
+    int rc;
+
+    if (!pysqlite_check_connection(self)) {
+        return NULL;
+    }
+    if (!PyArg_ParseTuple(args, "s*", &key_buffer)) {
+        return NULL;
+    }
+
+    rc = sqlite3_rekey(self->db, key_buffer.buf, key_buffer.len);
+    PyBuffer_Release(&key_buffer);
+
+    if (rc != SQLITE_OK) {
+        PyErr_SetString(pysqlite_OperationalError, sqlite3_errstr(rc));
+        return NULL;
+    }
+    Py_RETURN_NONE;
+}
+#endif
+
 static const char connection_doc[] =
 PyDoc_STR("SQLite database connection object.");
 
@@ -2033,6 +2083,12 @@ static PyMethodDef connection_methods[] = {
         PyDoc_STR("Creates a collation function. Non-standard.")},
     {"interrupt", (PyCFunction)pysqlite_connection_interrupt, METH_NOARGS,
         PyDoc_STR("Abort any pending database operation. Non-standard.")},
+#ifdef HAVE_ENCRYPTION
+    {"set_key", (PyCFunction)(void(*)(void))pysqlite_connection_key, METH_VARARGS,
+        PyDoc_STR("Set encryption key for database. Non-standard.")},
+    {"reset_key", (PyCFunction)(void(*)(void))pysqlite_connection_rekey, METH_VARARGS,
+        PyDoc_STR("Set encryption key for database. Non-standard.")},
+#endif
     #ifdef HAVE_BACKUP_API
     {"backup", (PyCFunction)(void(*)(void))pysqlite_connection_backup, METH_VARARGS | METH_KEYWORDS,
         PyDoc_STR("Makes a backup of the database. Non-standard.")},