From 4a9c6f0894f65d3cdd9ea89f3da28988ec5df6fd Mon Sep 17 00:00:00 2001
From: Eric Renfro <psi-jack@linux-help.org>
Date: Mon, 11 Nov 2024 11:40:00 -0500
Subject: [PATCH] Added backup of partition tables for use in restore

---
 scripts/functions.sh | 84 ++++++++++++++++++++++++++++++++++++++++++++
 scripts/os           |  2 ++
 2 files changed, 86 insertions(+)
 create mode 100644 scripts/functions.sh

diff --git a/scripts/functions.sh b/scripts/functions.sh
new file mode 100644
index 0000000..ee99ee5
--- /dev/null
+++ b/scripts/functions.sh
@@ -0,0 +1,84 @@
+#!/bin/bash 
+
+###########################################################
+### Utility & Common Functions
+###########################################################
+
+function is_bin_in_path {
+    builtin type -P "$1" &>/dev/null
+}
+
+###########################################################
+### BACKUP and RESTORE FUNCTIONS
+###########################################################
+
+#dumpPartitions() {
+backup_partitions() {
+    local rootPart rootDisk
+
+    rootPart="$(findmnt -no SOURCE / | sed -E 's/\[.*\]$//')"
+    rootDisk="/dev/$(lsblk -no pkname "$rootPart" | head -n1)"
+
+    sfdisk --dump "$rootDisk" > "${restoreDir}/sfdisk.dump"
+    blkid -o export "$rootPart" > "${restoreDir}/blkid.dump"
+    btrfs subvolume list -p / > "${restoreDir}/btrfs.dump"
+    mount | grep "$rootPart" > "${restoreDir}/mounts.dump"
+}
+
+#restorePartitions() {
+restore_partitions() {
+    local rootPart rootDisk
+
+    rootPart="$1"
+    rootDisk="$2"
+
+    if [[ -z "$rootDisk" || -z "$rootPart" ]]; then
+        echo "ERROR, restorePartitions not supplied with rootDisk and/or rootPartition"
+        exit 200
+    fi
+
+    # Restore Partition Table (without modification)
+    sfdisk "$rootDisk" < "${restoreDir}/sfdisk.dump"
+
+    # Restore UUIDs
+    while read -r line; do
+        if [[ $line == UUID=* ]]; then
+            eval "$line"
+            tune2fs "$rootPart" -U "$UUID"
+        fi
+    done < "${restoreDir}/blkid.dump"
+
+    echo "Partition table and UUID information have been restored."
+}
+
+#restoreBtrFSSubvolumes() {
+restore_btrfs_subvolumes() {
+    local rootBase subvolID subvolPath
+
+    rootBase="$1"
+
+    while read -r line; do
+        # Extract the subvolume ID and Path
+        subvolID="$(echo "$line" | awk '{print $2}')"
+        subvolPath="$(echo "$line" | awk '{print $NF}')"
+
+        # Restore the subvolume
+        btrfs subvolume create "${rootBase}/${subvolPath}"
+    done < "${restoreDir}/btrfs.dump"
+}
+
+#restoreBtrFSMounts() {
+restore_btfs_mounts() {
+    local rootBase mountSource mountDest
+
+    rootBase="$1"
+
+    while read -r line; do
+        # Extract mount source and destination
+        mountSource="$(echo "$line" | awk '{print $1}')"
+        mountDest="$(echo "$line" | awk '{print $3}')"
+
+        # Mount the subvolume
+        mount "$mountSource" "${rootBase}/${mountDest}"
+    done < "$restoreDir/mounts.dump"
+}
diff --git a/scripts/os b/scripts/os
index 9df90ad..2e0285a 100755
--- a/scripts/os
+++ b/scripts/os
@@ -5,6 +5,7 @@ restoreDir="/etc/restore"
 kernel="$(uname -s)"
 scriptPath="$(dirname "$(readlink -f "$0")")"
 
+source "${scriptPath}/functions.sh"
 
 if [[ "$kernel" == "Darwin" ]]; then
     restoreDir="${HOME}/restore"
@@ -78,6 +79,7 @@ hook_fail() {
 
 hook_after() {
     checkOS || exit 200
+    backup_partitions || exit 200
     runOsHook after
 }