#!/usr/bin/env bash
# Usage: ssh-wrap user@host [ssh-options]
# Wrapper to handle SSH host key changes automatically

set -uo pipefail

if [[ $# -eq 0 ]]; then
  echo "Usage: ssh-wrap user@host [ssh-options]" >&2
  exit 1
fi

# Capture SSH output
output=$(ssh "$@" 2>&1)
exit_code=$?

# Print the SSH output so user sees what happened
echo "$output"

# If SSH succeeded, we're done
if [[ $exit_code -eq 0 ]]; then
  exit 0
fi

# Check if the known_hosts warning appears
if echo "$output" | grep -q "REMOTE HOST IDENTIFICATION HAS CHANGED"; then
  echo ""
  echo "[WARNING] Host key has changed - possible man-in-the-middle attack or host reinstall."

  # Extract the known_hosts file and line number from the "Offending RSA key in ..." line
  # The line format typically is: "Offending RSA key in /path/to/known_hosts:line"
  if offending_info=$(echo "$output" | grep "Offending.*key in"); then
    KNOWN_HOSTS_FILE=$(echo "$offending_info" | awk '{print $5}' | cut -d: -f1)
    LINE_NUMBER=$(echo "$offending_info" | awk -F: '{print $NF}')

    if [[ -z "$KNOWN_HOSTS_FILE" || -z "$LINE_NUMBER" || ! -f "$KNOWN_HOSTS_FILE" ]]; then
      echo "[ERROR] Could not extract offending key information or file doesn't exist." >&2
      exit 1
    fi

    echo "[INFO] Offending key detected in: $KNOWN_HOSTS_FILE on line: $LINE_NUMBER"
    read -rp "Remove offending key and retry SSH connection? [y/N]: " RESPONSE
    
    if [[ "$RESPONSE" =~ ^[Yy]$ ]]; then
      # Backup known_hosts
      if cp "$KNOWN_HOSTS_FILE" "$KNOWN_HOSTS_FILE.bak"; then
        echo "[INFO] Backup created: $KNOWN_HOSTS_FILE.bak"
      else
        echo "[ERROR] Failed to create backup." >&2
        exit 1
      fi
      
      # Remove offending line
      if sed -i "${LINE_NUMBER}d" "$KNOWN_HOSTS_FILE"; then
        echo "[INFO] Offending key removed. Retrying SSH connection..."
        ssh "$@"
      else
        echo "[ERROR] Failed to remove offending key." >&2
        exit 1
      fi
    else
      echo "[INFO] Key was not removed. Exiting."
      exit 1
    fi
  else
    echo "[ERROR] Could not extract offending key information. Remove it manually if needed." >&2
    exit 1
  fi
else
  # SSH failed for another reason
  exit $exit_code
fi