add ability to remove ssh keys

update_ssh_keys.sh

@@ -5,23 +5,49 @@
 
 # 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"
 
+# 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/"
@@ -48,6 +74,7 @@ fi
     echo "Found SSH public key:"
     echo "$PUBLIC_KEY"
     echo ""
+fi
 
 # 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
+    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"
+    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 ""
+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!"
-echo "Backup files are saved with timestamp suffixes."
+fi