#! /bin/sh -x

# Copyright (C) 2024-2025 NVIDIA Corporation. ALL RIGHTS RESERVED.

# This service takes config-backup of running partition and restore in
# the new target
#
# This is only done on reboot, shutdown, halt, and poweroff, otherwise
# do nothing

# This service will config-save and config-restore the backup on new 
# target partition

# First mount the partition to which next switch is
#
#
CL_UPG_PATH=/var/lib/cumulus/image-upgrade
UPGRADE_OR_ROLLBACK_FILE="$CL_UPG_PATH/upgrade-or-rollback"
EXIT_CODE=0

if [ -f "${UPGRADE_OR_ROLLBACK_FILE}" ]; then
    # This file's contents should be "upgrade" or "rollback" set
    # by cl-image-upgrade .
    upgrade_or_rollback=$(cat ${UPGRADE_OR_ROLLBACK_FILE})
else
    # Default upgrade in case of 5.11 whose cl-image-upgrade
    # does not create this file.
    upgrade_or_rollback="upgrade"
fi

echo "upgrade_or_rollback=$upgrade_or_rollback"

if [ "$upgrade_or_rollback" != "upgrade" -a "$upgrade_or_rollback" != "rollback" ]; then
    echo "Warning: ${UPGRADE_OR_ROLLBACK_FILE} does not contain upgrade or rollback (contains $upgrade_or_rollback)"
    EXIT_CODE=1
else
    rm -f "${UPGRADE_OR_ROLLBACK_FILE}"
fi

#Is it Upgrade scenario or only activate ?
cl_upgrade_file=$CL_UPG_PATH/upgrade_valid


rectify_uuid()
{

    if [ ! -f "$CL_UPG_PATH/cl-etc-dev" ] || [ ! -f "$CL_UPG_PATH/cl-efi-dev" ] || [ ! -f "$CL_UPG_PATH/new_uuid" ] || [ ! -f "$CL_UPG_PATH/fstab_base" ]; then
        echo "Error: input files missing "
        return 1
    fi
    
    #Get cl_etc_dev and cl_efi_dev 
    cl_etc_dev=$(cat $CL_UPG_PATH/cl-etc-dev)
    cl_efi_dev=$(cat $CL_UPG_PATH/cl-efi-dev)
    new_uuid=$(cat $CL_UPG_PATH/new_uuid)

    #verify these values are not zero
    if [ -z "$cl_etc_dev" ] || [ -z "$cl_efi_dev" ] || [ -z "$new_uuid" ]; then
        echo "Error: one or more of $CL_UPG_PATH/{cl-etc-dev,cl-efi-dev,new_uuid} is be empty"
        return 1
    fi

    if [ ! -b "$cl_etc_dev" ] || [ ! -b "$cl_efi_dev" ]; then
        echo "Error: invalid dev $cl_etc_dev/$cl_efi_dev"
        return 1
    fi

    # cl-etc fstab file should have correct updated uuid
    # lets also copy this to CL-ETC etc/fstab
    
    echo "Overwriting CL-ETC fstab"

    cl_etc_mnt_pt=/mnt/cl_etc_mnt
    mkdir -p $cl_etc_mnt_pt

    mount $cl_etc_dev $cl_etc_mnt_pt || {
        echo "Could not mount $cl_etc_dev  on $cl_etc_mnt_pt"
        return 1
    }

    cp "$CL_UPG_PATH"/fstab_base $cl_etc_mnt_pt/fstab || {
        log_error "Copy of fstab to Cl-ETC fstab failed"
        return 1
    }

    umount $cl_etc_mnt_pt || {
        log_error "Could not umount $cl_etc_mnt_pt"
        return 1
    }

    ###
    # grub.cfg in EFI has uuid entry, which needs to be corrected
    # to reflect new one
    if [ -z "$cl_efi_dev" ]; then
        echo "EFI device variable is empty, no UUID update to do"
    elif [ ! -b "$cl_efi_dev" ]; then
        echo "EFI device $cl_efi_dev not valid, no UUID update to do"
    else
        echo "Correcting uuid of $dev in $cl_efi_dev grub.cfg "

    # Mount point
    cl_efi_mnt_pt=/mnt/cl_efi_mnt
    mkdir -p "$cl_efi_mnt_pt"

    # Files to process
    files="EFI/cumulus-linux/grub.cfg EFI/debian/grub.cfg"

    # Mount the device
    mount $cl_efi_dev $cl_efi_mnt_pt || {
        log_error "Could not mount $cl_efi_dev  on $cl_efi_mnt_pt"
        return 1
    }

    echo "Correcting grub.cfg in EFI System dev "
    # Process each file
    for file in $files; do
        # Extract old UUID from file
        old_uuid=$(grep "search.fs_uuid" "$cl_efi_mnt_pt/$file" | cut -d' ' -f2)

        if [ -z "$old_uuid" ]; then
            log_warning "No UUID found in $file"
            continue
        fi

        # Replace UUID in file
        if sed -i "s/$old_uuid/$new_uuid/g" "$cl_efi_mnt_pt/$file"; then
            echo "Updated UUID in $file from $old_uuid to $new_uuid"
        else
            log_error "Failed to update UUID in $file"
            return 1
        fi
    done

    # Unmount the device
    umount $cl_efi_mnt_pt || {
         log_error "Could not umount $cl_efi_mnt_pt"
         return 1
    }

    fi

    return 0
}


config_save_restore()
{
    target_dev=$(cat $CL_UPG_PATH/boot_device)
    # Sanity checks
    # target_dev should be valid

    if [ -z "$target_dev" ]; then
        echo "Error: target_dev is empty"
        return 1
    fi

    if [ ! -b "$target_dev" ]; then
        echo "Error: $target_dev is not a valid block device"
        return 1
    fi

    mount_pt=/mnt/target

    mkdir -p "$mount_pt"
    if [ $? -ne 0 ] ; then
        echo "Problems creating mount point for new device $target_dev"
        return 1
    fi

    if mount | awk '{print $3}' | grep -q "${mount_pt}"; then
        echo "Device is already mounted"
    else
        echo "Device is not mounted, mounting now"
        mount "${target_dev}" "${mount_pt}"
        if [ $? -eq 0 ]; then
            echo "Device $target_dev mounted successfully"
        else
            echo "Not able to mount dev $target_dev"
        fi
    fi

    # Because dnsmasq 2.89-1 to 2.90-4~deb12u1 changed from using
    # adduser to useradd to add its dnsmasq user, the uid of that
    # user and users added afterward may differ when upgrading from
    # <= 5.12.0.0011 to >= 5.12.0.0012.  So we need to save a copy of
    # the current partition /etc/passwd and edit it with the new partition
    # uids to copy it to the new partition after doing the config-backup
    # and config-restore.  See #4258009.
    passwd_file=$(mktemp /var/tmp/passwd.XXXXXXXX )
    echo "Making copy of /etc/passwd with fixed uids for dnsmasq and others in $passwd_file"
    chmod 600 $passwd_file
    cp /etc/passwd $passwd_file
    shadow_file=$(mktemp /var/tmp/shadow.XXXXXXXX )
    echo "Making copy of /etc/shadow in $shadow_file"
    chmod 600 $shadow_file
    cp /etc/shadow $shadow_file
    for user in dnsmasq frr Debian-snmp _lldpd ntpsec nvue ; do
        uid=$(grep $user $mount_pt/etc/passwd | awk -F: '{ print $3 }' )
        if [ "$uid" != "" ]; then
            echo "User $user uid changed to $uid in next boot filesystem"
            sed -i -e "s/^${user}:x:[0-9]*:/${user}:x:${uid}:/" $passwd_file
        fi
    done
    # Now add any system users (from /etc/login.defs, default uid 100 to 999)
    # currently in /etc/passwd to $passwd_file with accompanying entry
    # into $shadow_file
    sys_uid_min=$(grep -v "^#" /etc/login.defs | grep SYS_UID_MIN | awk '{ print $2 }')
    sys_uid_min=${sys_uid_max:-100}
    sys_uid_max=$(grep -v "^#" /etc/login.defs | grep SYS_UID_MAX | awk '{ print $2 }')
    sys_uid_max=${sys_uid_max:-999}
    while read -r userline ; do
        user=$(echo "$userline" | sed -e 's,:.*$,,')
        uid=$(echo "$userline" | awk -F: '{ print $3 }')
        if [ "0$uid" -ge "$sys_uid_min" -a "0$uid" -le "$sys_uid_max" ]; then
            if ! grep -q "$user" $passwd_file ; then
                echo "User $user added to passwd in next boot filesystem"
                echo "$userline" >> $passwd_file
                grep "^${user}:" /etc/shadow >> $shadow_file
            fi
        fi
    done < /etc/passwd

    # First do config save of the running image
    # Create a backup of current config, copy to etc,
    # and save name as backup_file_name

    # new_dirs_to_backup is a space separated list of directories
    # that need to be backed up that may be missing from the
    # config_backup.conf file in older versions of Cumulus Linux.
    # These directories need to be added to the config_backup.conf
    # file before we run config-backup .

    new_dirs_to_backup="/usr/local/share/ca-certificates/nvue/"

    config_backup_conf="/etc/cumulus/config_backup.conf"
    add_dirs_to_backup=""
    include_line=$(grep -E "^(\s*)include_file_dirs(\s*):" $config_backup_conf | tail -n 1)
    for d in $new_dirs_to_backup ; do
        if ! echo "$include_line" | grep -q "$d" ; then
            if [ -d "$d" -o -f "$d" ]; then
                add_dirs_to_backup="$add_dirs_to_backup $d"
            fi
        fi
    done
    if [ ! -z "$add_dirs_to_backup" ]; then
        if [ -z "$include_line" ]; then
            echo "include_file_dirs: $add_dirs_to_backup" >> $config_backup_conf
        else
            sed -i -e "s#$include_line#$include_line $add_dirs_to_backup#" $config_backup_conf
        fi
    fi

    config_backup_output=$(/usr/lib/cumulus/config-backup)
    echo "config_backup_output: $config_backup_output"
    backup_file=$(echo $config_backup_output | tail -1 | rev | cut -d' ' -f1 | rev)
    echo "backup_file: $backup_file"

    if [ -z "$backup_file" ]; then
        echo "Problems creating backup of running image"
        return 1
    fi

    #check if this dir exists 
    if [ ! -d "${mount_pt}/var/lib/config-backup/backups" ]; then
        mkdir -p "${mount_pt}/var/lib/config-backup/backups"
    fi

    #copy the backupfile in the target area

    cp ${backup_file} ${mount_pt}/var/lib/config-backup/backups/
    if [ $? -ne 0 ] ; then
        echo "Problems copying backup file ${backup_file} to new target ${mount_pt}/var/lib/config-backup/backups/"
        return 1
    fi

    # Restore the config in the target image
    
    cd $mount_pt
    /usr/cumulus/bin/in-target $mount_pt /usr/lib/cumulus/config-restore -f -B -b $backup_file

    # Delete the new files and directories in current configuration
    /usr/cumulus/bin/in-target $mount_pt /usr/lib/cumulus/config-restore -D -b $backup_file | grep "Only in current file-system" | awk -F ": " '{print $2}' | xargs chroot $mount_pt rm -r

    if [ $? -ne 0 ] ; then
        echo "Warning: Problems restoring current backup to new device "
    fi

    #now delete this config file
    # not deleting for debugging
    #rm ${mount_pt}/var/lib/config-backup/backups/${backup_file}

    # Now copy the saved edited $passwd_file and $shadow_file for
    # uid fixes and missing users to new partition if necessary.
    echo "Checking to see if copying changed $passwd_file is needed"
    if [ "$passwd_file" != "" ]; then
        if ! cmp -s "$passwd_file" ${mount_pt}/etc/passwd ; then
            echo "Copying fixed $passwd_file to ${mount_pt}/etc/passwd"
            cp "$passwd_file" ${mount_pt}/etc/passwd
        else
            echo "$passwd_file is the same as ${mount_pt}/etc/passwd, not copying"
        fi
        rm "$passwd_file"
    fi
    if [ "$shadow_file" != "" ]; then
        if ! cmp -s "$shadow_file" ${mount_pt}/etc/shadow ; then
            echo "Copying fixed $shadow_file to ${mount_pt}/etc/shadow"
            cp "$shadow_file" ${mount_pt}/etc/shadow
        else
            echo "$shadow_file is the same as ${mount_pt}/etc/shadow, not copying"
        fi
        rm "$shadow_file"
    fi

    # Some files and directories may be created with the incorrect uid
    echo "Fixing ownership of /var/lib/nvue /etc/nvue/certificates/{private,public} /usr/local/share/ca-certificates/nvue"
    uid=$(grep nvue $mount_pt/etc/passwd | awk -F: '{ print $3 }' )
    chown -R ${uid} ${mount_pt}/var/lib/nvue
    chown -R ${uid} ${mount_pt}/etc/nvue/certificates/private
    chown -R ${uid} ${mount_pt}/etc/nvue/certificates/public
    chown -R ${uid} ${mount_pt}/usr/local/share/ca-certificates/nvue

    echo "Fixing ownership of /var/lib/ntpsec"
    uid=$(grep ntpsec $mount_pt/etc/passwd | awk -F: '{ print $3 }' )
    chown -R ${uid} ntpsec ${mount_pt}/var/lib/ntpsec

    echo "Fixing ownership of /var/lib/snmp"
    uid=$(grep Debian-snmp $mount_pt/etc/passwd | awk -F: '{ print $3 }' )
    chown -R ${uid} ${mount_pt}/var/lib/snmp


    # Add upgrade-firstboot-valid artifact in target etc
    #
    if [ "$cl_upgrade" = "true" ] ; then 
        # mount current system partition somewhere else so we can
        # access its /var
        current_dev=$(df / | tail -1 | awk '{ print $1 }')
        current_mount_pt=/mnt/current
        mkdir -p "$current_mount_pt"
        if [ $? -ne 0 ] ; then
            echo "Problems creating mount point for current device $current_dev"
            return 1
        fi
        mount "${current_dev}" "${current_mount_pt}"
        if [ $? -eq 0 ]; then
            echo "Device $current_dev mounted successfully"
        else
            echo "Not able to mount dev $current_dev"
        fi
        # These /var directories need to go with each system partition:
        # lib/dpkg lib/apt cache/apt lib/cumulus/cumulus-local-apt-archive
        for d in lib/dpkg lib/apt cache/apt \
                 lib/cumulus/cumulus-local-apt-archive ; do
            # copy /var/$d current system partition
            echo "Copying /var/$d to $current_mount_pt/var"
            ( cd /var ; tar cf - $d ) | (cd $current_mount_pt/var ; tar xf - )
            # copy new system partition /var/$d to /var
            echo "Copying $mount_pt/var/$d to /var"
            ( cd $mount_pt/var ; tar cf - $d ) | (cd /var ; tar xf - )
        done
        # These /var directories need to be saved with each system partition
        # in case of future rollback: lib/nvue
        for d in lib/nvue ; do
            # copy /var/$d current system partition
            echo "Copying /var/$d to $current_mount_pt/var"
            ( cd /var ; tar cf - $d ) | (cd $current_mount_pt/var ; tar xf - )
        done
        touch $mount_pt/etc/cumulus-upgrade-firstboot-valid
    fi 
    umount -Rl $current_mount_pt
    umount -Rl $mount_pt

    return 0
}

rollback_restore_var() {
    target_dev=$(cat $CL_UPG_PATH/boot_device)
    # Sanity checks
    # target_dev should be valid

    if [ -z "$target_dev" ]; then
        echo "Error: target_dev is empty"
        return 1
    fi

    if [ ! -b "$target_dev" ]; then
        echo "Error: $target_dev is not a valid block device"
        return 1
    fi

    mount_pt=/mnt/target

    mkdir -p "$mount_pt"
    if [ $? -ne 0 ] ; then
        echo "Problems creating mount point for new device $target_dev"
        return 1
    fi

    if mount | awk '{print $3}' | grep -q "${mount_pt}"; then
        echo "Device is already mounted"
    else
        echo "Device is not mounted, mounting now"
        mount "${target_dev}" "${mount_pt}"
        if [ $? -eq 0 ]; then
            echo "Device $target_dev mounted successfully"
        else
            echo "Not able to mount dev $target_dev"
        fi
    fi

    # These /var directories need be saved with each system partition:
    # lib/dpkg lib/apt cache/apt lib/nvue
    # lib/cumulus/cumulus-local-apt-archive
    for d in lib/dpkg lib/apt cache/apt lib/nvue \
             lib/cumulus/cumulus-local-apt-archive ; do
        # copy /var/$d current system partition
        echo "Copying /var/$d to $current_mount_pt/var"
        ( cd /var ; tar cf - $d ) | (cd $current_mount_pt/var ; tar xf - )
    done

    # These /var directories need to be restored for rollback:
    # lib/nvue
    for d in lib/nvue ; do
        # copy new system partition /var/$d to /var
        echo "Copying $mount_pt/var/$d to /var"
        ( cd $mount_pt/var ; tar cf - $d ) | (cd /var ; tar xf - )
    done
    
    umount -Rl $current_mount_pt
    umount -Rl $mount_pt

    return 0
}

echo "CL upgrade-on-shutdown entry"

if [ ! -d "$CL_UPG_PATH" ]; then
    echo "Error: $CL_UPG_PATH not available"
    exit 1
fi

if [ -f $cl_upgrade_file ]; then
    cl_upgrade=true
    echo "CL Upgrade True"
fi

if [ -f $CL_UPG_PATH/boot_device ]; then
    cl_switch_part=true
    echo "CL Switch partition True"
fi

if [ "$cl_switch_part" = "true" ]; then
    if [ "$upgrade_or_rollback" != "rollback" ]; then
        # Do config save in running image, copy to target
        # and do in-target restore

        if [ "$upgrade_or_rollback" != "upgrade" ]; then
            echo "Warning: Unexpected value of upgrade_or_rollback: $upgrade_or_rollback"
        fi
        echo "executing Config Save/Restore"
        config_save_restore
        if [ $? -eq 0 ]; then
            echo "Config Save/Restore Successful"
        else
            echo "Config Save/Restore failed"
        fi
    else
        echo "Rollback, not doing Config Save/Restore"
        echo "Rollback, restoring needed /var directories"
        rollback_restore_var
    fi

    # in case of upgrade need to update uuid in in CL-ETC
    # and in grub.cfg in EFI dev

    if [ "$cl_upgrade" = "true" ] ; then 
	echo "CL upgrade flow, need to update uuid"
	rectify_uuid
	if [ $? -eq 0 ] ; then
	    echo "UUID updation succeeded"
	else
	    echo "UUID updation failed"
	fi

	rm $cl_upgrade_file
    fi
    rm $CL_UPG_PATH/boot_device
fi

# Disable our systemd services so they are not executed again
# without being explicity enabled by cl-image-upgrade -r or -a.
#systemctl stop cumulus-rollback-on-shutdown
#systemctl disable cumulus-rollback-on-shutdown
#systemctl stop cumulus-upgrade-on-shutdown
#systemctl disable cumulus-upgrade-on-shutdown

echo "CL upgrade-on-shutdown finished"

exit $EXIT_CODE
