format2.sh

#!/usr/bin/env bash
set -euo pipefail

############################################################
# ZERO-INTERACTION SUDO SYSTEM
############################################################
PASSWORD="${PASSWORD:-}"
if [[ -z "$PASSWORD" ]]; then
    echo "ERROR: PASSWORD variable not set."
    echo "Usage: PASSWORD=\"yourpass\" ./format.sh ..."
    exit 1
fi

sudo_exec() {
    echo "$PASSWORD" | sudo -S "$@"
}
unmount_device() {
    disk="$1"

    # Example: /dev/sda → list /dev/sda1 /dev/sda2 ...
    for p in $(lsblk -ln -o NAME "$disk" | tail -n +2); do
        part="/dev/$p"
        if mountpoint -q "/dev/$p"; then
            umount "$part"
        fi
    done

    # Also unmount by mountpoint paths
    for mp in $(lsblk -ln -o MOUNTPOINT "$disk" | grep /); do
        umount "$mp"
    done
}

resolve_device() {
    arg="$1"

    if [[ "$arg" == /media/* ]]; then
        findmnt -n -o SOURCE "$arg"
    else
        echo "$arg"
    fi
}

get_disk_from_partition() {
    part="$1"
    lsblk -no PKNAME "$part" | xargs -I{} echo "/dev/{}"
}

############################################################
# HELP MENU
############################################################
usage() {
    cat <<EOF
USB Toolkit — Zero Interaction Mode

Operations:
  --eraseDisk <disk> <fs>      Full wipe + new partition + mkfs + mount
  --erase, -e <part> <fs>      Format a single partition
  --zero <disk> <fs>           Full dd wipe → GPT/MBR → mkfs → mount
  --addPartition, -a <disk> <size> <fs> <name>
                               Add a new partition of <size>
  --name, -n <label>           Set filesystem label
  --mbr                        Use MBR partition scheme (default GPT)
  --help, -h                   Show help

Filesystem support:
  ext4, ext3, btrfs, swap, fat32, exfat, ntfs, apfs, hfs+

Example:
  PASSWORD="pass" ./format.sh --eraseDisk /dev/sdb ext4 -n MYUSB
EOF
    exit 1
}

############################################################
# FLAG STORAGE
############################################################
ERASE_DISK=false
ERASE_PART=false
ZERO_DISK=false
ADDPART=false
USE_MBR=false

DISK=""
PART=""
FS_TYPE_DISK=""
FS_TYPE_PART=""
ZERO_TARGET=""
ZERO_FS=""
ADDPART_DISK=""
ADDPART_SIZE_RAW=""
ADDPART_FS=""
ADDPART_NAME=""
LABEL="USB"
INVOKER_UID=$(id -u)
INVOKER_GID=$(id -g)
INVOKER_USER="${SUDO_USER:-$(logname 2>/dev/null || echo "$USER" || whoami 2>/dev/null)}"

############################################################
# DEVICE RESOLUTION: /media/... → /dev/...
############################################################
resolve_device() {
    local input="$1"

    # Already a block device?
    if [[ -b "$input" ]]; then
        echo "$input"
        return 0
    fi

    # Mountpoint resolution
    if [[ -d "$input" ]]; then
        local src
        src=$(findmnt -n -o SOURCE --target "$input" 2>/dev/null || true)
        if [[ -b "$src" ]]; then
            echo "$src"
            return 0
        fi
    fi

    # df fallback
    if df --output=source "$input" >/dev/null 2>&1; then
        local s
        s=$(df --output=source "$input" | tail -1)
        if [[ -b "$s" ]]; then
            echo "$s"
            return 0
        fi
    fi

    echo "ERROR: Unable to resolve device for: $input"
    exit 1
}

get_parent_disk() {
    local part="$1"
    local parent
    parent=$(lsblk -no pkname "$part" 2>/dev/null || true)
    if [[ -n "$parent" ]]; then
        echo "/dev/$parent"
    else
        echo "$part"
    fi
}

############################################################
# ARGUMENT PARSING
############################################################
while [[ $# -gt 0 ]]; do
    case "$1" in
        --eraseDisk)
            ERASE_DISK=true
            DISK=$(resolve_device "$2")
            DISK=$(get_parent_disk "$DISK")
            FS_TYPE_DISK="$3"
            shift 3
            ;;

        --erase|-e)
            ERASE_PART=true
            PART=$(resolve_device "$2")
            FS_TYPE_PART="$3"
            shift 3
            ;;

        --zero)
            ZERO_DISK=true
            ZERO_TARGET=$(resolve_device "$2")
            ZERO_TARGET=$(get_parent_disk "$ZERO_TARGET")
            ZERO_FS="$3"
            shift 3
            ;;

        --addPartition)
        
      dev=$(resolve_device "$2")
      disk=$(get_disk_from_partition "$dev")
      echo $disk
      unmount_device "$disk"

fdisk "$disk" <<EOF
n
p


w
EOF

      partprobe "$dev"
      shift
      shift
      ;;

        --name|-n)
            LABEL="$2"
            shift 2
            ;;

        --mbr)
            USE_MBR=true
            shift
            ;;

        --help|-h)
            usage
            ;;

        *)
            echo "Unknown flag: $1"
            usage
            ;;
    esac
done

if ! $ERASE_DISK && ! $ERASE_PART && ! $ZERO_DISK && ! $ADDPART; then
    echo "ERROR: No operation specified."
    usage
fi

############################################################
# FILESYSTEM CREATION
############################################################
mkfs_with_label() {
    local fs="$1"
    local target="$2"

    case "${fs,,}" in
        ext4) sudo_exec mkfs.ext4 -F -L "$LABEL" "$target" ;;
        ext3) sudo_exec mkfs.ext3 -F -L "$LABEL" "$target" ;;
        btrfs) sudo_exec mkfs.btrfs -f -L "$LABEL" "$target" ;;
        swap) sudo_exec mkswap -L "$LABEL" "$target" ;;

        exfat|exfat*)
            if command -v mkfs.exfat >/dev/null 2>&1; then
                sudo_exec mkfs.exfat -n "$LABEL" "$target"
            else
                sudo_exec mkfs.exfat.py -n "$LABEL" "$target"
            fi
            ;;

        fat32|vfat|fat)
            sudo_exec mkfs.vfat -F32 -n "$LABEL" "$target"
            ;;

        ntfs)
            sudo_exec mkfs.ntfs -F -L "$LABEL" "$target"
            ;;

        apfs)
            if ! command -v mkfs.apfs >/dev/null 2>&1; then
                echo "ERROR: mkfs.apfs not found (install apfsprogs)."
                exit 1
            fi
            sudo_exec mkfs.apfs -v "$LABEL" "$target"
            ;;

        hfs+|hfsplus)
            if ! command -v mkfs.hfsplus >/dev/null 2>&1; then
                echo "ERROR: mkfs.hfsplus not found (install hfsprogs)."
                exit 1
            fi
            sudo_exec mkfs.hfsplus -v "$LABEL" "$target"
            ;;

        *)
            echo "ERROR: Unsupported filesystem: $fs"
            exit 1
            ;;
    esac
}

############################################################
# MOUNTING HANDLER
############################################################
# Then replace mount_fs() with:
mount_fs() {
    local fs="$1"
    local dev="$2"
    local mp="$3"

    # numeric UID/GID captured at script start
    local uid_val="$INVOKER_UID"
    local gid_val="$INVOKER_GID"

    # ensure mountpoint exists (run as root)
    sudo_exec mkdir -p "$mp"

    case "${fs,,}" in
        swap)
            sudo_exec swapon "$dev"
            echo "Swap enabled on $dev"
            return
            ;;

        apfs)
            if ! command -v apfs-fuse >/dev/null 2>&1; then
                echo "ERROR: apfs-fuse missing."
                exit 1
            fi
            sudo_exec apfs-fuse "$dev" "$mp"
            # try to fix ownership (may be limited by FUSE)
            sudo_exec chown "$uid_val":"$gid_val" "$mp" 2>/dev/null || true
            echo "APFS mounted via FUSE at $mp"
            ;;

        hfs+|hfsplus)
            sudo_exec mount -t hfsplus "$dev" "$mp"
            sudo_exec chown "$uid_val":"$gid_val" "$mp" || true
            sudo_exec chmod 775 "$mp" || true
            echo "HFS+ mounted at $mp (ownership set)"
            ;;

        ntfs)
            if command -v ntfs-3g >/dev/null 2>&1; then
                sudo_exec umount "$mp" 2>/dev/null || true
                sudo_exec ntfs-3g -o uid="$uid_val",gid="$gid_val",umask=0022 "$dev" "$mp"
            else
                sudo_exec mount -t ntfs "$dev" "$mp" 2>/dev/null || sudo_exec mount "$dev" "$mp"
                sudo_exec chown "$uid_val":"$gid_val" "$mp" 2>/dev/null || true
            fi
            echo "NTFS mounted at $mp"
            ;;

        exfat|fat32|vfat|fat)
            sudo_exec umount "$mp" 2>/dev/null || true
            # use numeric uid/gid for vfat/exfat so files are owned by invoking user
            sudo_exec mount -o uid="$uid_val",gid="$gid_val",umask=0022 "$dev" "$mp"
            echo "FAT/exFAT mounted at $mp (uid=$uid_val gid=$gid_val)"
            ;;

        btrfs|ext4|ext3)
            sudo_exec mount "$dev" "$mp"
            # change ownership of the root of the FS to allow the user to write
            sudo_exec chown "$uid_val":"$gid_val" "$mp" || true
            sudo_exec chmod 775 "$mp" || true
            echo "$fs mounted at $mp (owned by invoking user)"
            ;;

        *)
            sudo_exec mount "$dev" "$mp" 2>/dev/null || true
            sudo_exec chown "$uid_val":"$gid_val" "$mp" 2>/dev/null || true
            echo "$dev mounted at $mp"
            ;;
    esac
}

############################################################
# ERASE PARTITION
############################################################
erase_partition() {
    sudo_exec umount "$PART" 2>/dev/null || true
    mkfs_with_label "$FS_TYPE_PART" "$PART"
    mount_fs "$FS_TYPE_PART" "$PART" "/mnt/$(basename "$PART")"
}

############################################################
# ERASE DISK
############################################################
erase_disk() {
    for p in $(lsblk -ln -o NAME "$DISK" | tail -n +2); do
        sudo_exec umount "/dev/$p" 2>/dev/null || true
    done

    sudo_exec wipefs -a "$DISK"

    if $USE_MBR; then
        sudo_exec parted -s "$DISK" mklabel msdos
        sudo_exec parted -s "$DISK" mkpart primary 1MiB 100%
    else
        sudo_exec parted -s "$DISK" mklabel gpt
        sudo_exec parted -s "$DISK" mkpart primary 0% 100%
    fi

    sleep 1
    local part="${DISK}1"
    mkfs_with_label "$FS_TYPE_DISK" "$part"
    mount_fs "$FS_TYPE_DISK" "$part" "/mnt/$(basename "$part")"
}

############################################################
# ZERO DISK
############################################################
zero_disk() {
    for p in $(lsblk -ln -o NAME "$ZERO_TARGET" | tail -n +2); do
        sudo_exec umount "/dev/$p" 2>/dev/null || true
    done

    sudo_exec dd if=/dev/zero of="$ZERO_TARGET" bs=1M status=progress

    if $USE_MBR; then
        sudo_exec parted -s "$ZERO_TARGET" mklabel msdos
        sudo_exec parted -s "$ZERO_TARGET" mkpart primary 1MiB 100%
    else
        sudo_exec parted -s "$ZERO_TARGET" mklabel gpt
        sudo_exec parted -s "$ZERO_TARGET" mkpart primary 0% 100%
    fi

    sleep 1
    local part="${ZERO_TARGET}1"
    mkfs_with_label "$ZERO_FS" "$part"
    mount_fs "$ZERO_FS" "$part" "/mnt/$(basename "$part")"
}

############################################################
# ADD PARTITION
############################################################
add_partition() {
    local disk="$ADDPART_DISK"
    local fs="$ADDPART_FS"
    local pname="$ADDPART_NAME"
    local size_raw="$ADDPART_SIZE_RAW"

    # unmount all partitions
    for p in $(lsblk -ln -o NAME "$disk" | tail -n +2); do
        sudo_exec umount "/dev/$p" 2>/dev/null || true
    done

    # get disk size in MiB
    local devbase=$(basename "$disk")
    local sectors=$(cat "/sys/block/$devbase/size")
    local disk_mib=$(( (sectors * 512) / 1024 / 1024 ))

    # parse size (MiB / GiB / %)
    parse_size() {
        local s="$1"
        if [[ $s =~ ^([0-9]+)G ]]; then echo $(( ${BASH_REMATCH[1]} * 1024 ))
        elif [[ $s =~ ^([0-9]+)M ]]; then echo "${BASH_REMATCH[1]}"
        elif [[ $s =~ ^([0-9]+)%$ ]]; then echo $(( disk_mib * ${BASH_REMATCH[1]} / 100 ))
        else
            echo "ERROR: Bad size format: $s"; exit 1
        fi
    }
    local size_mib=$(parse_size "$size_raw")

    # get all free ranges using parted -m print free
    local free_ranges
    free_ranges=$(sudo_exec parted -m -s "$disk" unit MiB print free | awk -F: '
        /free/ {
            start=$2; end=$3
            gsub("MiB","",start)
            gsub("MiB","",end)
            print start, end
        }
    ')

    # find largest free region
    local best_start=0 best_end=0 best_size=-1
    while read -r s e; do
        [[ -z "$s" || -z "$e" ]] && continue
        size=$(printf "%.0f" "$(echo "$e - $s" | bc)")
        if (( size > best_size )); then
            best_size=$size
            best_start=$(printf "%.0f" "$s")
            best_end=$(printf "%.0f" "$e")
        fi
    done <<< "$free_ranges"

    if (( best_size < size_mib )); then
        echo "ERROR: Not enough free space."
        exit 1
    fi

    local start_mib=$best_start
    local end_mib=$(( start_mib + size_mib ))

    # create partition
    sudo_exec parted -s "$disk" mkpart primary "$fs" "${start_mib}MiB" "${end_mib}MiB"
    sleep 1

    # find new partition name
    local newpart="/dev/$(lsblk -ln -o NAME "$disk" | tail -1)"

    # format with label
    local old_label="$LABEL"
    LABEL="$pname"
    mkfs_with_label "$fs" "$newpart"
    LABEL="$old_label"

    mount_fs "$fs" "$newpart" "/media/$INVOKER_USER/$pname"

    echo "Added partition: $newpart   Mounted at: /media/$INVOKER_USER/$pname"
}


############################################################
# DISPATCH
############################################################
$ERASE_PART && erase_partition
$ERASE_DISK && erase_disk
$ZERO_DISK && zero_disk
$ADDPART && add_partition