#!/usr/bin/env bash

#########################################################################################
#### systemd-hardening-auditor.sh — Audit systemd services for security hardening    ####
#### Checks PrivateTmp, NoNewPrivileges, ProtectSystem, and more                     ####
####                                                                                 ####
#### Author: Phil Connor                                                             ####
#### Contact: contact@mylinux.work                                                   ####
#### License: MIT                                                                    ####
#### Version 1.00                                                                    ####
####                                                                                 ####
#### Usage:                                                                          ####
####   ./systemd-hardening-auditor.sh                                                ####
####   ./systemd-hardening-auditor.sh --service nginx                                ####
####   ./systemd-hardening-auditor.sh --json                                         ####
####                                                                                 ####
#### See --help for all options.                                                     ####
#########################################################################################

set -euo pipefail

# ── Defaults ──────────────────────────────────────────────────────────
OUTPUT_FORMAT="text"
TARGET_SERVICE=""
ONLY_FAILING=false
MIN_SCORE=""
TEXTFILE_PATH=""
INCLUDE_SYSTEMD=false
NO_COLOR=false

# ── Hardening directives ─────────────────────────────────────────────
DIRECTIVES=(
    PrivateTmp PrivateDevices PrivateNetwork
    ProtectSystem ProtectHome NoNewPrivileges
    ProtectKernelTunables ProtectKernelModules ProtectControlGroups
    RestrictSUIDSGID ProtectClock ProtectKernelLogs
    ProtectHostname RestrictNamespaces RestrictRealtime
    MemoryDenyWriteExecute LockPersonality
)
TOTAL_DIRECTIVES=${#DIRECTIVES[@]}

# ── Systemd internal prefixes excluded by default ────────────────────
SYSTEMD_INTERNAL_PREFIXES=("systemd-" "init.scope" "dbus")

# ── Result storage ───────────────────────────────────────────────────
declare -a R_SERVICES=() R_SCORES=() R_PCTS=() R_DISABLED=()

# ── Colors ────────────────────────────────────────────────────────────
RED="" GREEN="" YELLOW="" BOLD="" RESET=""

setup_colors() {
    [[ "$NO_COLOR" == true ]] && return
    if [[ -t 1 ]]; then
        RED='\033[0;31m' GREEN='\033[0;32m' YELLOW='\033[0;33m'
        BOLD='\033[1m' RESET='\033[0m'
    fi
}

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

Audit running systemd services for security hardening directives.

Options:
  --service NAME        Audit a single service (e.g., nginx)
  --json                Output results as JSON array
  --csv                 Output results as CSV
  --only-failing        Show only services scoring below 50%
  --min-score N         Show only services scoring below N%
  --textfile PATH       Write Prometheus metrics to file
  --include-systemd     Include systemd internal services
  --no-color            Disable colored output
  --help                Show this help

Checks ${TOTAL_DIRECTIVES} hardening directives:
  $(printf '%s, ' "${DIRECTIVES[@]}" | sed 's/, $//')

Examples:
  $(basename "$0")
  $(basename "$0") --service nginx
  $(basename "$0") --json
  $(basename "$0") --only-failing
  $(basename "$0") --textfile /var/lib/node_exporter/textfile_collector/systemd_hardening.prom
EOF
    exit 0
}

# ── Argument parsing ─────────────────────────────────────────────────
parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --service)
                [[ -z "${2:-}" ]] && { echo "Error: --service requires a name" >&2; exit 1; }
                TARGET_SERVICE="$2"; shift 2 ;;
            --json)         OUTPUT_FORMAT="json"; NO_COLOR=true; shift ;;
            --csv)          OUTPUT_FORMAT="csv";  NO_COLOR=true; shift ;;
            --only-failing) ONLY_FAILING=true; shift ;;
            --min-score)
                [[ -z "${2:-}" ]] && { echo "Error: --min-score requires a number" >&2; exit 1; }
                MIN_SCORE="$2"; shift 2 ;;
            --textfile)
                [[ -z "${2:-}" ]] && { echo "Error: --textfile requires a path" >&2; exit 1; }
                TEXTFILE_PATH="$2"; shift 2 ;;
            --include-systemd) INCLUDE_SYSTEMD=true; shift ;;
            --no-color)     NO_COLOR=true; shift ;;
            --help|-h)      usage ;;
            *) echo "Error: unknown option: $1" >&2; exit 1 ;;
        esac
    done
}

# ── Helper functions ─────────────────────────────────────────────────

is_systemd_internal() {
    local svc="$1"
    for prefix in "${SYSTEMD_INTERNAL_PREFIXES[@]}"; do
        [[ "$svc" == ${prefix}* ]] && return 0
    done
    return 1
}

get_services() {
    if [[ -n "$TARGET_SERVICE" ]]; then
        local svc="$TARGET_SERVICE"
        [[ "$svc" != *.service ]] && svc="${svc}.service"
        if ! systemctl show "$svc" --property=LoadState 2>/dev/null | grep -q "LoadState=loaded"; then
            echo "Error: service '$svc' not found or not loaded" >&2
            exit 1
        fi
        echo "$svc"
        return
    fi
    systemctl list-units --type=service --state=running --no-legend --no-pager 2>/dev/null \
        | awk '{print $1}' | sort
}

# Check if a directive is enabled for a given service
check_directive() {
    local svc="$1" directive="$2"
    local value
    value=$(systemctl show "$svc" --property="$directive" 2>/dev/null | cut -d= -f2-)
    [[ -z "$value" ]] && return 1

    case "$directive" in
        ProtectSystem)
            case "$value" in strict|full|yes|true) return 0 ;; *) return 1 ;; esac ;;
        ProtectHome)
            case "$value" in yes|read-only|tmpfs|true) return 0 ;; *) return 1 ;; esac ;;
        RestrictNamespaces)
            case "$value" in false|no|"") return 1 ;; *) return 0 ;; esac ;;
        *)
            case "$value" in yes|true) return 0 ;; *) return 1 ;; esac ;;
    esac
}

progress_bar() {
    local score="$1" total="$2" bar="" i
    for ((i = 0; i < score; i++));           do bar+="█"; done
    for ((i = score; i < total; i++));       do bar+="░"; done
    echo "$bar"
}

score_color() {
    local pct="$1"
    if   [[ "$pct" -ge 70 ]]; then echo -ne "$GREEN"
    elif [[ "$pct" -ge 40 ]]; then echo -ne "$YELLOW"
    else                            echo -ne "$RED"
    fi
}

# ── Audit one service, store results directly (no subshell) ──────────
audit_and_collect() {
    local svc="$1" score=0
    local disabled_list=()

    for directive in "${DIRECTIVES[@]}"; do
        if check_directive "$svc" "$directive"; then
            ((score++)) || true
        else
            disabled_list+=("$directive")
        fi
    done

    local pct=$(( (score * 100) / TOTAL_DIRECTIVES ))

    # Apply filters — return silently if filtered out
    if [[ "$ONLY_FAILING" == true ]] && [[ "$pct" -ge 50 ]]; then return; fi
    if [[ -n "$MIN_SCORE" ]] && [[ "$pct" -ge "$MIN_SCORE" ]]; then return; fi

    R_SERVICES+=("$svc")
    R_SCORES+=("$score")
    R_PCTS+=("$pct")
    local disabled_str=""
    disabled_str=$(IFS='|'; echo "${disabled_list[*]+"${disabled_list[*]}"}")
    R_DISABLED+=("$disabled_str")
}

# ── Output: text ─────────────────────────────────────────────────────
output_text() {
    local hostname timestamp
    hostname=$(hostname 2>/dev/null || echo "unknown")
    timestamp=$(date -u +"%Y-%m-%dT%H:%M:%SZ")

    echo -e "${BOLD}Systemd Service Hardening Audit${RESET}"
    echo "Host: ${hostname}"
    echo "Time: ${timestamp}"
    echo ""
    printf "%-35s %-8s %-7s %s\n" "SERVICE" "SCORE" "GRADE" "DIRECTIVES"
    echo "──────────────────────────────────────────────────────────────────────────"

    local count="${#R_SERVICES[@]}"
    for ((i = 0; i < count; i++)); do
        local svc="${R_SERVICES[$i]}" score="${R_SCORES[$i]}" pct="${R_PCTS[$i]}"
        local bar color
        bar=$(progress_bar "$score" "$TOTAL_DIRECTIVES")
        color=$(score_color "$pct")
        printf "${color}%-35s %2d/%-4d %3d%%   %s${RESET}\n" "$svc" "$score" "$TOTAL_DIRECTIVES" "$pct" "$bar"

        # List missing directives for poorly hardened services
        if [[ "$pct" -lt 40 ]] && [[ -n "${R_DISABLED[$i]}" ]]; then
            local IFS='|' items line="" c=0
            read -ra items <<< "${R_DISABLED[$i]}"
            for d in "${items[@]}"; do
                if [[ $c -gt 0 ]] && [[ $((c % 3)) -eq 0 ]]; then
                    echo -e "  ${RED}${line}${RESET}"; line=""
                fi
                line+=$(printf "  ✗ %-28s" "$d")
                ((c++)) || true
            done
            [[ -n "$line" ]] && echo -e "  ${RED}${line}${RESET}"
        fi
    done

    # Summary
    local total="${#R_SERVICES[@]}" sum_pct=0 below_50=0
    for ((i = 0; i < total; i++)); do
        sum_pct=$((sum_pct + R_PCTS[i]))
        [[ "${R_PCTS[i]}" -lt 50 ]] && ((below_50++)) || true
    done
    local avg_pct=0
    [[ "$total" -gt 0 ]] && avg_pct=$((sum_pct / total))

    echo ""
    echo "────────────────────────────────────────"
    echo -e "${BOLD}Summary${RESET}"
    printf "  Total services:    %d\n" "$total"
    printf "  Average score:     %d%%\n" "$avg_pct"
    printf "  Below 50%%:         %d service(s)\n" "$below_50"
    echo "────────────────────────────────────────"
}

# ── Output: JSON ─────────────────────────────────────────────────────
output_json() {
    local total="${#R_SERVICES[@]}"
    echo "["
    for ((i = 0; i < total; i++)); do
        local svc="${R_SERVICES[$i]}" score="${R_SCORES[$i]}" pct="${R_PCTS[$i]}"
        local disabled="${R_DISABLED[$i]}"

        # Build missing array
        local missing_json="[]"
        if [[ -n "$disabled" ]]; then
            missing_json="["
            local IFS='|' first=true
            read -ra items <<< "$disabled"
            for d in "${items[@]}"; do
                [[ "$first" == true ]] && first=false || missing_json+=","
                missing_json+="\"${d}\""
            done
            missing_json+="]"
        fi

        local comma=","
        [[ $((i + 1)) -eq "$total" ]] && comma=""
        printf '  {"service":"%s","score":%d,"total":%d,"percentage":%d,"missing":%s}%s\n' \
            "$svc" "$score" "$TOTAL_DIRECTIVES" "$pct" "$missing_json" "$comma"
    done
    echo "]"
}

# ── Output: CSV ──────────────────────────────────────────────────────
output_csv() {
    echo "service,score,total,percentage,missing"
    local total="${#R_SERVICES[@]}"
    for ((i = 0; i < total; i++)); do
        local missing="${R_DISABLED[$i]//|/;}"
        printf '%s,%d,%d,%d,"%s"\n' \
            "${R_SERVICES[$i]}" "${R_SCORES[$i]}" "$TOTAL_DIRECTIVES" "${R_PCTS[$i]}" "$missing"
    done
}

# ── Output: Prometheus textfile ──────────────────────────────────────
output_prometheus() {
    local path="$1" total="${#R_SERVICES[@]}"
    local tmpfile
    tmpfile=$(mktemp "${path}.XXXXXX" 2>/dev/null) || {
        echo "Error: cannot write to ${path}" >&2; exit 1
    }

    {
        echo "# HELP systemd_hardening_score Hardening score for a systemd service (0-${TOTAL_DIRECTIVES})"
        echo "# TYPE systemd_hardening_score gauge"
        for ((i = 0; i < total; i++)); do
            echo "systemd_hardening_score{service=\"${R_SERVICES[$i]}\"} ${R_SCORES[$i]}"
        done
        echo "# HELP systemd_hardening_percentage Hardening percentage for a systemd service (0-100)"
        echo "# TYPE systemd_hardening_percentage gauge"
        for ((i = 0; i < total; i++)); do
            echo "systemd_hardening_percentage{service=\"${R_SERVICES[$i]}\"} ${R_PCTS[$i]}"
        done
        echo "# HELP systemd_hardening_total Total hardening directives checked"
        echo "# TYPE systemd_hardening_total gauge"
        echo "systemd_hardening_total ${TOTAL_DIRECTIVES}"
    } > "$tmpfile"

    mv "$tmpfile" "$path"
    echo "Wrote Prometheus metrics to ${path}" >&2
}

# ── Main ─────────────────────────────────────────────────────────────
main() {
    parse_args "$@"
    setup_colors

    if ! command -v systemctl &>/dev/null; then
        echo "Error: systemctl not found — this script requires systemd" >&2
        exit 1
    fi

    local services
    services=$(get_services)
    if [[ -z "$services" ]]; then
        echo "No running services found to audit" >&2
        exit 0
    fi

    # Audit each service
    while IFS= read -r svc; do
        [[ -z "$svc" ]] && continue
        if [[ "$INCLUDE_SYSTEMD" == false ]] && is_systemd_internal "$svc"; then
            continue
        fi
        audit_and_collect "$svc"
    done <<< "$services"

    # Produce output
    case "$OUTPUT_FORMAT" in
        json) output_json ;;
        csv)  output_csv  ;;
        text) output_text ;;
    esac

    # Write Prometheus metrics if requested
    [[ -n "$TEXTFILE_PATH" ]] && output_prometheus "$TEXTFILE_PATH"
}

main "$@"