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/04 10:36:20 UTC

[incubator-datalab] 02/04: [DATALAB-2674]: added block project ssh keys for gcp instances

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

commit 7e282286d8f98e70b2fff5e3ab7532bbdd1df163
Author: leonidfrolov <fr...@gmail.com>
AuthorDate: Tue Feb 1 14:19:00 2022 +0200

    [DATALAB-2674]: added block project ssh keys for gcp instances
---
 infrastructure-provisioning/scripts/deploy_datalab.py        |  2 ++
 .../src/general/lib/gcp/actions_lib.py                       |  6 +++++-
 .../src/general/scripts/gcp/common_create_instance.py        |  3 ++-
 .../src/general/scripts/gcp/common_prepare_notebook.py       |  8 ++++++--
 .../src/general/scripts/gcp/dataengine_prepare.py            | 12 ++++++++----
 .../src/general/scripts/gcp/project_prepare.py               |  9 +++++++--
 .../src/general/scripts/gcp/ssn_prepare.py                   |  9 +++++++--
 7 files changed, 37 insertions(+), 12 deletions(-)

diff --git a/infrastructure-provisioning/scripts/deploy_datalab.py b/infrastructure-provisioning/scripts/deploy_datalab.py
index c039834..40b233d 100644
--- a/infrastructure-provisioning/scripts/deploy_datalab.py
+++ b/infrastructure-provisioning/scripts/deploy_datalab.py
@@ -261,6 +261,8 @@ def build_parser():
                                    help='The SSN instance shape')
     gcp_parser.add_argument('--gcp_os_login_enabled', type=str, default='FALSE',
                             help='TRUE if os login enabled for gcp instances')
+    gcp_parser.add_argument('--gcp_block_project_ssh_keys', type=str, default='FALSE',
+                            help='TRUE to block project ssh keys for gcp instances')
 
     gcp_required_args = gcp_parser.add_argument_group('Required arguments')
     gcp_required_args.add_argument('--gcp_region', type=str, required=True, help='GCP region')
diff --git a/infrastructure-provisioning/src/general/lib/gcp/actions_lib.py b/infrastructure-provisioning/src/general/lib/gcp/actions_lib.py
index 37e8d30..025f1b5 100644
--- a/infrastructure-provisioning/src/general/lib/gcp/actions_lib.py
+++ b/infrastructure-provisioning/src/general/lib/gcp/actions_lib.py
@@ -320,7 +320,7 @@ class GCPActions:
                         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'):
+                        gpu_accelerator_type='None', gpu_accelerator_count='1', os_login_enabled='FALSE', block_project_ssh_keys='FALSE'):
         key = RSA.importKey(open(ssh_key_path, 'rb').read())
         ssh_key = key.publickey().exportKey("OpenSSH").decode('UTF-8')
         unique_index = datalab.meta_lib.GCPMeta().get_index_by_service_account_name(service_account_name)
@@ -428,6 +428,10 @@ class GCPActions:
                     {
                         "key": "enable-oslogin",
                         "value": "{}".format(os_login_enabled)
+                    },
+                    {
+                        "key": "block-project-ssh-keys",
+                        "value": "{}".format(block_project_ssh_keys)
                     }
                 ]
                 },
diff --git a/infrastructure-provisioning/src/general/scripts/gcp/common_create_instance.py b/infrastructure-provisioning/src/general/scripts/gcp/common_create_instance.py
index 246675d..adf2bf5 100644
--- a/infrastructure-provisioning/src/general/scripts/gcp/common_create_instance.py
+++ b/infrastructure-provisioning/src/general/scripts/gcp/common_create_instance.py
@@ -51,6 +51,7 @@ parser.add_argument('--network_tag', type=str, default='')
 parser.add_argument('--cluster_name', type=str, default='')
 parser.add_argument('--service_base_name', type=str, default='')
 parser.add_argument('--os_login_enabled', type=str, default='FALSE')
+parser.add_argument('--block_project_ssh_keys', type=str, default='FALSE')
 args = parser.parse_args()
 
 
@@ -66,7 +67,7 @@ if __name__ == "__main__":
                                          args.secondary_image_name, args.service_account_name, args.instance_class,
                                          args.network_tag, json.loads(args.labels), args.static_ip,
                                          args.primary_disk_size, args.secondary_disk_size, args.gpu_accelerator_type,
-                                         args.gpu_accelerator_count, args.os_login_enabled)
+                                         args.gpu_accelerator_count, args.os_login_enabled, args.block_project_ssh_keys)
     else:
         parser.print_help()
         sys.exit(2)
diff --git a/infrastructure-provisioning/src/general/scripts/gcp/common_prepare_notebook.py b/infrastructure-provisioning/src/general/scripts/gcp/common_prepare_notebook.py
index 9954592..5925480 100644
--- a/infrastructure-provisioning/src/general/scripts/gcp/common_prepare_notebook.py
+++ b/infrastructure-provisioning/src/general/scripts/gcp/common_prepare_notebook.py
@@ -151,6 +151,9 @@ if __name__ == "__main__":
         if os.environ['gcp_os_login_enabled'] != 'FALSE':
             notebook_config['gcp_os_login_enabled'] = 'TRUE'
 
+        if os.environ['gcp_block_project_ssh_keys'] != 'FALSE':
+            notebook_config['gcp_block_project_ssh_keys'] = 'TRUE'
+
         notebook_config['gpu_accelerator_type'] = 'None'
         notebook_config['gpu_accelerator_count'] = 'None'
 
@@ -196,7 +199,7 @@ if __name__ == "__main__":
                  "--ssh_key_path {6} --initial_user {7} --service_account_name {8} --image_name {9} " \
                  "--secondary_image_name {10} --instance_class {11} --primary_disk_size {12} " \
                  "--secondary_disk_size {13} --gpu_accelerator_type {14} --gpu_accelerator_count {15} --network_tag {16} --labels '{17}' " \
-                 "--service_base_name {18} --os_login_enabled {19}".\
+                 "--service_base_name {18} --os_login_enabled {19} --block_project_ssh_keys {20}".\
             format(notebook_config['instance_name'], notebook_config['region'], notebook_config['zone'],
                    notebook_config['vpc_name'], notebook_config['subnet_name'], notebook_config['instance_size'],
                    notebook_config['ssh_key_path'], notebook_config['initial_user'],
@@ -204,7 +207,8 @@ if __name__ == "__main__":
                    notebook_config['secondary_image_name'], 'notebook', notebook_config['primary_disk_size'],
                    notebook_config['secondary_disk_size'], notebook_config['gpu_accelerator_type'],
                    notebook_config['gpu_accelerator_count'], notebook_config['network_tag'],
-                   json.dumps(notebook_config['labels']), notebook_config['service_base_name'], notebook_config['gcp_os_login_enabled'])
+                   json.dumps(notebook_config['labels']), notebook_config['service_base_name'],
+                   notebook_config['gcp_os_login_enabled'], notebook_config['gcp_block_project_ssh_keys'])
         try:
             subprocess.run("~/scripts/{}.py {}".format('common_create_instance', params), shell=True, check=True)
         except:
diff --git a/infrastructure-provisioning/src/general/scripts/gcp/dataengine_prepare.py b/infrastructure-provisioning/src/general/scripts/gcp/dataengine_prepare.py
index 643354d..a549dee 100644
--- a/infrastructure-provisioning/src/general/scripts/gcp/dataengine_prepare.py
+++ b/infrastructure-provisioning/src/general/scripts/gcp/dataengine_prepare.py
@@ -101,6 +101,9 @@ if __name__ == "__main__":
         if os.environ['gcp_os_login_enabled'] != 'FALSE':
             data_engine['gcp_os_login_enabled'] = 'TRUE'
 
+        if os.environ['gcp_block_project_ssh_keys'] != 'FALSE':
+            data_engine['gcp_block_project_ssh_keys'] = 'TRUE'
+
         data_engine['cluster_name'] = "{}-{}-{}-de-{}".format(data_engine['service_base_name'],
                                                               data_engine['project_name'],
                                                               data_engine['endpoint_name'],
@@ -193,14 +196,15 @@ if __name__ == "__main__":
                  "--ssh_key_path {6} --initial_user {7} --service_account_name {8} --image_name {9} " \
                  "--secondary_image_name {10} --instance_class {11} --primary_disk_size {12} " \
                  "--secondary_disk_size {13} --gpu_accelerator_type {14} --gpu_accelerator_count {15} --network_tag {16} --cluster_name {17} " \
-                 "--labels '{18}' --service_base_name {19} --os_login_enabled {20}". \
+                 "--labels '{18}' --service_base_name {19} --os_login_enabled {20} --block_project_ssh_keys {21}". \
             format(data_engine['master_node_name'], data_engine['region'], data_engine['zone'], data_engine['vpc_name'],
                    data_engine['subnet_name'], data_engine['master_size'], data_engine['ssh_key_path'], initial_user,
                    data_engine['dataengine_service_account_name'], data_engine['primary_image_name'],
                    data_engine['secondary_image_name'], 'dataengine', data_engine['primary_disk_size'],
                    data_engine['secondary_disk_size'], data_engine['gpu_master_accelerator_type'],
                    data_engine['gpu_master_accelerator_count'], data_engine['network_tag'], data_engine['cluster_name'],
-                   json.dumps(data_engine['master_labels']), data_engine['service_base_name'], data_engine['gcp_os_login_enabled'])
+                   json.dumps(data_engine['master_labels']), data_engine['service_base_name'],
+                   data_engine['gcp_os_login_enabled'], data_engine['gcp_block_project_ssh_keys'])
         try:
             subprocess.run("~/scripts/{}.py {}".format('common_create_instance', params), shell=True, check=True)
         except:
@@ -219,7 +223,7 @@ if __name__ == "__main__":
                      "--instance_size {5} --ssh_key_path {6} --initial_user {7} --service_account_name {8} " \
                      "--image_name {9} --secondary_image_name {10} --instance_class {11} --primary_disk_size {12} " \
                      "--secondary_disk_size {13} --gpu_accelerator_type {14} --gpu_accelerator_count {15} --network_tag {16} --cluster_name {17} " \
-                     "--labels '{18}' --service_base_name {19} --os_login_enabled {20}". \
+                     "--labels '{18}' --service_base_name {19} --os_login_enabled {20} --block_project_ssh_keys {21}". \
                 format(slave_name, data_engine['region'], data_engine['zone'],
                        data_engine['vpc_name'], data_engine['subnet_name'], data_engine['slave_size'],
                        data_engine['ssh_key_path'], initial_user, data_engine['dataengine_service_account_name'],
@@ -228,7 +232,7 @@ if __name__ == "__main__":
                        data_engine['secondary_disk_size'], data_engine['gpu_slave_accelerator_type'],
                        data_engine['gpu_slave_accelerator_count'], data_engine['network_tag'],
                        data_engine['cluster_name'], json.dumps(data_engine['slave_labels']),
-                       data_engine['service_base_name'], data_engine['gcp_os_login_enabled'])
+                       data_engine['service_base_name'], data_engine['gcp_os_login_enabled'], data_engine['gcp_block_project_ssh_keys'])
             try:
                 subprocess.run("~/scripts/{}.py {}".format('common_create_instance', params), shell=True, check=True)
             except:
diff --git a/infrastructure-provisioning/src/general/scripts/gcp/project_prepare.py b/infrastructure-provisioning/src/general/scripts/gcp/project_prepare.py
index 02e9667..028087d 100644
--- a/infrastructure-provisioning/src/general/scripts/gcp/project_prepare.py
+++ b/infrastructure-provisioning/src/general/scripts/gcp/project_prepare.py
@@ -503,18 +503,23 @@ if __name__ == "__main__":
     if os.environ['gcp_os_login_enabled'] != 'FALSE':
         project_conf['gcp_os_login_enabled'] = 'TRUE'
 
+    if os.environ['gcp_block_project_ssh_keys'] != 'FALSE':
+        project_conf['gcp_block_project_ssh_keys'] = 'TRUE'
+
     try:
         project_conf['static_ip'] = \
             GCPMeta.get_static_address(project_conf['region'], project_conf['static_address_name'])['address']
         logging.info('[CREATE EDGE INSTANCE]')
         params = "--instance_name {} --region {} --zone {} --vpc_name {} --subnet_name {} --instance_size {} " \
                  "--ssh_key_path {} --initial_user {} --service_account_name {} --image_name {} --instance_class {} " \
-                 "--static_ip {} --network_tag {} --labels '{}' --service_base_name {} --os_login_enabled {}".format(
+                 "--static_ip {} --network_tag {} --labels '{}' --service_base_name {} --os_login_enabled {} " \
+                 "--block_project_ssh_keys {}".format(
                   project_conf['instance_name'], project_conf['region'], project_conf['zone'], project_conf['vpc_name'],
                   project_conf['subnet_name'], project_conf['instance_size'], project_conf['ssh_key_path'],
                   project_conf['initial_user'], project_conf['edge_service_account_name'], project_conf['image_name'],
                   'edge', project_conf['static_ip'], project_conf['network_tag'],
-                  json.dumps(project_conf['instance_labels']), project_conf['service_base_name'], project_conf['gcp_os_login_enabled'])
+                  json.dumps(project_conf['instance_labels']), project_conf['service_base_name'],
+                  project_conf['gcp_os_login_enabled'], project_conf['gcp_block_project_ssh_keys'])
         try:
             subprocess.run("~/scripts/{}.py {}".format('common_create_instance', params), shell=True, check=True)
         except:
diff --git a/infrastructure-provisioning/src/general/scripts/gcp/ssn_prepare.py b/infrastructure-provisioning/src/general/scripts/gcp/ssn_prepare.py
index 9346d15..f4fba87 100644
--- a/infrastructure-provisioning/src/general/scripts/gcp/ssn_prepare.py
+++ b/infrastructure-provisioning/src/general/scripts/gcp/ssn_prepare.py
@@ -254,6 +254,9 @@ if __name__ == "__main__":
     if os.environ['gcp_os_login_enabled'] != 'FALSE':
         ssn_conf['gcp_os_login_enabled'] = 'TRUE'
 
+    if os.environ['gcp_block_project_ssh_keys'] != 'FALSE':
+        ssn_conf['gcp_block_project_ssh_keys'] = 'TRUE'
+
     if os.environ['conf_os_family'] == 'debian':
         ssn_conf['initial_user'] = 'ubuntu'
         ssn_conf['sudo_group'] = 'sudo'
@@ -268,12 +271,14 @@ if __name__ == "__main__":
         params = "--instance_name {0} --region {1} --zone {2} --vpc_name {3} --subnet_name {4} --instance_size {5}"\
                  " --ssh_key_path {6} --initial_user {7} --service_account_name {8} --image_name {9}"\
                  " --instance_class {10} --static_ip {11} --network_tag {12} --labels '{13}' " \
-                 "--primary_disk_size {14} --service_base_name {15} --os_login_enabled {16}".\
+                 "--primary_disk_size {14} --service_base_name {15} --os_login_enabled {16} " \
+                 "--block_project_ssh_keys {17}".\
             format(ssn_conf['instance_name'], ssn_conf['region'], ssn_conf['zone'], ssn_conf['vpc_name'],
                    ssn_conf['subnet_name'], ssn_conf['instance_size'], ssn_conf['ssh_key_path'],
                    ssn_conf['initial_user'], ssn_conf['service_account_name'], ssn_conf['image_name'], 'ssn',
                    ssn_conf['static_ip'], ssn_conf['network_tag'], json.dumps(ssn_conf['instance_labels']), '20',
-                   ssn_conf['service_base_name'], ssn_conf['gcp_os_login_enabled'])
+                   ssn_conf['service_base_name'], ssn_conf['gcp_os_login_enabled'],
+                   ssn_conf['gcp_block_project_ssh_keys'])
         try:
             subprocess.run("~/scripts/{}.py {}".format('common_create_instance', params), shell=True, check=True)
         except:

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