#!/usr/bin/env bash

#########################################################################################
#### zero-trust-audit.sh — Audit Linux systems for zero-trust security posture       ####
#### Identity, network, encryption, least privilege, audit logging, supply chain     ####
#### Requires: bash 4+, root recommended                                             ####
####                                                                                 ####
#### Author: Phil Connor                                                             ####
#### Contact: contact@mylinux.work                                                   ####
#### License: MIT                                                                    ####
#### Version 1.01                                                                    ####
####                                                                                 ####
#### Usage:                                                                          ####
####   sudo ./zero-trust-audit.sh --audit                                            ####
####                                                                                 ####
#### See --help for all options.                                                     ####
#########################################################################################

set -euo pipefail

# ── Defaults ──────────────────────────────────────────────────────────
RUN_MODE=""
SECTION_FILTER="${ZT_SECTION:-all}"
OUTPUT_FORMAT="${ZT_FORMAT:-text}"
OUTPUT_FILE=""
VERBOSE="${VERBOSE:-false}"
COLOR="${COLOR:-auto}"

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

# ── Colors ────────────────────────────────────────────────────────────
RED="" GREEN="" YELLOW="" BLUE="" CYAN="" BOLD="" DIM="" RESET=""

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'
        CYAN='\033[0;36m'
        BOLD='\033[1m'
        DIM='\033[2m'
        RESET='\033[0m'
    else
        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; }

# ── Helpers ───────────────────────────────────────────────────────────
section_header() {
    echo ""
    echo -e "  ${BOLD}${CYAN}── $1 ──${RESET}"
    echo ""
}

field() {
    printf "  ${BOLD}%-22s${RESET} %s\n" "$1" "$2"
}

field_color() {
    printf "  ${BOLD}%-22s${RESET} %b\n" "$1" "$2"
}

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

# ── Check helpers ────────────────────────────────────────────────────
check_pass() {
    local name="$1" detail="${2:-}"
    ((PASS_COUNT++)) || true
    ((TOTAL_POSSIBLE++)) || true
    printf "  ${GREEN}✓ PASS${RESET}  %-38s %b\n" "$name" "${DIM}${detail}${RESET}"
}

check_fail() {
    local name="$1" detail="${2:-}" remediation="${3:-}"
    ((FAIL_COUNT++)) || true
    ((TOTAL_POSSIBLE++)) || true
    printf "  ${RED}✗ FAIL${RESET}  %-38s %b\n" "$name" "${DIM}${detail}${RESET}"
    if [[ -n "$remediation" ]]; then
        REMEDIATIONS+=("${name}|${remediation}")
    fi
}

check_warn() {
    local name="$1" detail="${2:-}" remediation="${3:-}"
    ((WARN_COUNT++)) || true
    ((TOTAL_POSSIBLE++)) || true
    printf "  ${YELLOW}! WARN${RESET}  %-38s %b\n" "$name" "${DIM}${detail}${RESET}"
    if [[ -n "$remediation" ]]; then
        REMEDIATIONS+=("${name}|${remediation}")
    fi
}

# ══════════════════════════════════════════════════════════════════════
#                     AUDIT: IDENTITY
# ══════════════════════════════════════════════════════════════════════
audit_identity() {
    section_header "Identity & Authentication"

    # SSH: PermitRootLogin
    if [[ -f /etc/ssh/sshd_config ]]; then
        local root_login
        root_login=$(grep -i "^[[:space:]]*PermitRootLogin" /etc/ssh/sshd_config 2>/dev/null | awk '{print $2}' || echo "")
        if [[ "${root_login,,}" == "no" ]]; then
            check_pass "SSH root login disabled" "PermitRootLogin no"
        else
            check_fail "SSH root login enabled" "Current: ${root_login:-not set}" \
                "Set 'PermitRootLogin no' in /etc/ssh/sshd_config"
        fi

        # SSH: PasswordAuthentication
        local pass_auth
        pass_auth=$(grep -i "^[[:space:]]*PasswordAuthentication" /etc/ssh/sshd_config 2>/dev/null | awk '{print $2}' || echo "")
        if [[ "${pass_auth,,}" == "no" ]]; then
            check_pass "SSH password auth disabled" "PasswordAuthentication no"
        else
            check_fail "SSH password auth enabled" "Current: ${pass_auth:-not set}" \
                "Set 'PasswordAuthentication no' in /etc/ssh/sshd_config"
        fi

        # SSH: PubkeyAuthentication
        local pubkey_auth
        pubkey_auth=$(grep -i "^[[:space:]]*PubkeyAuthentication" /etc/ssh/sshd_config 2>/dev/null | awk '{print $2}' || echo "")
        if [[ "${pubkey_auth,,}" != "no" ]]; then
            check_pass "SSH pubkey auth enabled" "PubkeyAuthentication yes"
        else
            check_fail "SSH pubkey auth disabled" "PubkeyAuthentication no" \
                "Set 'PubkeyAuthentication yes' in /etc/ssh/sshd_config"
        fi
    else
        check_warn "SSH config not found" "/etc/ssh/sshd_config missing"
    fi

    # Password policy: minlen
    if [[ -f /etc/security/pwquality.conf ]]; then
        local minlen
        minlen=$(grep -i "^[[:space:]]*minlen" /etc/security/pwquality.conf 2>/dev/null | awk -F= '{print $2}' | tr -d ' ' || echo "")
        if [[ -n "$minlen" ]] && [[ "$minlen" -ge 12 ]]; then
            check_pass "Password minimum length" "minlen=${minlen}"
        else
            check_fail "Password min length weak" "minlen=${minlen:-not set}" \
                "Set 'minlen = 14' in /etc/security/pwquality.conf"
        fi
    else
        check_warn "Password quality config missing" "/etc/security/pwquality.conf not found" \
            "Install libpam-pwquality and configure /etc/security/pwquality.conf"
    fi

    # Sudoers: NOPASSWD
    local nopasswd_count=0
    if [[ -f /etc/sudoers ]]; then
        nopasswd_count=$(grep -c "NOPASSWD" /etc/sudoers 2>/dev/null || true)
    fi
    if [[ -d /etc/sudoers.d ]]; then
        local extra
        extra=$(grep -r "NOPASSWD" /etc/sudoers.d/ 2>/dev/null | wc -l || true)
        nopasswd_count=$(( nopasswd_count + extra ))
    fi
    if [[ "$nopasswd_count" -eq 0 ]]; then
        check_pass "No NOPASSWD sudo rules" "0 entries"
    else
        check_warn "NOPASSWD sudo rules found" "${nopasswd_count} entries" \
            "Review sudoers for unnecessary NOPASSWD rules"
    fi
}

# ══════════════════════════════════════════════════════════════════════
#                     AUDIT: NETWORK
# ══════════════════════════════════════════════════════════════════════
audit_network() {
    section_header "Network Security"

    # Firewall active
    local fw_active=false
    if command -v ufw &>/dev/null && ufw status 2>/dev/null | grep -q "Status: active"; then
        fw_active=true
        check_pass "Firewall active" "ufw enabled"
    elif command -v nft &>/dev/null && nft list ruleset 2>/dev/null | grep -q "chain"; then
        fw_active=true
        check_pass "Firewall active" "nftables loaded"
    elif iptables -L -n 2>/dev/null | grep -qv "^Chain .* (policy" | grep -q .; then
        fw_active=true
        check_pass "Firewall active" "iptables rules present"
    fi
    if [[ "$fw_active" == "false" ]]; then
        check_fail "No firewall detected" "No ufw/nftables/iptables rules" \
            "Install and enable ufw or nftables"
    fi

    # Open ports
    local listen_count=0
    if command -v ss &>/dev/null; then
        listen_count=$(ss -tulnH 2>/dev/null | wc -l || true)
    fi
    if [[ "$listen_count" -le 10 ]]; then
        check_pass "Listening ports minimal" "${listen_count} ports"
    else
        check_warn "Many listening ports" "${listen_count} ports" \
            "Review listening services and disable unnecessary ones"
    fi

    # IP forwarding
    local ipv4_fwd
    ipv4_fwd=$(cat /proc/sys/net/ipv4/ip_forward 2>/dev/null || echo "0")
    if [[ "$ipv4_fwd" == "0" ]]; then
        check_pass "IP forwarding disabled" "net.ipv4.ip_forward=0"
    else
        check_fail "IP forwarding enabled" "net.ipv4.ip_forward=1" \
            "Set 'net.ipv4.ip_forward=0' in /etc/sysctl.conf"
    fi

    # ICMP redirect
    local icmp_redirect
    icmp_redirect=$(cat /proc/sys/net/ipv4/conf/all/accept_redirects 2>/dev/null || echo "0")
    if [[ "$icmp_redirect" == "0" ]]; then
        check_pass "ICMP redirects disabled" "accept_redirects=0"
    else
        check_warn "ICMP redirects accepted" "accept_redirects=1" \
            "Set 'net.ipv4.conf.all.accept_redirects=0' in /etc/sysctl.conf"
    fi
}

# ══════════════════════════════════════════════════════════════════════
#                     AUDIT: TRANSIT ENCRYPTION
# ══════════════════════════════════════════════════════════════════════
audit_transit() {
    section_header "Transit Encryption"

    # SSH protocol
    if [[ -f /etc/ssh/sshd_config ]]; then
        local protocol
        protocol=$(grep -i "^[[:space:]]*Protocol" /etc/ssh/sshd_config 2>/dev/null | awk '{print $2}' || echo "")
        if [[ -z "$protocol" ]] || [[ "$protocol" == "2" ]]; then
            check_pass "SSH protocol version 2" "Protocol 2 (default)"
        else
            check_fail "SSH protocol includes v1" "Protocol ${protocol}" \
                "Remove Protocol 1 from /etc/ssh/sshd_config"
        fi

        # SSH ciphers restricted
        local ciphers
        ciphers=$(grep -i "^[[:space:]]*Ciphers" /etc/ssh/sshd_config 2>/dev/null | awk '{print $2}' || echo "")
        if [[ -n "$ciphers" ]]; then
            check_pass "SSH ciphers restricted" "Custom cipher list configured"
        else
            check_warn "SSH ciphers use defaults" "No explicit cipher restriction" \
                "Set 'Ciphers' in /etc/ssh/sshd_config to restrict to modern ciphers"
        fi

        # SSH MACs restricted
        local macs
        macs=$(grep -i "^[[:space:]]*MACs" /etc/ssh/sshd_config 2>/dev/null | awk '{print $2}' || echo "")
        if [[ -n "$macs" ]]; then
            check_pass "SSH MACs restricted" "Custom MAC list configured"
        else
            check_warn "SSH MACs use defaults" "No explicit MAC restriction" \
                "Set 'MACs' in /etc/ssh/sshd_config to restrict to etm variants"
        fi
    else
        check_warn "SSH config not found" "Cannot verify SSH encryption settings"
    fi
}

# ══════════════════════════════════════════════════════════════════════
#                     AUDIT: ENCRYPTION AT REST
# ══════════════════════════════════════════════════════════════════════
audit_encryption() {
    section_header "Encryption at Rest"

    # LUKS volumes
    local luks_count=0
    if command -v lsblk &>/dev/null; then
        luks_count=$(lsblk --fs 2>/dev/null | grep -c "crypto_LUKS" || true)
    fi
    if [[ "$luks_count" -gt 0 ]]; then
        check_pass "LUKS encrypted volumes" "${luks_count} volume(s)"
    else
        check_warn "No LUKS encryption detected" "0 encrypted volumes" \
            "Consider encrypting data volumes with LUKS"
    fi

    # Swap encryption
    local swap_active
    swap_active=$(swapon --show --noheadings 2>/dev/null | wc -l || true)
    if [[ "$swap_active" -eq 0 ]]; then
        check_pass "No unencrypted swap" "Swap disabled or absent"
    else
        local swap_encrypted=false
        if swapon --show --noheadings 2>/dev/null | awk '{print $1}' | while read -r dev; do
            if [[ "$dev" == /dev/dm-* ]] || [[ "$dev" == /dev/mapper/* ]]; then
                return 0
            fi
        done; then
            swap_encrypted=true
        fi
        if [[ "$swap_encrypted" == "true" ]]; then
            check_pass "Swap appears encrypted" "Swap on dm/mapper device"
        else
            check_warn "Swap may be unencrypted" "${swap_active} swap device(s)" \
                "Configure encrypted swap or disable swap"
        fi
    fi

    # /tmp on tmpfs
    if mount | grep -q "tmpfs on /tmp"; then
        check_pass "/tmp on tmpfs" "In-memory tmpfs"
    else
        check_warn "/tmp not on tmpfs" "Persistent /tmp filesystem" \
            "Mount /tmp as tmpfs in /etc/fstab"
    fi
}

# ══════════════════════════════════════════════════════════════════════
#                     AUDIT: LEAST PRIVILEGE
# ══════════════════════════════════════════════════════════════════════
audit_least_privilege() {
    section_header "Least Privilege"

    # SUID binaries
    local suid_count=0
    suid_count=$(find / -perm -4000 -type f 2>/dev/null | wc -l || true)
    if [[ "$suid_count" -le 20 ]]; then
        check_pass "SUID binaries minimal" "${suid_count} found"
    elif [[ "$suid_count" -le 40 ]]; then
        check_warn "Elevated SUID count" "${suid_count} found" \
            "Review SUID binaries: find / -perm -4000 -type f"
    else
        check_fail "High SUID binary count" "${suid_count} found" \
            "Audit and remove unnecessary SUID bits"
    fi

    # SELinux / AppArmor
    local mac_active=false
    if command -v getenforce &>/dev/null; then
        local se_status
        se_status=$(getenforce 2>/dev/null || echo "Disabled")
        if [[ "$se_status" == "Enforcing" ]]; then
            mac_active=true
            check_pass "SELinux enforcing" "Mode: ${se_status}"
        elif [[ "$se_status" == "Permissive" ]]; then
            mac_active=true
            check_warn "SELinux permissive" "Not enforcing" \
                "Set SELINUX=enforcing in /etc/selinux/config"
        fi
    fi
    if [[ "$mac_active" == "false" ]] && command -v aa-status &>/dev/null; then
        local aa_profiles
        aa_profiles=$(aa-status --profiled 2>/dev/null || true)
        if [[ "$aa_profiles" -gt 0 ]]; then
            mac_active=true
            check_pass "AppArmor active" "${aa_profiles} profiles loaded"
        fi
    fi
    if [[ "$mac_active" == "false" ]]; then
        check_fail "No MAC framework active" "Neither SELinux nor AppArmor" \
            "Enable SELinux or AppArmor for mandatory access control"
    fi

    # Default umask
    local umask_val
    umask_val=$(umask 2>/dev/null || echo "0022")
    if [[ "$umask_val" == "0027" ]] || [[ "$umask_val" == "0077" ]]; then
        check_pass "Restrictive umask" "umask ${umask_val}"
    elif [[ "$umask_val" == "0022" ]]; then
        check_warn "Standard umask" "umask ${umask_val}" \
            "Set umask to 0027 in /etc/profile or /etc/login.defs"
    else
        check_warn "Non-standard umask" "umask ${umask_val}" \
            "Review umask setting for appropriateness"
    fi

    # World-writable directories (outside /tmp, /var/tmp)
    local ww_count=0
    ww_count=$(find / -maxdepth 3 -type d -perm -0002 \
        ! -path "/tmp*" ! -path "/var/tmp*" ! -path "/dev*" \
        ! -path "/proc*" ! -path "/sys*" ! -path "/run*" \
        2>/dev/null | wc -l || true)
    if [[ "$ww_count" -eq 0 ]]; then
        check_pass "No world-writable dirs" "Outside /tmp, /var/tmp"
    else
        check_warn "World-writable dirs found" "${ww_count} directories" \
            "Review: find / -maxdepth 3 -type d -perm -0002"
    fi
}

# ══════════════════════════════════════════════════════════════════════
#                     AUDIT: AUDIT LOGGING
# ══════════════════════════════════════════════════════════════════════
audit_logging() {
    section_header "Audit & Logging"

    # auditd
    if command -v auditctl &>/dev/null; then
        if systemctl is-active auditd &>/dev/null; then
            check_pass "auditd running" "Active and enabled"
            # Audit rules
            local rule_count
            rule_count=$(auditctl -l 2>/dev/null | grep -cv "^No rules" || true)
            if [[ "$rule_count" -gt 0 ]]; then
                check_pass "Audit rules configured" "${rule_count} rules"
            else
                check_warn "No audit rules defined" "auditd has no rules" \
                    "Add audit rules in /etc/audit/rules.d/"
            fi
        else
            check_fail "auditd not running" "Service inactive" \
                "Enable with: systemctl enable --now auditd"
        fi
    else
        check_fail "auditd not installed" "No audit framework" \
            "Install auditd: apt install auditd / yum install audit"
    fi

    # Journal persistence
    if [[ -d /var/log/journal ]]; then
        check_pass "Journal persistence" "Logs stored in /var/log/journal"
    else
        check_warn "Journal volatile only" "Logs lost on reboot" \
            "Create /var/log/journal and restart systemd-journald"
    fi

    # Syslog active
    if systemctl is-active rsyslog &>/dev/null || systemctl is-active syslog-ng &>/dev/null; then
        check_pass "Syslog daemon active" "rsyslog or syslog-ng running"
    else
        check_warn "No syslog daemon detected" "Relying on journal only" \
            "Install rsyslog for remote logging capability"
    fi

    # Log rotation
    if [[ -f /etc/logrotate.conf ]] || [[ -d /etc/logrotate.d ]]; then
        check_pass "Log rotation configured" "logrotate present"
    else
        check_warn "No log rotation" "logrotate not found" \
            "Install logrotate to prevent disk exhaustion"
    fi
}

# ══════════════════════════════════════════════════════════════════════
#                     AUDIT: SUPPLY CHAIN
# ══════════════════════════════════════════════════════════════════════
audit_supply_chain() {
    section_header "Supply Chain"

    # Package signing
    if command -v apt-get &>/dev/null; then
        local unsigned_repos=0
        if [[ -d /etc/apt/sources.list.d ]]; then
            unsigned_repos=$(grep -r "^\s*deb\s" /etc/apt/sources.list.d/ /etc/apt/sources.list 2>/dev/null \
                | grep -c "\[.*trusted=yes.*\]" || true)
        fi
        if [[ "$unsigned_repos" -eq 0 ]]; then
            check_pass "APT repos signed" "No trusted=yes overrides"
        else
            check_fail "Unsigned APT repos" "${unsigned_repos} repos skip GPG" \
                "Remove trusted=yes from APT source entries"
        fi
    elif command -v dnf &>/dev/null || command -v yum &>/dev/null; then
        local gpg_disabled=0
        gpg_disabled=$(grep -r "gpgcheck=0" /etc/yum.repos.d/ 2>/dev/null | wc -l || true)
        if [[ "$gpg_disabled" -eq 0 ]]; then
            check_pass "RPM repos GPG enabled" "All repos have gpgcheck=1"
        else
            check_fail "GPG check disabled" "${gpg_disabled} repos without gpgcheck" \
                "Set gpgcheck=1 in all repo files under /etc/yum.repos.d/"
        fi
    else
        check_warn "Unknown package manager" "Cannot verify signing"
    fi

    # Automatic updates
    if command -v unattended-upgrade &>/dev/null || \
       dpkg -l unattended-upgrades &>/dev/null 2>&1; then
        check_pass "Automatic security updates" "unattended-upgrades installed"
    elif command -v dnf &>/dev/null && rpm -q dnf-automatic &>/dev/null 2>&1; then
        check_pass "Automatic security updates" "dnf-automatic installed"
    else
        check_warn "No automatic updates" "Manual patching required" \
            "Install unattended-upgrades (Debian) or dnf-automatic (RHEL)"
    fi

    # Package integrity
    if command -v debsums &>/dev/null; then
        check_pass "Package integrity tool" "debsums available"
    elif command -v rpm &>/dev/null; then
        check_pass "Package integrity tool" "rpm -V available"
    else
        check_warn "No integrity checker" "Cannot verify package files" \
            "Install debsums (Debian) for package integrity checks"
    fi
}

# ══════════════════════════════════════════════════════════════════════
#                     RUN SECTION DISPATCHER
# ══════════════════════════════════════════════════════════════════════
run_section() {
    local name="$1"
    case "$name" in
        identity)       audit_identity ;;
        network)        audit_network ;;
        transit)        audit_transit ;;
        encryption)     audit_encryption ;;
        least-privilege) audit_least_privilege ;;
        logging)        audit_logging ;;
        supply-chain)   audit_supply_chain ;;
        *)              die "Unknown section: ${name}" ;;
    esac
}

# ══════════════════════════════════════════════════════════════════════
#                     SUMMARY
# ══════════════════════════════════════════════════════════════════════
print_summary() {
    local score=0 rating=""
    if [[ "$TOTAL_POSSIBLE" -gt 0 ]]; then
        score=$(( PASS_COUNT * 100 / TOTAL_POSSIBLE ))
    fi

    if [[ "$score" -ge 90 ]]; then
        rating="${GREEN}STRONG${RESET}"
    elif [[ "$score" -ge 70 ]]; then
        rating="${BLUE}GOOD${RESET}"
    elif [[ "$score" -ge 50 ]]; then
        rating="${YELLOW}MODERATE${RESET}"
    else
        rating="${RED}WEAK${RESET}"
    fi

    section_header "Zero Trust Score"
    field "Total checks:" "$TOTAL_POSSIBLE"
    field_color "Passed:" "${GREEN}${PASS_COUNT}${RESET}"
    field_color "Failed:" "${RED}${FAIL_COUNT}${RESET}"
    field_color "Warnings:" "${YELLOW}${WARN_COUNT}${RESET}"
    field "Score:" "${score}% (${PASS_COUNT}/${TOTAL_POSSIBLE})"
    field_color "Rating:" "$rating"
    echo ""
    field "Duration:" "$(elapsed)"
}

# ══════════════════════════════════════════════════════════════════════
#                     MODES
# ══════════════════════════════════════════════════════════════════════
do_audit() {
    log "Running zero-trust security audit..."

    local sections=("identity" "network" "transit" "encryption" "least-privilege" "logging" "supply-chain")

    if [[ "$SECTION_FILTER" != "all" ]]; then
        IFS=',' read -ra sections <<< "$SECTION_FILTER"
    fi

    for s in "${sections[@]}"; do
        run_section "$s"
    done

    print_summary
}

do_remediate() {
    log "Running audit with remediation suggestions..."

    local sections=("identity" "network" "transit" "encryption" "least-privilege" "logging" "supply-chain")

    if [[ "$SECTION_FILTER" != "all" ]]; then
        IFS=',' read -ra sections <<< "$SECTION_FILTER"
    fi

    for s in "${sections[@]}"; do
        run_section "$s"
    done

    print_summary

    if [[ ${#REMEDIATIONS[@]} -gt 0 ]]; then
        section_header "Remediation Steps"
        printf "  ${BOLD}%-40s %s${RESET}\n" "FINDING" "REMEDIATION"
        printf "  %s\n" "$(printf '%.0s─' {1..72})"
        for entry in "${REMEDIATIONS[@]}"; do
            local finding remediation
            finding="${entry%%|*}"
            remediation="${entry#*|}"
            printf "  %-40s %s\n" "$finding" "$remediation"
        done
    else
        log "No remediation steps needed — all checks passed"
    fi
}

do_export() {
    [[ -z "$OUTPUT_FILE" ]] && die "No output file specified (--output-file)"

    log "Running audit and exporting results..."

    local sections=("identity" "network" "transit" "encryption" "least-privilege" "logging" "supply-chain")
    for s in "${sections[@]}"; do
        run_section "$s"
    done

    local score=0
    if [[ "$TOTAL_POSSIBLE" -gt 0 ]]; then
        score=$(( PASS_COUNT * 100 / TOTAL_POSSIBLE ))
    fi

    case "$OUTPUT_FORMAT" in
        json)
            if command -v jq &>/dev/null; then
                jq -n \
                    --arg host "$(hostname -f 2>/dev/null || hostname)" \
                    --arg date "$(date -u '+%Y-%m-%dT%H:%M:%SZ')" \
                    --argjson pass "$PASS_COUNT" \
                    --argjson fail "$FAIL_COUNT" \
                    --argjson warn "$WARN_COUNT" \
                    --argjson total "$TOTAL_POSSIBLE" \
                    --argjson score "$score" \
                    '{host:$host, date:$date, pass:$pass, fail:$fail, warn:$warn, total:$total, score:$score}' \
                    > "$OUTPUT_FILE"
            else
                cat > "$OUTPUT_FILE" <<EOF
{"host":"$(hostname -f 2>/dev/null || hostname)","date":"$(date -u '+%Y-%m-%dT%H:%M:%SZ')","pass":${PASS_COUNT},"fail":${FAIL_COUNT},"warn":${WARN_COUNT},"total":${TOTAL_POSSIBLE},"score":${score}}
EOF
            fi
            ;;
        *)
            {
                echo "Zero Trust Audit Report"
                echo "======================"
                echo "Host:    $(hostname -f 2>/dev/null || hostname)"
                echo "Date:    $(date -u '+%Y-%m-%dT%H:%M:%SZ')"
                echo ""
                echo "Pass:    ${PASS_COUNT}"
                echo "Fail:    ${FAIL_COUNT}"
                echo "Warn:    ${WARN_COUNT}"
                echo "Total:   ${TOTAL_POSSIBLE}"
                echo "Score:   ${score}%"
            } > "$OUTPUT_FILE"
            ;;
    esac

    print_summary
    log "Results exported to ${OUTPUT_FILE}"
}

# ══════════════════════════════════════════════════════════════════════
#                     HELP
# ══════════════════════════════════════════════════════════════════════
show_help() {
    cat <<EOF
${BOLD}${SCRIPT_NAME}${RESET} — Zero Trust Security Audit

Audit Linux systems for zero-trust security posture across identity,
network, encryption, least privilege, logging, and supply chain.

${BOLD}MODES${RESET}
  --audit              Run full zero-trust audit (default: all sections)
  --section NAMES      Audit specific sections (comma-separated)
  --remediate          Audit with remediation suggestions
  --export             Export results to file (requires --output-file)

${BOLD}SECTIONS${RESET}
  identity             SSH, PAM, passwords, sudoers
  network              Firewall, open ports, IP forwarding
  transit              TLS, SSH protocol, cipher restrictions
  encryption           LUKS, swap encryption, tmpfs
  least-privilege      SUID, capabilities, SELinux/AppArmor, umask
  logging              auditd, journal, syslog, log rotation
  supply-chain         Package signing, auto-updates, integrity

${BOLD}OPTIONS${RESET}
  --format FORMAT      Output format for export: text, json (default: text)
  --output-file FILE   Write export results to file
  --verbose            Debug output
  --no-color           Disable colored output
  --help               Show this help message

${BOLD}ENVIRONMENT VARIABLES${RESET}
  ZT_SECTION           Default section filter (default: all)
  ZT_FORMAT            Default output format (default: text)
  VERBOSE              Enable verbose output (true/false)
  COLOR                Color mode: auto, always, never

${BOLD}EXAMPLES${RESET}
  # Full audit
  sudo ${SCRIPT_NAME} --audit

  # Audit specific sections
  sudo ${SCRIPT_NAME} --audit --section identity,network

  # Audit with fix suggestions
  sudo ${SCRIPT_NAME} --remediate

  # Export results as JSON
  sudo ${SCRIPT_NAME} --export --format json --output-file report.json

  # CI pipeline check
  sudo ${SCRIPT_NAME} --audit --no-color

${BOLD}EXIT CODES${RESET}
  0   Audit completed (any score)
  1   Runtime error
  2   Audit score below 50% (with --strict, not yet implemented)
EOF
}

# ══════════════════════════════════════════════════════════════════════
#                     PARSE ARGS
# ══════════════════════════════════════════════════════════════════════
parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --audit)        RUN_MODE="audit"; shift ;;
            --section)      RUN_MODE="audit"; SECTION_FILTER="${2:?--section requires names}"; shift 2 ;;
            --remediate)    RUN_MODE="remediate"; shift ;;
            --export)       RUN_MODE="export"; shift ;;
            --format)       OUTPUT_FORMAT="${2:?--format requires a value}"; shift 2 ;;
            --output-file)  OUTPUT_FILE="${2:?--output-file requires a path}"; shift 2 ;;
            --verbose)      VERBOSE="true"; shift ;;
            --no-color)     COLOR="never"; shift ;;
            --help|-h)      setup_colors; show_help; exit 0 ;;
            *)              die "Unknown option: $1 (see --help)" ;;
        esac
    done
}

# ══════════════════════════════════════════════════════════════════════
#                     MAIN
# ══════════════════════════════════════════════════════════════════════
main() {
    parse_args "$@"
    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}Zero Trust Audit${RESET}"
    echo "Host:   $(hostname -f 2>/dev/null || hostname)"
    echo "Mode:   ${RUN_MODE}"
    echo "Time:   $(date -u '+%Y-%m-%dT%H:%M:%SZ')"

    if [[ $EUID -ne 0 ]]; then
        warn "Running without root — some checks may be incomplete"
    fi

    case "$RUN_MODE" in
        audit)      do_audit ;;
        remediate)  do_remediate ;;
        export)     do_export ;;
        *)          die "Unknown mode: ${RUN_MODE}" ;;
    esac
}

main "$@"
