#!/bin/bash

set -euo pipefail

readonly FILENAME="${HOME}/.ssh/known_hosts"
readonly BACKUP_FILE="${FILENAME}.sshrm.bak"

check_known_hosts() {
	if [ ! -f "${FILENAME}" ]; then
		error "known_hosts not found."
		exit 1
	fi
}

usage() {
	echo "Usage: sshrm [-h|--help] [--no-backup] host|line_number"
}

error() {
	echo "Error: $1"
}

is_number() {
	case "$1" in
		''|*[!0-9]*) return 1 ;;
		*) return 0 ;;
	esac
}

get_host_from_line() {
	awk -v line="$1" 'NR == line {print $1; exit}' "${FILENAME}"
}

create_backup() {
	cp "${FILENAME}" "${BACKUP_FILE}"
}

remove_host() {
	ssh-keygen -R "$1"
}

confirm() {
	local prompt="$1"
	read -r -p "Remove ${prompt}? [y/N] " answer
	case "$answer" in
		y|Y|yes|YES)
			return 0
			;;
		*)
			error "cancelled."
			return 1
			;;
	esac
}

if [ "${1-}" = "" ]; then
	error "missing argument."
	usage
	exit 1
fi

if [ "$1" = "-h" ] || [ "$1" = "--help" ]; then
	usage
	exit 0
fi

backup=true
if [ "$1" = "--no-backup" ]; then
	backup=false
	shift
fi

if [ "${1-}" = "" ]; then
	error "missing argument."
	usage
	exit 1
fi

if [ "${2-}" != "" ]; then
	error "too many arguments."
	usage
	exit 1
fi

check_known_hosts

if is_number "$1"; then
	if [ "$1" -eq 0 ]; then
		error "invalid line number."
		exit 1
	fi
	HOST="$(get_host_from_line "$1")"
	if [ "${HOST}" = "" ]; then
		error "invalid line number."
		exit 1
	fi
else
	HOST="$1"
fi

if ! confirm "$HOST"; then
	exit 1
fi

if [ "${backup}" = true ]; then
	create_backup
fi

if ! remove_host "$HOST"; then
	error "removal failed."
	exit 1
fi
