From a959f71208e78f943299c3d3f76d7a687e3b91f8 Mon Sep 17 00:00:00 2001
From: Alex Vandiver <alexmv@zulip.com>
Date: Fri, 9 May 2025 01:19:10 -0400
Subject: [PATCH] kandra: Support installing multiple ssh keys from multiple
 secrets.

---
 puppet/kandra/files/install-ssh-keys | 44 ++++++++++++++++++++--------
 puppet/kandra/manifests/ssh_keys.pp  |  4 ++-
 2 files changed, 34 insertions(+), 14 deletions(-)

diff --git a/puppet/kandra/files/install-ssh-keys b/puppet/kandra/files/install-ssh-keys
index a93fbccf46..b7723f53a3 100755
--- a/puppet/kandra/files/install-ssh-keys
+++ b/puppet/kandra/files/install-ssh-keys
@@ -1,8 +1,24 @@
 #!/usr/bin/env bash
 set -euo pipefail
 
+args="$(getopt -o '' --long check -- "$@")"
+eval "set -- $args"
+check=false
+while true; do
+    case "$1" in
+        --check)
+            check=true
+            shift
+            ;;
+        --)
+            shift
+            break
+            ;;
+    esac
+done
+
 username="$1"
-ssh_secret_name="$2"
+shift
 
 homedir="$(getent passwd "$username" | cut -d: -f6)"
 sshdir="$homedir/.ssh"
@@ -15,20 +31,22 @@ trap cleanup EXIT
 
 umask 033
 
-keydata="$(/srv/zulip-aws-tools/bin/aws --output text \
-    secretsmanager get-secret-value \
-    --secret-id "$ssh_secret_name" \
-    --query SecretString)"
-for keyfile in $(jq -r 'keys[]' <<<"$keydata"); do
-    touch "$workdir/$keyfile"
-    if [[ "$keyfile" != *".pub" ]]; then
-        chmod 600 "$workdir/$keyfile"
-    fi
-    jq -r ".[\"$keyfile\"]" <<<"$keydata" | base64 -d >"$workdir/$keyfile"
-    chown "$username:$username" "$workdir/$keyfile"
+for ssh_secret_name in "$@"; do
+    keydata="$(/srv/zulip-aws-tools/bin/aws --output text \
+        secretsmanager get-secret-value \
+        --secret-id "$ssh_secret_name" \
+        --query SecretString)"
+    for keyfile in $(jq -r 'keys[]' <<<"$keydata"); do
+        touch "$workdir/$keyfile"
+        if [[ "$keyfile" != *".pub" ]]; then
+            chmod 600 "$workdir/$keyfile"
+        fi
+        jq -r ".[\"$keyfile\"]" <<<"$keydata" | base64 -d >"$workdir/$keyfile"
+        chown "$username:$username" "$workdir/$keyfile"
+    done
 done
 
-if [ "$#" -gt 2 ]; then
+if [ "$check" = "true" ]; then
     diff -rN -x config -x authorized_keys -x known_hosts \
         "$workdir/" "$sshdir/"
     exit 0
diff --git a/puppet/kandra/manifests/ssh_keys.pp b/puppet/kandra/manifests/ssh_keys.pp
index 2318430e1f..8e1697ea06 100644
--- a/puppet/kandra/manifests/ssh_keys.pp
+++ b/puppet/kandra/manifests/ssh_keys.pp
@@ -4,12 +4,14 @@ define kandra::ssh_keys(
   $user = $name
   if $keys == true {
     $keypath = "prod/ssh/keys/${user}"
+  } elsif $keys.is_a(Array) {
+    $keypath = join($keys.map |$k| {"prod/ssh/keys/${k}"}, ' ')
   } else {
     $keypath = "prod/ssh/keys/${keys}"
   }
   exec { "ssh_keys ${user}":
     require => File['/usr/local/bin/install-ssh-keys'],
     command => "/usr/local/bin/install-ssh-keys ${user} ${keypath}",
-    unless  => "[ -f /usr/local/bin/install-ssh-keys ] && /usr/local/bin/install-ssh-keys ${user} ${keypath} check",
+    unless  => "[ -f /usr/local/bin/install-ssh-keys ] && /usr/local/bin/install-ssh-keys --check ${user} ${keypath}",
   }
 }