#!/usr/bin/env bash

#########################################################################################
#### bastion-hardener.sh — Harden SSH bastion/jump hosts with audit and rollback     ####
#### Disables password auth, restricts ciphers, sets idle timeout, fail2ban config   ####
#### Requires: bash 4+, root privileges                                              ####
####                                                                                 ####
#### Author: Phil Connor                                                             ####
#### Contact: contact@mylinux.work                                                   ####
#### License: MIT                                                                    ####
#### Version 1.01                                                                    ####
####                                                                                 ####
#### Usage:                                                                          ####
####   sudo ./bastion-hardener.sh --audit                                            ####
####                                                                                 ####
#### See --help for all options.                                                     ####
#########################################################################################
# v1.01 changes:
#   - Fixed: ((0++)) returns 1 under set -e; added || true guards
#   - Fixed: grep in pipeline crashes under set -euo pipefail when no matches found. Added || true guard
#########################################################################################

set -euo pipefail

# ── Defaults ──────────────────────────────────────────────────────────
RUN_MODE=""
SSHD_CONFIG="${SSHD_CONFIG:-/etc/ssh/sshd_config}"
BACKUP_ROOT="${BACKUP_ROOT:-/etc/ssh}"
ALLOW_USERS="${ALLOW_USERS:-}"
ALLOW_GROUPS="${ALLOW_GROUPS:-}"
IDLE_TIMEOUT="${IDLE_TIMEOUT:-300}"
MAX_AUTH_TRIES="${MAX_AUTH_TRIES:-3}"
MAX_SESSIONS="${MAX_SESSIONS:-2}"
SESSION_LOG_DIR="${SESSION_LOG_DIR:-/var/log/bastion-sessions}"
FAIL2BAN_BANTIME="${FAIL2BAN_BANTIME:-3600}"
FAIL2BAN_MAXRETRY="${FAIL2BAN_MAXRETRY:-3}"
DRY_RUN="${DRY_RUN:-false}"
VERBOSE="${VERBOSE:-false}"
COLOR="${COLOR:-auto}"
ENABLE_SESSION_LOGGING="false"
CONFIGURE_FAIL2BAN="false"
ROLLBACK_DIR=""

# ── State ─────────────────────────────────────────────────────────────
SCRIPT_NAME="$(basename "$0")"
readonly SCRIPT_NAME
START_TIME=""
PASS_COUNT=0
FAIL_COUNT=0
WARN_COUNT=0
CHANGES=0

# ── Hardening settings ───────────────────────────────────────────────
readonly RECOMMENDED_CIPHERS="chacha20-poly1305@openssh.com,aes256-gcm@openssh.com,aes128-gcm@openssh.com"
readonly RECOMMENDED_MACS="hmac-sha2-512-etm@openssh.com,hmac-sha2-256-etm@openssh.com"
readonly RECOMMENDED_KEX="curve25519-sha256,curve25519-sha256@libssh.org"

# ── Colors ────────────────────────────────────────────────────────────
setup_colors() {
    if [[ "$COLOR" == "never" ]]; then
        RED="" GREEN="" YELLOW="" BLUE="" CYAN="" BOLD="" DIM="" RESET=""
        return
    fi
    if [[ "$COLOR" == "always" ]] || [[ -t 1 ]]; then
        RED='\033[0;31m'
        GREEN='\033[0;32m'
        YELLOW='\033[0;33m'
        BLUE='\033[0;34m'
        # shellcheck disable=SC2034
        CYAN='\033[0;36m'
        BOLD='\033[1m'
        DIM='\033[2m'
        RESET='\033[0m'
    else
        # shellcheck disable=SC2034
        RED="" GREEN="" YELLOW="" BLUE="" CYAN="" BOLD="" DIM="" RESET=""
    fi
}

# ── Logging ───────────────────────────────────────────────────────────
log()     { echo -e "${BLUE}[INFO]${RESET} $*"; }
warn()    { echo -e "${YELLOW}[WARN]${RESET} $*" >&2; }
err()     { echo -e "${RED}[ERROR]${RESET} $*" >&2; }
verbose() { if [[ "$VERBOSE" == "true" ]]; then echo -e "${DIM}[DEBUG]${RESET} $*"; fi; }

die() { err "$*"; exit 1; }

elapsed() {
    local end_time
    end_time=$(date +%s)
    echo "$(( end_time - START_TIME ))s"
}

require_root() {
    if [[ $EUID -ne 0 ]]; then
        die "This operation requires root privileges. Run with sudo."
    fi
}

# ── SSHD config helpers ──────────────────────────────────────────────
get_sshd_setting() {
    local key="$1"
    local val
    val=$({ grep -i "^[[:space:]]*${key}[[:space:]]" "$SSHD_CONFIG" 2>/dev/null || true; } | tail -1 | awk '{print $2}')
    if [[ -z "$val" ]]; then
        echo "(not set)"
    else
        echo "$val"
    fi
}

set_sshd_config() {
    local key="$1"
    local value="$2"
    local file="$3"

    if grep -qi "^[[:space:]]*${key}[[:space:]]" "$file" 2>/dev/null; then
        sed -i "s|^[[:space:]]*${key}[[:space:]].*|${key} ${value}|i" "$file"
    elif grep -qi "^[[:space:]]*#[[:space:]]*${key}[[:space:]]" "$file" 2>/dev/null; then
        sed -i "s|^[[:space:]]*#[[:space:]]*${key}[[:space:]].*|${key} ${value}|i" "$file"
    else
        echo "${key} ${value}" >> "$file"
    fi
    verbose "Set ${key} = ${value}"
}

# ── Audit check helper ───────────────────────────────────────────────
check_setting() {
    local name="$1"
    local current="$2"
    local recommended="$3"
    local is_warn="${4:-false}"

    local status_icon
    if [[ "${current,,}" == "${recommended,,}" ]]; then
        status_icon="${GREEN}✓ PASS${RESET}"
        ((PASS_COUNT++)) || true
    elif [[ "$is_warn" == "true" ]]; then
        status_icon="${YELLOW}! WARN${RESET}"
        ((WARN_COUNT++)) || true
    else
        status_icon="${RED}✗ FAIL${RESET}"
        ((FAIL_COUNT++)) || true
    fi

    printf "  %-34s %-16s %-16s %b\n" "$name" "$current" "$recommended" "$status_icon"
}

# ══════════════════════════════════════════════════════════════════════
#                          AUDIT MODE
# ══════════════════════════════════════════════════════════════════════

do_audit() {
    if [[ ! -f "$SSHD_CONFIG" ]]; then
        die "sshd_config not found at ${SSHD_CONFIG}"
    fi

    log "Auditing SSH configuration..."
    echo ""
    echo -e "  ${BOLD}SSH Configuration Audit${RESET}"
    printf "  ${BOLD}%-34s %-16s %-16s %s${RESET}\n" "SETTING" "CURRENT" "RECOMMENDED" "STATUS"
    printf "  %s\n" "$(printf '%.0s─' {1..78})"

    # Core auth settings
    check_setting "PermitRootLogin" "$(get_sshd_setting PermitRootLogin)" "no"
    check_setting "PasswordAuthentication" "$(get_sshd_setting PasswordAuthentication)" "no"
    check_setting "ChallengeResponseAuthentication" "$(get_sshd_setting ChallengeResponseAuthentication)" "no"
    check_setting "PubkeyAuthentication" "$(get_sshd_setting PubkeyAuthentication)" "yes"

    # Limits
    check_setting "MaxAuthTries" "$(get_sshd_setting MaxAuthTries)" "$MAX_AUTH_TRIES"
    check_setting "MaxSessions" "$(get_sshd_setting MaxSessions)" "$MAX_SESSIONS"

    # Timeouts
    check_setting "ClientAliveInterval" "$(get_sshd_setting ClientAliveInterval)" "$IDLE_TIMEOUT"

    local cac_current
    cac_current=$(get_sshd_setting ClientAliveCountMax)
    if [[ "$cac_current" == "3" ]]; then
        check_setting "ClientAliveCountMax" "$cac_current" "2" "true"
    else
        check_setting "ClientAliveCountMax" "$cac_current" "2"
    fi

    # Forwarding
    check_setting "X11Forwarding" "$(get_sshd_setting X11Forwarding)" "no"
    check_setting "AllowTcpForwarding" "$(get_sshd_setting AllowTcpForwarding)" "no"
    check_setting "AllowAgentForwarding" "$(get_sshd_setting AllowAgentForwarding)" "no"
    check_setting "PermitTunnel" "$(get_sshd_setting PermitTunnel)" "no"

    # Crypto
    local ciphers_current
    ciphers_current=$(get_sshd_setting Ciphers)
    if [[ "$ciphers_current" == "(not set)" ]]; then
        check_setting "Ciphers" "(default)" "(restricted)"
    elif [[ "$ciphers_current" == "$RECOMMENDED_CIPHERS" ]]; then
        check_setting "Ciphers" "(restricted)" "(restricted)"
    else
        check_setting "Ciphers" "(custom)" "(restricted)"
    fi

    local macs_current
    macs_current=$(get_sshd_setting MACs)
    if [[ "$macs_current" == "(not set)" ]]; then
        check_setting "MACs" "(default)" "(restricted)"
    elif [[ "$macs_current" == "$RECOMMENDED_MACS" ]]; then
        check_setting "MACs" "(restricted)" "(restricted)"
    else
        check_setting "MACs" "(custom)" "(restricted)"
    fi

    local kex_current
    kex_current=$(get_sshd_setting KexAlgorithms)
    if [[ "$kex_current" == "(not set)" ]]; then
        check_setting "KexAlgorithms" "(default)" "(restricted)"
    elif [[ "$kex_current" == "$RECOMMENDED_KEX" ]]; then
        check_setting "KexAlgorithms" "(restricted)" "(restricted)"
    else
        check_setting "KexAlgorithms" "(custom)" "(restricted)"
    fi

    # Logging and misc
    check_setting "LogLevel" "$(get_sshd_setting LogLevel)" "VERBOSE"
    check_setting "LoginGraceTime" "$(get_sshd_setting LoginGraceTime)" "30"

    # AllowUsers / AllowGroups (warn if not set)
    local au_current ag_current
    au_current=$(get_sshd_setting AllowUsers)
    ag_current=$(get_sshd_setting AllowGroups)
    if [[ "$au_current" == "(not set)" ]]; then
        check_setting "AllowUsers" "(not set)" "(recommended)" "true"
    else
        check_setting "AllowUsers" "(configured)" "(recommended)"
    fi
    if [[ "$ag_current" == "(not set)" ]]; then
        check_setting "AllowGroups" "(not set)" "(recommended)" "true"
    else
        check_setting "AllowGroups" "(configured)" "(recommended)"
    fi

    # Summary
    local total_checks=$((PASS_COUNT + FAIL_COUNT + WARN_COUNT))
    local score=0
    if [[ "$total_checks" -gt 0 ]]; then
        score=$(( PASS_COUNT * 100 / total_checks ))
    fi

    echo ""
    echo -e "  ${BOLD}Summary${RESET}"
    echo "    Total checks:     ${total_checks}"
    echo -e "    Passed:           ${GREEN}${PASS_COUNT}${RESET}"
    echo -e "    Failed:           ${RED}${FAIL_COUNT}${RESET}"
    echo -e "    Warnings:         ${YELLOW}${WARN_COUNT}${RESET}"
    echo "    Score:            ${score} / 100"

    # Extra warnings
    echo ""
    if ! command -v fail2ban-client &>/dev/null; then
        warn "Fail2ban not installed — brute-force protection unavailable"
    fi
    if [[ "$au_current" == "(not set)" && "$ag_current" == "(not set)" ]]; then
        warn "No AllowUsers/AllowGroups configured — all users can SSH in"
    fi

    log "Run with --apply to harden this host"
    log "Completed in $(elapsed)"
}

# ══════════════════════════════════════════════════════════════════════
#                          APPLY MODE
# ══════════════════════════════════════════════════════════════════════

do_apply() {
    require_root

    if [[ ! -f "$SSHD_CONFIG" ]]; then
        die "sshd_config not found at ${SSHD_CONFIG}"
    fi

    # Create backup
    local backup_dir
    backup_dir="${BACKUP_ROOT}/bastion-hardener-backup-$(date +%Y%m%d-%H%M%S)"
    log "Backing up ${SSHD_CONFIG} → ${backup_dir}/sshd_config"
    mkdir -p "$backup_dir"
    cp -p "$SSHD_CONFIG" "${backup_dir}/sshd_config"
    [[ -f /etc/ssh/banner.txt ]] && cp -p /etc/ssh/banner.txt "${backup_dir}/banner.txt"

    if [[ "$DRY_RUN" == "true" ]]; then
        log "${YELLOW}DRY RUN${RESET} — previewing changes (no files will be modified)"
        local tmp_config
        tmp_config=$(mktemp)
        cp "$SSHD_CONFIG" "$tmp_config"
        apply_settings "$tmp_config"
        echo ""
        log "Diff preview:"
        diff "$SSHD_CONFIG" "$tmp_config" || true
        rm -f "$tmp_config"
        log "Run without --dry-run to apply changes"
        return
    fi

    log "Applying SSH hardening..."
    apply_settings "$SSHD_CONFIG"

    # Create banner
    if [[ ! -f /etc/ssh/banner.txt ]]; then
        cat > /etc/ssh/banner.txt <<'BANNER'
***************************************************************************
*                        AUTHORIZED ACCESS ONLY                           *
*                                                                         *
*  This system is restricted to authorized users. All activities are       *
*  monitored and logged. Unauthorized access is prohibited and subject     *
*  to prosecution under applicable law.                                   *
*                                                                         *
*  By proceeding, you acknowledge that you have read and agree to the     *
*  organization's acceptable use policies.                                *
***************************************************************************
BANNER
        log "Created warning banner at /etc/ssh/banner.txt"
    fi

    # Session logging directory
    if [[ "$ENABLE_SESSION_LOGGING" == "true" ]]; then
        mkdir -p "$SESSION_LOG_DIR"
        chmod 700 "$SESSION_LOG_DIR"
        log "Session log directory: ${SESSION_LOG_DIR}"
    fi

    # Fail2ban configuration
    if [[ "$CONFIGURE_FAIL2BAN" == "true" ]]; then
        configure_fail2ban
    fi

    # Validate config
    log "Validating sshd configuration..."
    if sshd -t -f "$SSHD_CONFIG" 2>/dev/null; then
        echo -e "  ${GREEN}✓${RESET} sshd -t passed"
    else
        err "sshd -t validation failed — restoring backup"
        cp -p "${backup_dir}/sshd_config" "$SSHD_CONFIG"
        die "Config validation failed. Original config restored."
    fi

    # Restart sshd
    log "Restarting sshd..."
    if systemctl restart sshd 2>/dev/null || systemctl restart ssh 2>/dev/null; then
        echo -e "  ${GREEN}✓${RESET} sshd restarted successfully"
    else
        warn "Could not restart sshd — restart manually"
    fi

    # Write audit report
    local report_file
    report_file="/var/log/bastion-hardener-$(date +%Y%m%d-%H%M%S).log"
    {
        echo "Bastion Hardener — Apply Report"
        echo "Time: $(date -u '+%Y-%m-%dT%H:%M:%SZ')"
        echo "Host: $(hostname -f 2>/dev/null || hostname)"
        echo "Changes: ${CHANGES}"
        echo "Backup: ${backup_dir}"
    } > "$report_file" 2>/dev/null || true
    log "Writing audit report → ${report_file}"

    log "Changes applied: ${CHANGES}, skipped: 0"
    log "Backup directory: ${backup_dir}"
    log "To rollback: ./${SCRIPT_NAME} --rollback"
    log "Completed in $(elapsed)"
}

apply_settings() {
    local config_file="$1"

    local settings=(
        "PermitRootLogin no"
        "PasswordAuthentication no"
        "ChallengeResponseAuthentication no"
        "KbdInteractiveAuthentication no"
        "PubkeyAuthentication yes"
        "MaxAuthTries ${MAX_AUTH_TRIES}"
        "MaxSessions ${MAX_SESSIONS}"
        "ClientAliveInterval ${IDLE_TIMEOUT}"
        "ClientAliveCountMax 2"
        "X11Forwarding no"
        "AllowTcpForwarding no"
        "AllowAgentForwarding no"
        "PermitTunnel no"
        "Ciphers ${RECOMMENDED_CIPHERS}"
        "MACs ${RECOMMENDED_MACS}"
        "KexAlgorithms ${RECOMMENDED_KEX}"
        "LoginGraceTime 30"
        "LogLevel VERBOSE"
        "Banner /etc/ssh/banner.txt"
    )

    for setting in "${settings[@]}"; do
        local key value
        key="${setting%% *}"
        value="${setting#* }"
        set_sshd_config "$key" "$value" "$config_file"
        echo -e "  ${GREEN}✓${RESET} ${key} → ${value}"
        ((CHANGES++)) || true
    done

    # AllowUsers
    if [[ -n "$ALLOW_USERS" ]]; then
        local users_val="${ALLOW_USERS//,/ }"
        set_sshd_config "AllowUsers" "$users_val" "$config_file"
        echo -e "  ${GREEN}✓${RESET} AllowUsers → ${users_val}"
        ((CHANGES++)) || true
    fi

    # AllowGroups
    if [[ -n "$ALLOW_GROUPS" ]]; then
        local groups_val="${ALLOW_GROUPS//,/ }"
        set_sshd_config "AllowGroups" "$groups_val" "$config_file"
        echo -e "  ${GREEN}✓${RESET} AllowGroups → ${groups_val}"
        ((CHANGES++)) || true
    fi
}

configure_fail2ban() {
    if ! command -v fail2ban-client &>/dev/null; then
        warn "fail2ban not installed — skipping jail configuration"
        return
    fi

    local jail_file="/etc/fail2ban/jail.d/bastion-ssh.conf"
    log "Configuring fail2ban SSH jail → ${jail_file}"

    cat > "$jail_file" <<EOF
# Bastion Hardener — SSH jail configuration
# Generated: $(date -u '+%Y-%m-%dT%H:%M:%SZ')

[sshd]
enabled  = true
port     = ssh
filter   = sshd
logpath  = /var/log/auth.log
maxretry = ${FAIL2BAN_MAXRETRY}
bantime  = ${FAIL2BAN_BANTIME}
findtime = 600
action   = %(action_)s
EOF

    if systemctl restart fail2ban 2>/dev/null; then
        echo -e "  ${GREEN}✓${RESET} fail2ban SSH jail configured and restarted"
    else
        warn "Could not restart fail2ban"
    fi
}

# ══════════════════════════════════════════════════════════════════════
#                          ROLLBACK MODE
# ══════════════════════════════════════════════════════════════════════

do_rollback() {
    require_root

    local target_dir="$ROLLBACK_DIR"

    if [[ -z "$target_dir" ]]; then
        # Find most recent backup
        target_dir=$(find "$BACKUP_ROOT" -maxdepth 1 -type d -name "bastion-hardener-backup-*" 2>/dev/null | sort -r | head -1)
        if [[ -z "$target_dir" ]]; then
            die "No backup directories found in ${BACKUP_ROOT}"
        fi
    fi

    if [[ ! -d "$target_dir" ]]; then
        die "Backup directory not found: ${target_dir}"
    fi

    log "Restoring from ${target_dir}..."

    if [[ -f "${target_dir}/sshd_config" ]]; then
        cp -p "${target_dir}/sshd_config" "$SSHD_CONFIG"
        echo -e "  ${GREEN}✓${RESET} Restored sshd_config"
    else
        die "No sshd_config found in backup directory"
    fi

    if [[ -f "${target_dir}/banner.txt" ]]; then
        cp -p "${target_dir}/banner.txt" /etc/ssh/banner.txt
        echo -e "  ${GREEN}✓${RESET} Restored banner.txt"
    fi

    # Validate
    log "Validating restored configuration..."
    if sshd -t -f "$SSHD_CONFIG" 2>/dev/null; then
        echo -e "  ${GREEN}✓${RESET} sshd -t passed"
    else
        die "Restored config failed validation"
    fi

    # Restart
    log "Restarting sshd..."
    if systemctl restart sshd 2>/dev/null || systemctl restart ssh 2>/dev/null; then
        echo -e "  ${GREEN}✓${RESET} sshd restarted successfully"
    else
        warn "Could not restart sshd — restart manually"
    fi

    log "Rollback complete from ${target_dir}"
    log "Completed in $(elapsed)"
}

# ══════════════════════════════════════════════════════════════════════
#                          HELP
# ══════════════════════════════════════════════════════════════════════

show_help() {
    cat <<EOF
Usage: $SCRIPT_NAME [MODE] [OPTIONS]

Harden SSH bastion/jump hosts with audit, apply, and rollback.

MODES:
    --audit                 Read-only check of current SSH configuration
    --apply                 Apply hardening changes (backs up first)
    --rollback              Restore configuration from most recent backup

OPTIONS:
    --dry-run               Preview changes without writing to disk
    --allow-users USERS     Comma-separated list of allowed SSH users
    --allow-groups GROUPS   Comma-separated list of allowed SSH groups
    --idle-timeout SECS     Client alive interval in seconds (default: $IDLE_TIMEOUT)
    --max-auth-tries N      Maximum authentication attempts (default: $MAX_AUTH_TRIES)
    --max-sessions N        Maximum sessions per connection (default: $MAX_SESSIONS)
    --session-logging       Enable ForceCommand session recording
    --configure-fail2ban    Configure fail2ban SSH jail
    --backup-dir DIR        Specific backup directory for rollback
    --verbose               Debug output
    --no-color              Disable colored output
    --help, -h              Show this help

ENVIRONMENT VARIABLES:
    SSHD_CONFIG             Path to sshd_config (default: /etc/ssh/sshd_config)
    BACKUP_ROOT             Backup storage root (default: /etc/ssh)
    ALLOW_USERS             Comma-separated allowed users
    ALLOW_GROUPS            Comma-separated allowed groups
    IDLE_TIMEOUT            ClientAliveInterval in seconds (default: 300)
    MAX_AUTH_TRIES          Maximum auth attempts (default: 3)
    MAX_SESSIONS            Maximum sessions (default: 2)
    SESSION_LOG_DIR         Session recording directory
    FAIL2BAN_BANTIME        Ban duration in seconds (default: 3600)
    FAIL2BAN_MAXRETRY       Max retries before ban (default: 3)
    DRY_RUN                 Preview without writing (default: false)
    VERBOSE                 Debug output (default: false)

EXAMPLES:
    # Audit current SSH configuration
    sudo $SCRIPT_NAME --audit

    # Preview what would change
    sudo $SCRIPT_NAME --apply --dry-run

    # Apply hardening with specific allowed users
    sudo $SCRIPT_NAME --apply --allow-users "deployer,jumpuser"

    # Apply with fail2ban
    sudo $SCRIPT_NAME --apply --configure-fail2ban

    # Rollback to previous config
    sudo $SCRIPT_NAME --rollback
EOF
}

# ══════════════════════════════════════════════════════════════════════
#                          MAIN
# ══════════════════════════════════════════════════════════════════════

main() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --audit)              RUN_MODE="audit"; shift ;;
            --apply)              RUN_MODE="apply"; shift ;;
            --rollback)           RUN_MODE="rollback"; shift ;;
            --dry-run)            DRY_RUN="true"; shift ;;
            --allow-users)        ALLOW_USERS="$2"; shift 2 ;;
            --allow-groups)       ALLOW_GROUPS="$2"; shift 2 ;;
            --idle-timeout)       IDLE_TIMEOUT="$2"; shift 2 ;;
            --max-auth-tries)     MAX_AUTH_TRIES="$2"; shift 2 ;;
            --max-sessions)       MAX_SESSIONS="$2"; shift 2 ;;
            --session-logging)    ENABLE_SESSION_LOGGING="true"; shift ;;
            --configure-fail2ban) CONFIGURE_FAIL2BAN="true"; shift ;;
            --backup-dir)         ROLLBACK_DIR="$2"; shift 2 ;;
            --verbose)            VERBOSE="true"; shift ;;
            --no-color)           COLOR="never"; shift ;;
            --help|-h)            show_help; exit 0 ;;
            *)                    die "Unknown option: $1 (see --help)" ;;
        esac
    done

    setup_colors

    if [[ -z "$RUN_MODE" ]]; then err "No mode specified"; echo ""; show_help; exit 1; fi

    START_TIME=$(date +%s)

    echo ""
    echo -e "${BOLD}Bastion Hardener${RESET}"
    echo "Host:   $(hostname -f 2>/dev/null || hostname)"
    echo "Mode:   ${RUN_MODE}"
    echo "Time:   $(date -u '+%Y-%m-%dT%H:%M:%SZ')"
    echo ""

    case "$RUN_MODE" in
        audit)    do_audit ;;
        apply)    do_apply ;;
        rollback) do_rollback ;;
    esac
}

main "$@"