From 265e5c6f20eb8502a0d21015b73b1bb223d33579 Mon Sep 17 00:00:00 2001
From: Kevin Veen-Birkenbach <kevin@veen.world>
Date: Sat, 20 Jul 2024 22:16:51 +0200
Subject: [PATCH] Implemented count for blocks to overwrite;
 https://chatgpt.com/share/ac2b59af-c0f0-486f-89fe-0301d4915837

---
 scripts/base.sh                               | 31 +++++++++++++------
 .../encryption/storage/single_drive/setup.sh  |  2 +-
 scripts/image/setup.sh                        |  2 +-
 3 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/scripts/base.sh b/scripts/base.sh
index aa5242a..c4dfba0 100644
--- a/scripts/base.sh
+++ b/scripts/base.sh
@@ -95,15 +95,28 @@ set_device_path(){
   info "Optimal blocksize set to: $OPTIMAL_BLOCKSIZE" || error
 }
 
-overwritte_device_with_zeros(){
-  question "Should $device_path be overwritten with zeros before copying?(y/N)" && read -r copy_zeros_to_device
-  if [ "$copy_zeros_to_device" = "y" ]
-    then
-      info "Overwritting..." &&
-      dd if=/dev/zero of="$device_path" bs="$OPTIMAL_BLOCKSIZE" status=progress || error "Overwritting $device_path failed."
-    else
-      info "Skipping Overwritting..."
-  fi
+overwrite_device() {
+  question "Should $device_path be overwritten with zeros before copying? (y/N/block count)" && read -r copy_zeros_to_device
+  case "$copy_zeros_to_device" in
+    y)
+      info "Overwriting entire device..." &&
+      dd if=/dev/zero of="$device_path" bs="$OPTIMAL_BLOCKSIZE" status=progress || error "Overwriting $device_path failed."
+      ;;
+    N)
+      info "Skipping Overwriting..."
+      ;;
+    ''|*[!0-9]*)
+      error "Invalid input."
+      ;;
+    *)
+      if [[ "$copy_zeros_to_device" =~ ^[0-9]+$ ]]; then
+        info "Overwriting $copy_zeros_to_device blocks..." &&
+        dd if=/dev/zero of="$device_path" bs="$OPTIMAL_BLOCKSIZE" count="$copy_zeros_to_device" status=progress || error "Overwriting $device_path failed."
+      else
+        error "Invalid input. Block count must be a number."
+      fi
+      ;;
+  esac
 }
 
 get_packages(){
diff --git a/scripts/encryption/storage/single_drive/setup.sh b/scripts/encryption/storage/single_drive/setup.sh
index 7071f5d..95e1cb0 100644
--- a/scripts/encryption/storage/single_drive/setup.sh
+++ b/scripts/encryption/storage/single_drive/setup.sh
@@ -7,7 +7,7 @@ echo "Setups disk encryption"
 
 set_device_mount_partition_and_mapper_paths
 
-overwritte_device_with_zeros
+overwrite_device
 
 info "Creating new GPT partition table..."
 (	echo "g"	# create a new empty GPT partition table
diff --git a/scripts/image/setup.sh b/scripts/image/setup.sh
index 0e34052..f218867 100644
--- a/scripts/image/setup.sh
+++ b/scripts/image/setup.sh
@@ -263,7 +263,7 @@ if [ "$transfer_image" = "y" ]
         info "Skipping partition table deletion..."
     fi
 
-    overwritte_device_with_zeros
+    overwrite_device
 
     info "Starting image transfer..."
     if [ "$distribution" = "arch" ]