#!/bin/bash

################################################
#### Salt Key Manager                       ####
#### Automate salt-key operations           ####
####                                        ####
#### Author: Phil Connor                    ####
#### Contact: contact@mylinux.work          ####
#### Version: 1.00-030526                   ####
################################################

set -o pipefail

SCRIPT_NAME=$(basename "$0")
readonly SCRIPT_NAME

# Default configuration
readonly DEFAULT_STALE_DAYS=30
readonly DEFAULT_CACHE_DIR="/var/cache/salt/master/minions"

# Configuration variables
DEBUG=${DEBUG:-}

# Runtime flags
ACTION=""
TARGET_MINION=""
STALE_DAYS=$DEFAULT_STALE_DAYS
EXPORT_PATH=""
BULK_FILE=""
AUTO_YES=false
USE_COLOR=true

# Colors
C_GREEN=""
C_YELLOW=""
C_RED=""
C_CYAN=""
C_RESET=""

handle_error() {
    local exit_code=$1
    local line_number=$2
    echo "Error: $SCRIPT_NAME failed at line $line_number with exit code $exit_code" >&2
    exit "$exit_code"
}

trap 'handle_error $? $LINENO' ERR

debug_echo() {
    if [[ -n "$DEBUG" ]]; then
        echo "[DEBUG] $*" >&2
    fi
}

log_info() {
    echo "[INFO] $*"
}

log_warn() {
    echo "[WARN] $*" >&2
}

log_error() {
    echo "[ERROR] $*" >&2
}

setup_colors() {
    if [[ "$USE_COLOR" == true ]] && [[ -t 1 ]]; then
        C_GREEN='\033[0;32m'
        C_YELLOW='\033[0;33m'
        C_RED='\033[0;31m'
        C_CYAN='\033[0;36m'
        C_RESET='\033[0m'
    fi
}

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

Manage Salt minion keys — accept, reject, delete, verify, rotate, and
clean up stale keys.

ACTIONS:
    --list                  List all keys by status with counts
    --verify                Show pending keys with fingerprints for verification
    --accept-all            Accept all pending keys
    --accept MINION         Accept a specific minion key
    --reject MINION         Reject a specific minion key
    --delete MINION         Delete a specific minion key
    --rotate MINION         Rotate a minion key (delete, re-accept on reconnect)
    --cleanup-stale [DAYS]  Delete keys for minions not seen in DAYS days (default: $DEFAULT_STALE_DAYS)
    --export PATH           Export all accepted key fingerprints to a file
    --bulk-accept FILE      Accept minions listed in a file (one per line)

OPTIONS:
    --yes                   Skip confirmation prompts
    --no-color              Disable colored output
    --help, -h              Show this help message

ENVIRONMENT VARIABLES:
    DEBUG                   Enable debug output

EXAMPLES:
    # List all keys with status
    sudo $SCRIPT_NAME --list

    # Show pending keys for verification
    sudo $SCRIPT_NAME --verify

    # Accept all pending keys
    sudo $SCRIPT_NAME --accept-all --yes

    # Accept a specific minion
    sudo $SCRIPT_NAME --accept web01

    # Clean up minions not seen in 60 days
    sudo $SCRIPT_NAME --cleanup-stale 60

    # Export fingerprints for auditing
    sudo $SCRIPT_NAME --export /tmp/salt-keys.txt

    # Bulk accept from a file
    sudo $SCRIPT_NAME --bulk-accept /tmp/new-minions.txt --yes
EOF
}

count_keys() {
    local status="$1"
    salt-key --list "$status" 2>/dev/null | grep -cv "^$status\|^$" || echo 0
}

do_list() {
    echo "Salt Key Status"
    echo "==============="
    echo ""

    local accepted unaccepted denied rejected
    accepted=$(count_keys "accepted")
    unaccepted=$(count_keys "unaccepted")
    denied=$(count_keys "denied")
    rejected=$(count_keys "rejected")

    printf '  %bAccepted:%b   %d\n' "$C_GREEN" "$C_RESET" "$accepted"
    printf '  %bPending:%b    %d\n' "$C_YELLOW" "$C_RESET" "$unaccepted"
    printf '  %bDenied:%b     %d\n' "$C_RED" "$C_RESET" "$denied"
    printf '  %bRejected:%b   %d\n' "$C_RED" "$C_RESET" "$rejected"
    echo ""

    if ((accepted > 0)); then
        printf '%bAccepted Keys:%b\n' "$C_GREEN" "$C_RESET"
        salt-key --list accepted 2>/dev/null | grep -v "^Accepted Keys:" | sed 's/^/  /'
        echo ""
    fi

    if ((unaccepted > 0)); then
        printf '%bPending Keys:%b\n' "$C_YELLOW" "$C_RESET"
        salt-key --list unaccepted 2>/dev/null | grep -v "^Unaccepted Keys:" | sed 's/^/  /'
        echo ""
    fi

    if ((denied > 0)); then
        printf '%bDenied Keys:%b\n' "$C_RED" "$C_RESET"
        salt-key --list denied 2>/dev/null | grep -v "^Denied Keys:" | sed 's/^/  /'
        echo ""
    fi

    if ((rejected > 0)); then
        printf '%bRejected Keys:%b\n' "$C_RED" "$C_RESET"
        salt-key --list rejected 2>/dev/null | grep -v "^Rejected Keys:" | sed 's/^/  /'
        echo ""
    fi
}

do_verify() {
    local pending
    pending=$(salt-key --list unaccepted 2>/dev/null | grep -v "^Unaccepted Keys:$" | grep -v "^$")

    if [[ -z "$pending" ]]; then
        log_info "No pending keys to verify"
        return 0
    fi

    echo "Master Fingerprint:"
    printf '  %b' "$C_CYAN"
    salt-key -F master 2>/dev/null | grep -A1 "master.pub" | tail -1 | tr -d ' '
    printf '%b\n\n' "$C_RESET"

    echo "Pending Keys with Fingerprints:"
    echo ""

    while IFS= read -r minion; do
        [[ -z "$minion" ]] && continue
        minion=$(echo "$minion" | tr -d '[:space:]')
        local fingerprint
        fingerprint=$(salt-key -f "$minion" 2>/dev/null | grep -v "^Unaccepted Keys:" | awk '{print $2}' | head -1)
        printf '  %b%-30s%b %s\n' "$C_YELLOW" "$minion" "$C_RESET" "${fingerprint:-unknown}"
    done <<< "$pending"

    echo ""
    log_info "Verify each fingerprint matches the minion's local fingerprint:"
    log_info "  (on minion) salt-call --local key.finger"
}

do_accept_all() {
    local pending
    pending=$(count_keys "unaccepted")

    if ((pending == 0)); then
        log_info "No pending keys to accept"
        return 0
    fi

    log_info "Accepting $pending pending key(s)..."

    if [[ "$AUTO_YES" != true ]]; then
        echo "Accept all $pending pending keys? [y/N] "
        read -r confirm
        if [[ "$confirm" != "y" && "$confirm" != "Y" ]]; then
            log_info "Aborted"
            return 0
        fi
    fi

    salt-key -A -y 2>/dev/null
    log_info "All pending keys accepted"
}

do_accept() {
    local minion="$1"
    log_info "Accepting key for: $minion"

    if [[ "$AUTO_YES" != true ]]; then
        local fingerprint
        fingerprint=$(salt-key -f "$minion" 2>/dev/null | grep -v "^Unaccepted Keys:" | awk '{print $2}' | head -1)
        echo "Fingerprint: ${fingerprint:-unknown}"
        echo "Accept key for $minion? [y/N] "
        read -r confirm
        if [[ "$confirm" != "y" && "$confirm" != "Y" ]]; then
            log_info "Aborted"
            return 0
        fi
    fi

    salt-key -a "$minion" -y 2>/dev/null
    log_info "Key accepted for $minion"
}

do_reject() {
    local minion="$1"
    log_info "Rejecting key for: $minion"

    if [[ "$AUTO_YES" != true ]]; then
        echo "Reject key for $minion? [y/N] "
        read -r confirm
        if [[ "$confirm" != "y" && "$confirm" != "Y" ]]; then
            log_info "Aborted"
            return 0
        fi
    fi

    salt-key -r "$minion" -y 2>/dev/null
    log_info "Key rejected for $minion"
}

do_delete() {
    local minion="$1"
    log_info "Deleting key for: $minion"

    if [[ "$AUTO_YES" != true ]]; then
        echo "Delete key for $minion? This cannot be undone. [y/N] "
        read -r confirm
        if [[ "$confirm" != "y" && "$confirm" != "Y" ]]; then
            log_info "Aborted"
            return 0
        fi
    fi

    salt-key -d "$minion" -y 2>/dev/null
    log_info "Key deleted for $minion"
}

do_rotate() {
    local minion="$1"
    log_info "Rotating key for: $minion"
    log_info "This will delete the current key — the minion must reconnect to get a new key accepted"

    if [[ "$AUTO_YES" != true ]]; then
        echo "Rotate key for $minion? [y/N] "
        read -r confirm
        if [[ "$confirm" != "y" && "$confirm" != "Y" ]]; then
            log_info "Aborted"
            return 0
        fi
    fi

    salt-key -d "$minion" -y 2>/dev/null
    log_info "Key deleted for $minion — accept the new key when the minion reconnects"
    log_info "On the minion, restart salt-minion: systemctl restart salt-minion"
}

do_cleanup_stale() {
    local days="$1"
    log_info "Finding minions not seen in $days days..."

    if [[ ! -d "$DEFAULT_CACHE_DIR" ]]; then
        log_error "Minion cache directory not found: $DEFAULT_CACHE_DIR"
        return 1
    fi

    local stale_minions=()
    local cutoff
    cutoff=$(date -d "-${days} days" +%s 2>/dev/null) || cutoff=$(date -v-"${days}"d +%s 2>/dev/null)

    while IFS= read -r minion_dir; do
        local minion_name
        minion_name=$(basename "$minion_dir")
        local last_modified
        last_modified=$(stat -c %Y "$minion_dir" 2>/dev/null) || last_modified=$(stat -f %m "$minion_dir" 2>/dev/null) || continue

        if ((last_modified < cutoff)); then
            local days_ago=$(( ($(date +%s) - last_modified) / 86400 ))
            stale_minions+=("$minion_name")
            printf '  %b%-30s%b (last seen %d days ago)\n' "$C_RED" "$minion_name" "$C_RESET" "$days_ago"
        fi
    done < <(find "$DEFAULT_CACHE_DIR" -maxdepth 1 -mindepth 1 -type d 2>/dev/null)

    if [[ ${#stale_minions[@]} -eq 0 ]]; then
        log_info "No stale minions found"
        return 0
    fi

    echo ""
    log_info "Found ${#stale_minions[@]} stale minion(s)"

    if [[ "$AUTO_YES" != true ]]; then
        echo "Delete keys for all ${#stale_minions[@]} stale minions? [y/N] "
        read -r confirm
        if [[ "$confirm" != "y" && "$confirm" != "Y" ]]; then
            log_info "Aborted"
            return 0
        fi
    fi

    for minion in "${stale_minions[@]}"; do
        salt-key -d "$minion" -y 2>/dev/null && log_info "Deleted key: $minion"
    done

    log_info "Stale key cleanup complete"
}

do_export() {
    local output_path="$1"
    log_info "Exporting accepted key fingerprints to $output_path..."

    {
        echo "# Salt Key Fingerprint Export"
        echo "# Generated: $(date -u '+%Y-%m-%d %H:%M:%S UTC')"
        echo "# Master: $(hostname -f 2>/dev/null || hostname)"
        echo "#"
        echo "# Format: minion_id fingerprint"
        echo ""
        salt-key -F accepted 2>/dev/null | grep -v "^Accepted Keys:" | while IFS= read -r line; do
            [[ -z "$line" ]] && continue
            echo "$line"
        done
    } > "$output_path"

    local count
    count=$(grep -cv "^#\|^$" "$output_path" 2>/dev/null) || count=0
    log_info "Exported $count key fingerprint(s) to $output_path"
}

do_bulk_accept() {
    local input_file="$1"

    if [[ ! -f "$input_file" ]]; then
        log_error "File not found: $input_file"
        return 1
    fi

    local count=0
    local failed=0

    while IFS= read -r line; do
        [[ -z "$line" || "$line" == \#* ]] && continue
        local minion_id="${line%%:*}"
        minion_id=$(echo "$minion_id" | tr -d '[:space:]')

        if salt-key -a "$minion_id" -y 2>/dev/null; then
            log_info "Accepted: $minion_id"
            count=$((count + 1))
        else
            log_error "Failed to accept: $minion_id"
            failed=$((failed + 1))
        fi
    done < "$input_file"

    log_info "Bulk accept complete: $count accepted, $failed failed"
}

parse_arguments() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            --list)
                ACTION="list"
                shift
                ;;
            --verify)
                ACTION="verify"
                shift
                ;;
            --accept-all)
                ACTION="accept-all"
                shift
                ;;
            --accept)
                ACTION="accept"
                TARGET_MINION="$2"
                shift 2
                ;;
            --reject)
                ACTION="reject"
                TARGET_MINION="$2"
                shift 2
                ;;
            --delete)
                ACTION="delete"
                TARGET_MINION="$2"
                shift 2
                ;;
            --rotate)
                ACTION="rotate"
                TARGET_MINION="$2"
                shift 2
                ;;
            --cleanup-stale)
                ACTION="cleanup-stale"
                if [[ -n "${2:-}" && "$2" =~ ^[0-9]+$ ]]; then
                    STALE_DAYS="$2"
                    shift 2
                else
                    shift
                fi
                ;;
            --export)
                ACTION="export"
                EXPORT_PATH="$2"
                shift 2
                ;;
            --bulk-accept)
                ACTION="bulk-accept"
                BULK_FILE="$2"
                shift 2
                ;;
            --yes)
                AUTO_YES=true
                shift
                ;;
            --no-color)
                USE_COLOR=false
                shift
                ;;
            --help|-h)
                show_help
                exit 0
                ;;
            *)
                log_error "Unknown option: $1"
                show_help >&2
                exit 1
                ;;
        esac
    done
}

validate_requirements() {
    if [[ $EUID -ne 0 ]]; then
        log_error "This script must be run as root (use sudo)"
        exit 1
    fi

    if [[ -z "$ACTION" ]]; then
        log_error "An action is required"
        show_help >&2
        exit 1
    fi

    if ! command -v salt-key >/dev/null 2>&1; then
        log_error "salt-key not found — is salt-master installed?"
        exit 1
    fi

    if [[ "$ACTION" == "accept" || "$ACTION" == "reject" || "$ACTION" == "delete" || "$ACTION" == "rotate" ]]; then
        if [[ -z "$TARGET_MINION" ]]; then
            log_error "Minion name is required for --$ACTION"
            exit 1
        fi
    fi

    if [[ "$ACTION" == "export" && -z "$EXPORT_PATH" ]]; then
        log_error "Output path is required for --export"
        exit 1
    fi

    if [[ "$ACTION" == "bulk-accept" && -z "$BULK_FILE" ]]; then
        log_error "Input file is required for --bulk-accept"
        exit 1
    fi
}

main() {
    parse_arguments "$@"
    validate_requirements
    setup_colors

    case "$ACTION" in
        list)           do_list ;;
        verify)         do_verify ;;
        accept-all)     do_accept_all ;;
        accept)         do_accept "$TARGET_MINION" ;;
        reject)         do_reject "$TARGET_MINION" ;;
        delete)         do_delete "$TARGET_MINION" ;;
        rotate)         do_rotate "$TARGET_MINION" ;;
        cleanup-stale)  do_cleanup_stale "$STALE_DAYS" ;;
        export)         do_export "$EXPORT_PATH" ;;
        bulk-accept)    do_bulk_accept "$BULK_FILE" ;;
    esac

    debug_echo "Script completed successfully"
}

if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi