#!/bin/bash
################################################################################
# Script Name: openscap-audit.sh
# Description: OpenSCAP / CIS Benchmark security audit script with HTML and
#              terminal reporting on Ubuntu/Debian and RHEL/Rocky/Alma/Fedora
#
# Author:  Phil Connor
# Contact: contact@mylinux.work
# Website: https://mylinux.work
# License: MIT
# Version: 1.0
#
# Usage:
#   sudo ./openscap-audit.sh
#   sudo ./openscap-audit.sh --profile cis_server_l1
#   sudo ./openscap-audit.sh --report /tmp/audit-report.html
#   sudo ./openscap-audit.sh --dry-run
#
################################################################################

set -euo pipefail

# ============================================================================
# DEFAULTS
# ============================================================================

PROFILE=""
REPORT_DIR="/var/lib/openscap/reports"
REPORT_FILE=""
RESULTS_FILE=""
TAILORING_FILE=""
REMEDIATE=false
DRY_RUN=false
LIST_PROFILES=false

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
NC='\033[0m'

# ============================================================================
# HELPER FUNCTIONS
# ============================================================================

log_info()  { echo -e "${GREEN}[INFO]${NC} $*"; }
log_warn()  { echo -e "${YELLOW}[WARN]${NC} $*"; }
log_error() { echo -e "${RED}[ERROR]${NC} $*" >&2; }
log_step()  { echo -e "${CYAN}[STEP]${NC} $*"; }

show_usage() {
    cat <<EOF
Usage: $0 [OPTIONS]

Run an OpenSCAP security audit against CIS or DISA STIG benchmarks.

OPTIONS:
    --profile PROFILE       SCAP profile to evaluate (auto-detected if omitted)
    --list-profiles         List available profiles for this system
    --report PATH           HTML report output path (default: auto-generated in $REPORT_DIR)
    --results PATH          XCCDF results XML output path
    --tailoring PATH        XCCDF tailoring file for custom profile adjustments
    --remediate             Apply automatic remediation (USE WITH CAUTION)
    --dry-run               Show what would be done without executing
    -h, --help              Show this help message

PROFILES (common):
    cis_server_l1           CIS Level 1 - Server
    cis_server_l2           CIS Level 2 - Server
    cis_workstation_l1      CIS Level 1 - Workstation
    cis_workstation_l2      CIS Level 2 - Workstation
    stig                    DISA STIG
    standard                Standard System Security Profile
    pci-dss                 PCI-DSS v3.2.1

EXAMPLES:
    $0                                        # Auto-detect profile and run
    $0 --profile cis_server_l1                # CIS Level 1 Server
    $0 --profile stig --report /tmp/stig.html # STIG with custom report path
    $0 --list-profiles                        # Show available profiles
    $0 --remediate --profile cis_server_l1    # Auto-fix (careful!)
    $0 --dry-run                              # Preview

EOF
    exit 0
}

parse_args() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            -h|--help)          show_usage ;;
            --profile)          PROFILE="$2"; shift 2 ;;
            --list-profiles)    LIST_PROFILES=true; shift ;;
            --report)           REPORT_FILE="$2"; shift 2 ;;
            --results)          RESULTS_FILE="$2"; shift 2 ;;
            --tailoring)        TAILORING_FILE="$2"; shift 2 ;;
            --remediate)        REMEDIATE=true; shift ;;
            --dry-run)          DRY_RUN=true; shift ;;
            *) log_error "Unknown option: $1"; exit 1 ;;
        esac
    done
}

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

detect_os() {
    if [[ -f /etc/os-release ]]; then
        . /etc/os-release
        OS_ID="$ID"
        OS_VERSION="${VERSION_ID%%.*}"
        OS_FAMILY=""
        case "$OS_ID" in
            ubuntu|debian)          OS_FAMILY="debian" ;;
            rhel|centos|rocky|almalinux|fedora) OS_FAMILY="rhel" ;;
            *) log_error "Unsupported OS: $OS_ID"; exit 1 ;;
        esac
        log_info "Detected OS: $PRETTY_NAME"
    else
        log_error "Cannot detect OS"
        exit 1
    fi
}

# ============================================================================
# INSTALLATION
# ============================================================================

install_openscap() {
    if command -v oscap &>/dev/null; then
        log_info "OpenSCAP already installed: $(oscap --version | head -1)"
        return 0
    fi

    log_step "Installing OpenSCAP..."

    case "$OS_FAMILY" in
        debian)
            apt-get update -qq
            apt-get install -y -qq libopenscap8 openscap-utils ssg-debian ssg-base 2>/dev/null || \
                apt-get install -y -qq libopenscap8 openscap-utils scap-security-guide 2>/dev/null || \
                apt-get install -y -qq openscap-scanner openscap-utils scap-security-guide
            ;;
        rhel)
            dnf install -y -q openscap-scanner openscap-utils scap-security-guide
            ;;
    esac

    log_info "OpenSCAP installed: $(oscap --version | head -1)"
}

# ============================================================================
# CONTENT DETECTION
# ============================================================================

find_content() {
    log_step "Finding SCAP content for $PRETTY_NAME..."

    local content_paths=(
        "/usr/share/xml/scap/ssg/content"
        "/usr/share/scap-security-guide"
        "/usr/share/openscap"
    )

    SCAP_CONTENT=""

    # Build expected filename patterns based on OS
    local os_patterns=()
    case "$OS_ID" in
        ubuntu)     os_patterns=("ssg-ubuntu${OS_VERSION}04-ds.xml" "ssg-ubuntu${OS_VERSION}-ds.xml" "ssg-ubuntu-ds.xml") ;;
        debian)     os_patterns=("ssg-debian${OS_VERSION}-ds.xml" "ssg-debian-ds.xml") ;;
        rhel)       os_patterns=("ssg-rhel${OS_VERSION}-ds.xml" "ssg-rhel-ds.xml") ;;
        centos)     os_patterns=("ssg-centos${OS_VERSION}-ds.xml" "ssg-rhel${OS_VERSION}-ds.xml") ;;
        rocky)      os_patterns=("ssg-rl${OS_VERSION}-ds.xml" "ssg-rhel${OS_VERSION}-ds.xml") ;;
        almalinux)  os_patterns=("ssg-almalinux${OS_VERSION}-ds.xml" "ssg-rhel${OS_VERSION}-ds.xml") ;;
        fedora)     os_patterns=("ssg-fedora-ds.xml") ;;
    esac

    for dir in "${content_paths[@]}"; do
        if [[ -d "$dir" ]]; then
            for pattern in "${os_patterns[@]}"; do
                if [[ -f "$dir/$pattern" ]]; then
                    SCAP_CONTENT="$dir/$pattern"
                    log_info "Found SCAP content: $SCAP_CONTENT"
                    return 0
                fi
            done
        fi
    done

    # Fallback: find any matching content
    SCAP_CONTENT=$(find /usr/share -name "ssg-*-ds.xml" -type f 2>/dev/null | head -1)

    if [[ -z "$SCAP_CONTENT" ]]; then
        log_error "No SCAP content found for $PRETTY_NAME"
        log_error "Install scap-security-guide package"
        exit 1
    fi

    log_info "Using SCAP content: $SCAP_CONTENT"
}

list_available_profiles() {
    log_step "Available profiles for $PRETTY_NAME:"
    echo ""
    oscap info --profiles "$SCAP_CONTENT" 2>/dev/null | while IFS=: read -r id title; do
        printf "  %-50s %s\n" "$id" "$title"
    done
    echo ""
}

select_profile() {
    if [[ -n "$PROFILE" ]]; then
        # Match partial profile names
        local full_profile
        full_profile=$(oscap info --profiles "$SCAP_CONTENT" 2>/dev/null | \
            grep -i "$PROFILE" | head -1 | cut -d: -f1)

        if [[ -n "$full_profile" ]]; then
            PROFILE="$full_profile"
            log_info "Selected profile: $PROFILE"
        else
            log_error "Profile '$PROFILE' not found"
            log_info "Available profiles:"
            list_available_profiles
            exit 1
        fi
    else
        # Auto-select: prefer CIS Level 1 Server, then STIG, then standard
        for pattern in "cis.*server.*l1\|cis_server_l1" "stig\b" "standard"; do
            PROFILE=$(oscap info --profiles "$SCAP_CONTENT" 2>/dev/null | \
                grep -i "$pattern" | head -1 | cut -d: -f1) || true
            if [[ -n "$PROFILE" ]]; then
                break
            fi
        done

        if [[ -z "$PROFILE" ]]; then
            # Use first available profile
            PROFILE=$(oscap info --profiles "$SCAP_CONTENT" 2>/dev/null | head -1 | cut -d: -f1)
        fi

        log_info "Auto-selected profile: $PROFILE"
    fi
}

# ============================================================================
# AUDIT EXECUTION
# ============================================================================

run_audit() {
    local timestamp
    timestamp=$(date +%Y%m%d_%H%M%S)

    mkdir -p "$REPORT_DIR"

    if [[ -z "$REPORT_FILE" ]]; then
        REPORT_FILE="$REPORT_DIR/openscap-report-${timestamp}.html"
    fi

    if [[ -z "$RESULTS_FILE" ]]; then
        RESULTS_FILE="$REPORT_DIR/openscap-results-${timestamp}.xml"
    fi

    log_step "Running OpenSCAP audit..."
    log_info "Profile:  $PROFILE"
    log_info "Content:  $SCAP_CONTENT"
    log_info "Report:   $REPORT_FILE"
    log_info "Results:  $RESULTS_FILE"
    echo ""

    local oscap_cmd=(
        oscap xccdf eval
        --profile "$PROFILE"
        --report "$REPORT_FILE"
        --results "$RESULTS_FILE"
    )

    if [[ -n "$TAILORING_FILE" ]] && [[ -f "$TAILORING_FILE" ]]; then
        oscap_cmd+=(--tailoring-file "$TAILORING_FILE")
    fi

    if [[ "$REMEDIATE" == true ]]; then
        oscap_cmd+=(--remediate)
        log_warn "REMEDIATION ENABLED — changes will be applied!"
    fi

    oscap_cmd+=("$SCAP_CONTENT")

    # oscap returns exit code 2 for "some checks failed" — that's normal
    local exit_code=0
    "${oscap_cmd[@]}" 2>&1 || exit_code=$?

    if [[ $exit_code -eq 0 ]]; then
        log_info "All checks passed"
    elif [[ $exit_code -eq 2 ]]; then
        log_warn "Some checks failed (this is normal for a first audit)"
    else
        log_error "OpenSCAP encountered an error (exit code: $exit_code)"
    fi
}

# ============================================================================
# RESULTS SUMMARY
# ============================================================================

show_summary() {
    log_step "Audit Summary"
    echo ""

    if [[ -f "$RESULTS_FILE" ]]; then
        local pass fail error notapplicable notchecked total score
        pass=$(grep -c 'result="pass"' "$RESULTS_FILE" 2>/dev/null) || pass=0
        fail=$(grep -c 'result="fail"' "$RESULTS_FILE" 2>/dev/null) || fail=0
        error=$(grep -c 'result="error"' "$RESULTS_FILE" 2>/dev/null) || error=0
        notapplicable=$(grep -c 'result="notapplicable"' "$RESULTS_FILE" 2>/dev/null) || notapplicable=0
        notchecked=$(grep -c 'result="notchecked"' "$RESULTS_FILE" 2>/dev/null) || notchecked=0
        total=$((pass + fail))

        if [[ $total -gt 0 ]]; then
            score=$(echo "scale=1; $pass * 100 / $total" | bc)
        else
            score=0
        fi

        echo -e "  ${GREEN}Pass:${NC}            $pass"
        echo -e "  ${RED}Fail:${NC}            $fail"
        echo -e "  ${YELLOW}Error:${NC}           $error"
        echo -e "  Not Applicable:  $notapplicable"
        echo -e "  Not Checked:     $notchecked"
        echo ""
        echo -e "  ${CYAN}Compliance Score:${NC} ${score}%"
        echo ""

        # Show top failures
        if [[ $fail -gt 0 ]]; then
            log_info "Top failed rules (first 10):"
            grep -B2 'result="fail"' "$RESULTS_FILE" 2>/dev/null | \
                grep 'idref=' | sed 's/.*idref="\([^"]*\)".*/  - \1/' | head -10
            echo ""
        fi
    fi

    log_info "HTML report: $REPORT_FILE"
    log_info "XML results: $RESULTS_FILE"

    if [[ $fail -gt 0 ]]; then
        echo ""
        log_info "To auto-remediate failed checks (USE WITH CAUTION):"
        echo "  $0 --profile '$PROFILE' --remediate"
    fi
}

# ============================================================================
# DRY RUN
# ============================================================================

dry_run() {
    echo ""
    log_info "===== DRY RUN — No changes will be made ====="
    echo ""
    log_info "OS:         $PRETTY_NAME"
    log_info "Content:    $SCAP_CONTENT"
    log_info "Profile:    $PROFILE"
    log_info "Report dir: $REPORT_DIR"
    log_info "Remediate:  $REMEDIATE"
    echo ""
    log_info "Actions that would be performed:"
    echo "  1. Install OpenSCAP (if not present)"
    echo "  2. Find SCAP content for $OS_ID $OS_VERSION"
    echo "  3. Evaluate profile: $PROFILE"
    echo "  4. Generate HTML report in $REPORT_DIR/"
    echo "  5. Generate XCCDF results XML"
    if [[ "$REMEDIATE" == true ]]; then
        echo "  6. APPLY REMEDIATION for failed checks"
    fi
    echo ""
}

# ============================================================================
# MAIN
# ============================================================================

main() {
    parse_args "$@"
    check_root
    detect_os
    install_openscap
    find_content

    if [[ "$LIST_PROFILES" == true ]]; then
        list_available_profiles
        exit 0
    fi

    select_profile

    if [[ "$DRY_RUN" == true ]]; then
        dry_run
        exit 0
    fi

    echo ""
    log_info "===== OpenSCAP Security Audit ====="
    echo ""

    run_audit
    show_summary

    echo ""
    log_info "===== Audit Complete ====="
    echo ""
}

main "$@"