From 2446e11806edb9d1b130c649f768f99740c89451 Mon Sep 17 00:00:00 2001
From: hamdykhader <hamdy.khader@gmail.com>
Date: Fri, 19 Jun 2026 00:40:32 +0300
Subject: [PATCH] fix: add backup encryption key management to KMS and
 integrate with backup workflows

---
 .../controllers/backup_controller.py          | 30 +++++++++++++++++--
 simplyblock_core/kms/_base.py                 |  9 ++++++
 simplyblock_core/models/backup.py             |  1 +
 3 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/simplyblock_core/controllers/backup_controller.py b/simplyblock_core/controllers/backup_controller.py
index 52410ce2ac..df9755bcad 100644
--- a/simplyblock_core/controllers/backup_controller.py
+++ b/simplyblock_core/controllers/backup_controller.py
@@ -9,6 +9,7 @@
 
 from simplyblock_core.controllers import backup_events, tasks_controller
 from simplyblock_core.db_controller import DBController
+from simplyblock_core.kms import create_kms_connection, KMSException
 from simplyblock_core.models.backup import Backup, BackupPolicy, BackupPolicyAttachment
 from simplyblock_core.models.storage_node import StorageNode
 
@@ -283,6 +284,19 @@ def _create_single_backup(snapshot, lvol, node_id, cluster_id, prev_backup):
     backup.allowed_hosts = lvol.allowed_hosts
     backup.created_at = int(time.time())
     backup.status = Backup.STATUS_PENDING
+    backup.is_encrypted = bool(lvol.crypto_bdev)
+    if lvol.crypto_bdev:
+        cl = db_controller.get_cluster_by_id(cluster_id)
+        with create_kms_connection(cl) as kms:
+            try:
+                crypto_key = kms.get_backup_encryption_keys(backup)
+                if not crypto_key:
+                    kms.create_backup_encryption_keys(backup)
+                    logger.info("Created backup lvol keys")
+            except KMSException as e:
+                logger.error("Failed to create backup lvol keys")
+                raise e
+
     backup.write_to_db()
 
     _write_s3_metadata(None, backup)
@@ -447,13 +461,23 @@ def restore_backup(backup_id, lvol_name, pool_id_or_name, cluster_id=None,
     if not target_node.lvstore:
         return None, f"Target node {restore_node_id} has no lvstore (S3 bdev requires lvstore)"
 
-    original_lvol = db_controller.get_lvol_by_id(backup.lvol_id)
+    crypto_key = ()
+    if backup.is_encrypted:
+        cluster = db_controller.get_cluster_by_id(backup.cluster_id)
+        with create_kms_connection(cluster) as kms:
+            try:
+                crypto_key = kms.get_backup_encryption_keys(backup)
+            except KMSException:
+                msg = "Failed retrieve backup encryption keys"
+                logger.exception(msg)
+                return False, msg
+
     logger.info(f"Backup allowed hosts: {backup.allowed_hosts}")
     lvol_id, error = lvol_controller.add_lvol_ha(
         name=lvol_name,
         size=size,
         pool_id_or_name=pool_id_or_name,
-        use_crypto=bool(original_lvol.crypto_bdev),
+        use_crypto=bool(crypto_key),
         max_size=0,
         max_rw_iops=0,
         max_rw_mbytes=0,
@@ -461,7 +485,7 @@ def restore_backup(backup_id, lvol_name, pool_id_or_name, cluster_id=None,
         max_w_mbytes=0,
         host_id_or_name=restore_node_id,
         ha_type="default",
-        crypto_key=(original_lvol.crypto_key1, original_lvol.crypto_key2),
+        crypto_key=crypto_key,
         use_comp=False,
         distr_vuid=0,
         lvol_priority_class=0,
diff --git a/simplyblock_core/kms/_base.py b/simplyblock_core/kms/_base.py
index bb915d5829..e704b8334a 100644
--- a/simplyblock_core/kms/_base.py
+++ b/simplyblock_core/kms/_base.py
@@ -7,6 +7,7 @@
 
 if TYPE_CHECKING:
     from simplyblock_core.models.lvol_model import LVol
+    from simplyblock_core.models.backup import Backup
 
 
 class KMS(AbstractContextManager):
@@ -38,3 +39,11 @@ def create_key_encryption_key(self, name: str) -> None: ...
 
     @abstractmethod
     def delete_key_encryption_key(self, name: str) -> None: ...
+
+    @abstractmethod
+    def create_backup_encryption_keys(self, backup: "Backup") -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_backup_encryption_keys(self, backup: "Backup") -> tuple[str, str]:
+        pass
diff --git a/simplyblock_core/models/backup.py b/simplyblock_core/models/backup.py
index 54d2419e2b..a8d800c8e0 100644
--- a/simplyblock_core/models/backup.py
+++ b/simplyblock_core/models/backup.py
@@ -43,6 +43,7 @@ class Backup(BaseModel):
     allowed_hosts: List[dict] = []
     # S3 metadata written to metadata bucket
     s3_metadata: dict = {}
+    is_encrypted: bool = False
 
     def get_id(self):
         return "%s/%s" % (self.cluster_id, self.uuid)