add ability to remove ssh keys

update_ssh_keys.sh

@@ -5,50 +5,77 @@
 
 # Function to display usage
 usage() {
-    echo "Usage: $0 [ssh_key_name]"
+    echo "Usage: $0 [options] [ssh_key_name]"
+    echo "  options:"
+    echo "    -r, --remove   Remove SSH key from main.tf files"
+    echo "    -h, --help     Display this help message"
+    echo ""
     echo "  ssh_key_name: Name of the SSH key pair (default: terraform_key)"
     echo ""
     echo "Example:"
-    echo "  $0                    # Uses default 'terraform_key'"
-    echo "  $0 my_custom_key      # Uses 'my_custom_key' and 'my_custom_key.pub'"
+    echo "  $0                    # Updates with default 'terraform_key'"
+    echo "  $0 my_custom_key      # Updates with 'my_custom_key'"
+    echo "  $0 -r                 # Remove SSH key from files"
+    echo "  $0 -r my_custom_key   # Remove SSH key from files"
     exit 1
 }
 
-# Set the SSH key name (default to terraform_key)
-SSH_KEY_NAME="${1:-terraform_key}"
+# Parse command line arguments
+REMOVE_KEY=false
+SSH_KEY_NAME="terraform_key"
+
+# Check if any arguments are provided
+if [ $# -eq 0 ]; then
+    # No arguments - use default behavior (update)
+    :
+elif [ "$1" = "-h" ] || [ "$1" = "--help" ]; then
+    usage
+elif [ "$1" = "-r" ] || [ "$1" = "--remove" ]; then
+    # Remove mode enabled
+    REMOVE_KEY=true
+    if [ $# -gt 1 ]; then
+        SSH_KEY_NAME="$2"
+    fi
+else
+    # Normal update mode with key name provided as argument
+    SSH_KEY_NAME="$1"
+fi
 
 # Expand the home directory properly  
 HOME_DIR="${HOME:-/home/$(whoami)}"
 SSH_KEY_PATH="$HOME_DIR/.ssh/$SSH_KEY_NAME"
 SSH_KEY_PUB_PATH="$HOME_DIR/.ssh/$SSH_KEY_NAME.pub"
 
-# Check if SSH key exists
-if [ ! -f "$SSH_KEY_PATH" ] && [ ! -f "$SSH_KEY_PUB_PATH" ]; then
-    echo "Error: SSH key '$SSH_KEY_NAME' not found in $HOME_DIR/.ssh/"
-    echo "Please generate your SSH key first:"
-    echo "  ssh-keygen -t rsa -b 4096 -f $HOME_DIR/.ssh/$SSH_KEY_NAME"
-    exit 1
+# If not removing keys, validate SSH key exists
+if [ "$REMOVE_KEY" = false ]; then
+    # Check if SSH key exists
+    if [ ! -f "$SSH_KEY_PATH" ] && [ ! -f "$SSH_KEY_PUB_PATH" ]; then
+        echo "Error: SSH key '$SSH_KEY_NAME' not found in $HOME_DIR/.ssh/"
+        echo "Please generate your SSH key first:"
+        echo "  ssh-keygen -t rsa -b 4096 -f $HOME_DIR/.ssh/$SSH_KEY_NAME"
+        exit 1
+    fi
+
+    # Check if public key exists specifically (required for reading)
+    if [ ! -f "$SSH_KEY_PUB_PATH" ]; then
+        echo "Error: SSH public key '$SSH_KEY_NAME.pub' not found in $HOME_DIR/.ssh/"
+        exit 1
+    fi
+
+    # Get the public key content (remove any trailing whitespace)
+    PUBLIC_KEY=$(cat "$SSH_KEY_PUB_PATH" | tr -d '\n')
+
+    # Validate that we got a valid SSH key
+    if [[ ! "$PUBLIC_KEY" =~ ^ssh-[a-z]+[[:space:]]+[A-Za-z0-9+/]*[=]{0,3} ]]; then
+        echo "Error: Invalid SSH public key format detected"
+        exit 1
+    fi
+
+    echo "Found SSH public key:"
+    echo "$PUBLIC_KEY"
+    echo ""
 fi
 
-# Check if public key exists specifically (required for reading)
-if [ ! -f "$SSH_KEY_PUB_PATH" ]; then
-    echo "Error: SSH public key '$SSH_KEY_NAME.pub' not found in $HOME_DIR/.ssh/"
-    exit 1
-fi
-
-# Get the public key content (remove any trailing whitespace)
-PUBLIC_KEY=$(cat "$SSH_KEY_PUB_PATH" | tr -d '\n')
-
-# Validate that we got a valid SSH key
-if [[ ! "$PUBLIC_KEY" =~ ^ssh-[a-z]+[[:space:]]+[A-Za-z0-9+/]*[=]{0,3} ]]; then
-    echo "Error: Invalid SSH public key format detected"
-    exit 1
-fi
-
-echo "Found SSH public key:"
-echo "$PUBLIC_KEY"
-echo ""
-
 # Find all main.tf files and update them
 MAIN_TF_FILES=$(find . -name "main.tf" -type f)
 
@@ -61,10 +88,17 @@ echo "Updating SSH key in the following files:"
 echo "$MAIN_TF_FILES"
 echo ""
 
-# Replace the ssh_key line in all main.tf files using # as delimiter
-echo "Replacing SSH key in all main.tf files..."
+# Process each file based on remove mode
 for file in $MAIN_TF_FILES; do
-    sed -i "s#ssh_key = \".*\"#ssh_key = \"$PUBLIC_KEY\"#g" "$file"
+    if [ "$REMOVE_KEY" = true ]; then
+        echo "Removing SSH key from $file..."
+        # Set ssh_key to empty string for idempotent removal
+        sed -i "s/^[[:space:]]*ssh_key[[:space:]]*=[[:space:]]*\"[^\"]*\"/  ssh_key = \"\"/" "$file"
+    else
+        echo "Updating SSH key in $file..."
+        # Update the ssh_key line with new value
+        sed -i "s#ssh_key = \".*\"#ssh_key = \"$PUBLIC_KEY\"#g" "$file"
+    fi
 done
 
 # Verify the replacement worked
@@ -72,9 +106,18 @@ echo ""
 echo "Verification:"
 for file in $MAIN_TF_FILES; do
     echo "File: $file"
-    grep "ssh_key =" "$file" | head -1
+    if [ "$REMOVE_KEY" = true ]; then
+        # Show lines with empty ssh_key values
+        grep "ssh_key = \"\"" "$file" | head -1
+    else
+        # Show updated ssh_key lines
+        grep "ssh_key =" "$file" | head -1
+    fi
 done
 
 echo ""
-echo "SSH key has been successfully updated in all main.tf files!"
-echo "Backup files are saved with timestamp suffixes."
+if [ "$REMOVE_KEY" = true ]; then
+    echo "SSH key has been successfully removed (set to empty string) in all main.tf files!"
+else
+    echo "SSH key has been successfully updated in all main.tf files!"
+fi