You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datalab.apache.org by lf...@apache.org on 2022/02/10 12:44:38 UTC

[incubator-datalab] branch DATALAB-2674 updated: [DATALAB-2674]: fixed encrypted disk/image usage

This is an automated email from the ASF dual-hosted git repository.

lfrolov pushed a commit to branch DATALAB-2674
in repository https://gitbox.apache.org/repos/asf/incubator-datalab.git


The following commit(s) were added to refs/heads/DATALAB-2674 by this push:
     new a8e3c23  [DATALAB-2674]: fixed encrypted disk/image usage
a8e3c23 is described below

commit a8e3c236ea85d00ba381d536b269fca3a0f65b44
Author: leonidfrolov <fr...@gmail.com>
AuthorDate: Thu Feb 10 14:44:28 2022 +0200

    [DATALAB-2674]: fixed encrypted disk/image usage
---
 .../src/general/lib/gcp/actions_lib.py             | 38 +++++++++++++++-------
 .../general/scripts/gcp/common_start_notebook.py   |  7 +++-
 .../src/general/scripts/gcp/dataengine_start.py    |  6 +++-
 .../src/general/scripts/gcp/edge_start.py          |  6 +++-
 4 files changed, 42 insertions(+), 15 deletions(-)

diff --git a/infrastructure-provisioning/src/general/lib/gcp/actions_lib.py b/infrastructure-provisioning/src/general/lib/gcp/actions_lib.py
index 53df959..1ff8a4a 100644
--- a/infrastructure-provisioning/src/general/lib/gcp/actions_lib.py
+++ b/infrastructure-provisioning/src/general/lib/gcp/actions_lib.py
@@ -279,12 +279,16 @@ class GCPActions:
             if secondary_image_name == 'None':
                 params = {"sizeGb": size, "name": instance_name + '-secondary',
                           "type": "projects/{0}/zones/{1}/diskTypes/pd-ssd".format(self.project, zone)}
+                if rsa_encrypted_csek:
+                    params['diskEncryptionKey'] = {"rsaEncryptedKey": rsa_encrypted_csek}
+
             else:
                 params = {"sizeGb": size, "name": instance_name + '-secondary',
                           "type": "projects/{0}/zones/{1}/diskTypes/pd-ssd".format(self.project, zone),
                           "sourceImage": secondary_image_name}
-            if rsa_encrypted_csek:
-                params['diskEncryptionKey'] = {"rsaEncryptedKey": rsa_encrypted_csek}
+                if rsa_encrypted_csek:
+                    params["sourceImageEncryptionKey"] = {"rsaEncryptedKey": rsa_encrypted_csek}
+                    params['diskEncryptionKey'] = {"rsaEncryptedKey": rsa_encrypted_csek}
             request = self.service.disks().insert(project=self.project, zone=zone, body=params)
             result = request.execute()
             datalab.meta_lib.GCPMeta().wait_for_operation(result['name'], zone=zone)
@@ -320,10 +324,9 @@ class GCPActions:
                                    file=sys.stdout)}))
             traceback.print_exc(file=sys.stdout)
 
-    def create_instance(self, instance_name, service_base_name, cluster_name, region, zone, vpc_name, subnet_name, instance_size,
-                        ssh_key_path,
-                        initial_user, image_name, secondary_image_name, service_account_name, instance_class,
-                        network_tag, labels, static_ip='',
+    def create_instance(self, instance_name, service_base_name, cluster_name, region, zone, vpc_name, subnet_name,
+                        instance_size, ssh_key_path, initial_user, image_name, secondary_image_name,
+                        service_account_name, instance_class, network_tag, labels, static_ip='',
                         primary_disk_size='12', secondary_disk_size='30',
                         gpu_accelerator_type='None', gpu_accelerator_count='1',
                         os_login_enabled='FALSE', block_project_ssh_keys='FALSE', rsa_encrypted_csek=''):
@@ -416,7 +419,8 @@ class GCPActions:
 
         if service_base_name in image_name and rsa_encrypted_csek:
             for disk in disks:
-                disk["initializeParams"]["sourceImageEncryptionKey"] = {"rsaEncryptedKey": rsa_encrypted_csek}
+                if "initializeParams" in disk:
+                    disk["initializeParams"]["sourceImageEncryptionKey"] = {"rsaEncryptedKey": rsa_encrypted_csek}
                 disk["diskEncryptionKey"] = {"rsaEncryptedKey": rsa_encrypted_csek}
         elif rsa_encrypted_csek:
             for disk in disks:
@@ -560,8 +564,18 @@ class GCPActions:
                                    file=sys.stdout)}))
             traceback.print_exc(file=sys.stdout)
 
-    def start_instance(self, instance_name, zone):
-        request = self.service.instances().start(project=self.project, zone=zone, instance=instance_name)
+    def start_instance(self, instance_name, zone, rsa_encrypted_csek=''):
+        if rsa_encrypted_csek:
+            params = dict()
+            params['disks'] = list()
+            instance_data = datalab.meta_lib.GCPMeta().get_instance(instance_name)
+            for disk in instance_data['disks']:
+                params["disks"].append(
+                    {"diskEncryptionKey": {"rsaEncryptedKey": rsa_encrypted_csek}, "source": disk['source']})
+            request = self.service.instances().startWithEncryptionKey(project=self.project, zone=zone,
+                                                                      instance=instance_name, body=params)
+        else:
+            request = self.service.instances().start(project=self.project, zone=zone, instance=instance_name)
         try:
             result = request.execute()
             datalab.meta_lib.GCPMeta().wait_for_operation(result['name'], zone=zone)
@@ -831,12 +845,12 @@ class GCPActions:
             secondary_params["imageEncryptionKey"] = {"rsaEncryptedKey": rsa_encrypted_csek}
             secondary_params["sourceDiskEncryptionKey"] = {"rsaEncryptedKey": rsa_encrypted_csek}
         secondary_request = self.service.images().insert(project=self.project, body=secondary_params)
-        id_list=[]
+        id_list = []
         try:
             GCPActions().stop_instance(instance_name, zone)
             primary_image_check = datalab.meta_lib.GCPMeta().get_image_by_name(primary_image_name)
             if primary_image_check != '':
-                GCPActions().start_instance(instance_name, zone)
+                GCPActions().start_instance(instance_name, zone, rsa_encrypted_csek)
                 return ''
             primary_result = primary_request.execute()
             secondary_result = secondary_request.execute()
@@ -846,7 +860,7 @@ class GCPActions:
             datalab.meta_lib.GCPMeta().wait_for_operation(secondary_result['name'])
             print('Image {} has been created.'.format(secondary_image_name))
             id_list.append(secondary_result.get('id'))
-            GCPActions().start_instance(instance_name, zone)
+            GCPActions().start_instance(instance_name, zone, rsa_encrypted_csek)
             return id_list
         except Exception as err:
             logging.info(
diff --git a/infrastructure-provisioning/src/general/scripts/gcp/common_start_notebook.py b/infrastructure-provisioning/src/general/scripts/gcp/common_start_notebook.py
index f63d6fb..e854681 100644
--- a/infrastructure-provisioning/src/general/scripts/gcp/common_start_notebook.py
+++ b/infrastructure-provisioning/src/general/scripts/gcp/common_start_notebook.py
@@ -41,12 +41,17 @@ if __name__ == "__main__":
     notebook_config['service_base_name'] = (os.environ['conf_service_base_name'])
     notebook_config['notebook_name'] = os.environ['notebook_instance_name']
     notebook_config['zone'] = os.environ['gcp_zone']
+    if "gcp_wrapped_csek" in os.environ:
+        notebook_config['gcp_wrapped_csek'] = os.environ['gcp_wrapped_csek']
+    else:
+        notebook_config['gcp_wrapped_csek'] = ''
 
     try:
         logging.info('[START NOTEBOOK]')
         try:
             logging.info("Starting notebook")
-            GCPActions.start_instance(notebook_config['notebook_name'], notebook_config['zone'])
+            GCPActions.start_instance(notebook_config['notebook_name'], notebook_config['zone'],
+                                      notebook_config['gcp_wrapped_csek'])
         except Exception as err:
             traceback.print_exc()
             datalab.fab.append_result("Failed to start notebook.", str(err))
diff --git a/infrastructure-provisioning/src/general/scripts/gcp/dataengine_start.py b/infrastructure-provisioning/src/general/scripts/gcp/dataengine_start.py
index 865d846..daef908 100644
--- a/infrastructure-provisioning/src/general/scripts/gcp/dataengine_start.py
+++ b/infrastructure-provisioning/src/general/scripts/gcp/dataengine_start.py
@@ -39,7 +39,7 @@ def start_data_engine(zone, cluster_name):
         instances = GCPMeta.get_list_instances(zone, cluster_name)
         if 'items' in instances:
             for i in instances['items']:
-                GCPActions.start_instance(i['name'], zone)
+                GCPActions.start_instance(i['name'], zone, data_engine['gcp_wrapped_csek'])
     except Exception as err:
         datalab.fab.append_result("Failed to start dataengine", str(err))
         sys.exit(1)
@@ -68,6 +68,10 @@ if __name__ == "__main__":
                                                           data_engine['project_name'],
                                                           data_engine['endpoint_name'],
                                                           data_engine['computational_name'])
+    if "gcp_wrapped_csek" in os.environ:
+        data_engine['gcp_wrapped_csek'] = os.environ['gcp_wrapped_csek']
+    else:
+        data_engine['gcp_wrapped_csek'] = ''
     try:
         logging.info('[STARTING DATA ENGINE]')
         try:
diff --git a/infrastructure-provisioning/src/general/scripts/gcp/edge_start.py b/infrastructure-provisioning/src/general/scripts/gcp/edge_start.py
index 2d35732..8dea7b6 100644
--- a/infrastructure-provisioning/src/general/scripts/gcp/edge_start.py
+++ b/infrastructure-provisioning/src/general/scripts/gcp/edge_start.py
@@ -45,10 +45,14 @@ if __name__ == "__main__":
     edge_conf['static_address_name'] = '{0}-{1}-{2}-static-ip'.format(edge_conf['service_base_name'],
                                                                edge_conf['project_name'],
                                                                edge_conf['endpoint_name'])
+    if "gcp_wrapped_csek" in os.environ:
+        edge_conf['gcp_wrapped_csek'] = os.environ['gcp_wrapped_csek']
+    else:
+        edge_conf['gcp_wrapped_csek'] = ''
 
     logging.info('[START EDGE]')
     try:
-        GCPActions.start_instance(edge_conf['instance_name'], edge_conf['zone'])
+        GCPActions.start_instance(edge_conf['instance_name'], edge_conf['zone'], edge_conf['gcp_wrapped_csek'])
     except Exception as err:
         datalab.fab.append_result("Failed to start edge.", str(err))
         sys.exit(1)

---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datalab.apache.org
For additional commands, e-mail: commits-help@datalab.apache.org