From c8e3ce26a94b2d346e3140e310ad2f8155ff87c4 Mon Sep 17 00:00:00 2001
From: Michael Vines <mvines@gmail.com>
Date: Mon, 27 Aug 2018 20:37:35 -0700
Subject: [PATCH] Start of scripts/gcloud.sh

---
 ci/testnet-deploy.sh            |  19 ++--
 multinode-demo/gce_multinode.sh |  43 ++++++--
 scripts/gcloud.sh               | 190 ++++++++++++++++++++++++++++++++
 3 files changed, 229 insertions(+), 23 deletions(-)
 create mode 100644 scripts/gcloud.sh

diff --git a/ci/testnet-deploy.sh b/ci/testnet-deploy.sh
index da3bb79ba..8335b6c6e 100755
--- a/ci/testnet-deploy.sh
+++ b/ci/testnet-deploy.sh
@@ -5,9 +5,13 @@
 # This script must be run by a user/machine that has successfully authenticated
 # with GCP and has sufficient permission.
 #
+
 here=$(dirname "$0")
 metrics_write_datapoint="$here"/../scripts/metrics-write-datapoint.sh
 
+# shellcheck source=scripts/gcloud.sh
+source "$here"/../scripts/gcloud.sh
+
 # TODO: Switch over to rolling updates
 ROLLING_UPDATE=false
 #ROLLING_UPDATE=true
@@ -134,7 +138,7 @@ vm_foreach() {
   declare count=1
   for info in "${vmlist[@]}"; do
     declare vmClass vmName vmZone vmPublicIp
-    IFS=: read -r vmClass vmName vmZone vmPublicIp < <(echo "$info")
+    IFS=: read -r vmClass vmName vmZone vmPublicIp _ < <(echo "$info")
 
     eval "$cmd" "$vmName" "$vmZone" "$vmPublicIp" "$vmClass" "$count" "$@"
     count=$((count + 1))
@@ -179,16 +183,9 @@ vm_foreach_in_class() {
 findVms() {
   declare class="$1"
   declare filter="$2"
-  gcloud compute instances list --filter="$filter"
-  while read -r vmName vmZone vmPublicIp status; do
-    if [[ $status != RUNNING ]]; then
-      echo "Warning: $vmName is not RUNNING, ignoring it."
-      continue
-    fi
-    vmlist+=("$class:$vmName:$vmZone:$vmPublicIp")
-  done < <(gcloud compute instances list \
-             --filter="$filter" \
-             --format 'value(name,zone,networkInterfaces[0].accessConfigs[0].natIP,status)')
+
+  gcloud_FindInstances "$filter"
+  vmlist+=("${instances[@]/#/$class:}")
 }
 
 wait_for_pids() {
diff --git a/multinode-demo/gce_multinode.sh b/multinode-demo/gce_multinode.sh
index d9fd9b097..42fd00872 100755
--- a/multinode-demo/gce_multinode.sh
+++ b/multinode-demo/gce_multinode.sh
@@ -1,11 +1,15 @@
 #!/bin/bash
 
+here=$(dirname "$0")
+# shellcheck source=scripts/gcloud.sh
+source "$here"/../scripts/gcloud.sh
+
 command=$1
 prefix=
 num_nodes=
 out_file=
 image_name="ubuntu-16-04-cuda-9-2-new"
-ip_address_offset=5
+internalNetwork=false
 zone="us-west1-b"
 
 shift
@@ -43,7 +47,7 @@ while getopts "h?p:Pi:n:z:o:" opt; do
     prefix=$OPTARG
     ;;
   P)
-    ip_address_offset=4
+    internalNetwork=true
     ;;
   i)
     image_name=$OPTARG
@@ -69,22 +73,37 @@ set -e
 
 [[ -n $prefix ]] || usage "Need a prefix for GCE instance names"
 
-[[ -n $num_nodes ]] || usage "Need number of nodes"
-
-nodes=()
-for i in $(seq 1 "$num_nodes"); do
-  nodes+=("$prefix$i")
-done
 
 if [[ $command == "create" ]]; then
+  [[ -n $num_nodes ]] || usage "Need number of nodes"
   [[ -n $out_file ]] || usage "Need an outfile to store IP Addresses"
 
-  ip_addr_list=$(gcloud beta compute instances create "${nodes[@]}" --zone="$zone" --tags=testnet \
-    --image="$image_name" | awk -v offset="$ip_address_offset" '/RUNNING/ {print $offset}')
+  gcloud_CreateInstances "$prefix" "$num_nodes" "$zone" "$image_name"
+  gcloud_FindInstances "name~^$prefix"
 
-  echo "ip_addr_array=($ip_addr_list)" >"$out_file"
+  echo "ip_addr_array=()" > "$out_file"
+  recordPublicIp() {
+    declare name="$1"
+    declare publicIp="$3"
+    declare privateIp="$4"
+
+    if $internalNetwork; then
+      echo "ip_addr_array+=($privateIp) # $name" >> "$out_file"
+    else
+      echo "ip_addr_array+=($publicIp)  # $name" >> "$out_file"
+    fi
+  }
+  gcloud_ForEachInstance recordPublicIp
+
+  echo "Instance ip addresses recorded in $out_file"
 elif [[ $command == "delete" ]]; then
-  gcloud beta compute instances delete "${nodes[@]}"
+  gcloud_FindInstances "name~^$prefix"
+
+  if [[ ${#instances[@]} -eq 0 ]]; then
+    echo "No instances found matching '^$prefix'"
+    exit 0
+  fi
+  gcloud_DeleteInstances
 else
   usage "Unknown command: $command"
 fi
diff --git a/scripts/gcloud.sh b/scripts/gcloud.sh
new file mode 100644
index 000000000..59a1029fe
--- /dev/null
+++ b/scripts/gcloud.sh
@@ -0,0 +1,190 @@
+# |source| this file
+#
+# Utilities for working with gcloud
+#
+
+
+#
+# gcloud_FindInstances [filter]
+#
+# Find instances matching the specified pattern.
+#
+# For each matching instance, an entry in the `instances` array will be added with the
+# following information about the instance:
+#   "name:zone:public IP:private IP"
+#
+# filter   - The instances to filter on
+#
+# examples:
+#   $ gcloud_FindInstances "name=exact-machine-name"
+#   $ gcloud_FindInstances "name~^all-machines-with-a-common-machine-prefix"
+#
+gcloud_FindInstances() {
+  declare filter="$1"
+  #gcloud compute instances list --filter="$filter"
+  instances=()
+  while read -r name zone publicIp privateIp status; do
+    if [[ $status != RUNNING ]]; then
+      echo "Warning: $name is not RUNNING, ignoring it."
+      continue
+    fi
+  instances+=("$name:$zone:$publicIp:$privateIp")
+  done < <(gcloud compute instances list \
+             --filter="$filter" \
+             --format 'value(name,zone,networkInterfaces[0].accessConfigs[0].natIP,networkInterfaces[0].networkIP,status)')
+}
+
+#
+# gcloud_ForEachInstance [cmd] [extra args to cmd]
+#
+# Execute a command for each element in the `instances` array
+#
+#   cmd   - The command to execute on each instance
+#           The command will receive arguments followed by any
+#           additionl arguments supplied to gcloud_ForEachInstance:
+#               name     - name of the instance
+#               zone     - zone the instance is located in
+#               publicIp - The public IP address of this instance
+#               privateIp - The priate IP address of this instance
+#               count    - Monotonically increasing count for each
+#                          invocation of cmd, starting at 1
+#               ...      - Extra args to cmd..
+#
+#
+gcloud_ForEachInstance() {
+  declare cmd="$1"
+  shift
+  [[ -n $cmd ]] || { echo gcloud_ForEachInstance: cmd not specified; exit 1; }
+
+  declare count=1
+  for info in "${instances[@]}"; do
+    declare name zone publicIp privateIp
+    IFS=: read -r name zone publicIp privateIp < <(echo "$info")
+
+    eval "$cmd" "$name" "$zone" "$publicIp" "privateIp" "$count" "$@"
+    count=$((count + 1))
+  done
+}
+
+#
+# gcloud_CreateInstances [namePrefix] [numNodes] [zone] [imageName]
+#
+# Creates one more identical instances.
+#
+# namePrefix   - unique string to prefix all the instance names with
+# numNodes     - number of instances to create
+# zone         - zone to create the instances in
+# imageName    - Disk image for the instances
+#
+# Tip: use gcloud_FindInstances to locate the instances once this function
+#      returns
+gcloud_CreateInstances() {
+  declare namePrefix="$1"
+  declare numNodes="$2"
+  declare zone="$3"
+  declare imageName="$4"
+
+  declare nodes
+  read -ra nodes <<<$(seq -f "${namePrefix}%g" 1 "$numNodes")
+
+  (
+    set -x
+    gcloud beta compute instances create "${nodes[@]}" \
+      --zone="$zone" \
+      --tags=testnet \
+      --image="$imageName"
+  )
+}
+
+#
+# gcloud_DeleteInstances
+#
+# Deletes all the instances listed in the `instances` array
+#
+gcloud_DeleteInstances() {
+  declare names=("${instances[@]/:*/}")
+  (
+    set -x
+    gcloud beta compute instances delete "${names[@]}"
+  )
+}
+
+#
+# gcloud_FigureRemoteUsername [instanceInfo]
+#
+# The remote username when ssh-ing into GCP instances tends to not be the same
+# as the user's local username, but it needs to be discovered by ssh-ing into an
+# instance and examining the system.
+#
+# On success the gcloud_username global variable is updated
+#
+# instanceInfo  - an entry from the `instances` array
+#
+# example:
+#   gcloud_FigureRemoteUsername "name:zone:..."
+#
+gcloud_FigureRemoteUsername() {
+  if [[ -n $gcloud_username ]]; then
+    return
+  fi
+
+  declare instanceInfo="$1"
+  declare name zone
+  IFS=: read -r name zone _ < <(echo "$instanceInfo")
+
+  echo "Detecting remote username using $zone in $zone:"
+  # Figure the gcp ssh username
+  (
+    set -x
+    gcloud compute ssh "$name" --zone "$zone" -- "echo whoami \$(whoami)" | tee whoami
+  )
+
+  [[ "$(tr -dc '[:print:]' < whoami; rm -f whoami)" =~ ^whoami\ (.*)$ ]] || {
+    echo Unable to figure remote user name;
+    exit 1
+  }
+  gcloud_username="${BASH_REMATCH[1]}"
+  echo "Remote username: $gcloud_username--"
+}
+
+#
+# gcloud_PrepInstancesForSsh [username] [publicKey] [privateKey]
+#
+# Prepares all the instances in the `instances` array for ssh with the specified
+# keypair.  This eliminates the need to use the restrictive |gcloud compute ssh|,
+# use plain |ssh| instead.
+#
+# username    - gcp ssh username as computed by gcloud_FigureRemoteUsername
+# publicKey   - public key to install on all the instances
+# privateKey  - matching private key, used to verify ssh access
+#
+gcloud_PrepInstancesForSsh() {
+  declare username="$1"
+  declare publicKey="$2"
+  declare privateKey="$3"
+  [[ -r $publicKey ]] || {
+    echo "Unable to read public key: $publicKey"
+    exit 1
+  }
+
+  [[ -r $privateKey ]] || {
+    echo "Unable to read private key: $privateKey"
+    exit 1
+  }
+
+  for instanceInfo in "${instances[@]}"; do
+    declare name zone publicIp
+    IFS=: read -r name zone publicIp _ < <(echo "$instanceInfo")
+    (
+      set -x
+
+      # TODO: stomping on the authorized_keys isn't great, maybe do something
+      #       clever with |ssh-copy-id| one day
+      gcloud compute scp --zone "$zone" "$publicKey" "$name":.ssh/authorized_keys
+
+      # Confirm normal ssh now works
+      ssh -i "$privateKey" "$username@$publicIp" uptime
+    )
+  done
+}
+