#!/bin/bash
#
# STIG Compliance Checker
# 
# Per DoD/MilSpec requirements, checks STIG compliance across system components
# Supports: Kubernetes, PostgreSQL, Linux, Web Servers, Application Security
#

set -euo pipefail

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Logging
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[PASS]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[FAIL]${NC} $1"
}

# Check results
PASSED=0
FAILED=0
WARNINGS=0

# Check Kubernetes STIG compliance
check_kubernetes_stig() {
    log_info "Checking Kubernetes STIG compliance..."
    
    if ! command -v kubectl &> /dev/null; then
        log_warn "kubectl not found - skipping Kubernetes checks"
        return
    fi

    # Check if cluster is accessible
    if ! kubectl cluster-info &> /dev/null; then
        log_warn "Cannot access Kubernetes cluster - skipping checks"
        return
    fi

    # Check API server secure port
    local api_server=$(kubectl config view --minify -o jsonpath='{.clusters[0].cluster.server}')
    if [[ "$api_server" == https://* ]]; then
        log_success "Kubernetes API server uses HTTPS"
        ((PASSED++))
    else
        log_error "Kubernetes API server does not use HTTPS"
        ((FAILED++))
    fi

    # Check RBAC is enabled
    if kubectl get clusterrole &> /dev/null; then
        log_success "RBAC is enabled"
        ((PASSED++))
    else
        log_error "RBAC may not be enabled"
        ((FAILED++))
    fi

    # Check network policies
    local network_policies=$(kubectl get networkpolicies --all-namespaces --no-headers 2>/dev/null | wc -l)
    if [ "$network_policies" -gt 0 ]; then
        log_success "Network policies are configured ($network_policies found)"
        ((PASSED++))
    else
        log_warn "No network policies found - consider implementing network segmentation"
        ((WARNINGS++))
    fi
}

# Check PostgreSQL STIG compliance
check_postgresql_stig() {
    log_info "Checking PostgreSQL STIG compliance..."
    
    # Check if PostgreSQL is accessible
    if ! command -v psql &> /dev/null; then
        log_warn "psql not found - skipping PostgreSQL checks"
        return
    fi

    # Check SSL/TLS configuration
    local ssl_enabled=$(psql -h localhost -U postgres -t -c "SHOW ssl;" 2>/dev/null | tr -d ' ' || echo "off")
    if [ "$ssl_enabled" == "on" ]; then
        log_success "PostgreSQL SSL is enabled"
        ((PASSED++))
    else
        log_error "PostgreSQL SSL is not enabled"
        ((FAILED++))
    fi

    # Check password encryption
    local password_encryption=$(psql -h localhost -U postgres -t -c "SHOW password_encryption;" 2>/dev/null | tr -d ' ' || echo "md5")
    if [ "$password_encryption" == "scram-sha-256" ] || [ "$password_encryption" == "md5" ]; then
        log_success "PostgreSQL password encryption is configured"
        ((PASSED++))
    else
        log_warn "PostgreSQL password encryption may need review"
        ((WARNINGS++))
    fi
}

# Check Linux STIG compliance
check_linux_stig() {
    log_info "Checking Linux STIG compliance..."
    
    # Check SSH configuration
    if [ -f /etc/ssh/sshd_config ]; then
        # Check PermitRootLogin
        local permit_root=$(grep -i "^PermitRootLogin" /etc/ssh/sshd_config | awk '{print $2}' || echo "yes")
        if [ "$permit_root" == "no" ]; then
            log_success "SSH root login is disabled"
            ((PASSED++))
        else
            log_error "SSH root login should be disabled"
            ((FAILED++))
        fi

        # Check PasswordAuthentication
        local password_auth=$(grep -i "^PasswordAuthentication" /etc/ssh/sshd_config | awk '{print $2}' || echo "yes")
        if [ "$password_auth" == "no" ]; then
            log_success "SSH password authentication is disabled (key-based only)"
            ((PASSED++))
        else
            log_warn "SSH password authentication is enabled - consider disabling"
            ((WARNINGS++))
        fi
    fi

    # Check firewall
    if command -v ufw &> /dev/null; then
        local ufw_status=$(ufw status | head -1 | awk '{print $2}')
        if [ "$ufw_status" == "active" ]; then
            log_success "UFW firewall is active"
            ((PASSED++))
        else
            log_warn "UFW firewall is not active"
            ((WARNINGS++))
        fi
    elif command -v firewall-cmd &> /dev/null; then
        local firewalld_status=$(firewall-cmd --state 2>/dev/null || echo "not running")
        if [ "$firewalld_status" == "running" ]; then
            log_success "Firewalld is running"
            ((PASSED++))
        else
            log_warn "Firewalld is not running"
            ((WARNINGS++))
        fi
    fi
}

# Check application security STIG
check_application_stig() {
    log_info "Checking Application Security STIG compliance..."
    
    # Check for hardcoded secrets (basic check)
    if grep -r "password.*=.*['\"].*['\"]" api/src --include="*.ts" --include="*.js" 2>/dev/null | grep -v "CHANGE_ME" | grep -v "your-secret" | head -1; then
        log_error "Potential hardcoded passwords found in code"
        ((FAILED++))
    else
        log_success "No obvious hardcoded passwords found"
        ((PASSED++))
    fi

    # Check for default credentials
    if grep -r "postgres.*postgres\|admin.*admin\|root.*root" api/src --include="*.ts" --include="*.js" 2>/dev/null | head -1; then
        log_error "Potential default credentials found"
        ((FAILED++))
    else
        log_success "No obvious default credentials found"
        ((PASSED++))
    fi
}

# Main function
main() {
    echo "=========================================="
    echo "STIG Compliance Checker"
    echo "DoD/MilSpec Compliance Verification"
    echo "=========================================="
    echo ""

    check_kubernetes_stig
    echo ""
    check_postgresql_stig
    echo ""
    check_linux_stig
    echo ""
    check_application_stig
    echo ""

    echo "=========================================="
    echo "STIG Compliance Check Summary"
    echo "=========================================="
    echo -e "${GREEN}Passed:${NC} $PASSED"
    echo -e "${YELLOW}Warnings:${NC} $WARNINGS"
    echo -e "${RED}Failed:${NC} $FAILED"
    echo ""

    if [ $FAILED -eq 0 ]; then
        echo -e "${GREEN}All critical STIG checks passed!${NC}"
        exit 0
    else
        echo -e "${RED}Some STIG checks failed. Please review and remediate.${NC}"
        exit 1
    fi
}

main "$@"