#!/usr/bin/env bash
# ------------------------------------------------------------------------------
# Copyright (c) 2013-2026 Geoffray Levasseur <fatalerrors@geoffray-levasseur.org>
# Protected by the BSD3 license. Please read bellow for details.
#
# * Redistribution and use in source and binary forms,
# * with or without modification, are permitted provided
# * that the following conditions are met:
# *
# *   Redistributions of source code must retain the above
# *   copyright notice, this list of conditions and the
# *   following disclaimer.
# *
# *   Redistributions in binary form must reproduce the above
# *   copyright notice, this list of conditions and the following
# *   disclaimer in the documentation and/or other materials
# *   provided with the distribution.
# *
# *   Neither the name of the copyright holder nor the names
# *   of any other contributors may be used to endorse or
# *   promote products derived from this software without
# *   specific prior written permission.
# *
# * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND
# * CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
# * INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
# * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
# * OF SUCH DAMAGE.
# ------------------------------------------------------------------------------

load_conf "ssh"

# ------------------------------------------------------------------------------
# Internal helper: remove <key> (a name or an IP, in known_hosts notation, so
# also "[host]:port") from the known_hosts file <file> with ssh-keygen -R.
# ssh-keygen writes a new file: run by root on another user's file (--all-users)
# it would leave it owned by root with mode 0644, which keeps that user's ssh
# from adding host keys, or exposes a private file. The owner and mode of the
# file, and of its ".old" backup, are therefore put back.
# Returns ssh-keygen's status (non-zero when the key was not found).
# Usage: _rmhost_remove <key> <file>
_rmhost_remove()
{
    local key="$1" file="$2" owner="" mode="" rc=0
    if [[ ${EUID:-$(id -u)} -eq 0 ]]; then
        owner=$(stat -c '%u:%g' -- "$file" 2>/dev/null)
        mode=$(stat -c '%a' -- "$file" 2>/dev/null)
    fi

    ssh-keygen -R "$key" -f "$file" </dev/null >/dev/null 2>&1 || rc=$?

    if [[ -n $owner && -n $mode ]]; then
        chown "$owner" -- "$file" "${file}.old" 2>/dev/null
        chmod "$mode" -- "$file" "${file}.old" 2>/dev/null
    fi
    return "$rc"
}
export -f _rmhost_remove
# ------------------------------------------------------------------------------


# ------------------------------------------------------------------------------
# Remove host entries (name and IP) from ~/.ssh/known_hosts for the active user
# Usage: rmhost <hostname|ip> [hostname2|ip2 [...]]
rmhost()
{
    local PARSED
    local all_users=0
    local -a known_hosts_files=()

    PARSED=$(getopt -o ha --long help,all-users -n 'rmhost' -- "$@")
    # shellcheck disable=SC2181 # getopt return code is checked immediately after
    if [[ $? -ne 0 ]]; then
        disp E "Invalid options, use \"rmhost --help\" to display usage."
        return 1
    fi
    # Restore the getopt-normalised (shell-quoted) argument list
    eval set -- "$PARSED"

    while true; do
        case "$1" in
            -h|--help)
                printf "rmhost: Remove host/IP from known_hosts files.\n\n"
                printf "Usage: rmhost [--all-users] <pattern|ip> [pattern2|ip2 ...]\n\n"
                printf "Options:\n"
                printf "  -a, --all-users   Remove entries from all local users when run as root\n"
                printf "  -h, --help        Display this help screen\n\n"
                printf "A host name is removed together with all the IP addresses it resolves to.\n"
                printf "Wildcards:\n"
                printf "  Glob patterns (*, ?, [...]) are expanded against unhashed known_hosts entries.\n"
                printf "  Hashed entries (prefixed with |1|) are never matched by wildcards.\n"
                return 0
                ;;
            -a|--all-users)
                all_users=1
                shift
                ;;
            --)
                shift
                break
                ;;
            *)
                disp E "Invalid options, use \"rmhost --help\" to display usage."
                return 1
                ;;
        esac
    done

    [[ $# -eq 0 ]] && {
        disp E "Missing argument. Use 'rmhost --help' for usage."
        return 1
    }

    command -v ssh-keygen >/dev/null 2>&1 || {
        disp E "ssh-keygen is not installed."
        return 127
    }

    # A name starting with "-" would be read as an option by the tools below
    local target
    for target in "$@"; do
        if [[ "$target" == -* ]]; then
            disp E "Invalid host '$target'."
            return 1
        fi
    done

    if (( all_users )); then
        [[ ${EUID:-$(id -u)} -eq 0 ]] || {
            disp E "Option --all-users is only available when run as root."
            return 1
        }

        # Users may share a home directory: each file is listed once
        local -A _seen_homes=()
        local home
        while IFS=: read -r _ _ _ _ _ home _; do
            [[ -n $home && -f $home/.ssh/known_hosts && -z ${_seen_homes[$home]+x} ]] || continue
            _seen_homes[$home]=1
            known_hosts_files+=("$home/.ssh/known_hosts")
        done < /etc/passwd

        [[ -f /etc/ssh/ssh_known_hosts ]] && \
            known_hosts_files+=("/etc/ssh/ssh_known_hosts")

        [[ ${#known_hosts_files[@]} -gt 0 ]] || {
            disp W "No known_hosts files found for local users."
            return 0
        }
    else
        if [[ ! -f "${HOME}/.ssh/known_hosts" ]]; then
            disp W "No known_hosts file found (${HOME}/.ssh/known_hosts)."
            return 0
        fi
        known_hosts_files=("${HOME}/.ssh/known_hosts")
    fi

    for target in "$@"; do
        # Wildcard: expand the glob pattern against the unhashed entries of
        # each file, and remove the matches from that file only. A "[host]:port"
        # entry is not a pattern: its brackets are known_hosts notation.
        if [[ "$target" == *['*?']* ]] || \
           [[ "$target" == *'['* && ! "$target" =~ ^\[[^][]+\]:[0-9]+$ ]]; then
            local -a _matched=() _wentries=()
            local _wf _wl _wfield _whost _wmatch _key _found=0
            for _wf in "${known_hosts_files[@]}"; do
                _matched=()
                while IFS= read -r _wl; do
                    # Comments, hashed entries and @cert-authority/@revoked
                    # marker lines are left alone
                    [[ -z "$_wl" || "$_wl" == '#'* || "$_wl" == '|'* || "$_wl" == '@'* ]] && continue
                    _wfield="${_wl%% *}"
                    IFS=',' read -ra _wentries <<< "$_wfield"
                    for _whost in "${_wentries[@]}"; do
                        # Strip [host]:port notation to get the bare name for matching
                        if [[ "$_whost" == '['*']:'* ]]; then
                            _wmatch="${_whost#[}"
                            _wmatch="${_wmatch%%]:*}"
                        else
                            _wmatch="$_whost"
                        fi
                        # shellcheck disable=SC2053 # unquoted on purpose: glob match
                        [[ "$_wmatch" == $target ]] && _matched+=("$_whost")
                    done
                done < "$_wf"

                [[ ${#_matched[@]} -gt 0 ]] || continue
                _found=1
                while IFS= read -r _key; do
                    disp I "Removing '$_key' from $_wf..."
                    if ! _rmhost_remove "$_key" "$_wf"; then
                        disp W "No known_hosts entry found for '$_key' in '$_wf'."
                    fi
                done < <(printf '%s\n' "${_matched[@]}" | sort -u)
            done

            (( _found )) || disp W "No known_hosts entries match pattern '$target'."
            continue
        fi

        local hst="$target"
        local ip
        local -a ips=()

        if isipv4 "$hst" >/dev/null 2>&1 || isipv6 "$hst" >/dev/null 2>&1; then
            ips=("$hst")
            hst=""
        elif [[ "$hst" == '['*']:'* ]]; then
            # Explicit [host]:port entry: removed exactly as given
            :
        else
            # All the addresses of the name: ssh records the one it connected
            # to, which may be any of them (IPv4 and IPv6, several A records)
            if command -v host >/dev/null 2>&1; then
                mapfile -t ips < <(host "$hst" 2>/dev/null |
                    awk '/has address|has IPv6 address/ {print $NF}' | sort -u)
            elif command -v getent >/dev/null 2>&1; then
                # ahosts lists every address once per socket type
                mapfile -t ips < <(getent ahosts "$hst" 2>/dev/null |
                    awk '{print $1}' | sort -u)
            else
                disp W "No resolver tool found; removing hostname only for '$hst'."
            fi

            [[ ${#ips[@]} -eq 0 ]] && \
                disp W "Could not resolve IP for '$hst'; removing hostname only."
        fi

        local known_hosts_file=""
        for known_hosts_file in "${known_hosts_files[@]}"; do
            if [[ -n ${hst:-} ]]; then
                disp I "Removing host $hst from $known_hosts_file..."
                if ! _rmhost_remove "$hst" "$known_hosts_file"; then
                    disp W "No known_hosts entry found for '$hst' in '$known_hosts_file'."
                fi
            fi
            for ip in "${ips[@]}"; do
                disp I "Removing IP $ip from $known_hosts_file..."
                if ! _rmhost_remove "$ip" "$known_hosts_file"; then
                    disp W "No known_hosts entry found for '$ip' in '$known_hosts_file'."
                fi
            done
        done
    done
}
export -f rmhost
# ------------------------------------------------------------------------------


# ------------------------------------------------------------------------------
# Login root via SSH on the given machine
# Usage: ssr <server [ssh options]>
ssr()
{
    case "${1:-}" in
        -h|--help)
            printf "ssr: SSH into a server as root.\n\n"
            printf "Usage: ssr <server> [ssh_options...]\n\n"
            printf "Notes:\n"
            printf "  The first argument is the target server.\n"
            printf "  All remaining arguments are passed directly to ssh.\n\n"
            printf "Examples:\n"
            printf "  ssr srv01\n"
            printf "  ssr srv01 -p 2222\n"
            printf "  ssr srv01 -i ~/.ssh/id_ed25519 -J bastion\n"
            return 0
            ;;
    esac

    command -v ssh >/dev/null 2>&1 || {
        disp E "ssh is not installed."
        return 127
    }

    [[ $# -eq 0 || -z ${1:-} ]] && {
        disp E "Please specify the server you want to log in."
        return 1
    }

    local srv=$1
    shift

    # The first argument is the server: an option here ("ssr -p 22 host")
    # would end up in the root@... destination
    if [[ $srv == -* ]]; then
        disp E "'$srv' is not a server name: the server comes first, then the ssh options."
        return 1
    fi

    # Build the default options array from the configuration. An unset
    # SSH_DEFAULT_OPT means -Y (X11 forwarding); an empty one means no default
    # option at all, as the example configuration documents.
    local -a ssh_default_opts=()
    if [[ -n "${SSH_DEFAULT_OPT+x}" ]]; then
        read -ra ssh_default_opts <<< "${SSH_DEFAULT_OPT}"
    else
        ssh_default_opts=(-Y)
    fi

    # shellcheck disable=SC2029
    ssh "${ssh_default_opts[@]}" root@"$srv" "$@"
}
export -f ssr
# ------------------------------------------------------------------------------


# EOF