Set up Kubernetes container image security scanning with Trivy and admission controllers

#!/usr/bin/env bash

set -euo pipefail

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

# Global variables
TRIVY_VERSION="0.48.1"
SCRIPT_DIR="/tmp/k8s-trivy-install"
CLEANUP_PERFORMED=false

# Usage function
usage() {
    echo "Usage: $0 [OPTIONS]"
    echo "Install Kubernetes container image security scanning with Trivy and admission controllers"
    echo ""
    echo "Options:"
    echo "  --trivy-version VERSION    Trivy version to install (default: ${TRIVY_VERSION})"
    echo "  --skip-tests              Skip verification tests"
    echo "  -h, --help                Show this help message"
    echo ""
    echo "Prerequisites:"
    echo "  - Kubernetes cluster with kubectl configured"
    echo "  - Cluster admin privileges"
    echo "  - Internet connectivity"
}

# Parse arguments
SKIP_TESTS=false
while [[ $# -gt 0 ]]; do
    case $1 in
        --trivy-version)
            TRIVY_VERSION="$2"
            shift 2
            ;;
        --skip-tests)
            SKIP_TESTS=true
            shift
            ;;
        -h|--help)
            usage
            exit 0
            ;;
        *)
            echo -e "${RED}Error: Unknown option $1${NC}"
            usage
            exit 1
            ;;
    esac
done

# Cleanup function
cleanup() {
    if [ "$CLEANUP_PERFORMED" = false ]; then
        echo -e "${YELLOW}Cleaning up temporary files...${NC}"
        rm -rf "$SCRIPT_DIR"
        CLEANUP_PERFORMED=true
    fi
}

# Error handler
error_handler() {
    echo -e "${RED}Error occurred in script at line $1${NC}"
    cleanup
    exit 1
}

trap 'error_handler $LINENO' ERR
trap cleanup EXIT

# Log function
log() {
    echo -e "${GREEN}[$(date '+%Y-%m-%d %H:%M:%S')] $1${NC}"
}

warn() {
    echo -e "${YELLOW}[$(date '+%Y-%m-%d %H:%M:%S')] WARNING: $1${NC}"
}

error() {
    echo -e "${RED}[$(date '+%Y-%m-%d %H:%M:%S')] ERROR: $1${NC}"
}

# Check prerequisites
check_prerequisites() {
    echo -e "${BLUE}[1/10] Checking prerequisites...${NC}"
    
    # Check if running as root or with sudo
    if [[ $EUID -ne 0 ]] && ! sudo -n true 2>/dev/null; then
        error "This script requires root privileges or sudo access"
        exit 1
    fi
    
    # Check kubectl
    if ! command -v kubectl &> /dev/null; then
        error "kubectl is required but not installed"
        exit 1
    fi
    
    # Test kubernetes connectivity
    if ! kubectl cluster-info &> /dev/null; then
        error "Cannot connect to Kubernetes cluster. Please configure kubectl"
        exit 1
    fi
    
    # Check internet connectivity
    if ! curl -s --max-time 5 https://github.com &> /dev/null; then
        error "Internet connectivity required"
        exit 1
    fi
    
    log "Prerequisites check passed"
}

# Detect distribution
detect_distro() {
    if [ -f /etc/os-release ]; then
        . /etc/os-release
        case "$ID" in
            ubuntu|debian)
                PKG_MGR="apt"
                PKG_UPDATE="apt update -y"
                PKG_INSTALL="apt install -y"
                PKG_ARCH="Linux-64bit.deb"
                ;;
            almalinux|rocky|centos|rhel|ol|fedora)
                PKG_MGR="dnf"
                PKG_UPDATE="dnf update -y"
                PKG_INSTALL="dnf install -y"
                PKG_ARCH="Linux-64bit.rpm"
                # Fallback to yum for older systems
                if ! command -v dnf &> /dev/null; then
                    PKG_MGR="yum"
                    PKG_UPDATE="yum update -y"
                    PKG_INSTALL="yum install -y"
                fi
                ;;
            amzn)
                PKG_MGR="yum"
                PKG_UPDATE="yum update -y"
                PKG_INSTALL="yum install -y"
                PKG_ARCH="Linux-64bit.rpm"
                ;;
            *)
                error "Unsupported distribution: $ID"
                exit 1
                ;;
        esac
    else
        error "Cannot detect Linux distribution"
        exit 1
    fi
    log "Detected distribution: $ID"
}

# Install Trivy binary
install_trivy() {
    echo -e "${BLUE}[2/10] Installing Trivy binary...${NC}"
    
    mkdir -p "$SCRIPT_DIR"
    cd "$SCRIPT_DIR"
    
    # Download Trivy
    local trivy_url="https://github.com/aquasecurity/trivy/releases/download/v${TRIVY_VERSION}/trivy_${TRIVY_VERSION}_${PKG_ARCH}"
    wget -q "$trivy_url"
    
    # Install based on package type
    if [[ "$PKG_ARCH" == *".deb" ]]; then
        sudo dpkg -i "trivy_${TRIVY_VERSION}_${PKG_ARCH}" || true
        sudo $PKG_INSTALL -f
    else
        sudo rpm -ivh "trivy_${TRIVY_VERSION}_${PKG_ARCH}"
    fi
    
    log "Trivy binary installed successfully"
}

# Test Trivy standalone
test_trivy() {
    echo -e "${BLUE}[3/10] Testing Trivy standalone scanning...${NC}"
    
    if [ "$SKIP_TESTS" = false ]; then
        log "Testing Trivy with nginx:latest image..."
        trivy image --quiet --severity HIGH,CRITICAL nginx:latest || warn "Trivy test scan completed with findings"
    fi
    
    log "Trivy standalone test completed"
}

# Install Helm
install_helm() {
    echo -e "${BLUE}[4/10] Installing Helm...${NC}"
    
    if command -v helm &> /dev/null; then
        log "Helm already installed, skipping"
        return 0
    fi
    
    cd "$SCRIPT_DIR"
    
    if [[ "$PKG_MGR" == "apt" ]]; then
        curl -fsSL https://baltocdn.com/helm/signing.asc | gpg --dearmor | sudo tee /usr/share/keyrings/helm.gpg > /dev/null
        echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/helm.gpg] https://baltocdn.com/helm/stable/debian/ all main" | sudo tee /etc/apt/sources.list.d/helm-stable-debian.list
        sudo $PKG_UPDATE
        sudo $PKG_INSTALL helm
    else
        curl -fsSL -o get_helm.sh https://raw.githubusercontent.com/helm/helm/main/scripts/get-helm-3
        chmod 755 get_helm.sh
        ./get_helm.sh
    fi
    
    log "Helm installed successfully"
}

# Create Trivy namespace
create_namespace() {
    echo -e "${BLUE}[5/10] Creating Trivy namespace...${NC}"
    
    if kubectl get namespace trivy-system &> /dev/null; then
        log "trivy-system namespace already exists"
    else
        kubectl create namespace trivy-system
        kubectl label namespace trivy-system security=scanning
        log "Created trivy-system namespace"
    fi
}

# Deploy Trivy Operator
deploy_trivy_operator() {
    echo -e "${BLUE}[6/10] Deploying Trivy Operator...${NC}"
    
    helm repo add aqua https://aquasecurity.github.io/helm-charts/ || true
    helm repo update
    
    if helm list -n trivy-system | grep -q trivy-operator; then
        log "Trivy Operator already installed, upgrading..."
        helm upgrade trivy-operator aqua/trivy-operator \
            --namespace trivy-system \
            --set operator.scanJobsConcurrentLimit=3 \
            --set vulnerabilityReports.scanner=Trivy \
            --set compliance.failureThreshold=medium
    else
        helm install trivy-operator aqua/trivy-operator \
            --namespace trivy-system \
            --create-namespace \
            --set operator.scanJobsConcurrentLimit=3 \
            --set vulnerabilityReports.scanner=Trivy \
            --set compliance.failureThreshold=medium
    fi
    
    log "Trivy Operator deployed successfully"
}

# Configure Trivy policies
configure_trivy_policies() {
    echo -e "${BLUE}[7/10] Configuring Trivy scanning policies...${NC}"
    
    cat > "$SCRIPT_DIR/trivy-config.yaml" << 'EOF'
apiVersion: v1
kind: ConfigMap
metadata:
  name: trivy-operator-config
  namespace: trivy-system
data:
  trivy.severity: "HIGH,CRITICAL"
  trivy.ignoreUnfixed: "true"
  trivy.timeout: "5m0s"
  trivy.dbRepository: "ghcr.io/aquasecurity/trivy-db"
  vulnerabilityReports.scanner: "Trivy"
  configAuditReports.scanner: "Trivy"
EOF
    
    kubectl apply -f "$SCRIPT_DIR/trivy-config.yaml"
    log "Trivy policies configured"
}

# Install OPA Gatekeeper
install_gatekeeper() {
    echo -e "${BLUE}[8/10] Installing OPA Gatekeeper...${NC}"
    
    kubectl apply -f https://raw.githubusercontent.com/open-policy-agent/gatekeeper/release-3.14/deploy/gatekeeper.yaml
    
    log "Waiting for Gatekeeper to be ready..."
    kubectl wait --for=condition=Ready --timeout=300s -n gatekeeper-system pod -l control-plane=controller-manager || warn "Gatekeeper readiness check timed out"
    
    log "OPA Gatekeeper installed successfully"
}

# Create security policies
create_security_policies() {
    echo -e "${BLUE}[9/10] Creating security policy template...${NC}"
    
    cat > "$SCRIPT_DIR/constraint-template.yaml" << 'EOF'
apiVersion: templates.gatekeeper.sh/v1beta1
kind: ConstraintTemplate
metadata:
  name: k8srequiresecurityscan
spec:
  crd:
    spec:
      names:
        kind: K8sRequireSecurityScan
      validation:
        type: object
        properties:
          maxCritical:
            type: integer
            minimum: 0
          maxHigh:
            type: integer
            minimum: 0
  targets:
    - target: admission.k8s.gatekeeper.sh
      rego: |
        package k8srequiresecurityscan
        
        violation[{"msg": msg}] {
          input.review.object.kind == "Pod"
          container := input.review.object.spec.containers[_]
          msg := sprintf("Container image %v requires security scan", [container.image])
        }
EOF
    
    kubectl apply -f "$SCRIPT_DIR/constraint-template.yaml"
    
    cat > "$SCRIPT_DIR/constraint.yaml" << 'EOF'
apiVersion: constraints.gatekeeper.sh/v1beta1
kind: K8sRequireSecurityScan
metadata:
  name: require-security-scan
spec:
  match:
    kinds:
      - apiGroups: [""]
        kinds: ["Pod"]
  parameters:
    maxCritical: 0
    maxHigh: 5
EOF
    
    sleep 10 # Wait for CRD to be ready
    kubectl apply -f "$SCRIPT_DIR/constraint.yaml"
    
    log "Security policies created successfully"
}

# Verify installation
verify_installation() {
    echo -e "${BLUE}[10/10] Verifying installation...${NC}"
    
    # Check Trivy
    if command -v trivy &> /dev/null; then
        log "✓ Trivy binary: $(trivy --version | head -n1)"
    else
        warn "✗ Trivy binary not found"
    fi
    
    # Check Helm
    if command -v helm &> /dev/null; then
        log "✓ Helm: $(helm version --short)"
    else
        warn "✗ Helm not found"
    fi
    
    # Check namespace
    if kubectl get namespace trivy-system &> /dev/null; then
        log "✓ trivy-system namespace exists"
    else
        warn "✗ trivy-system namespace not found"
    fi
    
    # Check Trivy Operator
    if helm list -n trivy-system | grep -q trivy-operator; then
        log "✓ Trivy Operator deployed"
    else
        warn "✗ Trivy Operator not found"
    fi
    
    # Check Gatekeeper
    if kubectl get pods -n gatekeeper-system &> /dev/null; then
        log "✓ OPA Gatekeeper installed"
    else
        warn "✗ OPA Gatekeeper not found"
    fi
    
    echo -e "\n${GREEN}Installation completed successfully!${NC}"
    echo -e "${YELLOW}Next steps:${NC}"
    echo "1. Monitor vulnerability reports: kubectl get vulnerabilityreports -A"
    echo "2. Check Gatekeeper constraints: kubectl get constraints"
    echo "3. Test with a deployment to verify admission control"
}

# Main execution
main() {
    log "Starting Kubernetes Trivy security scanning setup"
    
    check_prerequisites
    detect_distro
    install_trivy
    test_trivy
    install_helm
    create_namespace
    deploy_trivy_operator
    configure_trivy_policies
    install_gatekeeper
    create_security_policies
    verify_installation
    
    log "Setup completed successfully!"
}

main "$@"