From eec0ebb00e57eb2e04c6ec629c7005d74a7776b6 Mon Sep 17 00:00:00 2001
From: PankajJoshi <pankajjoshi@microsoft.com>
Date: Wed, 29 Nov 2023 10:27:40 +0530
Subject: [PATCH 01/16] enabling azure_config_param.ini file for CVM. it was
 not created for CVM.

---
 VMEncryption/main/handle.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/VMEncryption/main/handle.py b/VMEncryption/main/handle.py
index 479fc00b3..a55b1a1d5 100644
--- a/VMEncryption/main/handle.py
+++ b/VMEncryption/main/handle.py
@@ -185,15 +185,16 @@ def disable_encryption():
 
 
 def stamp_disks_with_settings(items_to_encrypt, encryption_config, encryption_marker=None):
-    if security_Type == CommonVariables.ConfidentialVM:
-        logger.log(msg="Do not send vm setting to host for stamping.",level=CommonVariables.InfoLevel)
-        return
     disk_util = DiskUtil(hutil=hutil, patching=DistroPatcher, logger=logger, encryption_environment=encryption_environment)
     crypt_mount_config_util = CryptMountConfigUtil(logger=logger, encryption_environment=encryption_environment, disk_util=disk_util)
     bek_util = BekUtil(disk_util, logger,encryption_environment)
     current_passphrase_file = bek_util.get_bek_passphrase_file(encryption_config)
     public_settings = get_public_settings()
     extension_parameter = ExtensionParameter(hutil, logger, DistroPatcher, encryption_environment, get_protected_settings(), public_settings)
+    if security_Type == CommonVariables.ConfidentialVM:
+        logger.log(msg="Do not send vm setting to host for stamping.",level=CommonVariables.InfoLevel)
+        extension_parameter.commit()
+        return
     has_keystore_flag = CommonVariables.KeyStoreTypeKey in public_settings
 
     # post new encryption settings via wire server protocol

From 8397975ca05d283dba56be2fb01d965fa5cc4116 Mon Sep 17 00:00:00 2001
From: PankajJoshi <pankajjoshi@microsoft.com>
Date: Thu, 30 Nov 2023 15:25:14 +0530
Subject: [PATCH 02/16] update encryption settings with wrapped passphrase with
 new CMK for KEK rotation for CVM.

---
 VMEncryption/main/Common.py             |  3 +-
 VMEncryption/main/DiskUtil.py           | 22 +++++---
 VMEncryption/main/ExtensionParameter.py |  7 +++
 VMEncryption/main/handle.py             | 71 ++++++++++++++++++++++---
 4 files changed, 87 insertions(+), 16 deletions(-)

diff --git a/VMEncryption/main/Common.py b/VMEncryption/main/Common.py
index bae6e7fb5..a45753800 100644
--- a/VMEncryption/main/Common.py
+++ b/VMEncryption/main/Common.py
@@ -62,7 +62,8 @@ class CommonVariables:
     """
     cvm_ade_vm_encryption_token_id = 5
     ADEEncryptionVersionInLuksToken_1_0='1.0'
-    PassphraseNameValue = 'LUKSPasswordProtector'
+    PassphraseNameValueProtected = 'LUKSPasswordProtector'
+    PassphraseNameValueNotProtected = 'LUKSPasswordNotProtector'
     """
     IMDS IP:
     """
diff --git a/VMEncryption/main/DiskUtil.py b/VMEncryption/main/DiskUtil.py
index 5de5e9b69..bfb866359 100644
--- a/VMEncryption/main/DiskUtil.py
+++ b/VMEncryption/main/DiskUtil.py
@@ -193,21 +193,25 @@ def secure_key_release_operation(self,protectorbase64,kekUrl,operation,attestati
         self.logger.log("secure_key_release_operation {0} end.".format(operation))
         return process_comm.stdout.strip()
     
-    def import_token(self,device_path,passphrase_file,public_settings):
+    def import_token(self,device_path,passphrase_file,public_settings,PassphraseNameValue=CommonVariables.PassphraseNameValueProtected):
         '''this function reads passphrase from passphrase file, wrap it and update in token field of LUKS2 header.'''
         self.logger.log(msg="import_token for device: {0} started.".format(device_path))
-        protector = ""
+        Protector= ""
         with open(passphrase_file,"rb") as protector_file:
             #passphrase stored in keyfile is base64
-            protector = protector_file.read().decode('utf-8')
+            Protector = protector_file.read().decode('utf-8')
         KekVaultResourceId=public_settings.get(CommonVariables.KekVaultResourceIdKey)
         KeyEncryptionKeyUrl=public_settings.get(CommonVariables.KeyEncryptionKeyURLKey)
         AttestationUrl = public_settings.get(CommonVariables.AttestationURLKey)
-        wrappedProtector = self.secure_key_release_operation(protectorbase64=protector,
+        if PassphraseNameValue == CommonVariables.PassphraseNameValueProtected:
+            Protector = self.secure_key_release_operation(protectorbase64=Protector,
                                                         kekUrl=KeyEncryptionKeyUrl,
                                                         operation=CommonVariables.secure_key_release_wrap,
                                                         attestationUrl=AttestationUrl)
-        if not wrappedProtector:
+        else:
+            self.logger.log(msg="import_token passphrase is not wrapped, value of passphrase name key: {0}".format(PassphraseNameValue))
+
+        if not Protector:
             self.logger.log("import_token protector wrapping is unsuccessful for device {0}".format(device_path))
             return False
         data={
@@ -219,10 +223,9 @@ def import_token(self,device_path,passphrase_file,public_settings):
             CommonVariables.KeyVaultResourceIdKey:public_settings.get(CommonVariables.KeyVaultResourceIdKey),
             CommonVariables.KeyVaultURLKey:public_settings.get(CommonVariables.KeyVaultURLKey),
             CommonVariables.AttestationURLKey:AttestationUrl,
-            CommonVariables.PassphraseNameKey:CommonVariables.PassphraseNameValue,
-            CommonVariables.PassphraseKey:wrappedProtector
+            CommonVariables.PassphraseNameKey:PassphraseNameValue,
+            CommonVariables.PassphraseKey:Protector
         }
-        #TODO: needed to decide on temp path.
         custom_cmk = os.path.join("/var/lib/azure_disk_encryption_config/","custom_cmk.json")
         out_file = open(custom_cmk,"w")
         json.dump(data,out_file,indent=4)
@@ -254,6 +257,9 @@ def export_token(self,device_name):
         keyEncryptionKeyUrl=disk_encryption_setting[CommonVariables.KeyEncryptionKeyURLKey]
         wrappedProtector = disk_encryption_setting[CommonVariables.PassphraseKey]
         attestationUrl = disk_encryption_setting[CommonVariables.AttestationURLKey]
+        if disk_encryption_setting[CommonVariables.PassphraseNameKey] != CommonVariables.PassphraseNameValueProtected:
+            self.logger.log("passphrase is not Protectected. No need to do SKR.")
+            return wrappedProtector if wrappedProtector else None
         if wrappedProtector:
             #unwrap the protector.
             protector=self.secure_key_release_operation(attestationUrl=attestationUrl,
diff --git a/VMEncryption/main/ExtensionParameter.py b/VMEncryption/main/ExtensionParameter.py
index f1c4a0cf4..09e883bba 100644
--- a/VMEncryption/main/ExtensionParameter.py
+++ b/VMEncryption/main/ExtensionParameter.py
@@ -174,6 +174,13 @@ def _is_kv_equivalent(self, a, b):
             if b[-1] == '/': b = b[:-1]
         return a==b
 
+    def cmk_changed(self):
+        if (self.KeyEncryptionKeyURL or self.get_kek_url()) and \
+           (not self._is_kv_equivalent(self.KeyEncryptionKeyURL, self.get_kek_url())):
+           self.logger.log('Current config KeyEncryptionKeyURL {0} differs from effective config KeyEncryptionKeyURL {1}'.format(self.KeyEncryptionKeyURL, self.get_kek_url()))
+           return True
+        return False
+     
     def config_changed(self):
         if (self.command or self.get_command()) and \
            (self.command != self.get_command() and \
diff --git a/VMEncryption/main/handle.py b/VMEncryption/main/handle.py
index a55b1a1d5..6f82063eb 100644
--- a/VMEncryption/main/handle.py
+++ b/VMEncryption/main/handle.py
@@ -271,6 +271,56 @@ def get_protected_settings():
     else:
         return protected_settings_str
 
+def update_encryption_settings_luks2_header():
+    '''This function is used for CMK passphrse wrapping with new KEK URL and update metadata in LUKS2 header.'''
+    hutil.do_parse_context('UpdateEncryptionSettingsLuks2Header')
+    logger.log('Updating encryption settings LUKS-2 header')
+    # ensure cryptsetup package is still available in case it was for some reason removed after enable
+    try:
+        DistroPatcher.install_cryptsetup()
+    except Exception as e:
+        hutil.save_seq()
+        message = "Failed to update encryption settings with error: {0}, stack trace: {1}".format(e, traceback.format_exc())
+        hutil.do_exit(exit_code=CommonVariables.missing_dependency,
+                      operation='UpdateEncryptionSettingsLuks2Header',
+                      status=CommonVariables.extension_error_status,
+                      code=str(CommonVariables.missing_dependency),
+                      message=message)
+    try:
+        public_setting = get_public_settings()
+        encryption_config = EncryptionConfig(encryption_environment, logger)
+        extension_parameter = ExtensionParameter(hutil, logger, DistroPatcher, encryption_environment, get_protected_settings(), public_setting)
+        disk_util = DiskUtil(hutil=hutil, patching=DistroPatcher, logger=logger, encryption_environment=encryption_environment)
+        bek_util = BekUtil(disk_util, logger,encryption_environment)
+        device_items = disk_util.get_device_items(None)
+        for device_item in device_items:
+            device_item_path = disk_util.get_device_path(device_item.name)
+            if not disk_util.is_luks_device(device_item_path,None):
+                logger.log("Not a LUKS device, device path: {0}".format(device_item_path))
+                continue
+            logger.log("Reading passphrase from LUKS2 header, device name: {0}".format(device_item.name))
+            #get the unwrapped passphrase from LUKS2 header.
+            passphrase=disk_util.export_token(device_name=device_item.name)
+            if not passphrase:
+                logger.log("No passphrase in LUKS2 header, device name: {0}".format(device_item.name))
+                continue
+            logger.log("Updating wrapped passphrase to LUKS2 header with current public setting. device name {0}".format(device_item.name))
+            #protect passphrase before updating to LUKS2 is done in import_token
+            ret = disk_util.import_token(device_item_path,passphrase,public_setting,CommonVariables.PassphraseNameValueProtected)
+            if not ret:
+                logger.log("Udate passphrase with current public setting to LUKS2 header is not successful. device path {0}".format(device_item_path))
+        extension_parameter.commit()
+        bek_util.umount_azure_passhprase(encryption_config)
+    except Exception as e:
+        hutil.save_seq()
+        message = "Failed to update encryption settings Luks2 header with error: {0}, stack trace: {1}".format(e, traceback.format_exc())
+        logger.log(msg=message, level=CommonVariables.ErrorLevel)
+        bek_util.umount_azure_passhprase(encryption_config)
+        hutil.do_exit(exit_code=CommonVariables.unknown_error,
+                      operation='UpdateEncryptionSettingsLuks2Header',
+                      status=CommonVariables.extension_error_status,
+                      code=str(CommonVariables.unknown_error),
+                      message=message)
 
 def update_encryption_settings(extra_items_to_encrypt=[]):
     hutil.do_parse_context('UpdateEncryptionSettings')
@@ -946,14 +996,21 @@ def handle_encryption(public_settings, encryption_status, disk_util, bek_util, e
             logger.log("An operation already running. Cannot accept an update settings request.")
             hutil.reject_settings()
         are_devices_encrypted, items_to_encrypt = are_required_devices_encrypted(volume_type, encryption_status, disk_util, bek_util, encryption_operation)
-        if not are_devices_encrypted:
-            logger.log('Required devices not encrypted for volume type {0}. Calling update to stamp encryption settings.'.format(volume_type))
-            update_encryption_settings(items_to_encrypt)
-            logger.log('Encryption Settings stamped. Calling enable to encrypt new devices.')
-            enable_encryption()
+        if security_Type==CommonVariables.ConfidentialVM:
+            logger.log('Calling Update Encryption Setting in LUKS2 header.')
+            if extension_parameter.cmk_changed():
+                update_encryption_settings_luks2_header()
+            if not are_devices_encrypted:
+                enable_encryption()
         else:
-            logger.log('Calling Update Encryption Setting.')
-            update_encryption_settings()
+            if not are_devices_encrypted:
+                logger.log('Required devices not encrypted for volume type {0}. Calling update to stamp encryption settings.'.format(volume_type))
+                update_encryption_settings(items_to_encrypt)
+                logger.log('Encryption Settings stamped. Calling enable to encrypt new devices.')
+                enable_encryption()
+            else:
+                logger.log('Calling Update Encryption Setting.')
+                update_encryption_settings()
     else:
         logger.log("Config did not change or first call, enabling encryption")
         encryption_marker = EncryptionMarkConfig(logger, encryption_environment)

From 278b36992c85f60c26ad8d59d17ffff4bbfe1948 Mon Sep 17 00:00:00 2001
From: PankajJoshi <pankajjoshi@microsoft.com>
Date: Sat, 9 Dec 2023 11:12:53 +0530
Subject: [PATCH 03/16] commit1

---
 VMEncryption/main/Common.py   |  3 ++
 VMEncryption/main/DiskUtil.py | 71 +++++++++++++++++++++++++++++++++++
 VMEncryption/main/handle.py   | 18 ++++++++-
 3 files changed, 90 insertions(+), 2 deletions(-)

diff --git a/VMEncryption/main/Common.py b/VMEncryption/main/Common.py
index a45753800..4ca411a39 100644
--- a/VMEncryption/main/Common.py
+++ b/VMEncryption/main/Common.py
@@ -61,9 +61,12 @@ class CommonVariables:
     CVM LUKS2 header token
     """
     cvm_ade_vm_encryption_token_id = 5
+    cvm_ade_vm_encryption_backup_token_id = 6
     ADEEncryptionVersionInLuksToken_1_0='1.0'
     PassphraseNameValueProtected = 'LUKSPasswordProtector'
     PassphraseNameValueNotProtected = 'LUKSPasswordNotProtector'
+    AzureDiskEncryptionToken = 'Azure_Disk_Encryption'
+    AzureDiskEncryptionBackUpToken='Azure_Disk_Encryption_BackUp'
     """
     IMDS IP:
     """
diff --git a/VMEncryption/main/DiskUtil.py b/VMEncryption/main/DiskUtil.py
index bfb866359..a3544e888 100644
--- a/VMEncryption/main/DiskUtil.py
+++ b/VMEncryption/main/DiskUtil.py
@@ -24,6 +24,7 @@
 from subprocess import Popen
 import traceback
 import glob
+import tempfile
 
 from EncryptionConfig import EncryptionConfig
 from DecryptionMarkConfig import DecryptionMarkConfig
@@ -193,6 +194,18 @@ def secure_key_release_operation(self,protectorbase64,kekUrl,operation,attestati
         self.logger.log("secure_key_release_operation {0} end.".format(operation))
         return process_comm.stdout.strip()
     
+    def import_token(self,device_path,token_data,token_id):
+        self.logger.log(msg="import_token for device: {0} started.".format(device_path))
+        temp_file = tempfile.NamedTemporaryFile(delete=False)
+        temp_file.close()
+        json.dump(token_data,temp_file.name,indent=4)
+        cmd = "cryptsetup token import --json-file {0} --token-id {1} {2}".format(temp_file.name,token_id)
+        process_comm = ProcessCommunicator()
+        status = self.command_executor.Execute(cmd,communicator=process_comm)
+        self.logger.log(msg="import_token: device: {0} status: {1}".format(device_path,status))
+        os.unlink(temp_file.name)
+        return status==CommonVariables.process_success
+
     def import_token(self,device_path,passphrase_file,public_settings,PassphraseNameValue=CommonVariables.PassphraseNameValueProtected):
         '''this function reads passphrase from passphrase file, wrap it and update in token field of LUKS2 header.'''
         self.logger.log(msg="import_token for device: {0} started.".format(device_path))
@@ -226,6 +239,7 @@ def import_token(self,device_path,passphrase_file,public_settings,PassphraseName
             CommonVariables.PassphraseNameKey:PassphraseNameValue,
             CommonVariables.PassphraseKey:Protector
         }
+        #TODO handle with temp file.
         custom_cmk = os.path.join("/var/lib/azure_disk_encryption_config/","custom_cmk.json")
         out_file = open(custom_cmk,"w")
         json.dump(data,out_file,indent=4)
@@ -238,6 +252,29 @@ def import_token(self,device_path,passphrase_file,public_settings,PassphraseName
         self.logger.log(msg="import_token: device: {0} end.".format(device_path))
         return status==CommonVariables.process_success
     
+    def read_token(self,device_name,token_id):
+        '''this functions reads tokens from LUKS2 header.'''
+        device_path = os.path.join("/dev",device_name)
+        cmd = "cryptsetup token export --token-id {0} {1}".format(token_id,device_path)
+        process_comm = ProcessCommunicator()
+        status = self.command_executor.Execute(cmd, communicator=process_comm)
+        if status != 0:
+            self.logger.log("export_token token id {0} not found in device {1} LUKS header".format(CommonVariables.cvm_ade_vm_encryption_token_id,device_name))
+            return None
+        token = process_comm.stdout
+        return token
+
+    def remove_token(self,device_name,token_id):
+        '''this function remove the token'''
+        device_path = os.path.join("/dev",device_name)
+        cmd = "cryptsetup token remove --token-id {0} {1}".format(token_id,device_path)
+        process_comm = ProcessCommunicator()
+        status = self.command_executor.Execute(cmd, communicator=process_comm)
+        if status != 0:
+            self.logger.log("remove token id {0} not found in device {1} LUKS header".format(CommonVariables.cvm_ade_vm_encryption_token_id,device_name))
+            return False
+        return True
+
     def export_token(self,device_name):
         '''This function reads token id from luks2 header field and unwrap passphrase'''
         self.logger.log("export_token to device {0} started.".format(device_name))
@@ -403,6 +440,15 @@ def luks_get_uuid(self, header_or_dev_path):
                 return splits[1]
         return None
 
+    def get_token_id(self,header_or_dev_path,token_name):
+        '''if LUKS2 header has token name return the id else return none.'''
+        luks_dump_out = self._luks_get_header_dump(header_or_dev_path)
+        tokens = self._extract_luksv2_token(luks_dump_out)
+        for token in tokens:
+            if token[1] is token_name:
+                return token[0]
+        return None
+
     def _get_cryptsetup_version(self):
         # get version of currently installed cryptsetup
         cryptsetup_cmd = "{0} --version".format(self.distro_patcher.cryptsetup_path)
@@ -416,6 +462,31 @@ def _extract_luks_version_from_dump(self, luks_dump_out):
             if "version:" in line.lower():
                 return line.split()[-1]
 
+    def _extract_luksv2_token(self, luks_dump_out):
+        """
+        ...
+        Tokens:
+            1: Azure_Disk_Encryption_BackUp
+            5: Azure_Disk_Encryption
+        ...
+        """
+        lines = luks_dump_out.split("\n")
+        token_segment = False
+        token_lines = []
+        for line in lines:
+            parts = line.split(":")
+            if len(parts)<2:
+                continue
+            if token_segment and parts[1].strip() is '':
+                break
+            if "tokens" in parts[0].strip().lower():
+                token_segment = True
+                continue
+            if token_segment and self._isnumeric(parts[0].strip()):
+                token_lines.append(parts)
+                continue
+        return token_lines
+
     def _extract_luksv2_keyslot_lines(self, luks_dump_out):
         """
         A luks v2 luksheader looks kind of like this: (inessential stuff removed)
diff --git a/VMEncryption/main/handle.py b/VMEncryption/main/handle.py
index 6f82063eb..d611c1d74 100644
--- a/VMEncryption/main/handle.py
+++ b/VMEncryption/main/handle.py
@@ -299,16 +299,29 @@ def update_encryption_settings_luks2_header():
                 logger.log("Not a LUKS device, device path: {0}".format(device_item_path))
                 continue
             logger.log("Reading passphrase from LUKS2 header, device name: {0}".format(device_item.name))
+            #keep token copy for manual recovery
+            data = disk_util.read_token(device_name=device_item.name,token_id=CommonVariables.cvm_ade_vm_encryption_token_id)
+            data['type']='Azure_Disk_Encryption_Backup'
+            disk_util.import_token(device_path=device_item.name,token_data=data,token_id=1)
             #get the unwrapped passphrase from LUKS2 header.
             passphrase=disk_util.export_token(device_name=device_item.name)
+            #remove token
+            disk_util.remove_token(device_name=device_item.name,token_id=CommonVariables.cvm_ade_vm_encryption_token_id) 
             if not passphrase:
                 logger.log("No passphrase in LUKS2 header, device name: {0}".format(device_item.name))
                 continue
             logger.log("Updating wrapped passphrase to LUKS2 header with current public setting. device name {0}".format(device_item.name))
             #protect passphrase before updating to LUKS2 is done in import_token
-            ret = disk_util.import_token(device_item_path,passphrase,public_setting,CommonVariables.PassphraseNameValueProtected)
+            temp_keyfile = tempfile.NamedTemporaryFile(delete=False)
+            temp_keyfile.write(passphrase.encode("utf-8"))
+            temp_keyfile.close()
+            ret = disk_util.import_token(device_item_path,temp_keyfile.name,public_setting,CommonVariables.PassphraseNameValueProtected)
             if not ret:
-                logger.log("Udate passphrase with current public setting to LUKS2 header is not successful. device path {0}".format(device_item_path))
+                logger.log("Update passphrase with current public setting to LUKS2 header is not successful. device path {0}".format(device_item_path))
+                return None
+            os.unlink(temp_keyfile.name)
+            #removing backup token
+            disk_util.remove_token(device_name=device_item.name,token_id=1)
         extension_parameter.commit()
         bek_util.umount_azure_passhprase(encryption_config)
     except Exception as e:
@@ -839,6 +852,7 @@ def enable():
                                   encryption_config=encryption_config,
                                   passphrase_file=generated_passphrase_file)
         if security_Type == CommonVariables.ConfidentialVM:
+            #[TODO]token update/cleaning
             crypt_mount_config_util.device_unlock_using_luks2_header()
 
         encryption_status = json.loads(disk_util.get_encryption_status())

From 8ccd55e78e7011ddf76d568f07319a500dd4a973 Mon Sep 17 00:00:00 2001
From: PankajJoshi <pankajjoshi@microsoft.com>
Date: Sat, 9 Dec 2023 13:00:39 +0530
Subject: [PATCH 04/16] commit 2

---
 VMEncryption/main/CryptMountConfigUtil.py |  2 ++
 VMEncryption/main/DiskUtil.py             | 27 +++++++++++++++++++++--
 VMEncryption/main/handle.py               | 21 ++++++++++++------
 3 files changed, 41 insertions(+), 9 deletions(-)

diff --git a/VMEncryption/main/CryptMountConfigUtil.py b/VMEncryption/main/CryptMountConfigUtil.py
index 7ae6e27ce..d122436e0 100644
--- a/VMEncryption/main/CryptMountConfigUtil.py
+++ b/VMEncryption/main/CryptMountConfigUtil.py
@@ -266,6 +266,8 @@ def device_unlock_using_luks2_header(self):
         lock = threading.Lock()
         for device_item in device_items:
             if device_item.file_system == "crypto_LUKS": 
+                 #resotre Luks2 token from BackUp (if needed)
+                 self.disk_util.restore_luks2_token(device_name=device_item.name)
                  device_item_path = self.disk_util.get_device_path(device_item.name)
                  azure_item_path = azure_name_table[device_item_path] if device_item_path in azure_name_table else device_item_path
                  thread = threading.Thread(target=self._device_unlock_using_luks2_header,args=(device_item.name,device_item_path,azure_item_path,lock))
diff --git a/VMEncryption/main/DiskUtil.py b/VMEncryption/main/DiskUtil.py
index a3544e888..1bf8d248f 100644
--- a/VMEncryption/main/DiskUtil.py
+++ b/VMEncryption/main/DiskUtil.py
@@ -194,7 +194,7 @@ def secure_key_release_operation(self,protectorbase64,kekUrl,operation,attestati
         self.logger.log("secure_key_release_operation {0} end.".format(operation))
         return process_comm.stdout.strip()
     
-    def import_token(self,device_path,token_data,token_id):
+    def import_token_data(self,device_path,token_data,token_id):
         self.logger.log(msg="import_token for device: {0} started.".format(device_path))
         temp_file = tempfile.NamedTemporaryFile(delete=False)
         temp_file.close()
@@ -278,7 +278,7 @@ def remove_token(self,device_name,token_id):
     def export_token(self,device_name):
         '''This function reads token id from luks2 header field and unwrap passphrase'''
         self.logger.log("export_token to device {0} started.".format(device_name))
-        device_path = os.path.join("/dev",device_name)
+        device_path = self.get_device_path(device_name)
         protector = None
         cmd = "cryptsetup token export --token-id {0} {1}".format(CommonVariables.cvm_ade_vm_encryption_token_id,device_path)
         process_comm = ProcessCommunicator()
@@ -449,6 +449,29 @@ def get_token_id(self,header_or_dev_path,token_name):
                 return token[0]
         return None
 
+    def restore_luks2_token(self, device_name=None):
+        '''this function restoring token type Azure_Disk_Encryption_BackUp to Azure_Disk_Encryption'''
+        if not device_name:
+            return
+        device_path = self.get_device_path(device_name)
+        ade_token_id = self.get_token_id(header_or_dev_path=device_path,token_name=CommonVariables.AzureDiskEncryptionToken)
+        ade_token_id_backup = self.get_token_id(header_or_dev_path=device_path,token_name=CommonVariables.AzureDiskEncryptionBackUpToken)
+        if not ade_token_id_backup:
+            #do nothing
+            return
+        if ade_token_id:
+            #remove backup token id
+            self.remove_token(device_name=device_name,token_id=ade_token_id_backup)
+            return
+        self.logger.log("resotre luks2 token for device {0} is started.".format(device_name))
+        #read from backup and update AzureDiskEncryptionToken
+        data = self.read_token(device_name=device_name,token_id=ade_token_id_backup)
+        data['type']=CommonVariables.AzureDiskEncryptionBackUpToken
+        self.import_token_data(device_path=device_path,token_data=data,token_id=CommonVariables.AzureDiskEncryptionToken)
+        #remove backup
+        self.remove_token(device_name=device_name,token_id=ade_token_id_backup)
+        self.logger.log("resotre luks2 token for device {0} is successful.".format(device_name))
+
     def _get_cryptsetup_version(self):
         # get version of currently installed cryptsetup
         cryptsetup_cmd = "{0} --version".format(self.distro_patcher.cryptsetup_path)
diff --git a/VMEncryption/main/handle.py b/VMEncryption/main/handle.py
index d611c1d74..7344320ac 100644
--- a/VMEncryption/main/handle.py
+++ b/VMEncryption/main/handle.py
@@ -298,15 +298,22 @@ def update_encryption_settings_luks2_header():
             if not disk_util.is_luks_device(device_item_path,None):
                 logger.log("Not a LUKS device, device path: {0}".format(device_item_path))
                 continue
+            #estoring the token data to type Azure_Disk_Encryption
+            disk_util.restore_luks2_token(device_name=device_item.name)
             logger.log("Reading passphrase from LUKS2 header, device name: {0}".format(device_item.name))
             #keep token copy for manual recovery
-            data = disk_util.read_token(device_name=device_item.name,token_id=CommonVariables.cvm_ade_vm_encryption_token_id)
-            data['type']='Azure_Disk_Encryption_Backup'
-            disk_util.import_token(device_path=device_item.name,token_data=data,token_id=1)
+            ade_token_id = disk_util.get_token_id(header_or_dev_path=device_item_path)
+            if not ade_token_id:
+                logger.log("token type: Azure_Disk_Encryption not found for device {0}".format(device_item.name))
+                continue
+            data = disk_util.read_token(device_name=device_item.name,token_id=ade_token_id)
+            #writing to backup token
+            data['type']=CommonVariables.AzureDiskEncryptionBackUpToken
+            disk_util.import_token_data(device_path=device_item_path,token_data=data,token_id=CommonVariables.cvm_ade_vm_encryption_backup_token_id)
             #get the unwrapped passphrase from LUKS2 header.
             passphrase=disk_util.export_token(device_name=device_item.name)
             #remove token
-            disk_util.remove_token(device_name=device_item.name,token_id=CommonVariables.cvm_ade_vm_encryption_token_id) 
+            disk_util.remove_token(device_name=device_item.name,token_id=ade_token_id)
             if not passphrase:
                 logger.log("No passphrase in LUKS2 header, device name: {0}".format(device_item.name))
                 continue
@@ -315,13 +322,14 @@ def update_encryption_settings_luks2_header():
             temp_keyfile = tempfile.NamedTemporaryFile(delete=False)
             temp_keyfile.write(passphrase.encode("utf-8"))
             temp_keyfile.close()
-            ret = disk_util.import_token(device_item_path,temp_keyfile.name,public_setting,CommonVariables.PassphraseNameValueProtected)
+            #save passphrase to LUKS2 header with PassphraseNameValueProtected
+            ret = disk_util.import_token(device_path=device_item_path,passphrase_file=temp_keyfile.name,public_settings=public_setting,PassphraseNameValue=CommonVariables.PassphraseNameValueProtected)
             if not ret:
                 logger.log("Update passphrase with current public setting to LUKS2 header is not successful. device path {0}".format(device_item_path))
                 return None
             os.unlink(temp_keyfile.name)
             #removing backup token
-            disk_util.remove_token(device_name=device_item.name,token_id=1)
+            disk_util.remove_token(device_name=device_item.name,token_id=CommonVariables.cvm_ade_vm_encryption_backup_token_id)
         extension_parameter.commit()
         bek_util.umount_azure_passhprase(encryption_config)
     except Exception as e:
@@ -852,7 +860,6 @@ def enable():
                                   encryption_config=encryption_config,
                                   passphrase_file=generated_passphrase_file)
         if security_Type == CommonVariables.ConfidentialVM:
-            #[TODO]token update/cleaning
             crypt_mount_config_util.device_unlock_using_luks2_header()
 
         encryption_status = json.loads(disk_util.get_encryption_status())

From 081d4bbc6808c93941b30c5401245abeddbe73c4 Mon Sep 17 00:00:00 2001
From: PankajJoshi <pankajjoshi@microsoft.com>
Date: Sun, 10 Dec 2023 20:35:37 +0530
Subject: [PATCH 05/16] test outcome corrections

---
 VMEncryption/main/DiskUtil.py | 12 ++++++------
 VMEncryption/main/handle.py   | 20 ++++++++++++++++----
 2 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/VMEncryption/main/DiskUtil.py b/VMEncryption/main/DiskUtil.py
index 1bf8d248f..539f5d055 100644
--- a/VMEncryption/main/DiskUtil.py
+++ b/VMEncryption/main/DiskUtil.py
@@ -196,10 +196,10 @@ def secure_key_release_operation(self,protectorbase64,kekUrl,operation,attestati
     
     def import_token_data(self,device_path,token_data,token_id):
         self.logger.log(msg="import_token for device: {0} started.".format(device_path))
-        temp_file = tempfile.NamedTemporaryFile(delete=False)
+        temp_file = tempfile.NamedTemporaryFile(delete=False,mode='w+')
+        json.dump(token_data,temp_file,indent=4)
         temp_file.close()
-        json.dump(token_data,temp_file.name,indent=4)
-        cmd = "cryptsetup token import --json-file {0} --token-id {1} {2}".format(temp_file.name,token_id)
+        cmd = "cryptsetup token import --json-file {0} --token-id {1} {2}".format(temp_file.name,token_id,device_path)
         process_comm = ProcessCommunicator()
         status = self.command_executor.Execute(cmd,communicator=process_comm)
         self.logger.log(msg="import_token: device: {0} status: {1}".format(device_path,status))
@@ -262,7 +262,7 @@ def read_token(self,device_name,token_id):
             self.logger.log("export_token token id {0} not found in device {1} LUKS header".format(CommonVariables.cvm_ade_vm_encryption_token_id,device_name))
             return None
         token = process_comm.stdout
-        return token
+        return json.loads(token)
 
     def remove_token(self,device_name,token_id):
         '''this function remove the token'''
@@ -271,7 +271,7 @@ def remove_token(self,device_name,token_id):
         process_comm = ProcessCommunicator()
         status = self.command_executor.Execute(cmd, communicator=process_comm)
         if status != 0:
-            self.logger.log("remove token id {0} not found in device {1} LUKS header".format(CommonVariables.cvm_ade_vm_encryption_token_id,device_name))
+            self.logger.log("remove token id {0} not found in device {1} LUKS header".format(token_id,device_name))
             return False
         return True
 
@@ -506,7 +506,7 @@ def _extract_luksv2_token(self, luks_dump_out):
                 token_segment = True
                 continue
             if token_segment and self._isnumeric(parts[0].strip()):
-                token_lines.append(parts)
+                token_lines.append([int(parts[0].strip()),parts[1].strip()])
                 continue
         return token_lines
 
diff --git a/VMEncryption/main/handle.py b/VMEncryption/main/handle.py
index 7344320ac..315781a0c 100644
--- a/VMEncryption/main/handle.py
+++ b/VMEncryption/main/handle.py
@@ -271,7 +271,7 @@ def get_protected_settings():
     else:
         return protected_settings_str
 
-def update_encryption_settings_luks2_header():
+def update_encryption_settings_luks2_header(extra_items_to_encrypt=[]):
     '''This function is used for CMK passphrse wrapping with new KEK URL and update metadata in LUKS2 header.'''
     hutil.do_parse_context('UpdateEncryptionSettingsLuks2Header')
     logger.log('Updating encryption settings LUKS-2 header')
@@ -298,11 +298,11 @@ def update_encryption_settings_luks2_header():
             if not disk_util.is_luks_device(device_item_path,None):
                 logger.log("Not a LUKS device, device path: {0}".format(device_item_path))
                 continue
-            #estoring the token data to type Azure_Disk_Encryption
+            #restoring the token data to type Azure_Disk_Encryption
             disk_util.restore_luks2_token(device_name=device_item.name)
             logger.log("Reading passphrase from LUKS2 header, device name: {0}".format(device_item.name))
             #keep token copy for manual recovery
-            ade_token_id = disk_util.get_token_id(header_or_dev_path=device_item_path)
+            ade_token_id = disk_util.get_token_id(header_or_dev_path=device_item_path,token_name=CommonVariables.AzureDiskEncryptionToken)
             if not ade_token_id:
                 logger.log("token type: Azure_Disk_Encryption not found for device {0}".format(device_item.name))
                 continue
@@ -332,6 +332,18 @@ def update_encryption_settings_luks2_header():
             disk_util.remove_token(device_name=device_item.name,token_id=CommonVariables.cvm_ade_vm_encryption_backup_token_id)
         extension_parameter.commit()
         bek_util.umount_azure_passhprase(encryption_config)
+
+        if len(extra_items_to_encrypt) > 0:
+            hutil.do_status_report(operation='UpdateEncryptionSettingsLuks2Header',
+                                   status=CommonVariables.extension_success_status,
+                                   status_code=str(CommonVariables.success),
+                                   message='Encryption settings updated in LUKS2 header')
+        else:
+            hutil.do_exit(exit_code=0,
+                          operation='UpdateEncryptionSettingsLuks2Header',
+                          status=CommonVariables.extension_success_status,
+                          code=str(CommonVariables.success),
+                          message='Encryption settings updated in LUKS2 header')
     except Exception as e:
         hutil.save_seq()
         message = "Failed to update encryption settings Luks2 header with error: {0}, stack trace: {1}".format(e, traceback.format_exc())
@@ -1020,7 +1032,7 @@ def handle_encryption(public_settings, encryption_status, disk_util, bek_util, e
         if security_Type==CommonVariables.ConfidentialVM:
             logger.log('Calling Update Encryption Setting in LUKS2 header.')
             if extension_parameter.cmk_changed():
-                update_encryption_settings_luks2_header()
+                update_encryption_settings_luks2_header(items_to_encrypt)
             if not are_devices_encrypted:
                 enable_encryption()
         else:

From 46c7592871c298bd9cd2abad8d078ad109e11c02 Mon Sep 17 00:00:00 2001
From: PankajJoshi <pankajjoshi@microsoft.com>
Date: Mon, 11 Dec 2023 00:50:58 +0530
Subject: [PATCH 06/16] some correction, imporvement to import token on file
 creation.

---
 VMEncryption/main/DiskUtil.py | 15 +++++++--------
 1 file changed, 7 insertions(+), 8 deletions(-)

diff --git a/VMEncryption/main/DiskUtil.py b/VMEncryption/main/DiskUtil.py
index 539f5d055..b6c333560 100644
--- a/VMEncryption/main/DiskUtil.py
+++ b/VMEncryption/main/DiskUtil.py
@@ -240,15 +240,14 @@ def import_token(self,device_path,passphrase_file,public_settings,PassphraseName
             CommonVariables.PassphraseKey:Protector
         }
         #TODO handle with temp file.
-        custom_cmk = os.path.join("/var/lib/azure_disk_encryption_config/","custom_cmk.json")
-        out_file = open(custom_cmk,"w")
-        json.dump(data,out_file,indent=4)
-        out_file.close()
-        cmd = "cryptsetup token import --json-file {0} --token-id {1} {2}".format(custom_cmk,CommonVariables.cvm_ade_vm_encryption_token_id,device_path)
+        temp_file = tempfile.NamedTemporaryFile(delete=False,mode='w+')
+        json.dump(data,temp_file,indent=4)
+        temp_file.close()
+        cmd = "cryptsetup token import --json-file {0} --token-id {1} {2}".format(temp_file.name,CommonVariables.cvm_ade_vm_encryption_token_id,device_path)
         process_comm = ProcessCommunicator()
         status = self.command_executor.Execute(cmd,communicator=process_comm)
         self.logger.log(msg="import_token: device: {0} status: {1}".format(device_path,status))
-        os.remove(custom_cmk)
+        os.remove(temp_file.name)
         self.logger.log(msg="import_token: device: {0} end.".format(device_path))
         return status==CommonVariables.process_success
     
@@ -445,7 +444,7 @@ def get_token_id(self,header_or_dev_path,token_name):
         luks_dump_out = self._luks_get_header_dump(header_or_dev_path)
         tokens = self._extract_luksv2_token(luks_dump_out)
         for token in tokens:
-            if token[1] is token_name:
+            if token[1] == token_name:
                 return token[0]
         return None
 
@@ -500,7 +499,7 @@ def _extract_luksv2_token(self, luks_dump_out):
             parts = line.split(":")
             if len(parts)<2:
                 continue
-            if token_segment and parts[1].strip() is '':
+            if token_segment and parts[1].strip() == '':
                 break
             if "tokens" in parts[0].strip().lower():
                 token_segment = True

From 2181fc8effe3387d76e82d2ccbe95b87f91a5902 Mon Sep 17 00:00:00 2001
From: PankajJoshi <pankajjoshi@microsoft.com>
Date: Thu, 14 Dec 2023 13:09:39 +0530
Subject: [PATCH 07/16] updated the comments and description in functions.

---
 VMEncryption/main/Common.py               |  2 ++
 VMEncryption/main/CryptMountConfigUtil.py |  2 +-
 VMEncryption/main/DiskUtil.py             | 16 +++++++++++-----
 3 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/VMEncryption/main/Common.py b/VMEncryption/main/Common.py
index 4ca411a39..badf8387f 100644
--- a/VMEncryption/main/Common.py
+++ b/VMEncryption/main/Common.py
@@ -60,7 +60,9 @@ class CommonVariables:
     """
     CVM LUKS2 header token
     """
+    #used to store ADE (encryption setting + wrapped passphrase) token. Type: Azure_Disk_Encryption 
     cvm_ade_vm_encryption_token_id = 5
+    #used to store backup of token type Azure_Disk_Encryption token during CMK rotation. Type: Azure_Disk_Encryption_BackUp
     cvm_ade_vm_encryption_backup_token_id = 6
     ADEEncryptionVersionInLuksToken_1_0='1.0'
     PassphraseNameValueProtected = 'LUKSPasswordProtector'
diff --git a/VMEncryption/main/CryptMountConfigUtil.py b/VMEncryption/main/CryptMountConfigUtil.py
index d122436e0..6d76b0f5c 100644
--- a/VMEncryption/main/CryptMountConfigUtil.py
+++ b/VMEncryption/main/CryptMountConfigUtil.py
@@ -266,7 +266,7 @@ def device_unlock_using_luks2_header(self):
         lock = threading.Lock()
         for device_item in device_items:
             if device_item.file_system == "crypto_LUKS": 
-                 #resotre Luks2 token from BackUp (if needed)
+                 #resotre Luks2 token using BackUp.
                  self.disk_util.restore_luks2_token(device_name=device_item.name)
                  device_item_path = self.disk_util.get_device_path(device_item.name)
                  azure_item_path = azure_name_table[device_item_path] if device_item_path in azure_name_table else device_item_path
diff --git a/VMEncryption/main/DiskUtil.py b/VMEncryption/main/DiskUtil.py
index b6c333560..95305886a 100644
--- a/VMEncryption/main/DiskUtil.py
+++ b/VMEncryption/main/DiskUtil.py
@@ -195,20 +195,22 @@ def secure_key_release_operation(self,protectorbase64,kekUrl,operation,attestati
         return process_comm.stdout.strip()
     
     def import_token_data(self,device_path,token_data,token_id):
-        self.logger.log(msg="import_token for device: {0} started.".format(device_path))
+        '''Updating token_data json object to LUKS2 header's Tokens field.'''
+        self.logger.log(msg="import_token_data for device: {0} started.".format(device_path))
         temp_file = tempfile.NamedTemporaryFile(delete=False,mode='w+')
         json.dump(token_data,temp_file,indent=4)
         temp_file.close()
         cmd = "cryptsetup token import --json-file {0} --token-id {1} {2}".format(temp_file.name,token_id,device_path)
         process_comm = ProcessCommunicator()
         status = self.command_executor.Execute(cmd,communicator=process_comm)
-        self.logger.log(msg="import_token: device: {0} status: {1}".format(device_path,status))
+        self.logger.log(msg="import_token_data: device: {0} status: {1}".format(device_path,status))
         os.unlink(temp_file.name)
         return status==CommonVariables.process_success
 
     def import_token(self,device_path,passphrase_file,public_settings,PassphraseNameValue=CommonVariables.PassphraseNameValueProtected):
         '''this function reads passphrase from passphrase file, wrap it and update in token field of LUKS2 header.'''
         self.logger.log(msg="import_token for device: {0} started.".format(device_path))
+        self.logger.log(msg="import_token for passphrase file path: {0}.".format(passphrase_file))
         Protector= ""
         with open(passphrase_file,"rb") as protector_file:
             #passphrase stored in keyfile is base64
@@ -239,7 +241,6 @@ def import_token(self,device_path,passphrase_file,public_settings,PassphraseName
             CommonVariables.PassphraseNameKey:PassphraseNameValue,
             CommonVariables.PassphraseKey:Protector
         }
-        #TODO handle with temp file.
         temp_file = tempfile.NamedTemporaryFile(delete=False,mode='w+')
         json.dump(data,temp_file,indent=4)
         temp_file.close()
@@ -247,7 +248,7 @@ def import_token(self,device_path,passphrase_file,public_settings,PassphraseName
         process_comm = ProcessCommunicator()
         status = self.command_executor.Execute(cmd,communicator=process_comm)
         self.logger.log(msg="import_token: device: {0} status: {1}".format(device_path,status))
-        os.remove(temp_file.name)
+        os.unlink(temp_file.name)
         self.logger.log(msg="import_token: device: {0} end.".format(device_path))
         return status==CommonVariables.process_success
     
@@ -449,7 +450,12 @@ def get_token_id(self,header_or_dev_path,token_name):
         return None
 
     def restore_luks2_token(self, device_name=None):
-        '''this function restoring token type Azure_Disk_Encryption_BackUp to Azure_Disk_Encryption'''
+        '''this function restoring token type Azure_Disk_Encryption_BackUp to Azure_Disk_Encryption, 
+        this function acts on 4 secenarios. [token id Azure_Disk_Encryption, token id Azure_Disk_Encryption_BackUp, resote_action)'''
+        '''[Y,Y,remove token type Azure_Disk_Encryption_BackUp]
+           [Y,N,do nothing]
+           [N,N,do nothing]
+           [N,Y,move token type Azure_Disk_Encryption_BackUp to Azure_Disk_Encryption]'''
         if not device_name:
             return
         device_path = self.get_device_path(device_name)

From 262500df6b61ea89f87fe2061391c4877274350a Mon Sep 17 00:00:00 2001
From: PankajJoshi <pankajjoshi@microsoft.com>
Date: Thu, 14 Dec 2023 18:22:13 +0530
Subject: [PATCH 08/16] adding test cases

---
 VMEncryption/main/DiskUtil.py            |  36 ++++++--
 VMEncryption/main/test/test_disk_util.py | 112 +++++++++++++++++++++++
 2 files changed, 141 insertions(+), 7 deletions(-)

diff --git a/VMEncryption/main/DiskUtil.py b/VMEncryption/main/DiskUtil.py
index 95305886a..0f85c36ea 100644
--- a/VMEncryption/main/DiskUtil.py
+++ b/VMEncryption/main/DiskUtil.py
@@ -197,6 +197,12 @@ def secure_key_release_operation(self,protectorbase64,kekUrl,operation,attestati
     def import_token_data(self,device_path,token_data,token_id):
         '''Updating token_data json object to LUKS2 header's Tokens field.'''
         self.logger.log(msg="import_token_data for device: {0} started.".format(device_path))
+        if not token_data or not type(token_data) is dict:
+            self.logger.log(level=CommonVariables.WarningLevel, msg="import_token_data: token_data: {0} for device: {1} is not valid.".format(token_data,device_path))
+            return False
+        if not token_id:
+            self.logger.log(level= CommonVariables.WarningLevel, msg = "import_token_data: token_id: {0} for device: {1} is not valid.".format(token_id,device_path) )
+            return False
         temp_file = tempfile.NamedTemporaryFile(delete=False,mode='w+')
         json.dump(token_data,temp_file,indent=4)
         temp_file.close()
@@ -211,6 +217,9 @@ def import_token(self,device_path,passphrase_file,public_settings,PassphraseName
         '''this function reads passphrase from passphrase file, wrap it and update in token field of LUKS2 header.'''
         self.logger.log(msg="import_token for device: {0} started.".format(device_path))
         self.logger.log(msg="import_token for passphrase file path: {0}.".format(passphrase_file))
+        if not passphrase_file or not os.path.exists(passphrase_file):
+            self.logger.log(level=CommonVariables.WarningLevel,msg="import_token for passphrase file path: {0} not exists.".format(passphrase_file))
+            return False
         Protector= ""
         with open(passphrase_file,"rb") as protector_file:
             #passphrase stored in keyfile is base64
@@ -255,11 +264,14 @@ def import_token(self,device_path,passphrase_file,public_settings,PassphraseName
     def read_token(self,device_name,token_id):
         '''this functions reads tokens from LUKS2 header.'''
         device_path = os.path.join("/dev",device_name)
+        if not os.path.exists(device_path) or not token_id:
+            self.logger.log(level=CommonVariables.WarningLevel,msg="read_token: Inputs not valid. device_name: {0}, token id: {1}".format(device_name,token_id))
+            return None
         cmd = "cryptsetup token export --token-id {0} {1}".format(token_id,device_path)
         process_comm = ProcessCommunicator()
         status = self.command_executor.Execute(cmd, communicator=process_comm)
         if status != 0:
-            self.logger.log("export_token token id {0} not found in device {1} LUKS header".format(CommonVariables.cvm_ade_vm_encryption_token_id,device_name))
+            self.logger.log("read_token token id {0} not found in device {1} LUKS header".format(CommonVariables.cvm_ade_vm_encryption_token_id,device_name))
             return None
         token = process_comm.stdout
         return json.loads(token)
@@ -278,14 +290,18 @@ def remove_token(self,device_name,token_id):
     def export_token(self,device_name):
         '''This function reads token id from luks2 header field and unwrap passphrase'''
         self.logger.log("export_token to device {0} started.".format(device_name))
+        if not device_name or not os.path.exists(self.get_device_path(device_name)):
+            self.logger.log(level= CommonVariables.WarningLevel, msg="export_token Input is not valid. device name: {0}".format(device_name))
+            return None
         device_path = self.get_device_path(device_name)
         protector = None
-        cmd = "cryptsetup token export --token-id {0} {1}".format(CommonVariables.cvm_ade_vm_encryption_token_id,device_path)
-        process_comm = ProcessCommunicator()
-        status = self.command_executor.Execute(cmd, communicator=process_comm)
-        if status != 0:
-            self.logger.log("export_token token id {0} not found in device {1} LUKS header".format(CommonVariables.cvm_ade_vm_encryption_token_id,device_name))
+        cvm_ade_vm_encryption_token_id = self.get_token_id(header_or_dev_path=device_path,token_name=CommonVariables.AzureDiskEncryptionToken)
+        if not cvm_ade_vm_encryption_token_id:
+            self.logger.log("export_token token id {0} not found in device {1} LUKS header".format(cvm_ade_vm_encryption_token_id,device_name))
             return None
+        cmd = "cryptsetup token export --token-id {0} {1}".format(cvm_ade_vm_encryption_token_id,device_path)
+        process_comm = ProcessCommunicator()
+        self.command_executor.Execute(cmd, communicator=process_comm)
         token = process_comm.stdout
         disk_encryption_setting=json.loads(token)
         if disk_encryption_setting['version'] != CommonVariables.ADEEncryptionVersionInLuksToken_1_0:
@@ -442,6 +458,9 @@ def luks_get_uuid(self, header_or_dev_path):
 
     def get_token_id(self,header_or_dev_path,token_name):
         '''if LUKS2 header has token name return the id else return none.'''
+        if not header_or_dev_path or not os.path.exists(header_or_dev_path) or not token_name:
+            self.logger.log("get_token_id: invalid input, header_or_dev_path:{0} token_name:{1}".format(header_or_dev_path,token_name))
+            return None
         luks_dump_out = self._luks_get_header_dump(header_or_dev_path)
         tokens = self._extract_luksv2_token(luks_dump_out)
         for token in tokens:
@@ -456,7 +475,8 @@ def restore_luks2_token(self, device_name=None):
            [Y,N,do nothing]
            [N,N,do nothing]
            [N,Y,move token type Azure_Disk_Encryption_BackUp to Azure_Disk_Encryption]'''
-        if not device_name:
+        if not device_name or not os.path.exists(self.get_device_path(device_name)):
+            self.logger.log(level=CommonVariables.WarningLevel,msg="restore_luks2_token invalid input. device_name = {0}".format(device_name))
             return
         device_path = self.get_device_path(device_name)
         ade_token_id = self.get_token_id(header_or_dev_path=device_path,token_name=CommonVariables.AzureDiskEncryptionToken)
@@ -498,6 +518,8 @@ def _extract_luksv2_token(self, luks_dump_out):
             5: Azure_Disk_Encryption
         ...
         """
+        if not luks_dump_out:
+            return []
         lines = luks_dump_out.split("\n")
         token_segment = False
         token_lines = []
diff --git a/VMEncryption/main/test/test_disk_util.py b/VMEncryption/main/test/test_disk_util.py
index 0a6f339f8..59ace6f10 100644
--- a/VMEncryption/main/test/test_disk_util.py
+++ b/VMEncryption/main/test/test_disk_util.py
@@ -416,3 +416,115 @@ def test_get_luks_header_size_luks2_badoffset(self, ver_mock, dump_mock):
         dump_mock.return_value = ""
         header_size = self.disk_util.get_luks_header_size("/mocked/device/path")
         self.assertEqual(header_size, None)
+
+    @mock.patch("CommandExecutor.CommandExecutor.Execute", return_value=0)
+    def test_import_token_data(self,cmd_exc_mock):
+        '''update dict type data to luks2 header'''
+        cmd_exc_mock.return_value = 0
+        data = {}
+        result = self.disk_util.import_token_data(device_path="/dev/sda",token_data=data,token_id=1)
+        self.assertEqual(result,False)
+        data = None
+        result = self.disk_util.import_token_data(device_path="/dev/sda",token_data=data,token_id=1)
+        self.assertEqual(result,False)
+        data = [1,3]
+        result = self.disk_util.import_token_data(device_path="/dev/sda",token_data=data,token_id=1)
+        self.assertEqual(result,False)
+        data = "test data"
+        result = self.disk_util.import_token_data(device_path="/dev/sda",token_data=data,token_id=1)
+        self.assertEqual(result,False)
+        data = {'version':'1.0','type':'ADE'}
+        result = self.disk_util.import_token_data(device_path="/dev/sda",token_data=data,token_id=1)
+        self.assertEqual(result,True)
+
+    @mock.patch("os.path.exists", return_value=True)
+    @mock.patch("CommandExecutor.CommandExecutor.Execute", return_value=0)
+    def test_import_token(self,cmd_exc_mock,path_exit):
+        '''read passphrase from Passphrase_file, and update to LUKS2 header'''
+        cmd_exc_mock.return_value = 0
+        path_exit.return_value=False
+        result = self.disk_util.import_token(device_path="/dev/sdc",
+                                             passphrase_file=None,
+                                             public_settings=None)
+        self.assertEqual(result,False)
+        path_exit.return_value=False
+        result = self.disk_util.import_token(device_path="/dev/sdc",
+                                             passphrase_file="/var/lib/file_path",
+                                             public_settings=None)
+        self.assertEqual(result,False)
+
+    @mock.patch("os.path.exists")
+    @mock.patch("CommandExecutor.CommandExecutor.Execute", return_value=0)
+    def test_read_token(self,cmd_exc_mock,path_exists):
+        '''read token from LUKS2 header token'''
+        cmd_exc_mock.return_value = 1
+        path_exists.return_value = False
+        result = self.disk_util.read_token(device_name="",token_id=None)
+        self.assertEqual(result,None)
+        cmd_exc_mock.return_value = 1
+        path_exists.return_value = True
+        result = self.disk_util.read_token(device_name="sda",token_id=None)
+        self.assertEqual(result,None)
+        cmd_exc_mock.return_value = 1
+        path_exists.return_value = True
+        result = self.disk_util.read_token(device_name="sda",token_id=1)
+        self.assertEqual(result,None)
+
+    @mock.patch("DiskUtil.DiskUtil.get_token_id")
+    @mock.patch("os.path.exists")
+    @mock.patch("CommandExecutor.CommandExecutor.Execute", return_value=0)
+    def test_export_token(self,cmd_exc_mock,path_exists,ade_encryption_token):
+        '''read passphrase from token type Azure_Disk_Encryption'''
+        ade_encryption_token.return_value = 5
+        cmd_exc_mock.return_value = 1
+        path_exists.return_value = False
+        protector = self.disk_util.export_token("sda")
+        self.assertEqual(protector, None)
+        path_exists.return_value = True
+        ade_encryption_token.return_value = None
+        protector = self.disk_util.export_token("sda")
+        self.assertEqual(protector, None)
+
+    def test_extract_luksv2_token(self):
+        '''extracting Tokens field to get token from LUKS2 header. returns list'''
+        luks_dump_out = None
+        tokens = self.disk_util._extract_luksv2_token(luks_dump_out)
+        self.assertEqual(len(tokens), 0)
+        luks_dump_out = ""
+        tokens = self.disk_util._extract_luksv2_token(luks_dump_out)
+        self.assertEqual(len(tokens), 0)
+        luks_dump_out = "Tokens:\n\
+            1: Azure_Disk_Encryption_BackUp\n\
+            5: Azure_Disk_Encryption"
+        tokens = self.disk_util._extract_luksv2_token(luks_dump_out)
+        self.assertEqual(len(tokens), 2)
+        self.assertEqual(tokens[0][0], 1)
+        self.assertEqual(tokens[0][1], "Azure_Disk_Encryption_BackUp")
+        self.assertEqual(tokens[1][0], 5)
+        self.assertEqual(tokens[1][1], "Azure_Disk_Encryption")
+
+    @mock.patch("os.path.exists")
+    @mock.patch("DiskUtil.DiskUtil._luks_get_header_dump")
+    def test_get_token_id(self,luks_dump,header_or_dev_path_exist):
+        '''getting the token id from LUKS2 header, if not present return None'''
+        header_or_dev_path_exist.return_value = True
+        luks_dump.return_value="Tokens:\n\
+            1: Azure_Disk_Encryption_BackUp\n\
+            5: Azure_Disk_Encryption"
+        token_id = self.disk_util.get_token_id("/dev/sda","Azure_Disk_Encryption_BackUp")
+        self.assertEqual(token_id,1)
+        token_id = self.disk_util.get_token_id("/dev/sda","Azure_Disk_Encryption")
+        self.assertEqual(token_id,5)
+        luks_dump.return_value="Tokens:\n\
+            1: Azure_Disk_Encryption_BackUp"
+        token_id = self.disk_util.get_token_id("/dev/sda","Azure_Disk_Encryption_BackUp")
+        self.assertEqual(token_id,1)
+        token_id = self.disk_util.get_token_id("/dev/sda","Azure_Disk_Encryption")
+        self.assertEqual(token_id,None)
+        token_id = self.disk_util.get_token_id("","Azure_Disk_Encryption_BackUp")
+        self.assertEqual(token_id,None)
+        token_id = self.disk_util.get_token_id("/dev/sda","")
+        self.assertEqual(token_id,None)
+        header_or_dev_path_exist.return_value = False
+        token_id = self.disk_util.get_token_id("/dev/sda","Azure_Disk_Encryption")
+        self.assertEqual(token_id,None)
\ No newline at end of file

From 94cd02361db7b10eca6c94e6f2454446eea5b267 Mon Sep 17 00:00:00 2001
From: PankajJoshi <pankajjoshi@microsoft.com>
Date: Mon, 8 Jan 2024 16:52:17 +0530
Subject: [PATCH 09/16] code updation based on review comments

---
 VMEncryption/main/DiskUtil.py | 50 +++++++++++++++++++++++++----------
 VMEncryption/main/handle.py   |  2 +-
 2 files changed, 37 insertions(+), 15 deletions(-)

diff --git a/VMEncryption/main/DiskUtil.py b/VMEncryption/main/DiskUtil.py
index 0f85c36ea..2f6060d4f 100644
--- a/VMEncryption/main/DiskUtil.py
+++ b/VMEncryption/main/DiskUtil.py
@@ -195,7 +195,19 @@ def secure_key_release_operation(self,protectorbase64,kekUrl,operation,attestati
         return process_comm.stdout.strip()
     
     def import_token_data(self,device_path,token_data,token_id):
-        '''Updating token_data json object to LUKS2 header's Tokens field.'''
+        '''Updating token_data json object to LUKS2 header's Tokens field.
+        token data is as follow for version 1.0.
+        "version": "1.0",
+        "type": "Azure_Disk_Encryption",
+        "keyslots": [],
+        "KekVaultResourceId": "<kek_res_id>",
+        "KeyEncryptionKeyURL": "<kek_url>",
+        "KeyVaultResourceId": "<kv_res_id>",
+        "KeyVaultURL": "https://<vault_name>.vault.azure.net/",
+        "AttestationURL": null,
+        "PassphraseName": "LUKSPasswordProtector",
+        "Passphrase": "M53XE09n7O9r2AdKa7FYRYe..."
+        '''
         self.logger.log(msg="import_token_data for device: {0} started.".format(device_path))
         if not token_data or not type(token_data) is dict:
             self.logger.log(level=CommonVariables.WarningLevel, msg="import_token_data: token_data: {0} for device: {1} is not valid.".format(token_data,device_path))
@@ -299,11 +311,7 @@ def export_token(self,device_name):
         if not cvm_ade_vm_encryption_token_id:
             self.logger.log("export_token token id {0} not found in device {1} LUKS header".format(cvm_ade_vm_encryption_token_id,device_name))
             return None
-        cmd = "cryptsetup token export --token-id {0} {1}".format(cvm_ade_vm_encryption_token_id,device_path)
-        process_comm = ProcessCommunicator()
-        self.command_executor.Execute(cmd, communicator=process_comm)
-        token = process_comm.stdout
-        disk_encryption_setting=json.loads(token)
+        disk_encryption_setting=self.read_token(device_name=device_name,token_id=cvm_ade_vm_encryption_token_id)
         if disk_encryption_setting['version'] != CommonVariables.ADEEncryptionVersionInLuksToken_1_0:
             self.logger.log("export_token token version {0} is not a vaild version.".format(disk_encryption_setting['version']))
             return None
@@ -464,17 +472,17 @@ def get_token_id(self,header_or_dev_path,token_name):
         luks_dump_out = self._luks_get_header_dump(header_or_dev_path)
         tokens = self._extract_luksv2_token(luks_dump_out)
         for token in tokens:
-            if token[1] == token_name:
+            if len(token) == 2 and token[1] == token_name:
                 return token[0]
         return None
 
     def restore_luks2_token(self, device_name=None):
         '''this function restoring token type Azure_Disk_Encryption_BackUp to Azure_Disk_Encryption, 
-        this function acts on 4 secenarios. [token id Azure_Disk_Encryption, token id Azure_Disk_Encryption_BackUp, resote_action)'''
-        '''[Y,Y,remove token type Azure_Disk_Encryption_BackUp]
-           [Y,N,do nothing]
-           [N,N,do nothing]
-           [N,Y,move token type Azure_Disk_Encryption_BackUp to Azure_Disk_Encryption]'''
+        this function acts on 4 secenarios. [token id Azure_Disk_Encryption, token id Azure_Disk_Encryption_BackUp, resote_action, Scenario]'''
+        '''[Y,Y,remove token type Azure_Disk_Encryption_BackUp,VM restart/failure]
+           [Y,N,do nothing,Success scenario]
+           [N,N,do nothing,No token present]
+           [N,Y,move token type Azure_Disk_Encryption_BackUp to Azure_Disk_Encryption,VM restart/failure]'''
         if not device_name or not os.path.exists(self.get_device_path(device_name)):
             self.logger.log(level=CommonVariables.WarningLevel,msg="restore_luks2_token invalid input. device_name = {0}".format(device_name))
             return
@@ -512,9 +520,23 @@ def _extract_luks_version_from_dump(self, luks_dump_out):
 
     def _extract_luksv2_token(self, luks_dump_out):
         """
-        ...
+        A luks v2 luksheader looks kind of like this: (inessential stuff removed)
+
+        LUKS header information
+        Version:        2
+        Data segments:
+            0: crypt
+                offset: 0 [bytes]
+                length: 5539430400 [bytes]
+                cipher: aes-xts-plain64
+                sector: 512 [bytes]
+        Keyslots:
+            1: luks2
+                    Key:        512 bits
+            3: reencrypt (unbound)
+                    Key:        8 bits
         Tokens:
-            1: Azure_Disk_Encryption_BackUp
+            6: Azure_Disk_Encryption_BackUp
             5: Azure_Disk_Encryption
         ...
         """
diff --git a/VMEncryption/main/handle.py b/VMEncryption/main/handle.py
index 315781a0c..a3acbcf06 100644
--- a/VMEncryption/main/handle.py
+++ b/VMEncryption/main/handle.py
@@ -2189,7 +2189,7 @@ def mapped_device_item_match(device_item):
     return None
 
 
-def daemon_encrypt(): 
+def daemon_encrypt():
     logger.log("daemon_encrypt security type is {0}".format(security_Type))
     public_settings = get_public_settings()
     if security_Type == CommonVariables.ConfidentialVM:

From 3717d22eeee6d98e2d5554085124fc2c213d0e4e Mon Sep 17 00:00:00 2001
From: PankajJoshi <pankajjoshi@microsoft.com>
Date: Wed, 10 Jan 2024 17:03:09 +0530
Subject: [PATCH 10/16] correction as per review feedback

---
 VMEncryption/main/Common.py               |  8 +++--
 VMEncryption/main/CryptMountConfigUtil.py | 10 +++---
 VMEncryption/main/DiskUtil.py             | 44 +++++++++++------------
 VMEncryption/main/ExtensionParameter.py   |  1 +
 VMEncryption/main/handle.py               | 16 ++++++---
 5 files changed, 44 insertions(+), 35 deletions(-)

diff --git a/VMEncryption/main/Common.py b/VMEncryption/main/Common.py
index badf8387f..5e0cf8714 100644
--- a/VMEncryption/main/Common.py
+++ b/VMEncryption/main/Common.py
@@ -60,14 +60,16 @@ class CommonVariables:
     """
     CVM LUKS2 header token
     """
-    #used to store ADE (encryption setting + wrapped passphrase) token. Type: Azure_Disk_Encryption 
+    #this token id is used to store Primary ADE (encryption setting + wrapped passphrase) token.
     cvm_ade_vm_encryption_token_id = 5
-    #used to store backup of token type Azure_Disk_Encryption token during CMK rotation. Type: Azure_Disk_Encryption_BackUp
+    #this token id is used to store backup of token id 5, during CMK rotation.
     cvm_ade_vm_encryption_backup_token_id = 6
     ADEEncryptionVersionInLuksToken_1_0='1.0'
     PassphraseNameValueProtected = 'LUKSPasswordProtector'
-    PassphraseNameValueNotProtected = 'LUKSPasswordNotProtector'
+    #this token type is used to store Primary ADE token. Type id: 5
     AzureDiskEncryptionToken = 'Azure_Disk_Encryption'
+    #this token type is used to store backup ADE token. Type id: 6
+    #this token used for recovery to token 5 in case of reboot/interruption happened during reboot.
     AzureDiskEncryptionBackUpToken='Azure_Disk_Encryption_BackUp'
     """
     IMDS IP:
diff --git a/VMEncryption/main/CryptMountConfigUtil.py b/VMEncryption/main/CryptMountConfigUtil.py
index 6d76b0f5c..4d3de3440 100644
--- a/VMEncryption/main/CryptMountConfigUtil.py
+++ b/VMEncryption/main/CryptMountConfigUtil.py
@@ -142,7 +142,7 @@ def _restore_backup_crypttab_info(self,crypt_item,passphrase_file):
             return False
 
         if not os.path.exists(azure_crypt_mount_backup_location) and not os.path.exists(crypttab_backup_location):
-            self.logger.log(msg=("MountPoint info not found for" + device_item_real_path), level=CommonVariables.ErrorLevel)
+            self.logger.log(msg=("MountPoint info not found for" + crypt_item.dev_path), level=CommonVariables.ErrorLevel)
             # Not sure when this happens..
             # in this case also, just add an entry to the azure_crypt_mount without a mount point.
             self.add_crypt_item(crypt_item)
@@ -265,8 +265,8 @@ def device_unlock_using_luks2_header(self):
         threads = []
         lock = threading.Lock()
         for device_item in device_items:
-            if device_item.file_system == "crypto_LUKS": 
-                 #resotre Luks2 token using BackUp.
+            if device_item.file_system == "crypto_LUKS":
+                 #restore LUKS2 token using BackUp.
                  self.disk_util.restore_luks2_token(device_name=device_item.name)
                  device_item_path = self.disk_util.get_device_path(device_item.name)
                  azure_item_path = azure_name_table[device_item_path] if device_item_path in azure_name_table else device_item_path
@@ -274,12 +274,12 @@ def device_unlock_using_luks2_header(self):
                  threads.append(thread)
                  thread.start()
         for thread in threads:
-            thread.join()       
+            thread.join()
         self.logger.log("device_unlock_using_luks2_header End")
 
     def consolidate_azure_crypt_mount(self, passphrase_file):
         """
-        Reads the backup files from block devices that have a LUKS header and adds it to the cenral azure_crypt_mount file
+        Reads the backup files from block devices that have a LUKS2 header and adds it to the central azure_crypt_mount file
         """
         self.logger.log("Consolidating azure_crypt_mount")
 
diff --git a/VMEncryption/main/DiskUtil.py b/VMEncryption/main/DiskUtil.py
index 2f6060d4f..164698d7f 100644
--- a/VMEncryption/main/DiskUtil.py
+++ b/VMEncryption/main/DiskUtil.py
@@ -262,17 +262,12 @@ def import_token(self,device_path,passphrase_file,public_settings,PassphraseName
             CommonVariables.PassphraseNameKey:PassphraseNameValue,
             CommonVariables.PassphraseKey:Protector
         }
-        temp_file = tempfile.NamedTemporaryFile(delete=False,mode='w+')
-        json.dump(data,temp_file,indent=4)
-        temp_file.close()
-        cmd = "cryptsetup token import --json-file {0} --token-id {1} {2}".format(temp_file.name,CommonVariables.cvm_ade_vm_encryption_token_id,device_path)
-        process_comm = ProcessCommunicator()
-        status = self.command_executor.Execute(cmd,communicator=process_comm)
-        self.logger.log(msg="import_token: device: {0} status: {1}".format(device_path,status))
-        os.unlink(temp_file.name)
+        status = self.import_token_data(device_path=device_path,
+                                        token_data=data,
+                                        token_id=CommonVariables.cvm_ade_vm_encryption_token_id)
         self.logger.log(msg="import_token: device: {0} end.".format(device_path))
-        return status==CommonVariables.process_success
-    
+        return status
+
     def read_token(self,device_name,token_id):
         '''this functions reads tokens from LUKS2 header.'''
         device_path = os.path.join("/dev",device_name)
@@ -319,7 +314,7 @@ def export_token(self,device_name):
         wrappedProtector = disk_encryption_setting[CommonVariables.PassphraseKey]
         attestationUrl = disk_encryption_setting[CommonVariables.AttestationURLKey]
         if disk_encryption_setting[CommonVariables.PassphraseNameKey] != CommonVariables.PassphraseNameValueProtected:
-            self.logger.log("passphrase is not Protectected. No need to do SKR.")
+            self.logger.log(level=CommonVariables.WarningLevel, msg="passphrase is not Protected. No need to do SKR.")
             return wrappedProtector if wrappedProtector else None
         if wrappedProtector:
             #unwrap the protector.
@@ -464,7 +459,7 @@ def luks_get_uuid(self, header_or_dev_path):
                 return splits[1]
         return None
 
-    def get_token_id(self,header_or_dev_path,token_name):
+    def get_token_id(self, header_or_dev_path, token_name):
         '''if LUKS2 header has token name return the id else return none.'''
         if not header_or_dev_path or not os.path.exists(header_or_dev_path) or not token_name:
             self.logger.log("get_token_id: invalid input, header_or_dev_path:{0} token_name:{1}".format(header_or_dev_path,token_name))
@@ -477,33 +472,36 @@ def get_token_id(self,header_or_dev_path,token_name):
         return None
 
     def restore_luks2_token(self, device_name=None):
-        '''this function restoring token type Azure_Disk_Encryption_BackUp to Azure_Disk_Encryption, 
-        this function acts on 4 secenarios. [token id Azure_Disk_Encryption, token id Azure_Disk_Encryption_BackUp, resote_action, Scenario]'''
-        '''[Y,Y,remove token type Azure_Disk_Encryption_BackUp,VM restart/failure]
-           [Y,N,do nothing,Success scenario]
-           [N,N,do nothing,No token present]
-           [N,Y,move token type Azure_Disk_Encryption_BackUp to Azure_Disk_Encryption,VM restart/failure]'''
+        '''this function restores token
+        type:Azure_Disk_Encryption_BackUp, id:6 to type:Azure_Disk_Encryption id:5,
+        this function acts on 4 scenarios.
+        1. both token id: 5 and 6 present in LUKS2 Tokens filed, due to reboot/interrupt during
+        KEK rotation, such case remove token id 5 has latest data so remove token id 6.
+        2. token id 5 present but 6 is not present in LUKS2 Tokens field. do nothing.
+        3. token id 5 not present but 6 present in LUKS2 Tokens field, restore token id 5 using
+        token id 6, then remove token id 6.
+        4. no token ids 5 or 6 present in LUKS2 Tokens field, do nothing.'''
         if not device_name or not os.path.exists(self.get_device_path(device_name)):
             self.logger.log(level=CommonVariables.WarningLevel,msg="restore_luks2_token invalid input. device_name = {0}".format(device_name))
             return
         device_path = self.get_device_path(device_name)
-        ade_token_id = self.get_token_id(header_or_dev_path=device_path,token_name=CommonVariables.AzureDiskEncryptionToken)
+        ade_token_id_primary = self.get_token_id(header_or_dev_path=device_path,token_name=CommonVariables.AzureDiskEncryptionToken)
         ade_token_id_backup = self.get_token_id(header_or_dev_path=device_path,token_name=CommonVariables.AzureDiskEncryptionBackUpToken)
         if not ade_token_id_backup:
             #do nothing
             return
-        if ade_token_id:
+        if ade_token_id_primary:
             #remove backup token id
             self.remove_token(device_name=device_name,token_id=ade_token_id_backup)
             return
-        self.logger.log("resotre luks2 token for device {0} is started.".format(device_name))
+        self.logger.log("restore luks2 token for device {0} is started.".format(device_name))
         #read from backup and update AzureDiskEncryptionToken
         data = self.read_token(device_name=device_name,token_id=ade_token_id_backup)
-        data['type']=CommonVariables.AzureDiskEncryptionBackUpToken
+        data['type']=CommonVariables.AzureDiskEncryptionToken
         self.import_token_data(device_path=device_path,token_data=data,token_id=CommonVariables.AzureDiskEncryptionToken)
         #remove backup
         self.remove_token(device_name=device_name,token_id=ade_token_id_backup)
-        self.logger.log("resotre luks2 token for device {0} is successful.".format(device_name))
+        self.logger.log("restore luks2 token id {0} to {1} for device {2} is successful.".format(ade_token_id_backup,ade_token_id_primary,device_name))
 
     def _get_cryptsetup_version(self):
         # get version of currently installed cryptsetup
diff --git a/VMEncryption/main/ExtensionParameter.py b/VMEncryption/main/ExtensionParameter.py
index 09e883bba..2b682eb67 100644
--- a/VMEncryption/main/ExtensionParameter.py
+++ b/VMEncryption/main/ExtensionParameter.py
@@ -175,6 +175,7 @@ def _is_kv_equivalent(self, a, b):
         return a==b
 
     def cmk_changed(self):
+        '''current config CMK changed from effective config CMK.'''
         if (self.KeyEncryptionKeyURL or self.get_kek_url()) and \
            (not self._is_kv_equivalent(self.KeyEncryptionKeyURL, self.get_kek_url())):
            self.logger.log('Current config KeyEncryptionKeyURL {0} differs from effective config KeyEncryptionKeyURL {1}'.format(self.KeyEncryptionKeyURL, self.get_kek_url()))
diff --git a/VMEncryption/main/handle.py b/VMEncryption/main/handle.py
index a3acbcf06..6cb5e907c 100644
--- a/VMEncryption/main/handle.py
+++ b/VMEncryption/main/handle.py
@@ -192,7 +192,8 @@ def stamp_disks_with_settings(items_to_encrypt, encryption_config, encryption_ma
     public_settings = get_public_settings()
     extension_parameter = ExtensionParameter(hutil, logger, DistroPatcher, encryption_environment, get_protected_settings(), public_settings)
     if security_Type == CommonVariables.ConfidentialVM:
-        logger.log(msg="Do not send vm setting to host for stamping.",level=CommonVariables.InfoLevel)
+        logger.log(msg="Do not send CVM encryption setting to host for stamping.",
+                   level=CommonVariables.InfoLevel)
         extension_parameter.commit()
         return
     has_keystore_flag = CommonVariables.KeyStoreTypeKey in public_settings
@@ -299,6 +300,7 @@ def update_encryption_settings_luks2_header(extra_items_to_encrypt=[]):
                 logger.log("Not a LUKS device, device path: {0}".format(device_item_path))
                 continue
             #restoring the token data to type Azure_Disk_Encryption
+            #It is necessary to restore if we are resuming from previous attempt, otherwise its no-op.
             disk_util.restore_luks2_token(device_name=device_item.name)
             logger.log("Reading passphrase from LUKS2 header, device name: {0}".format(device_item.name))
             #keep token copy for manual recovery
@@ -323,13 +325,19 @@ def update_encryption_settings_luks2_header(extra_items_to_encrypt=[]):
             temp_keyfile.write(passphrase.encode("utf-8"))
             temp_keyfile.close()
             #save passphrase to LUKS2 header with PassphraseNameValueProtected
-            ret = disk_util.import_token(device_path=device_item_path,passphrase_file=temp_keyfile.name,public_settings=public_setting,PassphraseNameValue=CommonVariables.PassphraseNameValueProtected)
+            ret = disk_util.import_token(device_path=device_item_path,
+                                         passphrase_file=temp_keyfile.name,
+                                         public_settings=public_setting,
+                                         PassphraseNameValue=CommonVariables.PassphraseNameValueProtected)
             if not ret:
-                logger.log("Update passphrase with current public setting to LUKS2 header is not successful. device path {0}".format(device_item_path))
+                logger.log(level=CommonVariables.WarningLevel,
+                           msg="Update passphrase with current public setting to LUKS2 header is not successful. device path {0}".format(device_item_path))
                 return None
             os.unlink(temp_keyfile.name)
             #removing backup token
-            disk_util.remove_token(device_name=device_item.name,token_id=CommonVariables.cvm_ade_vm_encryption_backup_token_id)
+            disk_util.remove_token(device_name=device_item.name,
+                                   token_id=CommonVariables.cvm_ade_vm_encryption_backup_token_id)
+        #committing the extension parameter if KEK rotation is successful.
         extension_parameter.commit()
         bek_util.umount_azure_passhprase(encryption_config)
 

From 95969d485a82d8c0bfbea1c875014b244dbae121 Mon Sep 17 00:00:00 2001
From: PankajJoshi <pankajjoshi@microsoft.com>
Date: Wed, 10 Jan 2024 17:43:34 +0530
Subject: [PATCH 11/16] some minor fixes and comments corrections

---
 VMEncryption/main/Common.py   |  2 +-
 VMEncryption/main/DiskUtil.py | 23 +++++++++++++----------
 2 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/VMEncryption/main/Common.py b/VMEncryption/main/Common.py
index 5e0cf8714..5e7cf8c8f 100644
--- a/VMEncryption/main/Common.py
+++ b/VMEncryption/main/Common.py
@@ -69,7 +69,7 @@ class CommonVariables:
     #this token type is used to store Primary ADE token. Type id: 5
     AzureDiskEncryptionToken = 'Azure_Disk_Encryption'
     #this token type is used to store backup ADE token. Type id: 6
-    #this token used for recovery to token 5 in case of reboot/interruption happened during reboot.
+    #this token used for recovery to token 5 in case of reboot/interruption happened during KEK rotation.
     AzureDiskEncryptionBackUpToken='Azure_Disk_Encryption_BackUp'
     """
     IMDS IP:
diff --git a/VMEncryption/main/DiskUtil.py b/VMEncryption/main/DiskUtil.py
index 164698d7f..e97bf22c1 100644
--- a/VMEncryption/main/DiskUtil.py
+++ b/VMEncryption/main/DiskUtil.py
@@ -270,15 +270,17 @@ def import_token(self,device_path,passphrase_file,public_settings,PassphraseName
 
     def read_token(self,device_name,token_id):
         '''this functions reads tokens from LUKS2 header.'''
-        device_path = os.path.join("/dev",device_name)
-        if not os.path.exists(device_path) or not token_id:
-            self.logger.log(level=CommonVariables.WarningLevel,msg="read_token: Inputs not valid. device_name: {0}, token id: {1}".format(device_name,token_id))
+        device_path = self.get_device_path(dev_name=device_name)
+        if not device_path or not token_id:
+            self.logger.log(level=CommonVariables.WarningLevel,
+                            msg="read_token: Inputs not valid. device_name: {0}, token id: {1}".format(device_name,token_id))
             return None
         cmd = "cryptsetup token export --token-id {0} {1}".format(token_id,device_path)
         process_comm = ProcessCommunicator()
         status = self.command_executor.Execute(cmd, communicator=process_comm)
         if status != 0:
-            self.logger.log("read_token token id {0} not found in device {1} LUKS header".format(CommonVariables.cvm_ade_vm_encryption_token_id,device_name))
+            self.logger.log(level=CommonVariables.WarningLevel,
+                            msg="read_token token id {0} not found in device {1} LUKS header".format(CommonVariables.cvm_ade_vm_encryption_token_id,device_name))
             return None
         token = process_comm.stdout
         return json.loads(token)
@@ -290,17 +292,18 @@ def remove_token(self,device_name,token_id):
         process_comm = ProcessCommunicator()
         status = self.command_executor.Execute(cmd, communicator=process_comm)
         if status != 0:
-            self.logger.log("remove token id {0} not found in device {1} LUKS header".format(token_id,device_name))
+            self.logger.log(level=CommonVariables.WarningLevel,
+                            msg="remove token id {0} not found in device {1} LUKS header".format(token_id,device_name))
             return False
         return True
 
     def export_token(self,device_name):
         '''This function reads token id from luks2 header field and unwrap passphrase'''
         self.logger.log("export_token to device {0} started.".format(device_name))
-        if not device_name or not os.path.exists(self.get_device_path(device_name)):
+        device_path = self.get_device_path(device_name)
+        if not device_path:
             self.logger.log(level= CommonVariables.WarningLevel, msg="export_token Input is not valid. device name: {0}".format(device_name))
             return None
-        device_path = self.get_device_path(device_name)
         protector = None
         cvm_ade_vm_encryption_token_id = self.get_token_id(header_or_dev_path=device_path,token_name=CommonVariables.AzureDiskEncryptionToken)
         if not cvm_ade_vm_encryption_token_id:
@@ -481,10 +484,10 @@ def restore_luks2_token(self, device_name=None):
         3. token id 5 not present but 6 present in LUKS2 Tokens field, restore token id 5 using
         token id 6, then remove token id 6.
         4. no token ids 5 or 6 present in LUKS2 Tokens field, do nothing.'''
-        if not device_name or not os.path.exists(self.get_device_path(device_name)):
+        device_path = self.get_device_path(device_name)
+        if not device_path:
             self.logger.log(level=CommonVariables.WarningLevel,msg="restore_luks2_token invalid input. device_name = {0}".format(device_name))
             return
-        device_path = self.get_device_path(device_name)
         ade_token_id_primary = self.get_token_id(header_or_dev_path=device_path,token_name=CommonVariables.AzureDiskEncryptionToken)
         ade_token_id_backup = self.get_token_id(header_or_dev_path=device_path,token_name=CommonVariables.AzureDiskEncryptionBackUpToken)
         if not ade_token_id_backup:
@@ -498,7 +501,7 @@ def restore_luks2_token(self, device_name=None):
         #read from backup and update AzureDiskEncryptionToken
         data = self.read_token(device_name=device_name,token_id=ade_token_id_backup)
         data['type']=CommonVariables.AzureDiskEncryptionToken
-        self.import_token_data(device_path=device_path,token_data=data,token_id=CommonVariables.AzureDiskEncryptionToken)
+        self.import_token_data(device_path=device_path,token_data=data,token_id=ade_token_id_primary)
         #remove backup
         self.remove_token(device_name=device_name,token_id=ade_token_id_backup)
         self.logger.log("restore luks2 token id {0} to {1} for device {2} is successful.".format(ade_token_id_backup,ade_token_id_primary,device_name))

From 4e0bd4c10535a5a32cca723f2243379fd1c46337 Mon Sep 17 00:00:00 2001
From: PankajJoshi <pankajjoshi@microsoft.com>
Date: Wed, 10 Jan 2024 19:06:48 +0530
Subject: [PATCH 12/16] corrected pylint issues in newly added code, spell
 correction and update more concise comments.

---
 VMEncryption/main/Common.py   |  2 +-
 VMEncryption/main/DiskUtil.py | 83 +++++++++++++++++++----------------
 VMEncryption/main/handle.py   | 33 +++++++++-----
 3 files changed, 68 insertions(+), 50 deletions(-)

diff --git a/VMEncryption/main/Common.py b/VMEncryption/main/Common.py
index 5e7cf8c8f..ee3395440 100644
--- a/VMEncryption/main/Common.py
+++ b/VMEncryption/main/Common.py
@@ -62,7 +62,7 @@ class CommonVariables:
     """
     #this token id is used to store Primary ADE (encryption setting + wrapped passphrase) token.
     cvm_ade_vm_encryption_token_id = 5
-    #this token id is used to store backup of token id 5, during CMK rotation.
+    #this token id is used to store backup of token id 5, during KEK rotation.
     cvm_ade_vm_encryption_backup_token_id = 6
     ADEEncryptionVersionInLuksToken_1_0='1.0'
     PassphraseNameValueProtected = 'LUKSPasswordProtector'
diff --git a/VMEncryption/main/DiskUtil.py b/VMEncryption/main/DiskUtil.py
index e97bf22c1..1f64276f3 100644
--- a/VMEncryption/main/DiskUtil.py
+++ b/VMEncryption/main/DiskUtil.py
@@ -193,8 +193,8 @@ def secure_key_release_operation(self,protectorbase64,kekUrl,operation,attestati
             return None
         self.logger.log("secure_key_release_operation {0} end.".format(operation))
         return process_comm.stdout.strip()
-    
-    def import_token_data(self,device_path,token_data,token_id):
+
+    def import_token_data(self, device_path, token_data, token_id):
         '''Updating token_data json object to LUKS2 header's Tokens field.
         token data is as follow for version 1.0.
         "version": "1.0",
@@ -209,7 +209,7 @@ def import_token_data(self,device_path,token_data,token_id):
         "Passphrase": "M53XE09n7O9r2AdKa7FYRYe..."
         '''
         self.logger.log(msg="import_token_data for device: {0} started.".format(device_path))
-        if not token_data or not type(token_data) is dict:
+        if not token_data or not isinstance(token_data, dict):
             self.logger.log(level=CommonVariables.WarningLevel, msg="import_token_data: token_data: {0} for device: {1} is not valid.".format(token_data,device_path))
             return False
         if not token_id:
@@ -225,42 +225,45 @@ def import_token_data(self,device_path,token_data,token_id):
         os.unlink(temp_file.name)
         return status==CommonVariables.process_success
 
-    def import_token(self,device_path,passphrase_file,public_settings,PassphraseNameValue=CommonVariables.PassphraseNameValueProtected):
-        '''this function reads passphrase from passphrase file, wrap it and update in token field of LUKS2 header.'''
+    def import_token(self, device_path, passphrase_file, public_settings,
+                     passphrase_name_value=CommonVariables.PassphraseNameValueProtected):
+        '''This function reads passphrase from passphrase_file, do SKR and wrap passphrase with securely 
+        released key. Then it updates metadata (required encryption settings for SKR + wrapped passphrase) 
+        to primary token id: 5 type: Azure_Disk_Encryption in Tokens field of LUKS2 header.'''
         self.logger.log(msg="import_token for device: {0} started.".format(device_path))
         self.logger.log(msg="import_token for passphrase file path: {0}.".format(passphrase_file))
         if not passphrase_file or not os.path.exists(passphrase_file):
             self.logger.log(level=CommonVariables.WarningLevel,msg="import_token for passphrase file path: {0} not exists.".format(passphrase_file))
             return False
-        Protector= ""
+        protector= ""
         with open(passphrase_file,"rb") as protector_file:
             #passphrase stored in keyfile is base64
-            Protector = protector_file.read().decode('utf-8')
-        KekVaultResourceId=public_settings.get(CommonVariables.KekVaultResourceIdKey)
-        KeyEncryptionKeyUrl=public_settings.get(CommonVariables.KeyEncryptionKeyURLKey)
-        AttestationUrl = public_settings.get(CommonVariables.AttestationURLKey)
-        if PassphraseNameValue == CommonVariables.PassphraseNameValueProtected:
-            Protector = self.secure_key_release_operation(protectorbase64=Protector,
-                                                        kekUrl=KeyEncryptionKeyUrl,
+            protector = protector_file.read().decode('utf-8')
+        kek_vault_resource_id=public_settings.get(CommonVariables.KekVaultResourceIdKey)
+        key_encryption_key_url=public_settings.get(CommonVariables.KeyEncryptionKeyURLKey)
+        attestation_url = public_settings.get(CommonVariables.AttestationURLKey)
+        if passphrase_name_value == CommonVariables.PassphraseNameValueProtected:
+            protector = self.secure_key_release_operation(protectorbase64=protector,
+                                                        kekUrl=key_encryption_key_url,
                                                         operation=CommonVariables.secure_key_release_wrap,
-                                                        attestationUrl=AttestationUrl)
+                                                        attestationUrl=attestation_url)
         else:
-            self.logger.log(msg="import_token passphrase is not wrapped, value of passphrase name key: {0}".format(PassphraseNameValue))
+            self.logger.log(msg="import_token passphrase is not wrapped, value of passphrase name key: {0}".format(passphrase_name_value))
 
-        if not Protector:
+        if not protector:
             self.logger.log("import_token protector wrapping is unsuccessful for device {0}".format(device_path))
             return False
         data={
             "version":CommonVariables.ADEEncryptionVersionInLuksToken_1_0,
             "type":"Azure_Disk_Encryption",
             "keyslots":[],
-            CommonVariables.KekVaultResourceIdKey:KekVaultResourceId,
-            CommonVariables.KeyEncryptionKeyURLKey:KeyEncryptionKeyUrl,
+            CommonVariables.KekVaultResourceIdKey:kek_vault_resource_id,
+            CommonVariables.KeyEncryptionKeyURLKey:key_encryption_key_url,
             CommonVariables.KeyVaultResourceIdKey:public_settings.get(CommonVariables.KeyVaultResourceIdKey),
             CommonVariables.KeyVaultURLKey:public_settings.get(CommonVariables.KeyVaultURLKey),
-            CommonVariables.AttestationURLKey:AttestationUrl,
-            CommonVariables.PassphraseNameKey:PassphraseNameValue,
-            CommonVariables.PassphraseKey:Protector
+            CommonVariables.AttestationURLKey:attestation_url,
+            CommonVariables.PassphraseNameKey:passphrase_name_value,
+            CommonVariables.PassphraseKey:protector
         }
         status = self.import_token_data(device_path=device_path,
                                         token_data=data,
@@ -268,38 +271,44 @@ def import_token(self,device_path,passphrase_file,public_settings,PassphraseName
         self.logger.log(msg="import_token: device: {0} end.".format(device_path))
         return status
 
-    def read_token(self,device_name,token_id):
+    def read_token(self, device_name, token_id):
         '''this functions reads tokens from LUKS2 header.'''
         device_path = self.get_device_path(dev_name=device_name)
         if not device_path or not token_id:
             self.logger.log(level=CommonVariables.WarningLevel,
-                            msg="read_token: Inputs not valid. device_name: {0}, token id: {1}".format(device_name,token_id))
+                            msg="read_token: Inputs are not valid. device_name: {0}, token id: {1}".format(device_name,token_id))
             return None
         cmd = "cryptsetup token export --token-id {0} {1}".format(token_id,device_path)
         process_comm = ProcessCommunicator()
         status = self.command_executor.Execute(cmd, communicator=process_comm)
         if status != 0:
             self.logger.log(level=CommonVariables.WarningLevel,
-                            msg="read_token token id {0} not found in device {1} LUKS header".format(CommonVariables.cvm_ade_vm_encryption_token_id,device_name))
+                            msg="read_token: token id: {0} is not found for device_name: {1} in LUKS header".format(token_id,device_name))
             return None
         token = process_comm.stdout
         return json.loads(token)
 
-    def remove_token(self,device_name,token_id):
+    def remove_token(self, device_name, token_id):
         '''this function remove the token'''
-        device_path = os.path.join("/dev",device_name)
+        device_path = self.get_device_path(dev_name=device_name)
+        if not device_path or not token_id:
+            self.logger.log(level=CommonVariables.WarningLevel,
+                            msg="remove_token: Inputs are not valid. device name: {0}, token_id: {1}".format(device_name,token_id))
+            return False
         cmd = "cryptsetup token remove --token-id {0} {1}".format(token_id,device_path)
         process_comm = ProcessCommunicator()
         status = self.command_executor.Execute(cmd, communicator=process_comm)
         if status != 0:
             self.logger.log(level=CommonVariables.WarningLevel,
-                            msg="remove token id {0} not found in device {1} LUKS header".format(token_id,device_name))
+                            msg="remove_token: token id: {0} is not found for device_name: {1} in LUKS header".format(token_id,device_name))
             return False
         return True
 
     def export_token(self,device_name):
-        '''This function reads token id from luks2 header field and unwrap passphrase'''
-        self.logger.log("export_token to device {0} started.".format(device_name))
+        '''This function reads wrapped passphrase from LUKS2 Tokens for
+        token id:5, which belongs to primary token type: Azure_Disk_Encryption
+        and do SKR and returns unwrapped passphrase'''
+        self.logger.log("export_token: for device_name: {0} started.".format(device_name))
         device_path = self.get_device_path(device_name)
         if not device_path:
             self.logger.log(level= CommonVariables.WarningLevel, msg="export_token Input is not valid. device name: {0}".format(device_name))
@@ -313,17 +322,17 @@ def export_token(self,device_name):
         if disk_encryption_setting['version'] != CommonVariables.ADEEncryptionVersionInLuksToken_1_0:
             self.logger.log("export_token token version {0} is not a vaild version.".format(disk_encryption_setting['version']))
             return None
-        keyEncryptionKeyUrl=disk_encryption_setting[CommonVariables.KeyEncryptionKeyURLKey]
-        wrappedProtector = disk_encryption_setting[CommonVariables.PassphraseKey]
-        attestationUrl = disk_encryption_setting[CommonVariables.AttestationURLKey]
+        key_encryption_key_url=disk_encryption_setting[CommonVariables.KeyEncryptionKeyURLKey]
+        wrapped_protector = disk_encryption_setting[CommonVariables.PassphraseKey]
+        attestation_url = disk_encryption_setting[CommonVariables.AttestationURLKey]
         if disk_encryption_setting[CommonVariables.PassphraseNameKey] != CommonVariables.PassphraseNameValueProtected:
             self.logger.log(level=CommonVariables.WarningLevel, msg="passphrase is not Protected. No need to do SKR.")
-            return wrappedProtector if wrappedProtector else None
-        if wrappedProtector:
+            return wrapped_protector if wrapped_protector else None
+        if wrapped_protector:
             #unwrap the protector.
-            protector=self.secure_key_release_operation(attestationUrl=attestationUrl,
-                                                        kekUrl=keyEncryptionKeyUrl,
-                                                        protectorbase64=wrappedProtector,
+            protector=self.secure_key_release_operation(attestationUrl=attestation_url,
+                                                        kekUrl=key_encryption_key_url,
+                                                        protectorbase64=wrapped_protector,
                                                         operation=CommonVariables.secure_key_release_unwrap)
         self.logger.log("export_token to device {0} end.".format(device_name))
         return protector
diff --git a/VMEncryption/main/handle.py b/VMEncryption/main/handle.py
index 6cb5e907c..619394cbc 100644
--- a/VMEncryption/main/handle.py
+++ b/VMEncryption/main/handle.py
@@ -272,8 +272,11 @@ def get_protected_settings():
     else:
         return protected_settings_str
 
-def update_encryption_settings_luks2_header(extra_items_to_encrypt=[]):
-    '''This function is used for CMK passphrse wrapping with new KEK URL and update metadata in LUKS2 header.'''
+def update_encryption_settings_luks2_header(extra_items_to_encrypt=None):
+    '''This function is used for CMK passphrase wrapping with new KEK URL and update
+    metadata in LUKS2 header.'''
+    if extra_items_to_encrypt is None:
+        extra_items_to_encrypt=[]
     hutil.do_parse_context('UpdateEncryptionSettingsLuks2Header')
     logger.log('Updating encryption settings LUKS-2 header')
     # ensure cryptsetup package is still available in case it was for some reason removed after enable
@@ -304,21 +307,23 @@ def update_encryption_settings_luks2_header(extra_items_to_encrypt=[]):
             disk_util.restore_luks2_token(device_name=device_item.name)
             logger.log("Reading passphrase from LUKS2 header, device name: {0}".format(device_item.name))
             #keep token copy for manual recovery
-            ade_token_id = disk_util.get_token_id(header_or_dev_path=device_item_path,token_name=CommonVariables.AzureDiskEncryptionToken)
-            if not ade_token_id:
-                logger.log("token type: Azure_Disk_Encryption not found for device {0}".format(device_item.name))
+            ade_primary_token_id = disk_util.get_token_id(header_or_dev_path=device_item_path,token_name=CommonVariables.AzureDiskEncryptionToken)
+            if not ade_primary_token_id:
+                logger.log("primary token type: Azure_Disk_Encryption not found for device {0}".format(device_item.name))
                 continue
-            data = disk_util.read_token(device_name=device_item.name,token_id=ade_token_id)
-            #writing to backup token
+            data = disk_util.read_token(device_name=device_item.name,token_id=ade_primary_token_id)
+            #writing primary token data to backup token, update token type to back up.
             data['type']=CommonVariables.AzureDiskEncryptionBackUpToken
+            #update backup token data to backup token id:6.
             disk_util.import_token_data(device_path=device_item_path,token_data=data,token_id=CommonVariables.cvm_ade_vm_encryption_backup_token_id)
             #get the unwrapped passphrase from LUKS2 header.
             passphrase=disk_util.export_token(device_name=device_item.name)
-            #remove token
-            disk_util.remove_token(device_name=device_item.name,token_id=ade_token_id)
             if not passphrase:
-                logger.log("No passphrase in LUKS2 header, device name: {0}".format(device_item.name))
+                logger.log(level=CommonVariables.WarningLevel,
+                           msg="No passphrase found in LUKS2 header, device name: {0}".format(device_item.name))
                 continue
+            #remove primary token from Tokens field of LUKS2 header.
+            disk_util.remove_token(device_name=device_item.name,token_id=ade_primary_token_id)
             logger.log("Updating wrapped passphrase to LUKS2 header with current public setting. device name {0}".format(device_item.name))
             #protect passphrase before updating to LUKS2 is done in import_token
             temp_keyfile = tempfile.NamedTemporaryFile(delete=False)
@@ -334,7 +339,7 @@ def update_encryption_settings_luks2_header(extra_items_to_encrypt=[]):
                            msg="Update passphrase with current public setting to LUKS2 header is not successful. device path {0}".format(device_item_path))
                 return None
             os.unlink(temp_keyfile.name)
-            #removing backup token
+            #removing backup token as KEK rotation is successful here.
             disk_util.remove_token(device_name=device_item.name,
                                    token_id=CommonVariables.cvm_ade_vm_encryption_backup_token_id)
         #committing the extension parameter if KEK rotation is successful.
@@ -1036,7 +1041,11 @@ def handle_encryption(public_settings, encryption_status, disk_util, bek_util, e
         if is_daemon_running():
             logger.log("An operation already running. Cannot accept an update settings request.")
             hutil.reject_settings()
-        are_devices_encrypted, items_to_encrypt = are_required_devices_encrypted(volume_type, encryption_status, disk_util, bek_util, encryption_operation)
+        are_devices_encrypted, items_to_encrypt = are_required_devices_encrypted(volume_type=volume_type,
+                                                                                 encryption_status=encryption_status,
+                                                                                disk_util=disk_util,
+                                                                                bek_util=bek_util,
+                                                                                encryption_operation=encryption_operation)
         if security_Type==CommonVariables.ConfidentialVM:
             logger.log('Calling Update Encryption Setting in LUKS2 header.')
             if extension_parameter.cmk_changed():

From b0471b60e579bd1f58e49f323703634bfe912982 Mon Sep 17 00:00:00 2001
From: PankajJoshi <pankajjoshi@microsoft.com>
Date: Thu, 11 Jan 2024 00:40:45 +0530
Subject: [PATCH 13/16] minor correction in code

---
 VMEncryption/main/DiskUtil.py | 5 +++--
 VMEncryption/main/handle.py   | 2 +-
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/VMEncryption/main/DiskUtil.py b/VMEncryption/main/DiskUtil.py
index 1f64276f3..ee08269c6 100644
--- a/VMEncryption/main/DiskUtil.py
+++ b/VMEncryption/main/DiskUtil.py
@@ -506,14 +506,15 @@ def restore_luks2_token(self, device_name=None):
             #remove backup token id
             self.remove_token(device_name=device_name,token_id=ade_token_id_backup)
             return
+        #ade_token_id_backup having value but ade_token_id_primary is none
         self.logger.log("restore luks2 token for device {0} is started.".format(device_name))
         #read from backup and update AzureDiskEncryptionToken
         data = self.read_token(device_name=device_name,token_id=ade_token_id_backup)
         data['type']=CommonVariables.AzureDiskEncryptionToken
-        self.import_token_data(device_path=device_path,token_data=data,token_id=ade_token_id_primary)
+        self.import_token_data(device_path=device_path,token_data=data,token_id=CommonVariables.cvm_ade_vm_encryption_token_id)
         #remove backup
         self.remove_token(device_name=device_name,token_id=ade_token_id_backup)
-        self.logger.log("restore luks2 token id {0} to {1} for device {2} is successful.".format(ade_token_id_backup,ade_token_id_primary,device_name))
+        self.logger.log("restore luks2 token id {0} to {1} for device {2} is successful.".format(ade_token_id_backup,CommonVariables.cvm_ade_vm_encryption_token_id,device_name))
 
     def _get_cryptsetup_version(self):
         # get version of currently installed cryptsetup
diff --git a/VMEncryption/main/handle.py b/VMEncryption/main/handle.py
index 619394cbc..3255dd8ea 100644
--- a/VMEncryption/main/handle.py
+++ b/VMEncryption/main/handle.py
@@ -333,7 +333,7 @@ def update_encryption_settings_luks2_header(extra_items_to_encrypt=None):
             ret = disk_util.import_token(device_path=device_item_path,
                                          passphrase_file=temp_keyfile.name,
                                          public_settings=public_setting,
-                                         PassphraseNameValue=CommonVariables.PassphraseNameValueProtected)
+                                         passphrase_name_value=CommonVariables.PassphraseNameValueProtected)
             if not ret:
                 logger.log(level=CommonVariables.WarningLevel,
                            msg="Update passphrase with current public setting to LUKS2 header is not successful. device path {0}".format(device_item_path))

From a651ad0c35f8a9d94835576529bbbc324975a40f Mon Sep 17 00:00:00 2001
From: PankajJoshi <pankajjoshi@microsoft.com>
Date: Thu, 11 Jan 2024 14:07:44 +0530
Subject: [PATCH 14/16] removing some spell issues and redundant code

---
 VMEncryption/main/DiskUtil.py | 2 +-
 VMEncryption/main/handle.py   | 5 +----
 2 files changed, 2 insertions(+), 5 deletions(-)

diff --git a/VMEncryption/main/DiskUtil.py b/VMEncryption/main/DiskUtil.py
index ee08269c6..f15ff1eff 100644
--- a/VMEncryption/main/DiskUtil.py
+++ b/VMEncryption/main/DiskUtil.py
@@ -487,7 +487,7 @@ def restore_luks2_token(self, device_name=None):
         '''this function restores token
         type:Azure_Disk_Encryption_BackUp, id:6 to type:Azure_Disk_Encryption id:5,
         this function acts on 4 scenarios.
-        1. both token id: 5 and 6 present in LUKS2 Tokens filed, due to reboot/interrupt during
+        1. both token id: 5 and 6 present in LUKS2 Tokens field, due to reboot/interrupt during
         KEK rotation, such case remove token id 5 has latest data so remove token id 6.
         2. token id 5 present but 6 is not present in LUKS2 Tokens field. do nothing.
         3. token id 5 not present but 6 present in LUKS2 Tokens field, restore token id 5 using
diff --git a/VMEncryption/main/handle.py b/VMEncryption/main/handle.py
index 3255dd8ea..2e038289f 100644
--- a/VMEncryption/main/handle.py
+++ b/VMEncryption/main/handle.py
@@ -295,7 +295,6 @@ def update_encryption_settings_luks2_header(extra_items_to_encrypt=None):
         encryption_config = EncryptionConfig(encryption_environment, logger)
         extension_parameter = ExtensionParameter(hutil, logger, DistroPatcher, encryption_environment, get_protected_settings(), public_setting)
         disk_util = DiskUtil(hutil=hutil, patching=DistroPatcher, logger=logger, encryption_environment=encryption_environment)
-        bek_util = BekUtil(disk_util, logger,encryption_environment)
         device_items = disk_util.get_device_items(None)
         for device_item in device_items:
             device_item_path = disk_util.get_device_path(device_item.name)
@@ -306,7 +305,7 @@ def update_encryption_settings_luks2_header(extra_items_to_encrypt=None):
             #It is necessary to restore if we are resuming from previous attempt, otherwise its no-op.
             disk_util.restore_luks2_token(device_name=device_item.name)
             logger.log("Reading passphrase from LUKS2 header, device name: {0}".format(device_item.name))
-            #keep token copy for manual recovery
+            #copy primary token to backup token for recovery, if reboot or interrupt happened during KEK rotation.
             ade_primary_token_id = disk_util.get_token_id(header_or_dev_path=device_item_path,token_name=CommonVariables.AzureDiskEncryptionToken)
             if not ade_primary_token_id:
                 logger.log("primary token type: Azure_Disk_Encryption not found for device {0}".format(device_item.name))
@@ -344,7 +343,6 @@ def update_encryption_settings_luks2_header(extra_items_to_encrypt=None):
                                    token_id=CommonVariables.cvm_ade_vm_encryption_backup_token_id)
         #committing the extension parameter if KEK rotation is successful.
         extension_parameter.commit()
-        bek_util.umount_azure_passhprase(encryption_config)
 
         if len(extra_items_to_encrypt) > 0:
             hutil.do_status_report(operation='UpdateEncryptionSettingsLuks2Header',
@@ -361,7 +359,6 @@ def update_encryption_settings_luks2_header(extra_items_to_encrypt=None):
         hutil.save_seq()
         message = "Failed to update encryption settings Luks2 header with error: {0}, stack trace: {1}".format(e, traceback.format_exc())
         logger.log(msg=message, level=CommonVariables.ErrorLevel)
-        bek_util.umount_azure_passhprase(encryption_config)
         hutil.do_exit(exit_code=CommonVariables.unknown_error,
                       operation='UpdateEncryptionSettingsLuks2Header',
                       status=CommonVariables.extension_error_status,

From e9d939b3ce99333b0992230333b56d26427a2bdc Mon Sep 17 00:00:00 2001
From: PankajJoshi <pankajjoshi@microsoft.com>
Date: Thu, 29 Feb 2024 18:50:11 +0530
Subject: [PATCH 15/16] passphrase rotation on update encryption setting for
 CVM

---
 VMEncryption/main/handle.py | 99 +++++++++++++++++++++++++++++++++++--
 1 file changed, 94 insertions(+), 5 deletions(-)

diff --git a/VMEncryption/main/handle.py b/VMEncryption/main/handle.py
index 2e038289f..00aef61c1 100644
--- a/VMEncryption/main/handle.py
+++ b/VMEncryption/main/handle.py
@@ -272,6 +272,75 @@ def get_protected_settings():
     else:
         return protected_settings_str
 
+def create_temp_file(file_content):
+    '''This function is creating a temp file of file_content. returning a file name.'''
+    temp_keyfile = tempfile.NamedTemporaryFile(delete=False)
+    if isinstance(file_content, bytes):
+        temp_keyfile.write(file_content)
+    else:
+        temp_keyfile.write(file_content.encode("utf-8"))
+    temp_keyfile.close()
+    return temp_keyfile.name
+
+def add_new_passphrase_luks2_key_slot(disk_util,\
+                        existing_passphrase,\
+                        new_passphrase,\
+                        device_path,\
+                        luks_header_path):
+    '''This function is used to add a new passphrase in luks key slot.'''
+    logger.log("add_new_passphrase_luks2_key_slot: start!")
+    ret = False
+    try:
+        before_key_slots = disk_util.luks_dump_keyslots(device_path, luks_header_path)
+        logger.log("Before key addition, key slots for {0}: {1}".format(device_path, before_key_slots))
+        logger.log("Adding new key for {0}".format(device_path))
+        existing_passphrase_file_name = create_temp_file(existing_passphrase)
+        new_passphrase_file_name = create_temp_file(new_passphrase)
+        luks_add_result = disk_util.luks_add_key(passphrase_file=existing_passphrase_file_name,
+                                            dev_path=device_path,
+                                            mapper_name=None,
+                                            header_file=luks_header_path,
+                                            new_key_path=new_passphrase_file_name)
+        logger.log("luks add result is {0}".format(luks_add_result))
+        after_key_slots = disk_util.luks_dump_keyslots(device_path, luks_header_path)
+        logger.log("After key addition, key slots for {0}: {1}".format(device_path, after_key_slots))
+        new_key_slot = list([x[0] != x[1] for x in zip(before_key_slots, after_key_slots)]).index(True)
+        logger.log("New key was added in key slot {0}".format(new_key_slot))
+        os.unlink(existing_passphrase_file_name)
+        os.unlink(new_passphrase_file_name)
+        ret = True
+    except Exception as e:
+        msg="add_new_passphrase_luks2_key_slot failed with error {0}, stack trace: {1}".format(e, traceback.format_exc())
+        logger.log(msg=msg,level=CommonVariables.WarningLevel)
+    logger.log("add_new_passphrase_luks2_key_slot: end!")
+    return ret
+
+def remove_passphrase_luks2_key_slot(disk_util,\
+                        passphrase,\
+                        device_path,\
+                        luks_header_path):
+    '''This function is used for removing a passphrase from luks key slot.'''
+    logger.log("remove_passphrase_luks2_key_slot: start!")
+    ret = False
+    try:
+        before_key_slots = disk_util.luks_dump_keyslots(device_path, luks_header_path)
+        logger.log("Before key removal, key slots for {0}: {1}".format(device_path, before_key_slots))
+        logger.log("Removing new key for {0}".format(device_path))
+        passphrase_file_name = create_temp_file(passphrase)
+        luks_remove_result = disk_util.luks_remove_key(passphrase_file=passphrase_file_name,
+                                                    dev_path=device_path,
+                                                    header_file=luks_header_path)
+        logger.log("luks remove result is {0}".format(luks_remove_result))
+        after_key_slots = disk_util.luks_dump_keyslots(device_path, luks_header_path)
+        logger.log("After key removal, key slots for {0}: {1}".format(device_path, after_key_slots))
+        os.unlink(passphrase_file_name)
+        ret = True
+    except Exception as e:
+        msg="remove_passphrase_luks2_key_slot failed with error {0}, stack trace: {1}".format(e, traceback.format_exc())
+        logger.log(msg=msg,level=CommonVariables.WarningLevel)
+    logger.log("remove_passphrase_luks2_key_slot: end!")
+    return ret
+
 def update_encryption_settings_luks2_header(extra_items_to_encrypt=None):
     '''This function is used for CMK passphrase wrapping with new KEK URL and update
     metadata in LUKS2 header.'''
@@ -295,7 +364,11 @@ def update_encryption_settings_luks2_header(extra_items_to_encrypt=None):
         encryption_config = EncryptionConfig(encryption_environment, logger)
         extension_parameter = ExtensionParameter(hutil, logger, DistroPatcher, encryption_environment, get_protected_settings(), public_setting)
         disk_util = DiskUtil(hutil=hutil, patching=DistroPatcher, logger=logger, encryption_environment=encryption_environment)
+        bek_util = BekUtil(disk_util, logger,encryption_environment)
         device_items = disk_util.get_device_items(None)
+        if extension_parameter.passphrase is None or extension_parameter.passphrase == "":
+            extension_parameter.passphrase = bek_util.generate_passphrase()
+
         for device_item in device_items:
             device_item_path = disk_util.get_device_path(device_item.name)
             if not disk_util.is_luks_device(device_item_path,None):
@@ -324,20 +397,36 @@ def update_encryption_settings_luks2_header(extra_items_to_encrypt=None):
             #remove primary token from Tokens field of LUKS2 header.
             disk_util.remove_token(device_name=device_item.name,token_id=ade_primary_token_id)
             logger.log("Updating wrapped passphrase to LUKS2 header with current public setting. device name {0}".format(device_item.name))
+            #add new slot with new passphrase.
+            is_added = add_new_passphrase_luks2_key_slot(disk_util=disk_util,
+                                existing_passphrase=passphrase,
+                                new_passphrase=extension_parameter.passphrase,
+                                device_path= device_item_path,
+                                luks_header_path=None)
+            if not is_added:
+                logger.log(level=CommonVariables.WarningLevel,
+                           msg="new passphrase is not added to LUKS2 slot. Skip operation for device: {0}".format(device_item.name))
+                continue
             #protect passphrase before updating to LUKS2 is done in import_token
-            temp_keyfile = tempfile.NamedTemporaryFile(delete=False)
-            temp_keyfile.write(passphrase.encode("utf-8"))
-            temp_keyfile.close()
+            new_passphrase_file = create_temp_file(extension_parameter.passphrase)
             #save passphrase to LUKS2 header with PassphraseNameValueProtected
             ret = disk_util.import_token(device_path=device_item_path,
-                                         passphrase_file=temp_keyfile.name,
+                                         passphrase_file=new_passphrase_file,
                                          public_settings=public_setting,
                                          passphrase_name_value=CommonVariables.PassphraseNameValueProtected)
             if not ret:
                 logger.log(level=CommonVariables.WarningLevel,
                            msg="Update passphrase with current public setting to LUKS2 header is not successful. device path {0}".format(device_item_path))
                 return None
-            os.unlink(temp_keyfile.name)
+            os.unlink(new_passphrase_file)
+            #removing old password form key slot.
+            is_removed = remove_passphrase_luks2_key_slot(disk_util=disk_util,
+                                                          passphrase=passphrase,
+                                                          device_path=device_item_path,
+                                                          luks_header_path=None)
+            if not is_removed:
+                logger.log(level=CommonVariables.WarningLevel,
+                           msg="old passphrase is not removed from LUKS2 slot. Skip operation for device: {0}".format(device_item.name))
             #removing backup token as KEK rotation is successful here.
             disk_util.remove_token(device_name=device_item.name,
                                    token_id=CommonVariables.cvm_ade_vm_encryption_backup_token_id)

From 4e856270425954b1997f0bcae9d3cfa4bae30ebb Mon Sep 17 00:00:00 2001
From: PankajJoshi <pankajjoshi@microsoft.com>
Date: Fri, 1 Mar 2024 14:21:05 +0530
Subject: [PATCH 16/16] storing passphrase to root-fs post to sucessful
 passphrase rotation.

---
 VMEncryption/main/BekUtil.py | 16 +++++++++++++++-
 VMEncryption/main/handle.py  |  9 +++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/VMEncryption/main/BekUtil.py b/VMEncryption/main/BekUtil.py
index 20336a3ca..91ca51c25 100644
--- a/VMEncryption/main/BekUtil.py
+++ b/VMEncryption/main/BekUtil.py
@@ -15,7 +15,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+import sys
+import os
 from Common import CommonVariables
 from IMDSUtil import IMDSStoredResults
 from BekUtilVolumeImpl import BekUtilVolumeImpl
@@ -43,6 +44,19 @@ def __init__(self, disk_util, logger, encryption_environment=None):
     def generate_passphrase(self):
         return self.bekUtilImpl.generate_passphrase()
 
+    def store_bek_passphrase_file_name(self,encryption_config, passphrase,key_file_name):
+        '''this function is used to store passphrase to specific file name'''
+        #update new passphrase to key_file_path.
+        key_file_dir = os.path.dirname(self.get_bek_passphrase_file(encryption_config))
+        key_file_path = os.path.join(key_file_dir,key_file_name)
+        if sys.version_info[0] < 3:
+            if isinstance(passphrase, str):
+                passphrase = passphrase.decode('utf-8')
+        with open(key_file_path, 'wb') as f:
+            f.write(passphrase)
+        #making sure the permissions are read only to root user.
+        os.chmod(key_file_path,0o400)
+
     def store_bek_passphrase(self, encryption_config, passphrase):
         return self.bekUtilImpl.store_bek_passphrase(encryption_config,passphrase)
 
diff --git a/VMEncryption/main/handle.py b/VMEncryption/main/handle.py
index 00aef61c1..a2bd07fed 100644
--- a/VMEncryption/main/handle.py
+++ b/VMEncryption/main/handle.py
@@ -430,6 +430,15 @@ def update_encryption_settings_luks2_header(extra_items_to_encrypt=None):
             #removing backup token as KEK rotation is successful here.
             disk_util.remove_token(device_name=device_item.name,
                                    token_id=CommonVariables.cvm_ade_vm_encryption_backup_token_id)
+            #update passphrase file for auto unlock
+            key_file_name = CommonVariables.encryption_key_file_name
+            scsi_lun_numbers = disk_util.get_azure_data_disk_controller_and_lun_numbers([os.path.realpath(device_item_path)])
+            if len(scsi_lun_numbers) != 0:
+                scsi_controller, lun_number = scsi_lun_numbers[0]
+                key_file_name = "{0}_{1}_{2}".format(key_file_name,str(scsi_controller),str(lun_number))
+            bek_util.store_bek_passphrase_file_name(encryption_config=encryption_config,
+                                                    passphrase=extension_parameter.passphrase,
+                                                    key_file_name=key_file_name)
         #committing the extension parameter if KEK rotation is successful.
         extension_parameter.commit()