#!/bin/bash
#
#################################################################################################
#
# SPDX-License-Identifier: GPL-3.0-or-later
#
# Copyright (C) 2026 snix
#
# snix luks zfs ssh mkinitfs Optimized script to update initramfs with SSH and encryption support
# 
#################################################################################################
#

banner() {
banner_filename=$SCRIPT_NAME
banner_author=snix
banner_license=GPL-3.0-or-later
banner_description='snix luks zfs ssh mkinitfs Optimized script to update initramfs with SSH and encryption support'
echo "############################################################################################################################################################################
# Filename:      $banner_filename
# Author:        $banner_author
# License:       $banner_license
# Description:   $banner_description
############################################################################################################################################################################
#                                               #
#     ▄████████ ███▄▄▄▄    ▄█  ▀████    ▐████▀  #
#    ███    ███ ███▀▀▀██▄ ███    ███▌   ████▀   #
#    ███    █▀  ███   ███ ███▌    ███  ▐███     #
#    ███        ███   ███ ███▌    ▀███▄███▀     #
#  ▀███████████ ███   ███ ███▌    ████▀██▄      #
#           ███ ███   ███ ███    ▐███  ▀███     #
#     ▄█    ███ ███   ███ ███   ▄███     ███▄   #
#   ▄████████▀   ▀█   █▀  █▀   ████       ███▄  #
#                                               #
############################################################################################################################################################################"
}

set -euo pipefail

readonly SCRIPT_NAME="${0##*/}"
BASENAME="$SCRIPT_NAME"
VERSION="2.0_r13"
DEFAULT_MKINITFS_PATH="/etc/mkinitfs"
DEFAULT_CUSTOM_PATH="$DEFAULT_MKINITFS_PATH/$BASENAME"

# Base configuration
readonly CONFIG_FILE="/etc/$BASENAME.conf"
readonly WORK_DIR="/tmp/initfs.$$"
readonly BACKUP_SUFFIX=".backup.$(date +%Y%m%d_%H%M%S)"

# Default values (can be overridden by configuration file)
DEFAULT_MKINITFS_CONF="$DEFAULT_MKINITFS_PATH/mkinitfs.conf"
DEFAULT_ENCRYPT_SCRIPT="/usr/sbin/slzs-enc"
DEFAULT_SCRIPTS_DIR="$DEFAULT_CUSTOM_PATH/scripts"
DEFAULT_SSH_KEY="/etc/ssh/ssh_host_ed25519_key"
DEFAULT_AUTHORIZED_KEYS="/etc/ssh/authorized_keys"
DEFAULT_COMPRESSION="zstd"
BACKUP_KEEP="n"
DEFAULT_AUTO_UNLOCK="no"

# Global variables
VERBOSE=false
ALL_KERNELS=false
MANUAL_KERNEL=""
MANUAL_INIT=""
MANUAL_COMPRESSION=""
MANUAL_MKINITFS_CONF=""
MANUAL_SSH_KEY=""
SAVE_PASSPHRASE=false
MANUAL_PASSPHRASE=""
PASSPHRASE=""
BACKUP_PATH=""

# Load configuration from file
load_config() {
    [[ -f "$CONFIG_FILE" ]] || return 0

    log "Loading configuration from: $CONFIG_FILE"

    # Read variables from configuration file
    while IFS='=' read -r key value; do
        # Ignore comments and empty lines
        [[ "$key" =~ ^[[:space:]]*# ]] && continue
        [[ -z "$key" ]] && continue

        # Remove spaces and quotes
        key=$(echo "$key"|tr -d '[:space:]')
        value=$(echo "$value"|sed 's/^[[:space:]]*//;s/[[:space:]]*$//;s/^["'\'']*//;s/["'\'']*$//')

        case "$key" in
            DEFAULT_MKINITFS_CONF)    DEFAULT_MKINITFS_CONF="$value"    ;;
            DEFAULT_ENCRYPT_SCRIPT)   DEFAULT_ENCRYPT_SCRIPT="$value"   ;;
            DEFAULT_SCRIPTS_DIR)      DEFAULT_SCRIPTS_DIR="$value"      ;;
            DEFAULT_SSH_KEY)          DEFAULT_SSH_KEY="$value"          ;;
            DEFAULT_AUTHORIZED_KEYS)  DEFAULT_AUTHORIZED_KEYS="$value"  ;;
            DEFAULT_COMPRESSION)      DEFAULT_COMPRESSION="$value"      ;;
            BACKUP_KEEP)              BACKUP_KEEP="$value"              ;;
            DEFAULT_AUTO_UNLOCK)      DEFAULT_AUTO_UNLOCK="$value"      ;;
        esac
    done < "$CONFIG_FILE"
}

# Utility functions
log() {
    [[ "$VERBOSE" == true ]] && echo "[$(date '+%Y-%m-%d %H:%M:%S')] $SCRIPT_NAME: $*" >&2
    return 0
}

log_always() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $SCRIPT_NAME: $*" >&2
}

error() {
    log_always "ERROR: $*"
    exit 1
}

get_kernels() {
    ls -1 /lib/modules|grep -E '^[0-9]+\.[0-9]+'|sort -V
}

get_kernel_name() {
    local kver="$1"
    local name="${kver##*-}"
    [[ "$name" == "$kver" || "$name" =~ ^[0-9]+$ ]] && { echo "ERROR: cannot detect kernel type from: $kver" >&2; return 1; }
    echo "$name"
}

get_initramfs_for_kernel() {
    local kver="$1"
    local name
    name=$(get_kernel_name "$kver") || return 1
    echo "/boot/initramfs-$name"
}

usage() {
    cat >&2 << EOF
Usage: $SCRIPT_NAME [-v] [-a] [-k KERNEL_VERSION] [-i INIT_FILE] [-c COMPRESSION] [-m MKINITFS_CONF] [-s SSH_KEY] [-p PASSPHRASE]

Options:
  -v                  - Verbose mode (show detailed logs)
  -a, --all           - Process all kernels found in /lib/modules
  --keep              - Keep initramfs backup after update (default: auto-removed)
  -k KERNEL_VERSION   - Manually specify kernel version
  -i INIT_FILE        - Manually specify initramfs-init script (default: auto-generated from /usr/share/mkinitfs/initramfs-init)
  -c COMPRESSION      - Specify compression type (default: $DEFAULT_COMPRESSION)
  -m MKINITFS_CONF    - Specify mkinitfs.conf file (default: $DEFAULT_MKINITFS_CONF)
  -s SSH_KEY          - Specify SSH key (default: $DEFAULT_SSH_KEY)
  -p PASSPHRASE       - Passphrase for keys.enc (default: read from $DEFAULT_CUSTOM_PATH/passphrase)
  --save              - Save passphrase to $DEFAULT_CUSTOM_PATH/passphrase when using -p
  -V                  - Show version
  -h                  - Show this help

Kernel name is auto-detected from /lib/modules (e.g. 6.12.56-0-lts -> initramfs-lts)
Without -a, only the latest kernel in /lib/modules is used.

Configuration file: $CONFIG_FILE
  May contain: DEFAULT_MKINITFS_CONF, DEFAULT_COMPRESSION, BACKUP_KEEP

Examples:
  $SCRIPT_NAME                       # Latest kernel only
  $SCRIPT_NAME -a                    # All kernels in /lib/modules
  $SCRIPT_NAME -v -a                 # All kernels with verbose
  $SCRIPT_NAME -k 6.12.56-0-lts      # Specific kernel version
EOF
    exit 1
}

cleanup() {
    local exit_code=$?
    local dst_dir="$DEFAULT_MKINITFS_PATH/features.d"
    for feat in "${LINKED_FEATURES[@]}"; do
        for ext in files modules; do
            local link="${dst_dir}/${feat}.${ext}"
            [[ -L "$link" ]] && rm -f "$link" && log "Removed feature link: ${feat}.${ext}"
        done
    done
    [[ -d "$WORK_DIR" ]] && { log "Cleaning temporary directory: $WORK_DIR"; rm -rf "$WORK_DIR"; }
    rm -f "$DEFAULT_CUSTOM_PATH/features.d/$BASENAME.modules"
    [[ -n "$BACKUP_PATH" && "$BACKUP_KEEP" != "y" && "$BACKUP_KEEP" != "Y" ]] && rm -f "$BACKUP_PATH"
    exit $exit_code
}

prepare_ssh_keys() {
    local ssh_key="${1:-$DEFAULT_SSH_KEY}"
    local auth_keys="${DEFAULT_AUTHORIZED_KEYS}"
    local dst_dir="${WORK_DIR}/etc/ssh"

    [[ ! -f "$ssh_key" ]] && { log_always "SSH key not found: $ssh_key, generating..."; ssh-keygen -t ed25519 -N "" -f "$ssh_key" || error "Failed to generate SSH key"; }

    mkdir -p "$dst_dir"
    cp "$ssh_key" "$dst_dir/" || error "Failed to copy SSH private key"

    [[ -f "$auth_keys" ]] && { cp "$auth_keys" "${dst_dir}/authorized_keys" || error "Failed to copy authorized_keys"; } || {
        log_always "authorized_keys not found: $auth_keys"
        printf 'Enter public key for initramfs SSH access: ' > /dev/tty
        IFS= read -r rawkey < /dev/tty
        [[ -n "$rawkey" ]] || error "No public key provided"
        local keytype="" tmp
        tmp=$(mktemp)
        for type in ssh-ed25519 ssh-rsa ecdsa-sha2-nistp256 ecdsa-sha2-nistp384 ecdsa-sha2-nistp521; do
            printf '%s %s\n' "$type" "$rawkey" > "$tmp"
            ssh-keygen -l -f "$tmp" >/dev/null 2>&1 && { keytype="$type"; break; }
        done
        rm -f "$tmp"
        [[ -n "$keytype" ]] || error "Unrecognized key type"
        printf '%s %s\n' "$keytype" "$rawkey" > "${dst_dir}/authorized_keys"
        printf '%s %s\n' "$keytype" "$rawkey" > "$auth_keys"
        log_always "authorized_keys saved to: $auth_keys"
    }

    log "SSH keys prepared in: $dst_dir"
}

prepare_init_script() {
    local src="/usr/share/mkinitfs/initramfs-init"
    local dst="${WORK_DIR}/initramfs-init-run"

    [[ -f "$src" ]] || error "file '$src' not found"

    local count
    count=$(grep -c '/bin/busybox sh' "$src")
    [[ "$count" -eq 1 ]] || error "$count occurrences of '/bin/busybox sh' in '$src', expected exactly 1"

    cp "$src" "$dst"
    sed -i '/\/bin\/busybox sh/i\\t\tstartup' "$dst"

    log "Init script patched: $dst"
    echo "$dst"
}

check_dependencies() {
    local deps=("mkinitfs" "zstdcat" "cpio")
    for dep in "${deps[@]}"; do
        command -v "$dep" >/dev/null || error "Required command not found: $dep"
    done
}

get_physical_iface() {
    local iface="$1" result member lower

    # Physical interface: has device/uevent
    [[ -f "/sys/class/net/$iface/device/uevent" ]] && { echo "$iface"; return 0; }

    # Bridge: iterate all members looking for a physical one
    [[ -d "/sys/class/net/$iface/bridge" ]] && {
        for member in /sys/class/net/$iface/brif/*; do
            member=$(basename "$member")
            [[ -d "/sys/class/net/$member" ]] || continue
            result=$(get_physical_iface "$member") && { echo "$result"; return 0; }
        done
        return 1
   }

    # Virtual interface (VLAN, macvlan, etc.): follow lower_* links
    for lower in /sys/class/net/$iface/lower_*; do
        [[ -e "$lower" ]] || continue
        lower=$(basename "$lower"); lower="${lower#lower_}"
        result=$(get_physical_iface "$lower") && { echo "$result"; return 0; }
    done

    return 1
}

get_nic_module() {
    local iface="$1"
    local phys
    phys=$(get_physical_iface "$iface") || return 1
    [[ -f "/sys/class/net/$phys/device/uevent" ]] || return 1
    awk -F= '/^DRIVER=/{print $2}' "/sys/class/net/$phys/device/uevent"
}

detect_network() {
    local ssh_key="${1:-$DEFAULT_SSH_KEY}"
    local iface gw ip phys

    iface=$(ip -4 route show default|awk 'NR==1{print $5}')
    [[ -n "$iface" ]] || error "Cannot detect default network interface"

    gw=$(ip -4 route show default|awk 'NR==1{print $3}')
    [[ -n "$gw" ]] || error "Cannot detect default gateway"

    ip=$({ ip addr show "$iface" || true; }|awk '/inet /{print $2; exit}')
    [[ -n "$ip" ]] || error "Cannot detect IP for interface: $iface"

    phys=$(get_physical_iface "$iface") || error "Cannot detect physical interface for: $iface"

    mkdir -p "${WORK_DIR}/sbin"
    cat > "${WORK_DIR}/sbin/startup" << EOF
#!/bin/sh
ifname="\${ifname:-$phys}"
ip="\${ip:-$ip}"
gw="\${gw:-$gw}"
port="\${port:-22}"
[ -d "/sys/class/net/\$ifname" ] || { echo "Warning: \$ifname not found, falling back to eth0"; ifname=eth0; }
[ -d "/sys/class/net/\$ifname" ] || { echo "interface \$ifname not found"; exit 1; }
ip link set dev "\$ifname" up
ip addr add "\$ip" dev "\$ifname"
ip route add default via "\$gw"
mkdir -p /var/empty /etc/ssh
grep -q "^/bin/ash" /etc/shells 2>/dev/null || echo /bin/ash >> /etc/shells
grep -q "^sshd:" /etc/passwd || echo 'sshd:x:22:22:sshd:/var/empty:/sbin/nologin' >> /etc/passwd
cat > /etc/ssh/sshd_config << SSHEOF
HostKey /etc/ssh/$(basename "$ssh_key")
PermitRootLogin prohibit-password
PubkeyAuthentication yes
AuthorizedKeysFile /etc/ssh/authorized_keys
MaxAuthTries 3
X11Forwarding no
SSHEOF
/usr/sbin/sshd -p "\$port"
$(  [[ "$DEFAULT_AUTO_UNLOCK" == "yes" || "$DEFAULT_AUTO_UNLOCK" == "YES" ]] && echo '/sbin/go || echo "Auto-unlock failed, SSH available for manual intervention"' )
EOF
    chmod +x "${WORK_DIR}/sbin/startup"

    log_always "Network detected: $phys ip=$ip gw=$gw"
}

LINKED_FEATURES=()

generate_features() {
    local src_dir="$DEFAULT_CUSTOM_PATH/features.d"
    local dst_dir="$DEFAULT_MKINITFS_PATH/features.d"

    [[ -d "$src_dir" ]] || error "Features source directory not found: $src_dir"

    # Genera custom.modules con il modulo NIC rilevato
    local iface module
    iface=$(ip -4 route show default|awk 'NR==1{print $5}')
    [[ -n "$iface" ]] || error "Cannot detect default network interface"
    module=$(get_nic_module "$iface") || error "Cannot detect NIC module for: $iface"
    printf '%s\n' "$module" > "${src_dir}/${BASENAME}.modules"
    log_always "NIC module detected: $module"

    # Linka tutti i file da $DEFAULT_CUSTOM_PATH/features.d/ in features.d/
    for f in "$src_dir"/*; do
        [[ -f "$f" ]] || continue
        local name dst
        name=$(basename "$f")
        dst="${dst_dir}/${name}"
        [[ -e "$dst" && ! -L "$dst" ]] && { log_always "Warning: $dst already exists and is not a symlink, skipping"; continue; }
        ln -sf "$f" "$dst"
        LINKED_FEATURES+=("${name%.*}")
        log "Linked feature: $name"
    done
}

ensure_features() {
    local conf="$1"
    local features
    features=$(grep '^features=' "$conf"|tail -1|sed 's/features=//;s/"//g')

    local updated=false

    # Features obbligatorie
    for feat in scsi network zfs cryptsetup; do
           echo "$features"|grep -qw "$feat" || { features="$features $feat"; updated=true; log_always "Feature '$feat' not found in mkinitfs.conf, adding at runtime"; }
    done

    # Features linkate da $DEFAULT_CUSTOM_PATH/features.d/
    for feat in $(printf '%s\n' "${LINKED_FEATURES[@]}"|sort -u); do
        echo "$features"|grep -qw "$feat" || { features="$features $feat"; updated=true; log "Adding linked feature to conf: $feat"; }
    done

    [[ "$updated" == true ]] && { local tmp_conf="${WORK_DIR}/mkinitfs.conf"; { grep -v '^features=' "$conf"; echo "features=\"$features\""; } > "$tmp_conf"; echo "$tmp_conf"; } || echo "$conf"
}

validate_kernel() {
    local kver="$1"
    [[ -d "/lib/modules/$kver" ]] || error "Kernel $kver not found in /lib/modules"
    get_kernel_name "$kver" > /dev/null || error "Cannot detect kernel type from: $kver"
}

check_files() {
    local init_script="$1"
    local mkinitfs_conf="$2"
    local encrypt_script="$3"
    local files=("$mkinitfs_conf" "$init_script" "$encrypt_script")

    for file in "${files[@]}"; do
        [[ -f "$file" ]] || error "Required file not found: $file"
    done
}

create_backup() {
    local initramfs_path="$1"
    local backup_path="${initramfs_path}${BACKUP_SUFFIX}"

    log "Creating backup: $backup_path"
    cp "$initramfs_path" "$backup_path" || error "Unable to create backup"
    echo "$backup_path"
}

build_initramfs() {
    local kernel="$1"
    local init_script="$2"
    local compression="$3"
    local mkinitfs_conf="$4"

    log "Building initramfs for kernel: $kernel"
    log "Init file used: $init_script"
    log "Compression used: $compression"
    log "Mkinitfs config used: $mkinitfs_conf"

    local cmd=(
        mkinitfs
        -c "$mkinitfs_conf"
        -i "$init_script"
        -C "$compression"
        -b /
        "$kernel"
    )

    # Echo command only in verbose mode
    log "Command: ${cmd[*]}"

    { "${cmd[@]}" 2>&1 1>&3|{ grep -v "not created: newer or same age" >&2 || true; }; } 3>&1 || error "Failed to build initramfs"
}

extract_initramfs() {
    local initramfs_path="$1"

    log "Extracting existing initramfs to: $WORK_DIR"

    mkdir -p "$WORK_DIR" || error "Unable to create work directory"
    cd "$WORK_DIR" || error "Unable to access work directory"

    { zstdcat "$initramfs_path"|cpio -idu 2>&1 1>&3|{ grep -v -E "(ld-musl-.*\.so\.[0-9]+|[0-9]+ blocks)" >&2 || true; }; } 3>&1 || error "Failed to extract initramfs"
}

get_passphrase() {
    if [[ -n "$MANUAL_PASSPHRASE" ]]; then
        PASSPHRASE="$MANUAL_PASSPHRASE"
        [[ "$SAVE_PASSPHRASE" == true ]] && { printf '%s' "$PASSPHRASE" > "$DEFAULT_CUSTOM_PATH/passphrase"; chmod 600 "$DEFAULT_CUSTOM_PATH/passphrase"; log_always "Passphrase saved to: $DEFAULT_CUSTOM_PATH/passphrase"; }
    elif [[ -f "$DEFAULT_CUSTOM_PATH/passphrase" ]]; then
        PASSPHRASE=$(cat "$DEFAULT_CUSTOM_PATH/passphrase")
    else
        printf 'Enter passphrase for keys.enc: ' > /dev/tty
        IFS= read -rs PASSPHRASE < /dev/tty
        echo > /dev/tty
        [[ -n "$PASSPHRASE" ]] || error "No passphrase provided"
        printf '%s' "$PASSPHRASE" > "$DEFAULT_CUSTOM_PATH/passphrase"
        chmod 600 "$DEFAULT_CUSTOM_PATH/passphrase"
        log_always "Passphrase saved to: $DEFAULT_CUSTOM_PATH/passphrase"
    fi
}

apply_encryption() {
    local encrypt_script="$1"
    get_passphrase
    log "Applying encryption configuration"
    "$encrypt_script" -p "$PASSPHRASE" -d "${WORK_DIR}/etc/keys.enc" || error "Failed encryption script"

    if [[ "$DEFAULT_AUTO_UNLOCK" == "yes" || "$DEFAULT_AUTO_UNLOCK" == "YES" ]]; then
        local autopass_file autopass_path
        autopass_file=$(cat /proc/sys/kernel/random/uuid | tr -d '-' | cut -c1-12)
        autopass_path="${WORK_DIR}/usr/lib/${autopass_file}"
        printf '%s' "$PASSPHRASE" | rev | base64 > "$autopass_path"
        chmod 600 "$autopass_path"
        log_always "Auto-unlock enabled: passphrase stored as /usr/lib/${autopass_file}"

        local safego_path="${WORK_DIR}/sbin/safego"
        [[ -f "$safego_path" ]] || error "safego not found in WORK_DIR (copy_scripts must run before apply_encryption)"
        sed -i "s|^AUTOPASS=.*|AUTOPASS=\"/usr/lib/${autopass_file}\"|" "$safego_path"
    fi
}

copy_scripts() {
    local scripts_dir="${1}"
    local dst="${WORK_DIR}/sbin"

    [[ -d "$scripts_dir" ]] || error "Scripts directory not found: $scripts_dir"
    mkdir -p "$dst"

    for f in "$scripts_dir"/*; do
        [[ -f "$f" ]] || continue
        log "Copying script: ${f##*/} -> sbin"
        cp "$f" "$dst/" || error "Failed to copy ${f##*/}"
        chmod +x "$dst/${f##*/}"
    done

    # Generate go
    cat > "${dst}/go" << 'EOF'
#!/bin/sh
/sbin/safego "$@" || { echo "safego failed, aborting"; exit 1; }
/sbin/gameover
EOF
    chmod +x "${dst}/go"

    # Generate gameover
    cat > "${dst}/gameover" << 'EOF'
#!/bin/sh
rm -f /dev/shm/keys
T=$(mktemp)
cat > "$T" << 'INNER'
#!/bin/sh
sleep 1
pkill -9 sshd
ip -o link show | awk -F': ' '$2 != "lo" {print $2}' | while read -r iface; do
    ip addr flush dev "$iface"
    ip link set dev "$iface" down
done
pkill -9 busybox
rm -f "$0"
INNER
chmod +x "$T"
setsid "$T" >/dev/null 2>&1 &
EOF
    chmod +x "${dst}/gameover"
    log "Generated: go, gameover"
}

repackage_initramfs() {
    local initramfs_path="$1"
    local temp_file="${WORK_DIR}/../initfs_new.zstd"

    log "Repackaging initramfs"
    cd "$WORK_DIR" || error "Unable to access work directory"
    find .|sort|cpio --quiet --renumber-inodes -o -H newc|zstd -T0 -19 > "$temp_file" || error "Failed repackaging"

    log "Replacing initramfs: $initramfs_path"
    mv "$temp_file" "$initramfs_path" || error "Unable to replace initramfs"
}

process_kernel() {
    local kernel_version="$1"
    local init_script="$2"
    local compression="$3"
    local mkinitfs_conf="$4"
    local ssh_key="${5:-$DEFAULT_SSH_KEY}"
    local initramfs_path
    initramfs_path=$(get_initramfs_for_kernel "$kernel_version") || return 1
    log_always "Kernel selected: $kernel_version"
    log_always "Initramfs target: $initramfs_path"
    check_files "$init_script" "$mkinitfs_conf" "$DEFAULT_ENCRYPT_SCRIPT"

    # Backup (only if file exists)
    [[ -f "$initramfs_path" ]] && { BACKUP_PATH=$(create_backup "$initramfs_path"); log_always "Backup saved to: $BACKUP_PATH"; } || log "No existing initramfs to backup"
    build_initramfs "$kernel_version" "$init_script" "$compression" "$mkinitfs_conf"
    extract_initramfs "$initramfs_path"
    prepare_ssh_keys "$ssh_key"
    detect_network "$ssh_key"
    copy_scripts "$DEFAULT_SCRIPTS_DIR"
    apply_encryption "$DEFAULT_ENCRYPT_SCRIPT"
    repackage_initramfs "$initramfs_path"
    log_always "Initramfs update completed: $kernel_version"
    [[ -n "$BACKUP_PATH" ]] && {
        [[ "$BACKUP_KEEP" == "y" || "$BACKUP_KEEP" == "Y" ]] && log_always "Backup kept at: $BACKUP_PATH" || { 
          rm -f "$BACKUP_PATH" || log "Warning: unable to remove backup"
          log_always "Backup automatically removed"
        }
        BACKUP_PATH=""; }
}

main() {
    # Argument parsing
    while [[ $# -gt 0 ]]; do
        case "$1" in
            "-v"|"--verbose")
                VERBOSE=true
                shift
                ;;
            "-a"|"--all")
                ALL_KERNELS=true
                shift
                ;;
            "--keep")
                BACKUP_KEEP="y"
                shift
                ;;
            "-k"|"--kernel")
                [[ $# -lt 2 ]] && error "Option -k requires an argument (kernel version)"
                MANUAL_KERNEL="$2"
                shift 2
                ;;
            "-i"|"--init")
                [[ $# -lt 2 ]] && error "Option -i requires an argument (init file)"
                MANUAL_INIT="$2"
                shift 2
                ;;
            "-c"|"--compression")
                [[ $# -lt 2 ]] && error "Option -c requires an argument (compression type)"
                MANUAL_COMPRESSION="$2"
                shift 2
                ;;
            "-m"|"--mkinitfs-conf")
                [[ $# -lt 2 ]] && error "Option -m requires an argument (mkinitfs.conf file)"
                MANUAL_MKINITFS_CONF="$2"
                shift 2
                ;;
            "-s"|"--ssh-key")
                [[ $# -lt 2 ]] && error "Option -s requires an argument (ssh key path)"
                MANUAL_SSH_KEY="$2"
                shift 2
                ;;
            "-p"|"--passphrase")
                [[ $# -lt 2 ]] && error "Option -p requires an argument (passphrase)"
                MANUAL_PASSPHRASE="$2"
                shift 2
                ;;
            "--save")
                SAVE_PASSPHRASE=true
                shift
                ;;
            "-V"|"--version")
                echo "$SCRIPT_NAME $VERSION"
                exit 0
                ;;
            "-h"|"--help")
                usage
                ;;
            -*)
                error "Unrecognized option: $1"
                ;;
            *)
                error "Unrecognized argument: $1"
                ;;
        esac
    done

    banner
    log_always "Starting initramfs update"

    trap cleanup EXIT INT TERM
    check_dependencies
    mkdir -p "$WORK_DIR"

    local init_script
    [[ -n "$MANUAL_INIT" ]] && init_script="$MANUAL_INIT" || init_script=$(prepare_init_script)
    local compression="${MANUAL_COMPRESSION:-$DEFAULT_COMPRESSION}"
    local mkinitfs_conf="${MANUAL_MKINITFS_CONF:-$DEFAULT_MKINITFS_CONF}"
    local ssh_key="${MANUAL_SSH_KEY:-$DEFAULT_SSH_KEY}"
    generate_features
    mkinitfs_conf=$(ensure_features "$mkinitfs_conf")

    # Determine kernels to process
    local kernels
    if [[ -n "$MANUAL_KERNEL" ]]; then
        validate_kernel "$MANUAL_KERNEL"
        kernels="$MANUAL_KERNEL"
    elif [[ "$ALL_KERNELS" == true ]]; then
        kernels=$(get_kernels)
        [[ -n "$kernels" ]] || error "No kernels found in /lib/modules"
    else
        kernels=$(get_kernels|tail -1)
        [[ -n "$kernels" ]] || error "No kernels found in /lib/modules"
    fi

    while read -r kver; do
        [[ -z "$kver" ]] && continue
        process_kernel "$kver" "$init_script" "$compression" "$mkinitfs_conf" "$ssh_key"
    done <<< "$kernels"

    echo "############################################################################################################################################################################"
    log_always "ATTENTION!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!"
    log_always "Authorized keys ($DEFAULT_AUTHORIZED_KEYS):"
    echo "############################################################################################################################################################################"
    cat "$DEFAULT_AUTHORIZED_KEYS" >&2
    echo "############################################################################################################################################################################"
    log_always "ATTENTION!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!"
    log_always "Decryption passphrase:"
    echo "############################################################################################################################################################################"
    echo "$PASSPHRASE"
    echo "############################################################################################################################################################################"
}

# Generate default config if missing
[[ -f "$CONFIG_FILE" ]] || cat > "$CONFIG_FILE" << EOF
# $SCRIPT_NAME configuration file
#
# DEFAULT_MKINITFS_CONF=$DEFAULT_MKINITFS_CONF
# DEFAULT_ENCRYPT_SCRIPT=$DEFAULT_ENCRYPT_SCRIPT
# DEFAULT_SCRIPTS_DIR=$DEFAULT_SCRIPTS_DIR
# DEFAULT_SSH_KEY=$DEFAULT_SSH_KEY
# DEFAULT_AUTHORIZED_KEYS=$DEFAULT_AUTHORIZED_KEYS
# DEFAULT_COMPRESSION=$DEFAULT_COMPRESSION
# BACKUP_KEEP=$BACKUP_KEEP
# DEFAULT_AUTO_UNLOCK=$DEFAULT_AUTO_UNLOCK
EOF

# Load configuration from file (if exists)
load_config

# Verify script is run as root
[[ $EUID -eq 0 ]] || error "This script must be run as root"

# Main execution
main "$@"
