From aa4fd5e1c0f3fc90af8d6268a1571d833c72a737 Mon Sep 17 00:00:00 2001
From: "C.J. Collier" <cjac@colliertech.org>
Date: Mon, 7 Oct 2024 16:19:10 -0700
Subject: [PATCH] delete install image after customization instance has been
 started (#93)

* deleting install image after customization instance has been started

* removed zone argument, reduced to single line, repositioned -q

* service account is per-purpose instead of per-cluster

* grant access to bucket before run ; revoke access to bucket, private key material and secret list access after the run

* reduced noise following successful image creation message
---
 custom_image_utils/shell_script_generator.py |  4 ++
 examples/secure-boot/cuda.sh                 | 45 +++++++++++++++-----
 examples/secure-boot/env.json.sample         |  2 +-
 3 files changed, 39 insertions(+), 12 deletions(-)

diff --git a/custom_image_utils/shell_script_generator.py b/custom_image_utils/shell_script_generator.py
index 9bedac6..d9f3140 100644
--- a/custom_image_utils/shell_script_generator.py
+++ b/custom_image_utils/shell_script_generator.py
@@ -120,6 +120,9 @@
 
   touch /tmp/{run_id}/vm_created
 
+  # clean up intermediate install image
+  gcloud compute images delete -q {image_name}-install --project={project_id}
+
   echo 'Waiting for customization script to finish and VM shutdown.'
   gcloud compute instances tail-serial-port-output {image_name}-install \
       --project={project_id} \
@@ -148,6 +151,7 @@
     --source-disk={image_name}-install \
     {storage_location_flag} \
     --family={family}
+
   touch /tmp/{run_id}/image_created
 }}
 
diff --git a/examples/secure-boot/cuda.sh b/examples/secure-boot/cuda.sh
index e9a09c5..ba04a5f 100644
--- a/examples/secure-boot/cuda.sh
+++ b/examples/secure-boot/cuda.sh
@@ -17,7 +17,7 @@
 set -ex
 
 export PROJECT_ID="$(jq    -r .PROJECT_ID    env.json)"
-export CLUSTER_NAME="$(jq  -r .CLUSTER_NAME  env.json)"
+export PURPOSE="$(jq       -r .PURPOSE       env.json)"
 export BUCKET="$(jq        -r .BUCKET        env.json)"
 export IMAGE_VERSION="$(jq -r .IMAGE_VERSION env.json)"
 export ZONE="$(jq          -r .ZONE          env.json)"
@@ -25,7 +25,7 @@ export ZONE="$(jq          -r .ZONE          env.json)"
 custom_image_zone="${ZONE}"
 disk_size_gb="50" # greater than or equal to 30
 
-SA_NAME="sa-${CLUSTER_NAME}"
+SA_NAME="sa-${PURPOSE}"
 GSA="${SA_NAME}@${PROJECT_ID}.iam.gserviceaccount.com"
 
 gcloud config set project ${PROJECT_ID}
@@ -40,23 +40,33 @@ metadata="${metadata},private_secret_name=${private_secret_name}"
 metadata="${metadata},secret_project=${secret_project}"
 metadata="${metadata},secret_version=${secret_version}"
 
-# Instructions for creating the service account can be found here:
-# https://github.com/LLC-Technologies-Collier/dataproc-repro/blob/78945b5954ab47aac56f55ac22b3c35569d154e0/shared-functions.sh#L759
+if gcloud iam service-accounts list --filter email="${GSA}" 2>&1 | grep 'Listed 0 items.' ; then
+  # Create service account for this purpose
+  echo "creating pre-init customization service account ${GSA}"
+  gcloud iam service-accounts create "${SA_NAME}" \
+    --description="Service account for pre-init customization" \
+    --display-name="${SA_NAME}"
+fi
+
+# Grant service account access to bucket
+gcloud storage buckets add-iam-policy-binding "gs://${BUCKET}" \
+  --member="serviceAccount:${GSA}" \
+  --role="roles/storage.objectViewer"
 
 # Grant the service account access to list secrets for the project
 gcloud projects add-iam-policy-binding "${PROJECT_ID}" \
   --member="serviceAccount:${GSA}" \
   --role="roles/secretmanager.viewer"
 
-# grant service account permission to access the private secret
+# Grant service account permission to access the private secret
 gcloud secrets add-iam-policy-binding "${private_secret_name}" \
-    --member="serviceAccount:${GSA}" \
-    --role="roles/secretmanager.secretAccessor"
+  --member="serviceAccount:${GSA}" \
+  --role="roles/secretmanager.secretAccessor"
 
-# grant service account permission to access the public secret
+# Grant service account permission to access the public secret
 gcloud secrets add-iam-policy-binding "${public_secret_name}" \
-    --member="serviceAccount:${GSA}" \
-    --role="roles/secretmanager.secretAccessor"
+  --member="serviceAccount:${GSA}" \
+  --role="roles/secretmanager.secretAccessor"
 
 # If no OS family specified, default to debian
 if [[ "${IMAGE_VERSION}" != *-* ]] ; then
@@ -95,5 +105,18 @@ python generate_custom_image.py \
     --no-smoke-test \
     --gcs-bucket "${BUCKET}" \
     --shutdown-instance-timer-sec=30
-
 set +x
+# Revoke permission to access the private secret
+gcloud secrets remove-iam-policy-binding "${private_secret_name}" \
+  --member="serviceAccount:${GSA}" \
+  --role="roles/secretmanager.secretAccessor" > /dev/null 2>&1
+
+# Revoke access to bucket
+gcloud storage buckets remove-iam-policy-binding "gs://${BUCKET}" \
+  --member="serviceAccount:${GSA}" \
+  --role="roles/storage.objectViewer" > /dev/null 2>&1
+
+# Revoke access to list secrets for the project
+gcloud projects remove-iam-policy-binding "${PROJECT_ID}" \
+  --member="serviceAccount:${GSA}" \
+  --role="roles/secretmanager.viewer" > /dev/null 2>&1
diff --git a/examples/secure-boot/env.json.sample b/examples/secure-boot/env.json.sample
index 295b2f7..9df37dc 100644
--- a/examples/secure-boot/env.json.sample
+++ b/examples/secure-boot/env.json.sample
@@ -1,6 +1,6 @@
 {
   "PROJECT_ID":"example-yyyy-nn",
-  "CLUSTER_NAME":"my-dataproc-cluster-name",
+  "PURPOSE":"pre-init",
   "BUCKET":"my-bucket-name",
   "IMAGE_VERSION":"2.2-debian12",
   "ZONE":"us-west4-a"