#!/bin/bash
# ---------------------------------------------------------------
# Company:              University of Rostock
# Engineer:             Tim Brockmann
# Create Date:          17.10.2023
# Project Name:         APU-testbed
# File Name:            apu_backup_plan.sh
# Target OS:            Tiny Core Linux 14.0
# Dependencies:         - init_backup.sh
#                                       - rsync v3.1.1
# Task:
# This script is downloaded during the backup process from the
# control pc (ctrlpc: 192.168.0.1), is called by the init_backup.sh
# script and performes the backup. Note that this script is started
# as sub console and uses global variables defined and initialized
# by the parent console (init_backup.sh).
# ---------------------------------------------------------------
# Define the hostname of the reference node and what have to be done
# when the script is running automatically
# BACKUP=true .. the reference node rsync its system partition to the ctrl pc
# RESTORE=true .. all nodes rsync the backup from ctrl pc to its system partition
# attention!!! having both switches initialized by the value true can cause
# weired problems, scince copying from and to the backup folder on
# the control pc is done at the same time
# ---------------------------------------------------------------
# ---------------------- switches -------------------------------
S_DEBUG=0 # use this to get full terminal output

# ---------------------- actions --------------------------------
A_UROM=0 # use this to update the BIOS to the version specified in V_ROM
A_SWOS=0 # use this to switch to another operating system while keeping all partition informations and the data from sda4 (data partition)
A_INIT=0 # forces a complete restore including full erasing drive
A_PULL=0 # all nodes (except D_REFNOD) fetching from P_REMOTE to local backup partition before restoring
# A_REST=0 # all nodes (except D_REFNOD) restore update from local backup partition
A_GRUB=0 # install grub only (e.g. after installing an os by hand)
A_PART=0 # init the ssd drive only (mostly used when installing a new OS from scratch)
A_SWKL=0 # switches the kernel to the version specified in V_KERNEL
A_PUSH=0 # !!! the reference node !!! updates the remote backup at defined path P_REMOTE

# ---------------------- defines --------------------------------
D_REFND="apu50"      # defines the reference node, which is only able to push its drive to remote
D_DRIVE="/dev/sda"   # defines the drive we want to inspect
# D_DRIVE="/dev/mmcblk0"

# ---------------------- versions and paths ---------------------
#V_BACKUP="kali2_v1"      # backup version
#V_BACKUP="kali2_6.6_rc7" # backup version
# V_BACKUP="kali2_6.6.0+_reg_mod" # backup version
# V_BACKUP="kali2_6.6.0+_blacklist" # backup version
V_BACKUP="kali2_6.6.0+_wiptp" # backup version
# V_BACKUP="kali2_6.6.0+_lora_chirpstack" # backup version

# V_KERNEL="6.6.0_rc7" # kernel version
# V_KERNEL="6.6.0_no_reg_restrictions" # kernel version
V_KERNEL="6.6.0+_x86" # kernel version
# V_KERNEL="6.6.0+_x86_mesh_blacklist" # added a blacklist for persitant peer link blocking
# V_KERNEL="6.6.0+_x86_wiptp"
# V_KERNEL="6.6.0+_x86_wiptp_automcs"
# V_KERNEL="6.6.0+_x86_wiptp_close_opposite_link"
# V_KERNEL="6.6.0+_x86_ptp+filter"
# V_KERNEL="6.6.0+_x86_test_branch"

V_ROM="apu2/apu2_v4.19.0.1.rom"

P_REMOTE="/home/apu/backup/full_backup/${V_BACKUP}/" # remote backup path
P_KERNEL="/home/apu/backup/apu_kernel/${V_KERNEL}/" # remote kernel path
P_ROM="/home/apu/roms/${V_ROM}"

echo " "
echo "==== Tiny Core 14.0 backup plan started ===="
echo "==== Hostname: ${APU_HOSTNAME} ===="
echo "==== Local log file: ${LOCAL_LOG_FILE} ===="
echo " "

# ---------------------------------------------------------------
# Specify the drive we want to inspect
# ---------------------------------------------------------------

# drive="/dev/sda"
# drive="/dev/mmcblk0p1"

# ---------------------------------------------------------------
# Define the predefined labels/types to compare against
# ---------------------------------------------------------------
# expected_labels=("system" "swap" "backup" "data")
# expected_types=("ext4" "swap" "ext4" "ext4")
expected_labels=("system" "swap")
expected_types=("ext4" "swap")

# ---------------------------------------------------------------
# Scrpt control flags
# ---------------------------------------------------------------
something_todo_flag=0
mount_flag=0
label_flag=0
type_flag=0
init_filesystem=0

# ---------------------------------------------------------------
# Initialize arrays to store partition information
# ---------------------------------------------------------------
partitions=()
discovered_labels=()
discovered_types=()

update_rom() {
    rsync -ahxL --stats --delete -e 'ssh -o StrictHostKeyChecking=no' apu@"${CTRLPC}":${P_ROM} /home/tc/

}

# ---------------------------------------------------------------
# Use the transfer function inside the backup_plan-script to update the remote log file
# ---------------------------------------------------------------
transfer_logfile() {
        echo " "
        echo "==== Transfering .log ===="
        # create directory
        rsync -ahx --delete -e 'ssh -o StrictHostKeyChecking=no' /home/tc/*.log apu@"${CTRLPC}":"${REMOTE_LOG_PATH}"
        wait
}

erase_everything() {
    device="$1"
    echo "==== Start erasing everything ===="
    # Use blkdiscard to discard all data on the drive
    blkdiscard -f "$device"

    # Check the exit status of the blkdiscard command
    if [ $? -eq 0 ]; then
        echo "Drive $device has been set to a blank state."
    else
        echo "Failed to set the drive $device to a blank state."
    fi
    dd if=/dev/zero of="$device" bs=512 count=1
}

format_partitions() {
    device="$1"
    echo "==== Start partitioning ssd ===="
    if [[ $device != "/dev/sda" ]]; then
        device="${1}p"
    fi
    mkfs.ext4 -F -L "system" "${device}1"
    sudo swapoff -av
    mkswap -f -L "swap" -U c44c66d3-6b51-4493-9f5d-7e9d9ae4abd6 "${device}2"

    # mkfs.ext4 -F -L "backup" "${device}3"
    # mkfs.ext4 -F -L "data" "${device}4"

    # updates uuid
    sudo e2fsck -fy "${device}1"
    unmount_all_partitions
    sudo tune2fs -U 61b58db2-59a0-4db6-8648-c6780c53df3a "${device}1"
}

# ---------------------------------------------------------------
# Everything to do if a drive is completely naked or needs to be reset
# ---------------------------------------------------------------
init_drive() {
    echo " "
    echo "==== Start initializing drive ===="
    transfer_logfile

    unmount_all_partitions
    # sudo bash /opt/scripts/delete_all.sh "$D_DRIVE"
    erase_everything "$D_DRIVE"
    sudo bash /home/tc/partitioning_ssd.sh "$D_DRIVE"
    # sudo bash /opt/scripts/partitioning_ssd.sh "$D_DRIVE"
    # yes y | sudo bash /opt/scripts/format_partitions.sh "$D_DRIVE"
    format_partitions "$D_DRIVE"
    wait

    echo " "
    echo "==== Initializing drive done ===="
    transfer_logfile
}

# ---------------------------------------------------------------
# scince we want to change the OS, we need to remove all the data
# from both backup and system partition first and than copy the
# new files from the control pc
# ---------------------------------------------------------------
rm_partition_data() {
    echo " "
    echo "==== Start wiping the partitions backup and system (sda1) ===="

    mount_needed_partitions
    # sudo rm -r /mnt/sda3/*
    sudo rm -r /mnt/sda1/*

        echo " "
    echo "==== Wiping partitions done ===="
    transfer_logfile
}

rm_ids_system() {
    mount_needed_partitions
    sudo echo "" > /mnt/sda1/etc/machine-id
    sudo echo "" > /mnt/sda1/var/lib/dbus/machine-id
    sudo echo "" > /mnt/sda1/etc/hostname
        echo " "
        echo "==== Done removing system id's ===="
}

# ---------------------------------------------------------------
# Updating the linux kernel
# variable KERNEL needs to be set accordingly (e.g.: v6.6-rc7)
# and needs to match with the corresponding folder on the ctrlpc
# ---------------------------------------------------------------
update_kernel() {
    echo " "
    echo "==== Start updating kernel to version ${V_KERNEL} ===="
    transfer_logfile

    mkdir -p /home/tc/kernel
    if [[ $S_DEBUG -eq 1 ]]; then
        rsync -ahx --stats --delete --info=progress2 -e 'ssh -o StrictHostKeyChecking=no' apu@"${CTRLPC}":"${P_KERNEL}linux*" /home/tc/kernel/
                wait
    else
        rsync -ahx --stats --delete -e 'ssh -o StrictHostKeyChecking=no' apu@"${CTRLPC}":"${P_KERNEL}linux*" /home/tc/kernel/
                wait
    fi
    mount_for_chroot

    sudo rm /mnt/sda1/boot/initrd*
    sudo rm /mnt/sda1/boot/vmlinuz*
    sudo rm /mnt/sda1/boot/config*
    sudo rm /mnt/sda1/boot/System*

    cd /home/tc/kernel/
    sudo tar -xf *.tar.xz -C /mnt/sda1
    sudo chroot /mnt/sda1 /usr/sbin/update-initramfs -k all -c
    sudo chroot /mnt/sda1 /usr/sbin/update-grub

        echo " "
    echo "==== Done updating kernel to version ${V_KERNEL} ===="
    transfer_logfile
}

# ---------------------------------------------------------------
# Updating the remote backup on the control pc is only be performed
# when its the reference node
# ---------------------------------------------------------------
update_remote_backup() {
    echo " "
    echo "==== Start updating remote backup of version ${V_BACKUP} ===="
    mount_needed_partitions
    transfer_logfile
    rm_ids_system
    ssh apu@${CTRLPC} "mkdir -p "${P_REMOTE}""

    if [[ $S_DEBUG -eq 1 ]]; then
        sudo rsync -ahx --stats --delete --info=progress2 -e 'ssh -o StrictHostKeyChecking=no' --exclude-from '/opt/backup_exclude_1.txt' /mnt/sda1/ root@192.168.0.1:${P_REMOTE}
        wait
    else
        sudo rsync -ahx --stats --delete -e 'ssh -o StrictHostKeyChecking=no' --exclude-from '/opt/backup_exclude_1.txt' /mnt/sda1/ root@192.168.0.1:${P_REMOTE}
        wait
        fi

    echo " "
    echo "==== Done updating remote backup of version ${V_BACKUP} ===="
    transfer_logfile
}

# ---------------------------------------------------------------
# all nodes except references node (apu00) gets the backup and saves it to the
# backup partition (/mnt/sda3)
# ..dry() for debugging purposes
# ---------------------------------------------------------------
get_backup() {
    echo " "
    echo "==== Start getting backup version ${V_BACKUP} from control pc ===="
    mount_needed_partitions
    wait
    transfer_logfile
    wait
    # attention: this overwrites all made changes on backup partition
    # better only use it when device is completely empty
    # sudo rsync -ahx --stats --delete --progress -e 'ssh -o StrictHostKeyChecking=no' root@192.168.0.1:/home/apu/backup/apu-backup/ /mnt/sda3/
    if [[ $S_DEBUG -eq 1 ]]; then
        sudo rsync -ahx --stats --delete --info=progress2 -e 'ssh -o StrictHostKeyChecking=no' root@192.168.0.1:${P_REMOTE} /mnt/sda1/
        wait
    else
        sudo rsync -ahx --stats --delete -e 'ssh -o StrictHostKeyChecking=no' root@192.168.0.1:${P_REMOTE} /mnt/sda1/
        wait
    fi

    # rm_ids_backup
    rm_ids_system
    wait
    echo " "
    echo "==== Done getting backup version ${V_BACKUP} from control pc ===="
    transfer_logfile
    wait
}

# ---------------------------------------------------------------
# since there are no mounting points when running the script during boot
# we first need to create the folders and can then mount the partitions
# unmounting -> self explaining
# ---------------------------------------------------------------
mount_needed_partitions() {
    if [[ $mount_flag -eq 0 ]]; then
        echo " "
        echo "==== Start mounting needed partitions ===="

        mount_flag=1
        sudo mkdir -p /mnt/sda1
        sudo mount /dev/sda1 /mnt/sda1
        echo " "
        echo "==== Done mounting needed partitions ===="
        transfer_logfile
    fi
}

unmount_all_partitions() {
    echo " "
    echo "==== Start unmounting all partitions ===="

    mount_flag=0
    # Disable swap partition so that we are able to unmount all partitions
    sudo swapoff -av
    sudo umount -f ${D_DRIVE}*

        echo " "
        echo "==== Done unmounting all partitions ===="
    transfer_logfile
}

mount_for_chroot() {
        echo " "
        echo "==== Start mounting for chroot ===="

        mount_needed_partitions
        sudo mount --bind /dev /mnt/sda1/dev
        sudo mount --bind /proc /mnt/sda1/proc
        sudo mount --bind /sys /mnt/sda1/sys

        echo " "
        echo "==== Done mounting for chroot ===="
        transfer_logfile
}

# ---------------------------------------------------------------
# this function is never needed until now, because its part of the format
# script, which comes with the "site.gz" package of the predefined
# tiny core linux
# ---------------------------------------------------------------
adjust_new_node_to_backup() {
    echo " "
    echo "==== Updating the system partition uuid ===="

        # updates the system partition uuid
    if [[ $D_DRIVE == "/dev/sda" ]]; then
        sudo e2fsck -fy /dev/sda1
        unmount_all_partitions
        sudo tune2fs -fU 61b58db2-59a0-4db6-8648-c6780c53df3a /dev/sda1
    else
        sudo e2fsck -fy "${D_DRIVE}p1"
        unmount_all_partitions
        sudo tune2fs -fU 61b58db2-59a0-4db6-8648-c6780c53df3a "${D_DRIVE}p1"
    fi

        echo " "
        echo "==== Done Updating the system partition uuid ===="
    transfer_logfile
}

# ---------------------------------------------------------------
# Installs the boot sector of the main partition (/dev/sda1)
# ---------------------------------------------------------------
install_grub() {
        echo " "
        echo "==== Start installing grub ===="

        # attention: the system drive sda1 needs to be mounted at /mnt/sda1 for this
        mount_needed_partitions
        sudo grub-install --target=i386-pc --boot-directory=/mnt/sda1/boot ${D_DRIVE}

        echo " "
        echo "==== Done installing grub ===="
        transfer_logfile
}

# ---------------------------------------------------------------
# Use blkid to get partition information
# ---------------------------------------------------------------
while read -r line; do
    if [[ $line == $D_DRIVE* ]]; then
        #echo $line
        partition_info=($line)
        partition_device="${partition_info[0]}"
        partition_label=""
        partition_type=""
        # Extract the label (if available)
        if [[ $line == *LABEL=* ]]; then
            partition_label=$(echo "$line" | sed -n 's/.*LABEL="\([^"]*\)".*/\1/p')
            discovered_labels+=("$partition_label")
                        #echo $partition_label
        fi
            # Extract the Type (if available)
            if [[ $line == *TYPE=* ]]; then
                partition_type=$(echo "$line" | sed -n 's/.*TYPE="\([^"]*\)".*/\1/p')
                discovered_types+=("$partition_type")
                #echo $partition_type
            fi
    fi
done < <(blkid $D_DRIVE*[0-9])

# ---------------------------------------------------------------
# Compare discovered labels with expected labels
# ---------------------------------------------------------------
for i in "${!expected_labels[@]}"; do
    if [[ "${discovered_labels[$i]}" != "${expected_labels[$i]}" ]]; then
        echo "Label $i: Discovered '${discovered_labels[$i]}' does not match Expected '${expected_labels[$i]}'"
        label_flag=1
        break
    else
        echo "Label $i: Discovered '${discovered_labels[$i]}' matches Expected '${expected_labels[$i]}'"
    fi
done

# ---------------------------------------------------------------
# Compare discovered types with expected types
# ---------------------------------------------------------------
for j in "${!expected_types[@]}"; do
    if [[ "${discovered_types[$j]}" != "${expected_types[$j]}" ]]; then
        echo "Type $j: Discovered '${discovered_types[$j]}' does not match Expected '${expected_types[$j]}'"
        type_flag=1
        break
    else
        echo "Type $j: Discovered '${discovered_types[$j]}' matches Expected '${expected_types[$j]}'"
    fi
done

# ---------------------------------------------------------------
# Check the label flag to see if all labels matched
# ---------------------------------------------------------------
if [ "$label_flag" -eq 1 ]; then
    echo " "
    echo "==== Not all labels matched predefined labels ===="
    init_filesystem=1
else
    echo " "
    echo "==== All labels matched predefined labels and are in the correct order ==="
fi

# ---------------------------------------------------------------
# Check the type flag to see if all types matched
# ---------------------------------------------------------------
if [ "$type_flag" -eq 1 ]; then
    echo " "
    echo "==== Not all types matched predefined types ===="
    init_filesystem=1
else
    echo " "
    echo "==== All types matched predefined types and are in the correct order ===="
fi

# ---------------------------------------------------------------
# Every node can update its kernel version if according flag is 1
# ---------------------------------------------------------------
if [ $A_SWKL -eq 1 ]; then
    something_todo_flag=1
    update_kernel
    echo " "
    echo "==== I'm node $APU_HOSTNAME and my new kernel version is $V_KERNEL ====="
fi

# ---------------------------------------------------------------
# Check if the kernel should get updated
# ---------------------------------------------------------------
if [[ $APU_HOSTNAME == $D_REFND ]] && [[ $init_filesystem -eq 0 ]]; then # its the reference node
    if [ $A_PUSH -eq 1 ]; then # pushing after the kernel update
        something_todo_flag=1
        update_remote_backup
        echo " "
        echo "==== I'm the reference node $D_REFND and I updated $V_BACKUP ===="
    fi
    if [ $something_todo_flag -eq 0 ]; then # everything is fine
        echo " "
        echo "==== I'm the reference node $D_REFND and ther's nothing to do for me ===="
    fi
else # its not the reference node
    if [ $A_PART -eq 1 ]; then # only partitioniing the drive
        echo " "
        echo "==== Only initializing the SSD ===="
        init_drive
                echo " "
        echo "==== SSD ready for OS ===="
    fi

    if [[ $A_INIT -eq 1 ]] || [[ $init_filesystem -eq 1 ]]; then # initializing the whole system to defined backup
        A_SWOS=0;
        init_drive
        A_PULL=1;
        # A_REST=1;
    fi

    if [[ $A_SWOS -eq 1 ]]; then # switching operating system to defined backup
        rm_partition_data
        A_PULL=1;
        # A_REST=1;
    fi

    if [[ $A_PULL -eq 1 ]]; then # fetching the defined backup from control pc to backup partition
        get_backup
    fi

    if [[ $A_INIT -eq 1 ]] || [[ $A_SWOS -eq 1 ]] || [[ $A_GRUB -eq 1 ]]; then # installing the grub bootloader to system drive
        install_grub
    fi

    echo " "
    echo "==== I'm node $APU_HOSTNAME and I'm done with everything ===="
fi

echo " "
echo "==== Backup plan done ===="