improve throughout.

README.md

@@ -39,9 +39,6 @@ sudo apt-get install exfat-fuse exfat-utils
 
 # For XFS support
 sudo apt-get install xfsprogs
-
-# For forensic analysis tools
-sudo apt-get install sleuthkit
 ```
 
 ## Preset Layouts
@@ -131,12 +128,18 @@ The script will:
 
 Checking filesystem tool availability...
 
-  ✓ NTFS    (mkfs.ntfs available)
+  ✓ FAT12/16 (mkfs.fat available)
   ✓ FAT32   (mkfs.vfat available)
   ✓ exFAT   (mkfs.exfat available)
-  ✓ ext2/3/4 (mkfs.ext4 available)
+  ✓ NTFS    (mkfs.ntfs available)
+  ✓ ext2    (mke2fs/mkfs.ext2 available)
+  ✓ ext3    (mke2fs/mkfs.ext3 available)
+  ✓ ext4    (mkfs.ext4 available)
   ✓ XFS     (mkfs.xfs available)
+  ✓ HFS+    (mkfs.hfsplus available)
+  ✓ APFS    (limited Linux support; creation typically requires macOS)
   ✓ swap    (mkswap available)
+  ✓ Unallocated (no mkfs required)
 
 Enter output filename (default: forensic_disk.dd): win11.dd
 

cleanup.sh

@@ -67,49 +67,166 @@ show_user_loop_devices() {
     return 0
 }
 
+# Helper to list processes blocking a device (if possible)
+list_blocking_processes() {
+    local dev="$1"
+    if command -v fuser >/dev/null 2>&1; then
+        fuser -v "$dev" 2>/dev/null || true
+    elif command -v lsof >/dev/null 2>&1; then
+        lsof "$dev" 2>/dev/null || true
+    else
+        print_info "Install 'psmisc' (provides fuser) or 'lsof' to list blocking processes"
+    fi
+}
+
 # Function to unmount all partitions of a loop device
 unmount_loop_partitions() {
     local loop_device=$1
     local unmounted=0
 
-    # Try both naming conventions
+    # Build a list of partition nodes for this loop device
-    for part in ${loop_device}p* ${loop_device}[0-9]*; do
+    local parts=()
-        if [ -e "$part" ]; then
+    for p in "${loop_device}p"* "${loop_device}"[0-9]*; do
-            local mount_point=$(findmnt -n -o TARGET "$part" 2>/dev/null || true)
+        # Only consider block device nodes that actually exist
+        if [ -b "$p" ]; then
+            parts+=("$p")
+        fi
+    done
+
+    if [ ${#parts[@]} -eq 0 ]; then
+        print_info "No partition block devices found for $loop_device"
+        return 0
+    fi
+
+    for part in "${parts[@]}"; do
+        # Determine if the partition is mounted
+        local mount_point
+        mount_point=$(findmnt -n -o TARGET --source "$part" 2>/dev/null || true)
         if [ -n "$mount_point" ]; then
             print_info "Unmounting $part from $mount_point"
-                if umount "$part"; then
+            if umount "$part" 2>/dev/null; then
                 print_success "Unmounted $part"
                 unmounted=$((unmounted + 1))
             else
-                    print_warning "Failed to unmount $part"
+                print_warning "Regular umount failed for $part - attempting lazy unmount"
+                if umount -l "$part" 2>/dev/null; then
+                    print_success "Lazy-unmounted $part"
+                    unmounted=$((unmounted + 1))
+                else
+                    print_warning "Failed to unmount $part even with lazy unmount"
                 fi
             fi
+        else
+            print_info "Partition $part does not appear to be mounted"
         fi
     done
 
+    # Attempt to remove partition mappings so kernel releases device nodes
+    if command -v partx >/dev/null 2>&1; then
+        partx -d "$loop_device" 2>/dev/null || true
+    elif command -v kpartx >/dev/null 2>&1; then
+        kpartx -d "$loop_device" 2>/dev/null || true
+    fi
+
+    # Give udev a moment to remove stale device nodes
+    if command -v udevadm >/dev/null 2>&1; then
+        udevadm settle 2>/dev/null || true
+    fi
+    sleep 1
+
     return $unmounted
 }
 
 # Function to detach loop device
 detach_loop_device() {
     local loop_device=$1
+    local force=${2:-false}  # second optional argument: "true" to forcibly kill blockers
 
     print_info "Detaching $loop_device"
 
-    # Check if device still exists in losetup output
+    # If the device is not present in losetup listing, consider it already detached
-    if ! losetup -l | grep -q "^$loop_device "; then
+    if ! losetup -l -n -O NAME 2>/dev/null | awk '{print $1}' | grep -qxF "$loop_device"; then
         print_success "Loop device already detached"
         return 0
     fi
 
-    if losetup -d "$loop_device"; then
+    # Try a normal detach first
+    if losetup -d "$loop_device" 2>/dev/null; then
         print_success "Detached $loop_device"
         return 0
-    else
-        print_error "Failed to detach $loop_device"
-        return 1
     fi
+
+    print_warning "Initial detach failed for $loop_device; attempting recovery steps"
+
+    # Try removing partition mappings and let udev settle
+    if command -v partx >/dev/null 2>&1; then
+        partx -d "$loop_device" 2>/dev/null || true
+    elif command -v kpartx >/dev/null 2>&1; then
+        kpartx -d "$loop_device" 2>/dev/null || true
+    fi
+    if command -v udevadm >/dev/null 2>&1; then
+        udevadm settle 2>/dev/null || true
+    fi
+    sleep 1
+
+    # Second detach attempt
+    if losetup -d "$loop_device" 2>/dev/null; then
+        print_success "Detached $loop_device"
+        return 0
+    fi
+
+    # If allowed, try to kill processes referencing the loop device (SIGTERM then SIGKILL)
+    if [ "$force" = "true" ]; then
+        if command -v fuser >/dev/null 2>&1; then
+            print_info "Killing processes using $loop_device (SIGTERM)"
+            fuser -k -TERM "$loop_device" 2>/dev/null || true
+            sleep 1
+
+            # Try detach again
+            if losetup -d "$loop_device" 2>/dev/null; then
+                print_success "Detached $loop_device after killing blockers"
+                return 0
+            fi
+
+            print_info "Killing any remaining processes using $loop_device (SIGKILL)"
+            fuser -k -KILL "$loop_device" 2>/dev/null || true
+            sleep 1
+
+            if losetup -d "$loop_device" 2>/dev/null; then
+                print_success "Detached $loop_device after force-killing blockers"
+                return 0
+            fi
+        elif command -v lsof >/dev/null 2>&1; then
+            print_info "Listing processes using $loop_device via lsof"
+            lsof "$loop_device" 2>/dev/null || true
+            local pids
+            pids=$(lsof -t "$loop_device" 2>/dev/null || true)
+            if [ -n "$pids" ]; then
+                print_info "Sending SIGTERM to: $pids"
+                kill -TERM $pids 2>/dev/null || true
+                sleep 1
+                if losetup -d "$loop_device" 2>/dev/null; then
+                    print_success "Detached $loop_device after killing blockers"
+                    return 0
+                fi
+                print_info "Sending SIGKILL to: $pids"
+                kill -KILL $pids 2>/dev/null || true
+                sleep 1
+                if losetup -d "$loop_device" 2>/dev/null; then
+                    print_success "Detached $loop_device after force-killing blockers"
+                    return 0
+                fi
+            fi
+        else
+            print_warning "No 'fuser' or 'lsof' available; cannot automatically kill blocking processes"
+        fi
+    fi
+
+    print_error "Failed to detach $loop_device. The device is likely still referenced by processes or the kernel"
+    print_info "Processes referencing $loop_device (if any):"
+    list_blocking_processes "$loop_device"
+
+    return 1
 }
 
 # Automatic mode
@@ -149,9 +266,10 @@ auto_cleanup() {
         local file=$(echo "$line" | awk '{$1=""; print $0}' | sed 's/^ *//')
 
         echo "Processing: $device"
-        unmount_loop_partitions "$device"
+        unmount_loop_partitions "$device" || true
 
-        if detach_loop_device "$device"; then
+        # In automatic cleanup assume we are allowed to force-kill blockers when necessary
+        if detach_loop_device "$device" "true"; then
             success=$((success+1))
         else
             failed=$((failed+1))
@@ -188,8 +306,16 @@ manual_cleanup() {
     print_info "Found loop device: $loop_device"
     echo ""
 
-    unmount_loop_partitions "$loop_device"
+    unmount_loop_partitions "$loop_device" || true
+
+    # Ask user whether to force if needed
+    read -p "Force-kill processes using $loop_device if necessary? (y/N): " force_ans
+    force_ans=${force_ans:-N}
+    if [ "$force_ans" = "y" ] || [ "$force_ans" = "Y" ]; then
+        detach_loop_device "$loop_device" "true"
+    else
         detach_loop_device "$loop_device"
+    fi
 
     echo ""
     print_success "Cleanup complete for: $target"
@@ -251,8 +377,15 @@ interactive_cleanup() {
     print_info "Cleaning up: $device -> $file"
     echo ""
 
-    unmount_loop_partitions "$device"
+    unmount_loop_partitions "$device" || true
+
+    read -p "Force-kill processes using $device if necessary? (y/N): " force_ans
+    force_ans=${force_ans:-N}
+    if [ "$force_ans" = "y" ] || [ "$force_ans" = "Y" ]; then
+        detach_loop_device "$device" "true"
+    else
         detach_loop_device "$device"
+    fi
 
     echo ""
     print_success "Cleanup complete"

pseudodisk.sh

@@ -4,7 +4,8 @@
 # Creates disk images with various filesystems for forensic analysis practice
 # Now with improved UX, sanity checks, and extended filesystem support
 
-#set -e  # Exit on error
+set -e
+set -o pipefail
 
 # Color codes for output
 RED='\033[0;31m'
@@ -15,6 +16,9 @@ CYAN='\033[0;36m'
 MAGENTA='\033[0;35m'
 NC='\033[0m' # No Color
 
+# Initialize mount points array to avoid bad-substitution errors when cleanup runs
+MOUNT_POINTS=()
+
 # Function to print colored messages
 print_info() {
     echo -e "${BLUE}[INFO]${NC} $1"
@@ -101,68 +105,48 @@ check_filesystem_tools() {
     echo "Checking filesystem tool availability..."
     echo ""
 
-    # FAT12/16
+    # Each entry: Display Name | space-separated candidate commands | install hint / note
-    if command -v mkfs.fat >/dev/null 2>&1 || command -v mkfs.vfat >/dev/null 2>&1; then
+    local checks=(
-        echo -e "  ${GREEN}✓${NC} FAT12/16 (mkfs.fat available)"
+        "FAT12/16|mkfs.fat mkfs.msdos mkfs.vfat|sudo apt-get install dosfstools"
-    else
+        "FAT32|mkfs.vfat mkfs.fat|sudo apt-get install dosfstools"
-        echo -e "  ${YELLOW}✗${NC} FAT12/16 (install: sudo apt-get install dosfstools)"
+        "exFAT|mkfs.exfat mkfs.exfatprogs mkfs.exfat-utils|sudo apt-get install exfatprogs or exfat-utils"
+        "NTFS|mkfs.ntfs mkntfs|sudo apt-get install ntfs-3g"
+        "ext2|mkfs.ext2 mke2fs|sudo apt-get install e2fsprogs"
+        "ext3|mkfs.ext3 mke2fs|sudo apt-get install e2fsprogs"
+        "ext4|mkfs.ext4 mke2fs|sudo apt-get install e2fsprogs"
+        "XFS|mkfs.xfs|sudo apt-get install xfsprogs"
+        "HFS+|mkfs.hfsplus newfs_hfs|sudo apt-get install hfsprogs"
+        "APFS|mkfs.apfs apfs-fuse|limited support — creation typically requires macOS or specialized tools"
+        "swap|mkswap|should be present in util-linux"
+        "Unallocated|:|no mkfs required"
+    )
+
+    local entry name cmds hint cmd found
+
+    for entry in "${checks[@]}"; do
+        # Split the packed entry into three fields
+        IFS='|' read -r name cmds hint <<< "$entry"
+        found=0
+
+        # If cmds is a single colon, treat it as 'no tool required'
+        if [ "$cmds" = ":" ]; then
+            printf "  %b %s (%s)\n" "${GREEN}✓${NC}" "$name" "$hint"
+            continue
         fi
 
-    # FAT32
+        # Check each candidate command for the filesystem
-    if command -v mkfs.vfat >/dev/null 2>&1; then
+        for cmd in $cmds; do
-        echo -e "  ${GREEN}✓${NC} FAT32    (mkfs.vfat available)"
+            if command -v "$cmd" >/dev/null 2>&1; then
-    else
+                printf "  %b %-12s (%s available)\n" "${GREEN}✓${NC}" "$name" "$cmd"
-        echo -e "  ${YELLOW}✗${NC} FAT32    (install: sudo apt-get install dosfstools)"
+                found=1
+                break
             fi
+        done
 
-    # exFAT
+        if [ "$found" -eq 0 ]; then
-    if command -v mkfs.exfat >/dev/null 2>&1; then
+            printf "  %b %-12s (install: %s)\n" "${YELLOW}✗${NC}" "$name" "$hint"
-        echo -e "  ${GREEN}✓${NC} exFAT    (mkfs.exfat available)"
-    else
-        echo -e "  ${YELLOW}✗${NC} exFAT    (install: sudo apt-get install exfatprogs)"
-    fi
-    
-    # NTFS
-    if command -v mkfs.ntfs >/dev/null 2>&1; then
-        echo -e "  ${GREEN}✓${NC} NTFS     (mkfs.ntfs available)"
-    else
-        echo -e "  ${YELLOW}✗${NC} NTFS     (install: sudo apt-get install ntfs-3g)"
-    fi
-    
-    # ext2/3/4
-    if command -v mkfs.ext4 >/dev/null 2>&1; then
-        echo -e "  ${GREEN}✓${NC} ext2/3/4 (mkfs.ext4 available)"
-    else
-        echo -e "  ${YELLOW}✗${NC} ext2/3/4 (install: sudo apt-get install e2fsprogs)"
-    fi
-    
-    # XFS
-    if command -v mkfs.xfs >/dev/null 2>&1; then
-        echo -e "  ${GREEN}✓${NC} XFS      (mkfs.xfs available)"
-    else
-        echo -e "  ${YELLOW}✗${NC} XFS      (install: sudo apt-get install xfsprogs)"
-    fi
-    
-    # HFS+
-    if command -v mkfs.hfsplus >/dev/null 2>&1; then
-        echo -e "  ${GREEN}✓${NC} HFS+     (mkfs.hfsplus available)"
-    else
-        echo -e "  ${YELLOW}✗${NC} HFS+     (install: sudo apt-get install hfsprogs)"
-    fi
-    
-    # APFS
-    if command -v mkfs.apfs >/dev/null 2>&1; then
-        echo -e "  ${GREEN}✓${NC} APFS     (mkfs.apfs available)"
-    else
-        echo -e "  ${YELLOW}✗${NC} APFS     (limited Linux support - not recommended)"
-    fi
-    
-    # swap
-    if command -v mkswap >/dev/null 2>&1; then
-        echo -e "  ${GREEN}✓${NC} swap     (mkswap available)"
-    else
-        echo -e "  ${YELLOW}✗${NC} swap     (should be in util-linux)"
         fi
+    done
 
     echo ""
 }
@@ -296,16 +280,30 @@ get_filename() {
     read -p "Enter output filename (default: forensic_disk.dd): " FILENAME
     FILENAME=${FILENAME:-forensic_disk.dd}
 
-    # Validate filename
+    # Strip whitespace
-    if [[ "$FILENAME" =~ [^a-zA-Z0-9._-] ]]; then
+    FILENAME=$(echo "$FILENAME" | xargs)
-        print_warning "Filename contains special characters. This may cause issues."
+
-        read -p "Continue with this filename? (y/n): " continue
+    # Reject invalid characters (allow only safe chars)
-        if [ "$continue" != "y" ]; then
+    if [[ ! "$FILENAME" =~ ^[a-zA-Z0-9._-]+$ ]]; then
+        print_error "Filename can only contain: letters, numbers, dots, underscores, hyphens"
         get_filename
         return
     fi
+
+    # Prevent path traversal and absolute paths
+    if [[ "$FILENAME" == *".."* ]] || [[ "$FILENAME" == /* ]]; then
+        print_error "Path traversal not allowed"
+        get_filename
+        return
     fi
 
+    # Ensure .dd extension
+    if [[ "$FILENAME" != *.dd ]]; then
+        FILENAME="${FILENAME}.dd"
+        print_info "Added .dd extension: $FILENAME"
+    fi
+
+    # Check if file exists
     if [ -f "$FILENAME" ]; then
         print_warning "File already exists: $FILENAME"
         read -p "Overwrite? (y/n): " OVERWRITE
@@ -337,20 +335,39 @@ get_disk_size() {
         4) DISK_SIZE_MB=5120 ;;
         5) DISK_SIZE_MB=10240 ;;
         6)
+            while true; do
                 read -p "Enter size in MB: " DISK_SIZE_MB
-            if ! [[ "$DISK_SIZE_MB" =~ ^[0-9]+$ ]] || [ "$DISK_SIZE_MB" -lt 10 ]; then
+
-                print_error "Invalid size. Must be at least 10 MB"
+                # Validate input
-                get_disk_size
+                if ! [[ "$DISK_SIZE_MB" =~ ^[0-9]+$ ]]; then
-                return
+                    print_error "Invalid input. Enter a number."
+                    continue
                 fi
-            if [ "$DISK_SIZE_MB" -gt 102400 ]; then
+
-                print_warning "Very large disk size (>100GB). This may take a while."
+                if [ "$DISK_SIZE_MB" -lt 10 ]; then
-                read -p "Continue? (y/n): " continue
+                    print_error "Minimum size is 10 MB"
-                if [ "$continue" != "y" ]; then
+                    continue
-                    get_disk_size
-                    return
                 fi
+
+                # Maximum 10TB (reasonable limit)
+                if [ "$DISK_SIZE_MB" -gt 10485760 ]; then
+                    print_error "Maximum size is 10TB (10485760 MB)"
+                    continue
                 fi
+
+                # Check available disk space
+                local available_kb
+                available_kb=$(df --output=avail -k "." | tail -1)
+                local required_kb=$((DISK_SIZE_MB * 1024))
+
+                if [ "$available_kb" -lt "$required_kb" ]; then
+                    local available_mb=$((available_kb / 1024))
+                    print_error "Not enough disk space. Available: ${available_mb}MB"
+                    continue
+                fi
+
+                break
+            done
             ;;
         *)
             print_error "Invalid choice"
@@ -601,12 +618,19 @@ apply_preset() {
             ;;
         9)  # Linux with /home
             local root_size=$((DISK_SIZE_MB / 4))
-            if [ "$root_size" -lt 5120 ]; then
+            local min_root=5120
-                root_size=5120  # Minimum 5GB for root
+
+            if [ "$root_size" -lt "$min_root" ]; then
+                # Check if disk is large enough for minimum
+                if [ "$DISK_SIZE_MB" -lt "$((min_root + 1024))" ]; then
+                    # Disk too small, use proportional sizing
+                    root_size=$((DISK_SIZE_MB * 2 / 3))
+                    print_warning "Disk too small for 5GB root, using ${root_size}MB"
+                else
+                    root_size=$min_root
                 fi
-            if [ "$root_size" -gt $((DISK_SIZE_MB - 1024)) ]; then
-                root_size=$((DISK_SIZE_MB / 2))  # If not enough space, use half
             fi
+
             PARTITION_CONFIGS+=("vfat|260|EFI")
             PARTITION_CONFIGS+=("ext4|${root_size}|rootfs")
             PARTITION_CONFIGS+=("ext4|remaining|home")
@@ -938,36 +962,65 @@ create_disk_image() {
     case $INIT_METHOD in
         fallocate)
             if command -v fallocate >/dev/null 2>&1; then
-                fallocate -l ${DISK_SIZE_MB}M "$FILENAME"
+                if ! fallocate -l ${DISK_SIZE_MB}M "$FILENAME"; then
+                    print_error "Failed to create disk image with fallocate"
+                    exit 1
+                fi
             else
                 print_warning "fallocate not available, falling back to /dev/zero"
-                dd if=/dev/zero of="$FILENAME" bs=1M count=$DISK_SIZE_MB status=progress 2>&1 | grep -v "records"
+                if ! dd if=/dev/zero of="$FILENAME" bs=1M count="$DISK_SIZE_MB" status=progress; then
+                    print_error "Failed to create disk image"
+                    exit 1
+                fi
             fi
             ;;
         zero)
-            dd if=/dev/zero of="$FILENAME" bs=1M count=$DISK_SIZE_MB status=progress 2>&1 | grep -v "records"
+            if ! dd if=/dev/zero of="$FILENAME" bs=1M count="$DISK_SIZE_MB" status=progress; then
+                print_error "Failed to create disk image"
+                exit 1
+            fi
             ;;
         random)
             print_warning "Using /dev/urandom - this will be SLOW!"
-            dd if=/dev/urandom of="$FILENAME" bs=1M count=$DISK_SIZE_MB status=progress 2>&1 | grep -v "records"
+            if ! dd if=/dev/urandom of="$FILENAME" bs=1M count="$DISK_SIZE_MB" status=progress; then
+                print_error "Failed to create disk image"
+                exit 1
+            fi
             ;;
     esac
 
-    print_success "Disk image created with $INIT_METHOD"
+    # Verify file was created with correct size
+    if [ ! -f "$FILENAME" ]; then
+        print_error "Disk image file was not created"
+        exit 1
+    fi
+
+    local actual_size
+    actual_size=$(stat -c%s "$FILENAME" 2>/dev/null || stat -f%z "$FILENAME" 2>/dev/null || echo 0)
+    local expected_size=$((DISK_SIZE_MB * 1024 * 1024))
+
+    if [ "$actual_size" -ne "$expected_size" ]; then
+        print_error "Disk image size mismatch (expected ${expected_size} bytes, got ${actual_size})"
+        rm -f "$FILENAME"
+        exit 1
+    fi
+
+    print_success "Disk image created and verified"
 }
 
 # Setup loop device
 setup_loop_device() {
     print_info "Setting up loop device..."
-    LOOP_DEVICE=$(losetup -f)
+
+    # Use atomic operation (find + attach in one command)
+    LOOP_DEVICE=$(losetup -f --show "$FILENAME" 2>/dev/null || true)
 
     if [ -z "$LOOP_DEVICE" ]; then
-        print_error "No free loop devices available"
+        print_error "Failed to create loop device"
         print_info "Try: sudo modprobe loop max_loop=16"
         exit 1
     fi
 
-    losetup "$LOOP_DEVICE" "$FILENAME"
     print_success "Loop device created: $LOOP_DEVICE"
 }
 
@@ -1037,76 +1090,162 @@ create_partitions() {
 
 # Format the partitions
 format_partitions() {
-    local part_num=1
+    local config_num=1
+    local actual_part_num=1
 
     for config in "${PARTITION_CONFIGS[@]}"; do
         IFS='|' read -r fs size label <<< "$config"
 
-        # Skip unallocated space - no partition to format
         if [ "$fs" = "unallocated" ]; then
-            print_info "Skipping unallocated space (no partition to format)"
+            print_info "Skipping unallocated space in config $config_num"
+            config_num=$((config_num + 1))
             continue
         fi
 
         # Determine partition device name
-        PARTITION="${LOOP_DEVICE}p${part_num}"
+        PARTITION="${LOOP_DEVICE}p${actual_part_num}"
         if [ ! -e "$PARTITION" ]; then
-            PARTITION="${LOOP_DEVICE}${part_num}"
+            PARTITION="${LOOP_DEVICE}${actual_part_num}"
         fi
 
         if [ ! -e "$PARTITION" ]; then
-            print_error "Cannot find partition device for partition $part_num"
+            print_error "Cannot find partition device for config $config_num (expected ${LOOP_DEVICE}p${actual_part_num} or ${LOOP_DEVICE}${actual_part_num})"
-            print_info "Expected: ${LOOP_DEVICE}p${part_num} or ${LOOP_DEVICE}${part_num}"
             cleanup
             exit 1
         fi
 
-        print_info "Formatting partition $part_num ($PARTITION) with $fs filesystem..."
+        print_info "Formatting config #${config_num} -> partition ${actual_part_num} ($PARTITION) with $fs filesystem..."
 
         case $fs in
             fat12)
-                # FAT12 requires specific cluster size
+                output=$(mkfs.fat -F 12 -n "$label" "$PARTITION" 2>&1)
-                mkfs.fat -F 12 -n "$label" "$PARTITION" 2>&1 | grep -v "^mkfs.fat"
+                ret=$?
+                if [ $ret -ne 0 ]; then
+                    print_error "Failed to format partition ${actual_part_num} as FAT12"
+                    echo "$output"
+                    cleanup
+                    exit 1
+                fi
+                echo "$output" | sed -n '1,50p'
                 ;;
             fat16)
-                # FAT16
+                output=$(mkfs.fat -F 16 -n "$label" "$PARTITION" 2>&1)
-                mkfs.fat -F 16 -n "$label" "$PARTITION" 2>&1 | grep -v "^mkfs.fat"
+                ret=$?
+                if [ $ret -ne 0 ]; then
+                    print_error "Failed to format partition ${actual_part_num} as FAT16"
+                    echo "$output"
+                    cleanup
+                    exit 1
+                fi
+                echo "$output" | sed -n '1,50p'
                 ;;
             vfat)
-                # FAT32
+                output=$(mkfs.vfat -F 32 -n "$label" "$PARTITION" 2>&1)
-                mkfs.vfat -F 32 -n "$label" "$PARTITION" 2>&1 | grep -v "^mkfs.fat"
+                ret=$?
+                if [ $ret -ne 0 ]; then
+                    print_error "Failed to format partition ${actual_part_num} as FAT32"
+                    echo "$output"
+                    cleanup
+                    exit 1
+                fi
+                echo "$output" | sed -n '1,50p'
                 ;;
             ntfs)
-                mkfs.ntfs -f -L "$label" "$PARTITION" 2>&1 | grep -E "^(Cluster|Creating|mkntfs completed)"
+                output=$(mkfs.ntfs -f -L "$label" "$PARTITION" 2>&1)
+                ret=$?
+                if [ $ret -ne 0 ]; then
+                    print_error "Failed to format partition ${actual_part_num} as NTFS"
+                    echo "$output"
+                    cleanup
+                    exit 1
+                fi
+                echo "$output" | grep -E "^(Cluster|Creating|mkntfs completed)" || true
                 ;;
             exfat)
-                mkfs.exfat -n "$label" "$PARTITION" 2>&1 | grep -v "^exfatprogs"
+                output=$(mkfs.exfat -n "$label" "$PARTITION" 2>&1)
+                ret=$?
+                if [ $ret -ne 0 ]; then
+                    print_error "Failed to format partition ${actual_part_num} as exFAT"
+                    echo "$output"
+                    cleanup
+                    exit 1
+                fi
+                echo "$output" | sed -n '1,50p'
                 ;;
             ext2|ext3|ext4)
-                mkfs."$fs" -L "$label" "$PARTITION" 2>&1 | grep -E "^(Creating|Writing|mke2fs)"
+                output=$(mkfs."$fs" -L "$label" "$PARTITION" 2>&1)
+                ret=$?
+                if [ $ret -ne 0 ]; then
+                    print_error "Failed to format partition ${actual_part_num} as $fs"
+                    echo "$output"
+                    cleanup
+                    exit 1
+                fi
+                echo "$output" | grep -E "^(Creating|Writing|mke2fs)" || true
                 ;;
             xfs)
-                mkfs.xfs -f -L "$label" "$PARTITION" 2>&1 | grep -v "^meta-data"
+                output=$(mkfs.xfs -f -L "$label" "$PARTITION" 2>&1)
+                ret=$?
+                if [ $ret -ne 0 ]; then
+                    print_error "Failed to format partition ${actual_part_num} as XFS"
+                    echo "$output"
+                    cleanup
+                    exit 1
+                fi
+                echo "$output" | sed -n '1,50p'
                 ;;
             hfsplus)
-                mkfs.hfsplus -v "$label" "$PARTITION" 2>&1
+                output=$(mkfs.hfsplus -v "$label" "$PARTITION" 2>&1)
+                ret=$?
+                if [ $ret -ne 0 ]; then
+                    print_error "Failed to format partition ${actual_part_num} as HFS+"
+                    echo "$output"
+                    cleanup
+                    exit 1
+                fi
+                echo "$output" | sed -n '1,50p'
                 ;;
             apfs)
                 print_warning "APFS formatting on Linux is not well supported"
                 print_info "Skipping format for APFS partition"
                 ;;
             swap)
-                mkswap -L "SWAP$part_num" "$PARTITION" 2>&1 | grep "^Setting up"
+                output=$(mkswap -L "SWAP${actual_part_num}" "$PARTITION" 2>&1)
+                ret=$?
+                if [ $ret -ne 0 ]; then
+                    print_error "Failed to set up swap on partition ${actual_part_num}"
+                    echo "$output"
+                    cleanup
+                    exit 1
+                fi
+                echo "$output" | grep -E "^Setting up" || true
+                ;;
+            *)
+                print_warning "Unknown filesystem type: $fs - skipping format"
                 ;;
         esac
 
-        print_success "Partition $part_num formatted"
+        print_success "Partition ${actual_part_num} formatted (config ${config_num})"
-        part_num=$((part_num + 1))
+        actual_part_num=$((actual_part_num + 1))
+        config_num=$((config_num + 1))
     done
 }
 
 # Cleanup function
 cleanup() {
+    # Unmount any mounted filesystems
+    if [ "${#MOUNT_POINTS[@]}" -gt 0 ]; then
+        print_info "Unmounting filesystems..."
+        for mp in "${MOUNT_POINTS[@]}"; do
+            if mountpoint -q "$mp" 2>/dev/null; then
+                umount "$mp" 2>/dev/null || umount -l "$mp" 2>/dev/null || true
+                print_info "Unmounted $mp"
+            fi
+            rmdir "$mp" 2>/dev/null || true
+        done
+    fi
+
+    # Detach loop device
     if [ -n "$LOOP_DEVICE" ]; then
         print_info "Cleaning up loop device..."
         losetup -d "$LOOP_DEVICE" 2>/dev/null || true
@@ -1208,7 +1347,7 @@ show_summary() {
         config_num=$((config_num + 1))
     done
     
-    if [ ${#MOUNT_POINTS[@]} -gt 0 ]; then
+    if [ "${#MOUNT_POINTS[@]}" -gt 0 ]; then
         echo ""
         echo "Mount Points:"
         for mp in "${MOUNT_POINTS[@]}"; do
@@ -1239,7 +1378,7 @@ show_summary() {
     echo "  xxd -s 0x1BE -l 64 $FILENAME      # MBR partition table"
     echo ""
     echo "Clean up (when done):"
-    if [ ${#MOUNT_POINTS[@]} -gt 0 ]; then
+    if [ "${#MOUNT_POINTS[@]}" -gt 0 ]; then
         for mp in "${MOUNT_POINTS[@]}"; do
             echo "  sudo umount $mp"
         done