From 7b32216684c6c5286ad2e585fbbdf3a6c03e3067 Mon Sep 17 00:00:00 2001
From: Noah Masur <7386960+nmasur@users.noreply.github.com>
Date: Mon, 7 Apr 2025 16:35:31 -0400
Subject: [PATCH] aws ssm ssh in profile

---
 .../presets/programs/aws-ssh/aws-ssh.nix      | 19 +++++
 .../aws-ssh/aws-ssm-ssh-proxy-command.sh      | 69 +++++++++++++++++++
 .../modules/nmasur/profiles/work.nix          |  3 +
 3 files changed, 91 insertions(+)
 create mode 100644 platforms/home-manager/modules/nmasur/presets/programs/aws-ssh/aws-ssh.nix
 create mode 100644 platforms/home-manager/modules/nmasur/presets/programs/aws-ssh/aws-ssm-ssh-proxy-command.sh

diff --git a/platforms/home-manager/modules/nmasur/presets/programs/aws-ssh/aws-ssh.nix b/platforms/home-manager/modules/nmasur/presets/programs/aws-ssh/aws-ssh.nix
new file mode 100644
index 0000000..cb5d7b1
--- /dev/null
+++ b/platforms/home-manager/modules/nmasur/presets/programs/aws-ssh/aws-ssh.nix
@@ -0,0 +1,19 @@
+{ config, lib, ... }:
+
+let
+  cfg = config.nmasur.presets.programs.aws-ssh;
+in
+
+{
+  options.nmasur.presets.programs.aws-ssh.enable = lib.mkEnableOption "AWS SSH tools";
+
+  config = lib.mkIf cfg.enable {
+
+    # Ignore wine directories in searches
+    home.file.".ssh/aws-ssm-ssh-proxy-command.sh" = {
+      text = builtins.readFile ./aws-ssm-ssh-proxy-command.sh;
+      executable = true;
+    };
+
+  };
+}
diff --git a/platforms/home-manager/modules/nmasur/presets/programs/aws-ssh/aws-ssm-ssh-proxy-command.sh b/platforms/home-manager/modules/nmasur/presets/programs/aws-ssh/aws-ssm-ssh-proxy-command.sh
new file mode 100644
index 0000000..60103c0
--- /dev/null
+++ b/platforms/home-manager/modules/nmasur/presets/programs/aws-ssh/aws-ssm-ssh-proxy-command.sh
@@ -0,0 +1,69 @@
+#!/usr/bin/env bash
+set -eu
+
+################################################################################
+#
+# For documentation see https://github.com/qoomon/aws-ssm-ssh-proxy-command
+#
+################################################################################
+
+getInstanceId() {
+  local instance_name="$1"
+  local instance_id=$(aws ec2 describe-instances --filters "Name=tag:Name,Values=${instance_name}" --query "Reservations[].Instances[?State.Name == 'running'].InstanceId" --output text)
+
+  echo "${instance_id}"
+}
+
+instance_name="$1"
+ssh_user="$2"
+ssh_port="$3"
+ssh_public_key_path="$4"
+
+ec2InstanceIdPattern='^m?i-[0-9a-f]{8,17}$'
+if [[ $instance_name =~ $ec2InstanceIdPattern ]]; then
+  instance_id=$instance_name
+else
+  instance_id=$(getInstanceId "$instance_name")
+
+  if [[ -z $instance_id ]]; then
+    echo "Found no running instances with name \"${instance_name}\"."
+    exit 1
+  else
+    echo "Instance ID for \"${instance_name}\": \"${instance_id}\""
+  fi
+fi
+
+REGION_SEPARATOR='--'
+if echo "$instance_id" | grep -q -e "${REGION_SEPARATOR}"; then
+  export AWS_REGION="${instance_id##*"${REGION_SEPARATOR}"}"
+  instance_id="${instance_id%%"$REGION_SEPARATOR"*}"
+fi
+
+>/dev/stderr echo "Add public key ${ssh_public_key_path} for ${ssh_user} at instance ${instance_id} for 10 seconds"
+ssh_public_key="$(cat "${ssh_public_key_path}")"
+aws ssm send-command \
+  --instance-ids "${instance_id}" \
+  --document-name 'AWS-RunShellScript' \
+  --comment "Add an SSH public key to authorized_keys for 10 seconds" \
+  --parameters commands="
+  \"
+    set -eu
+    
+    mkdir -p ~${ssh_user}/.ssh && cd ~${ssh_user}/.ssh
+
+    authorized_key='${ssh_public_key} ssm-session'
+    
+    echo \\\"\${authorized_key}\\\" >> authorized_keys
+    
+    sleep 10
+    
+    (grep -v -F \\\"\${authorized_key}\\\" authorized_keys || true) > authorized_keys~
+    mv authorized_keys~ authorized_keys
+  \"
+  "
+
+>/dev/stderr echo "Start ssm session to instance ${instance_id}"
+aws ssm start-session \
+  --target "${instance_id}" \
+  --document-name 'AWS-StartSSHSession' \
+  --parameters "portNumber=${ssh_port}"
diff --git a/platforms/home-manager/modules/nmasur/profiles/work.nix b/platforms/home-manager/modules/nmasur/profiles/work.nix
index f1eb697..e81eb8c 100644
--- a/platforms/home-manager/modules/nmasur/profiles/work.nix
+++ b/platforms/home-manager/modules/nmasur/profiles/work.nix
@@ -39,11 +39,14 @@ in
       pkgs.nmasur.terraform-init # Quick shortcut for initializing Terraform backend
     ];
 
+    programs.fish.shellAliases.ec2 = "aws-ec2";
+
     nmasur.presets = {
       fonts.enable = lib.mkDefault true;
       programs = {
         _1password.enable = lib.mkDefault true;
         atuin.enable = lib.mkDefault true;
+        aws-ssh.enable = lib.mkDefault true;
         bash.enable = lib.mkDefault true;
         bat.enable = lib.mkDefault true;
         direnv.enable = lib.mkDefault true;