#!/usr/bin/env bash

######################################################################################
#### ssh-key-auditor.sh — Audit SSH keys across user accounts for compliance      ####
#### Scans authorized_keys, flags weak keys, detects duplicates, checks perms.    ####
#### Requires: bash 4+, coreutils, ssh-keygen                                    ####
####                                                                              ####
#### Author: Phil Connor                                                          ####
#### Contact: contact@mylinux.work                                                ####
#### License: MIT                                                                 ####
#### Version 1.01                                                                 ####
####                                                                              ####
#### Usage:                                                                       ####
####   sudo ./ssh-key-auditor.sh                                                  ####
####   sudo ./ssh-key-auditor.sh --min-rsa-bits 4096 --format tap                 ####
####                                                                              ####
#### See --help for all options.                                                  ####
######################################################################################

set -euo pipefail

# ── Defaults ──────────────────────────────────────────────────────────
SCAN_PATH="${SCAN_PATH:-/home}"
SCAN_SYSTEM_USERS="${SCAN_SYSTEM_USERS:-false}"
MIN_RSA_BITS="${MIN_RSA_BITS:-2048}"
SKIP_DUPLICATES="${SKIP_DUPLICATES:-false}"
OUTPUT_FORMAT="${OUTPUT_FORMAT:-text}"  # text, tap, junit
JUNIT_FILE="${JUNIT_FILE:-ssh-audit-results.xml}"
VERBOSE="${VERBOSE:-false}"
COLOR="${COLOR:-auto}"

# ── State ─────────────────────────────────────────────────────────────
PASS=0
FAIL=0
SKIP=0
TOTAL=0
RESULTS=()
START_TIME=""

# Key tracking for duplicate detection and summary
declare -A KEY_FINGERPRINTS   # fingerprint -> "user1 user2 ..."
declare -A KEY_TYPE_COUNTS    # type -> count
TOTAL_KEYS=0
USERS_WITH_KEYS=0

# ── Colors ────────────────────────────────────────────────────────────
setup_colors() {
    if [[ "$COLOR" == "never" ]]; then
        RED="" GREEN="" YELLOW="" BLUE="" BOLD="" 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'
        BOLD='\033[1m'
        RESET='\033[0m'
    else
        RED="" GREEN="" YELLOW="" BLUE="" BOLD="" 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 "${BLUE}[DEBUG]${RESET} $*"; fi; }

# ── Test Result Recording ─────────────────────────────────────────────
record_pass() {
    local name="$1"
    local detail="${2:-}"
    ((PASS++)) || true
    ((TOTAL++)) || true
    RESULTS+=("PASS|${name}|${detail}")
    if [[ "$OUTPUT_FORMAT" == "tap" ]]; then
        echo "ok ${TOTAL} - ${name}"
    else
        echo -e "  ${GREEN}✓${RESET} ${name}${detail:+ — ${detail}}"
    fi
}

record_fail() {
    local name="$1"
    local detail="${2:-}"
    ((FAIL++)) || true
    ((TOTAL++)) || true
    RESULTS+=("FAIL|${name}|${detail}")
    if [[ "$OUTPUT_FORMAT" == "tap" ]]; then
        echo "not ok ${TOTAL} - ${name}"
        [[ -n "$detail" ]] && echo "  # ${detail}"
    else
        echo -e "  ${RED}✗${RESET} ${name}${detail:+ — ${detail}}"
    fi
}

record_skip() {
    local name="$1"
    local reason="${2:-}"
    ((SKIP++)) || true
    ((TOTAL++)) || true
    RESULTS+=("SKIP|${name}|${reason}")
    if [[ "$OUTPUT_FORMAT" == "tap" ]]; then
        echo "ok ${TOTAL} - ${name} # SKIP ${reason}"
    else
        echo -e "  ${YELLOW}⊘${RESET} ${name}${reason:+ — ${reason}}"
    fi
}

# ── JUnit XML Writer ──────────────────────────────────────────────────
write_junit() {
    local end_time
    end_time=$(date +%s)
    local duration=$(( end_time - START_TIME ))

    cat > "$JUNIT_FILE" <<JUNIT_EOF
<?xml version="1.0" encoding="UTF-8"?>
<testsuites tests="${TOTAL}" failures="${FAIL}" skipped="${SKIP}" time="${duration}">
  <testsuite name="ssh-key-auditor" tests="${TOTAL}" failures="${FAIL}" skipped="${SKIP}" time="${duration}">
JUNIT_EOF

    for result in "${RESULTS[@]}"; do
        local status name detail
        status=$(echo "$result" | cut -d'|' -f1)
        name=$(echo "$result" | cut -d'|' -f2)
        detail=$(echo "$result" | cut -d'|' -f3)

        # XML-escape
        name="${name//&/&amp;}"
        name="${name//</&lt;}"
        name="${name//>/&gt;}"
        name="${name//\"/&quot;}"
        detail="${detail//&/&amp;}"
        detail="${detail//</&lt;}"
        detail="${detail//>/&gt;}"
        detail="${detail//\"/&quot;}"

        case "$status" in
            PASS)
                echo "    <testcase name=\"${name}\" classname=\"ssh-audit\">" >> "$JUNIT_FILE"
                [[ -n "$detail" ]] && echo "      <system-out>${detail}</system-out>" >> "$JUNIT_FILE"
                echo "    </testcase>" >> "$JUNIT_FILE"
                ;;
            FAIL)
                echo "    <testcase name=\"${name}\" classname=\"ssh-audit\">" >> "$JUNIT_FILE"
                echo "      <failure message=\"${detail}\">FAILED: ${name} — ${detail}</failure>" >> "$JUNIT_FILE"
                echo "    </testcase>" >> "$JUNIT_FILE"
                ;;
            SKIP)
                echo "    <testcase name=\"${name}\" classname=\"ssh-audit\">" >> "$JUNIT_FILE"
                echo "      <skipped message=\"${detail}\"/>" >> "$JUNIT_FILE"
                echo "    </testcase>" >> "$JUNIT_FILE"
                ;;
        esac
    done

    echo "  </testsuite>" >> "$JUNIT_FILE"
    echo "</testsuites>" >> "$JUNIT_FILE"

    log "JUnit report written to ${JUNIT_FILE}"
}

# ══════════════════════════════════════════════════════════════════════
#                          AUDIT FUNCTIONS
# ══════════════════════════════════════════════════════════════════════

# ── Check .ssh directory permissions ──────────────────────────────────
check_ssh_dir_perms() {
    local user="$1"
    local ssh_dir="$2"

    if [[ ! -d "$ssh_dir" ]]; then
        verbose "No .ssh directory for ${user}"
        return
    fi

    local perms
    perms=$(stat -c '%a' "$ssh_dir" 2>/dev/null) || return

    if [[ "$perms" == "700" ]]; then
        record_pass "[${user}] .ssh directory permissions" "${perms}"
    else
        record_fail "[${user}] .ssh directory permissions" "${perms} (expected 700)"
    fi
}

# ── Check authorized_keys file permissions ────────────────────────────
check_authkeys_perms() {
    local user="$1"
    local authkeys_file="$2"

    if [[ ! -f "$authkeys_file" ]]; then
        return
    fi

    local perms
    perms=$(stat -c '%a' "$authkeys_file" 2>/dev/null) || return

    case "$perms" in
        600|644)
            record_pass "[${user}] authorized_keys permissions" "${perms}"
            ;;
        *)
            local world_bit="${perms:2:1}"
            if [[ "$world_bit" =~ [2367] ]]; then
                record_fail "[${user}] authorized_keys permissions" "${perms} (world-writable)"
            else
                record_fail "[${user}] authorized_keys permissions" "${perms} (expected 600 or 644)"
            fi
            ;;
    esac
}

# ── Check for deprecated authorized_keys2 ─────────────────────────────
check_authkeys2() {
    local user="$1"
    local ssh_dir="$2"

    local ak2="${ssh_dir}/authorized_keys2"
    if [[ -f "$ak2" ]]; then
        record_fail "[${user}] authorized_keys2 present" "deprecated, migrate to authorized_keys"
    fi
}

# ── Check private key permissions ─────────────────────────────────────
check_private_key_perms() {
    local user="$1"
    local ssh_dir="$2"

    local privkey_names=("id_rsa" "id_ed25519" "id_ecdsa" "id_dsa" "id_xmss")

    for keyname in "${privkey_names[@]}"; do
        local keyfile="${ssh_dir}/${keyname}"
        if [[ -f "$keyfile" ]]; then
            local perms
            perms=$(stat -c '%a' "$keyfile" 2>/dev/null) || continue
            local world_bit="${perms:2:1}"
            if [[ "$world_bit" != "0" ]]; then
                record_fail "[${user}] Private key ${keyname}" "world-readable (${perms})"
            else
                verbose "[${user}] Private key ${keyname} permissions OK (${perms})"
            fi
        fi
    done
}

# ── Audit a single key line ───────────────────────────────────────────
audit_key_line() {
    local user="$1"
    local key_num="$2"
    local line="$3"
    local authkeys_file="$4"

    # Skip empty lines and comments
    [[ -z "$line" || "$line" =~ ^[[:space:]]*# ]] && return

    # Check for command= or from= restrictions
    local restrictions=""
    if [[ "$line" =~ ^(.*[[:space:]])?(command=\"[^\"]*\") ]]; then
        restrictions="${BASH_REMATCH[2]}"
    fi
    if [[ "$line" =~ ^(.*[[:space:]])?(from=\"[^\"]*\") ]]; then
        restrictions="${restrictions:+${restrictions}, }${BASH_REMATCH[2]}"
    fi

    # Write key to temp file for ssh-keygen parsing
    local tmpkey
    tmpkey=$(mktemp)
    echo "$line" > "$tmpkey"

    local key_info
    key_info=$(ssh-keygen -l -f "$tmpkey" 2>/dev/null) || {
        rm -f "$tmpkey"
        verbose "[${user}] Key ${key_num}: could not parse"
        return
    }
    rm -f "$tmpkey"

    # Parse ssh-keygen output: "2048 SHA256:xxx comment (RSA)"
    local bits key_type comment fingerprint
    bits=$(echo "$key_info" | awk '{print $1}')
    fingerprint=$(echo "$key_info" | awk '{print $2}')
    key_type=$(echo "$key_info" | grep -oP '\(([A-Z0-9]+)\)' | tr -d '()')
    comment=$(echo "$key_info" | awk '{for(i=3;i<NF;i++) printf "%s ", $i; print ""}' | sed 's/ *$//')

    # Normalize key type to lowercase for counting
    local key_type_lower
    key_type_lower=$(echo "$key_type" | tr '[:upper:]' '[:lower:]')

    ((TOTAL_KEYS++)) || true
    KEY_TYPE_COUNTS["$key_type_lower"]=$(( ${KEY_TYPE_COUNTS["$key_type_lower"]:-0} + 1 ))

    # Track fingerprint for duplicate detection
    if [[ "$SKIP_DUPLICATES" != "true" && -n "$fingerprint" ]]; then
        KEY_FINGERPRINTS["$fingerprint"]="${KEY_FINGERPRINTS["$fingerprint"]:-} ${user}"
    fi

    # Key type checks
    case "$key_type" in
        DSA)
            record_fail "[${user}] Key ${key_num}: ${key_type_lower} (${bits} bits)" "DEPRECATED key type (DSA)"
            ;;
        RSA)
            if [[ "$bits" -lt "$MIN_RSA_BITS" ]]; then
                record_fail "[${user}] Key ${key_num}: ${key_type_lower} (${bits} bits)" "below minimum ${MIN_RSA_BITS} bits"
            else
                record_pass "[${user}] Key ${key_num}: ssh-${key_type_lower} (${bits} bits)" "${comment:-no comment}"
            fi
            ;;
        ECDSA|ED25519)
            record_pass "[${user}] Key ${key_num}: ssh-${key_type_lower} (${bits} bits)" "${comment:-no comment}"
            ;;
        *)
            record_pass "[${user}] Key ${key_num}: ${key_type_lower} (${bits} bits)" "${comment:-no comment}"
            ;;
    esac

    # Check for missing comment
    if [[ -z "$comment" || "$comment" == "$fingerprint" ]]; then
        record_fail "[${user}] Key ${key_num}: ${key_type_lower} (${bits} bits)" "no comment/identifier"
    fi

    # Note restrictions
    if [[ -n "$restrictions" ]]; then
        record_skip "[${user}] Key ${key_num} has restriction" "${restrictions}"
    fi
}

# ── Audit a single user ───────────────────────────────────────────────
audit_user() {
    local home_dir="$1"
    local user
    user=$(basename "$home_dir")

    local ssh_dir="${home_dir}/.ssh"
    local authkeys_file="${ssh_dir}/authorized_keys"

    # Skip if no .ssh directory at all
    if [[ ! -d "$ssh_dir" ]]; then
        verbose "Skipping ${user} — no .ssh directory"
        return
    fi

    if [[ "$OUTPUT_FORMAT" != "tap" ]]; then
        echo ""
        echo -e "${BOLD}User: ${user}${RESET}"
    fi

    # Permission checks
    check_ssh_dir_perms "$user" "$ssh_dir"
    check_authkeys_perms "$user" "$authkeys_file"
    check_authkeys2 "$user" "$ssh_dir"
    check_private_key_perms "$user" "$ssh_dir"

    # Key audit
    if [[ -f "$authkeys_file" && -r "$authkeys_file" ]]; then
        local key_num=0
        local user_has_keys=false
        local ak_contents
        ak_contents=$(< "$authkeys_file")
        while IFS= read -r line || [[ -n "$line" ]]; do
            # Skip empty lines and comments
            [[ -z "$line" || "$line" =~ ^[[:space:]]*# ]] && continue
            ((key_num++)) || true
            user_has_keys=true
            audit_key_line "$user" "$key_num" "$line" "$authkeys_file"
        done <<< "$ak_contents"
        if [[ "$user_has_keys" == "true" ]]; then
            ((USERS_WITH_KEYS++)) || true
        fi
    else
        verbose "[${user}] No authorized_keys file or not readable"
    fi
}

# ── Check for duplicate keys across users ─────────────────────────────
check_duplicates() {
    if [[ "$SKIP_DUPLICATES" == "true" ]]; then
        return
    fi

    local found_dupes=false

    for fp in "${!KEY_FINGERPRINTS[@]}"; do
        local users_str="${KEY_FINGERPRINTS[$fp]}"
        # Trim leading space, deduplicate
        users_str=$(echo "$users_str" | xargs -n1 | sort -u | xargs)
        local user_count
        user_count=$(echo "$users_str" | wc -w)

        if [[ "$user_count" -gt 1 ]]; then
            if [[ "$found_dupes" == "false" && "$OUTPUT_FORMAT" != "tap" ]]; then
                echo ""
                echo -e "${BOLD}Duplicate Keys${RESET}"
                found_dupes=true
            fi
            record_fail "Duplicate key found" "${fp} present in: ${users_str}"
        fi
    done
}

# ══════════════════════════════════════════════════════════════════════
#                          OUTPUT
# ══════════════════════════════════════════════════════════════════════

print_summary() {
    local end_time
    end_time=$(date +%s)
    local duration=$(( end_time - START_TIME ))

    # Build type breakdown string
    local type_breakdown=""
    for ktype in "${!KEY_TYPE_COUNTS[@]}"; do
        local count="${KEY_TYPE_COUNTS[$ktype]}"
        type_breakdown="${type_breakdown:+${type_breakdown}, }${count}× ${ktype}"
    done

    echo ""
    echo -e "${BOLD}────────────────────────────────────────${RESET}"
    echo -e "${BOLD}Summary${RESET}"
    echo -e "  Users scanned: ${USERS_WITH_KEYS}"
    echo -e "  Total keys:    ${TOTAL_KEYS}"
    if [[ -n "$type_breakdown" ]]; then
        echo -e "  Key types:     ${type_breakdown}"
    fi
    echo -e "  ${GREEN}${PASS} passed${RESET}  ${RED}${FAIL} failed${RESET}  ${YELLOW}${SKIP} skipped${RESET}  (${duration}s)"
    echo -e "${BOLD}────────────────────────────────────────${RESET}"

    if [[ $FAIL -eq 0 ]]; then
        echo -e "${GREEN}${BOLD}All checks passed.${RESET}"
    else
        echo -e "${RED}${BOLD}${FAIL} issue(s) found.${RESET}"
    fi
}

print_tap_header() {
    echo "TAP version 13"
}

print_tap_footer() {
    echo "1..${TOTAL}"
    echo "# pass ${PASS}"
    echo "# fail ${FAIL}"
    echo "# skip ${SKIP}"
}

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

usage() {
    cat <<EOF
Usage: $(basename "$0") [OPTIONS]

Audit SSH authorized_keys across user accounts. Flags weak key types, short
key lengths, duplicates, and permission issues. Read-only — never modifies keys.

Requires: bash 4+, coreutils, ssh-keygen

Options:
  --scan-path PATH    Base path to scan for user home dirs (default: /home)
  --min-rsa-bits N    Minimum RSA key length in bits (default: 2048)
  --format FORMAT     Output: text (default), tap, junit
  --junit-file FILE   JUnit output path (default: ssh-audit-results.xml)
  --verbose           Show debug output
  --no-color          Disable colored output
  --help              Show this help

Environment variables:
  SCAN_PATH           Base scan path (default: /home)
  SCAN_SYSTEM_USERS   Also scan /root and system accounts (default: false)
  MIN_RSA_BITS        Minimum RSA key bits (default: 2048)
  SKIP_DUPLICATES     Skip cross-user duplicate detection (default: false)
  OUTPUT_FORMAT       Output format (default: text)
  JUNIT_FILE          JUnit output path (default: ssh-audit-results.xml)
  VERBOSE             Debug output (default: false)
  COLOR               Color mode: auto, always, never (default: auto)

Examples:
  # Scan all users under /home
  sudo ./$(basename "$0")

  # Require 4096-bit RSA, JUnit output
  sudo ./$(basename "$0") --min-rsa-bits 4096 --format junit

  # TAP output for CI pipeline
  sudo ./$(basename "$0") --format tap

  # Scan specific path
  sudo ./$(basename "$0") --scan-path /srv/users
EOF
}

main() {
    # Parse arguments
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --scan-path)      SCAN_PATH="$2"; shift ;;
            --min-rsa-bits)   MIN_RSA_BITS="$2"; shift ;;
            --format)         OUTPUT_FORMAT="$2"; shift ;;
            --junit-file)     JUNIT_FILE="$2"; shift ;;
            --verbose)        VERBOSE=true ;;
            --no-color)       COLOR=never ;;
            --help|-h)        usage; exit 0 ;;
            *)                err "Unknown option: $1"; usage; exit 1 ;;
        esac
        shift
    done

    setup_colors

    # Validate scan path
    if [[ ! -d "$SCAN_PATH" ]]; then
        err "Scan path does not exist: ${SCAN_PATH}"
        exit 1
    fi

    START_TIME=$(date +%s)

    if [[ "$OUTPUT_FORMAT" == "tap" ]]; then
        print_tap_header
    else
        echo ""
        echo -e "${BOLD}SSH Key Auditor${RESET}"
        echo -e "Scan path: ${SCAN_PATH}"
        echo -e "Time:      $(date -u +%Y-%m-%dT%H:%M:%SZ)"
    fi

    # Build list of directories to scan
    local -a scan_dirs=()

    # Scan user home directories under SCAN_PATH
    if [[ -d "$SCAN_PATH" ]]; then
        for dir in "${SCAN_PATH}"/*/; do
            [[ -d "$dir" ]] && scan_dirs+=("$dir")
        done
    fi

    # Optionally scan system users
    if [[ "$SCAN_SYSTEM_USERS" == "true" ]]; then
        [[ -d "/root" ]] && scan_dirs+=("/root/")
        # Scan common system user home dirs from /etc/passwd
        while IFS=: read -r _ _ uid _ _ home _; do
            if [[ "$uid" -lt 1000 && "$uid" -gt 0 && -d "$home/.ssh" && "$home" != "/root" ]]; then
                scan_dirs+=("${home}/")
            fi
        done < /etc/passwd
    fi

    if [[ ${#scan_dirs[@]} -eq 0 ]]; then
        warn "No directories found to scan in ${SCAN_PATH}"
        if [[ "$OUTPUT_FORMAT" == "tap" ]]; then
            print_tap_footer
        fi
        exit 0
    fi

    verbose "Scanning ${#scan_dirs[@]} directories"

    # Audit each user
    for dir in "${scan_dirs[@]}"; do
        # Remove trailing slash for basename
        dir="${dir%/}"
        audit_user "$dir"
    done

    # Check for duplicates across users
    check_duplicates

    # Output
    if [[ "$OUTPUT_FORMAT" == "tap" ]]; then
        print_tap_footer
    elif [[ "$OUTPUT_FORMAT" == "junit" ]]; then
        print_summary
        write_junit
    else
        print_summary
    fi

    # Exit code
    [[ $FAIL -eq 0 ]] && exit 0 || exit 1
}

main "$@"