restore from memory

backend/scripts/restore.py

@@ -4,6 +4,7 @@ import sys
 import subprocess
 import datetime
 import sqlite3
+import tempfile
 from pathlib import Path
 from huggingface_hub import HfApi, hf_hub_download
 
@@ -12,21 +13,17 @@ SCRIPT_DIR = Path(__file__).parent.resolve()
 BACKEND_DIR = SCRIPT_DIR.parent
 sys.path.append(str(BACKEND_DIR))
 
-# Database paths
+# Database path (actual application database)
 DATA_DIR = os.environ.get("DATA_DIR", "/app/backend/data")
 DB_FILE_PATH = os.path.join(DATA_DIR, "webui.db")
 
-# Restore paths
-RESTORE_BACKUP_DIR = os.environ.get("BACKUP_DIR", "/tmp/open_webui/db_backup")
-DB_GPG_PATH = os.path.join(RESTORE_BACKUP_DIR, "webui.db.gpg")
-TIMESTAMP_FILE_PATH = os.path.join(RESTORE_BACKUP_DIR, "last_backup_time.txt")
-
-# Hugging Face repo paths
-REPO_TIMESTAMP_FILE = "db_backup/last_backup_time.txt"
+# Hugging Face repo paths (virtual paths in HF storage)
 REPO_DB_GPG_FILE = "db_backup/webui.db.gpg"
+REPO_TIMESTAMP_FILE = "db_backup/last_backup_time.txt"
 
 
-def check_requirements():
+def check_gpg():
+    """Verify GPG is available."""
     try:
         subprocess.run(["gpg", "--version"], check=True, capture_output=True)
         return True
@@ -35,33 +32,33 @@ def check_requirements():
         return False
 
 
-def validate_secrets():
-    required_vars = ["BACKUP_PASSPHRASE", "HF_TOKEN", "SPACE_ID"]
-    missing = [var for var in required_vars if not os.environ.get(var)]
-    
+def validate_environment():
+    """Verify all required environment variables are set."""
+    required = ["BACKUP_PASSPHRASE", "HF_TOKEN", "SPACE_ID"]
+    missing = [var for var in required if not os.environ.get(var)]
     if missing:
-        print(f"Error: Missing required environment variables: {', '.join(missing)}")
+        print(f"Error: Missing environment variables: {', '.join(missing)}")
         return False
     return True
 
 
-def ensure_directories():
+def ensure_data_dir():
+    """Ensure the database directory exists."""
     try:
         os.makedirs(DATA_DIR, mode=0o755, exist_ok=True)
-        os.makedirs(RESTORE_BACKUP_DIR, mode=0o755, exist_ok=True)
         return True
     except Exception as e:
-        print(f"Error creating directories: {e}")
+        print(f"Error creating data directory: {e}")
         return False
 
 
 def get_latest_backup_info(repo_id, hf_token):
+    """Check if backup exists and get its timestamp."""
     api = HfApi()
     try:
         files = api.list_repo_files(repo_id=repo_id, repo_type="space", token=hf_token)
         
-        backup_exists = (REPO_DB_GPG_FILE in files)
-        if not backup_exists:
+        if REPO_DB_GPG_FILE not in files:
             print("No backup file found in repository")
             return False, None
         
@@ -87,52 +84,59 @@ def get_latest_backup_info(repo_id, hf_token):
         return False, None
 
 
-def download_backup(repo_id, hf_token):
+def decrypt_database_from_memory(encrypted_data, passphrase):
+    """Decrypt database directly from memory."""
     try:
-        print("Downloading encrypted database backup...")
-        temp_file = hf_hub_download(
-            repo_id=repo_id,
-            repo_type="space",
-            filename=REPO_DB_GPG_FILE,
-            token=hf_token
-        )
-        os.replace(temp_file, DB_GPG_PATH)
-        print("Backup downloaded successfully")
-        return True
-    except Exception as e:
-        print(f"Error downloading backup: {e}")
-        return False
-
+        # Create a secure temporary directory for GPG operations
+        with tempfile.TemporaryDirectory(prefix='gpg_home_') as gpg_home:
+            os.chmod(gpg_home, 0o700)
+            
+            # Create a temporary file for the encrypted data
+            with tempfile.NamedTemporaryFile(mode='wb', suffix='.gpg', delete=False) as temp_encrypted:
+                temp_encrypted.write(encrypted_data)
+                temp_encrypted_path = temp_encrypted.name
 
-def decrypt_database(passphrase):
-    if not os.path.exists(DB_GPG_PATH):
-        print("No encrypted backup found locally. Starting with fresh DB")
-        return True
-    
-    try:
-        print("Decrypting database...")
-        os.makedirs("/root/.gnupg", mode=0o700, exist_ok=True)
-        
-        decrypt_cmd = [
-            "gpg",
-            "--batch",
-            "--yes",
-            "--passphrase", passphrase,
-            "-d",
-            "-o", DB_FILE_PATH,
-            DB_GPG_PATH
-        ]
-        subprocess.run(decrypt_cmd, check=True, stderr=subprocess.PIPE)
-        print(f"Database decrypted successfully to {DB_FILE_PATH}")
-        return True
+            try:
+                print(f"Decrypting database ({len(encrypted_data)} bytes)...")
+                decrypt_cmd = [
+                    "gpg",
+                    "--batch",
+                    "--yes",
+                    "--homedir", gpg_home,
+                    "--passphrase", passphrase,
+                    "--pinentry-mode", "loopback",
+                    "-d",
+                    "-o", DB_FILE_PATH,
+                    temp_encrypted_path
+                ]
+                
+                result = subprocess.run(decrypt_cmd, capture_output=True, check=True)
+                
+                if os.path.exists(DB_FILE_PATH) and os.path.getsize(DB_FILE_PATH) > 0:
+                    print(f"Database decrypted successfully ({os.path.getsize(DB_FILE_PATH)} bytes)")
+                    return True
+                else:
+                    print("Error: Decrypted database is missing or empty")
+                    return False
+                    
+            finally:
+                # Clean up the temporary encrypted file
+                if os.path.exists(temp_encrypted_path):
+                    os.unlink(temp_encrypted_path)
+                    
     except subprocess.CalledProcessError as e:
-        print(f"Decryption failed: {e.stderr.decode(errors='ignore')}")
+        print(f"Decryption failed: {e.stderr.decode()}")
+        return False
+    except Exception as e:
+        print(f"Decryption error: {e}")
         return False
 
 
 def verify_database():
+    """Verify the restored database integrity."""
     if not os.path.exists(DB_FILE_PATH):
-        return True
+        print(f"Error: Database file not found at {DB_FILE_PATH}")
+        return False
     
     try:
         print("Verifying database integrity...")
@@ -146,42 +150,60 @@ def verify_database():
             if result.lower() == "ok" and len(tables) > 0:
                 print(f"Database verified: {len(tables)} tables found")
                 return True
-            else:
-                print("Database verification failed")
-                return False
+            print("Database verification failed")
+            return False
     except Exception as e:
         print(f"Database verification error: {e}")
         return False
 
 
 def restore_db():
-    if not check_requirements() or not validate_secrets():
+    """Main restore function using in-memory approach."""
+    if not check_gpg() or not validate_environment() or not ensure_data_dir():
         return False
-    
-    if not ensure_directories():
-        return False
-    
+
     passphrase = os.environ["BACKUP_PASSPHRASE"]
     hf_token = os.environ["HF_TOKEN"]
     space_id = os.environ["SPACE_ID"]
     
     backup_exists, timestamp = get_latest_backup_info(space_id, hf_token)
-    if backup_exists:
-        if not download_backup(space_id, hf_token):
-            return False
+    if not backup_exists:
+        print("No backup found - starting with fresh database")
+        return True
+
+    try:
+        print("Downloading encrypted database...")
+        encrypted_file = hf_hub_download(
+            repo_id=space_id,
+            repo_type="space",
+            filename=REPO_DB_GPG_FILE,
+            token=hf_token
+        )
+        
+        # Read encrypted data into memory
+        with open(encrypted_file, 'rb') as f:
+            encrypted_data = f.read()
+        
+        print(f"Downloaded encrypted data: {len(encrypted_data)} bytes")
         
-        if not decrypt_database(passphrase):
+        if not decrypt_database_from_memory(encrypted_data, passphrase):
+            print("Failed to decrypt database")
             return False
         
         if not verify_database():
+            print("Failed to verify database")
             if os.path.exists(DB_FILE_PATH):
                 os.unlink(DB_FILE_PATH)
             return False
-    else:
-        print("No backup found - starting with fresh database")
-    
-    print("Database restore completed successfully!")
-    return True
+        
+        print("Database restore completed successfully!")
+        return True
+        
+    except Exception as e:
+        print(f"Restore error: {e}")
+        if os.path.exists(DB_FILE_PATH):
+            os.unlink(DB_FILE_PATH)
+        return False
 
 
 if __name__ == "__main__":