ソースを参照

Merge pull request #16132 from rndmcnlly/feature/sqlcipher-database-encryption

feat: Implement SQLCipher support for database encryption
Tim Jaeryang Baek 1 ヶ月 前
コミット
86fa564b44

+ 3 - 0
backend/open_webui/env.py

@@ -288,6 +288,9 @@ DB_VARS = {
 
 if all(DB_VARS.values()):
     DATABASE_URL = f"{DB_VARS['db_type']}://{DB_VARS['db_cred']}@{DB_VARS['db_host']}:{DB_VARS['db_port']}/{DB_VARS['db_name']}"
+elif DATABASE_TYPE == "sqlite+sqlcipher" and not os.environ.get("DATABASE_URL"):
+    # Handle SQLCipher with local file when DATABASE_URL wasn't explicitly set
+    DATABASE_URL = f"sqlite+sqlcipher:///{DATA_DIR}/webui.db"
 
 # Replace the postgres:// with postgresql://
 if "postgres://" in DATABASE_URL:

+ 29 - 1
backend/open_webui/internal/db.py

@@ -1,3 +1,4 @@
+import os
 import json
 import logging
 from contextlib import contextmanager
@@ -79,7 +80,34 @@ handle_peewee_migration(DATABASE_URL)
 
 
 SQLALCHEMY_DATABASE_URL = DATABASE_URL
-if "sqlite" in SQLALCHEMY_DATABASE_URL:
+
+# Handle SQLCipher URLs
+if SQLALCHEMY_DATABASE_URL.startswith('sqlite+sqlcipher://'):
+    database_password = os.environ.get("DATABASE_PASSWORD")
+    if not database_password or database_password.strip() == "":
+        raise ValueError("DATABASE_PASSWORD is required when using sqlite+sqlcipher:// URLs")
+    
+    # Extract database path from SQLCipher URL
+    db_path = SQLALCHEMY_DATABASE_URL.replace('sqlite+sqlcipher://', '')
+    if db_path.startswith('/'):
+        db_path = db_path[1:]  # Remove leading slash for relative paths
+    
+    # Create a custom creator function that uses sqlcipher3
+    def create_sqlcipher_connection():
+        import sqlcipher3
+        conn = sqlcipher3.connect(db_path, check_same_thread=False)
+        conn.execute(f"PRAGMA key = '{database_password}'")
+        return conn
+    
+    engine = create_engine(
+        "sqlite://",  # Dummy URL since we're using creator
+        creator=create_sqlcipher_connection,
+        echo=False
+    )
+    
+    log.info("Connected to encrypted SQLite database using SQLCipher")
+
+elif "sqlite" in SQLALCHEMY_DATABASE_URL:
     engine = create_engine(
         SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False}
     )

+ 38 - 16
backend/open_webui/internal/wrappers.py

@@ -1,4 +1,5 @@
 import logging
+import os
 from contextvars import ContextVar
 
 from open_webui.env import SRC_LOG_LEVELS
@@ -7,6 +8,7 @@ from peewee import InterfaceError as PeeWeeInterfaceError
 from peewee import PostgresqlDatabase
 from playhouse.db_url import connect, parse
 from playhouse.shortcuts import ReconnectMixin
+from playhouse.sqlcipher_ext import SqlCipherDatabase
 
 log = logging.getLogger(__name__)
 log.setLevel(SRC_LOG_LEVELS["DB"])
@@ -43,24 +45,44 @@ class ReconnectingPostgresqlDatabase(CustomReconnectMixin, PostgresqlDatabase):
 
 
 def register_connection(db_url):
-    db = connect(db_url, unquote_user=True, unquote_password=True)
-    if isinstance(db, PostgresqlDatabase):
-        # Enable autoconnect for SQLite databases, managed by Peewee
+    # Check if using SQLCipher protocol
+    if db_url.startswith('sqlite+sqlcipher://'):
+        database_password = os.environ.get("DATABASE_PASSWORD")
+        if not database_password or database_password.strip() == "":
+            raise ValueError("DATABASE_PASSWORD is required when using sqlite+sqlcipher:// URLs")
+        
+        # Parse the database path from SQLCipher URL
+        # Convert sqlite+sqlcipher:///path/to/db.sqlite to /path/to/db.sqlite
+        db_path = db_url.replace('sqlite+sqlcipher://', '')
+        if db_path.startswith('/'):
+            db_path = db_path[1:]  # Remove leading slash for relative paths
+        
+        # Use Peewee's native SqlCipherDatabase with encryption
+        db = SqlCipherDatabase(db_path, passphrase=database_password)
         db.autoconnect = True
         db.reuse_if_open = True
-        log.info("Connected to PostgreSQL database")
+        log.info("Connected to encrypted SQLite database using SQLCipher")
+        
+    else:
+        # Standard database connection (existing logic)
+        db = connect(db_url, unquote_user=True, unquote_password=True)
+        if isinstance(db, PostgresqlDatabase):
+            # Enable autoconnect for SQLite databases, managed by Peewee
+            db.autoconnect = True
+            db.reuse_if_open = True
+            log.info("Connected to PostgreSQL database")
 
-        # Get the connection details
-        connection = parse(db_url, unquote_user=True, unquote_password=True)
+            # Get the connection details
+            connection = parse(db_url, unquote_user=True, unquote_password=True)
 
-        # Use our custom database class that supports reconnection
-        db = ReconnectingPostgresqlDatabase(**connection)
-        db.connect(reuse_if_open=True)
-    elif isinstance(db, SqliteDatabase):
-        # Enable autoconnect for SQLite databases, managed by Peewee
-        db.autoconnect = True
-        db.reuse_if_open = True
-        log.info("Connected to SQLite database")
-    else:
-        raise ValueError("Unsupported database connection")
+            # Use our custom database class that supports reconnection
+            db = ReconnectingPostgresqlDatabase(**connection)
+            db.connect(reuse_if_open=True)
+        elif isinstance(db, SqliteDatabase):
+            # Enable autoconnect for SQLite databases, managed by Peewee
+            db.autoconnect = True
+            db.reuse_if_open = True
+            log.info("Connected to SQLite database")
+        else:
+            raise ValueError("Unsupported database connection")
     return db

+ 31 - 7
backend/open_webui/migrations/env.py

@@ -2,8 +2,8 @@ from logging.config import fileConfig
 
 from alembic import context
 from open_webui.models.auths import Auth
-from open_webui.env import DATABASE_URL
-from sqlalchemy import engine_from_config, pool
+from open_webui.env import DATABASE_URL, DATABASE_PASSWORD
+from sqlalchemy import engine_from_config, pool, create_engine
 
 # this is the Alembic Config object, which provides
 # access to the values within the .ini file in use.
@@ -62,11 +62,35 @@ def run_migrations_online() -> None:
     and associate a connection with the context.
 
     """
-    connectable = engine_from_config(
-        config.get_section(config.config_ini_section, {}),
-        prefix="sqlalchemy.",
-        poolclass=pool.NullPool,
-    )
+    # Handle SQLCipher URLs
+    if DB_URL and DB_URL.startswith('sqlite+sqlcipher://'):
+        if not DATABASE_PASSWORD or DATABASE_PASSWORD.strip() == "":
+            raise ValueError("DATABASE_PASSWORD is required when using sqlite+sqlcipher:// URLs")
+        
+        # Extract database path from SQLCipher URL
+        db_path = DB_URL.replace('sqlite+sqlcipher://', '')
+        if db_path.startswith('/'):
+            db_path = db_path[1:]  # Remove leading slash for relative paths
+        
+        # Create a custom creator function that uses sqlcipher3
+        def create_sqlcipher_connection():
+            import sqlcipher3
+            conn = sqlcipher3.connect(db_path, check_same_thread=False)
+            conn.execute(f"PRAGMA key = '{DATABASE_PASSWORD}'")
+            return conn
+        
+        connectable = create_engine(
+            "sqlite://",  # Dummy URL since we're using creator
+            creator=create_sqlcipher_connection,
+            echo=False
+        )
+    else:
+        # Standard database connection (existing logic)
+        connectable = engine_from_config(
+            config.get_section(config.config_ini_section, {}),
+            prefix="sqlalchemy.",
+            poolclass=pool.NullPool,
+        )
 
     with connectable.connect() as connection:
         context.configure(connection=connection, target_metadata=target_metadata)

+ 1 - 0
backend/requirements.txt

@@ -20,6 +20,7 @@ sqlalchemy==2.0.38
 alembic==1.14.0
 peewee==3.18.1
 peewee-migrate==1.12.2
+sqlcipher3-wheels==0.5.4
 psycopg2-binary==2.9.9
 pgvector==0.4.0
 PyMySQL==1.1.1