#!/bin/bash

## bash library for formatusb

##arguments format, label, device

partnum=""

device="$1"
format="$2"
label="$3"
part="$4"  #can be part, defaults, gpt, or msdos

case "${part,,}" in
     part)     part=part     ;;
     gpt)      part=gpt      ;;
     msdos)    part=msdos    ;;
     defaults) part=defaults ;;
     # Fail closed: an unrecognised mode must not silently fall through to
     # "defaults" (whole-disk repartitioning) on a destructive privileged helper.
        *) echo "Invalid partition mode: $part" >&2; exit 2 ;;
esac

# SECURITY: $device is interpolated (often unquoted) into many root commands
# below (df/dd/parted/blkid/sfdisk). This helper is directly pkexec-invocable,
# so GUI-side checks are bypassable — validate the device name here. A real
# block-device name (sdb1, mmcblk0p1, nvme0n1p1) is alphanumeric only; this
# rejects any shell metacharacter / path separator.
if [[ ! "$device" =~ ^[A-Za-z0-9]+$ ]]; then
    echo "Invalid device name: $device" >&2
    exit 2
fi

# Must be a real block device (the helper is directly pkexec-invocable, so this
# is not guaranteed by the GUI device list).
if [ ! -b "/dev/$device" ]; then
    echo "Not a block device: /dev/$device" >&2
    exit 2
fi

whole_disk_name()
{
    local node="${1#/dev/}" parent type
    while [ -n "$node" ] && [ -b "/dev/$node" ]; do
        type=$(lsblk -ndo TYPE "/dev/$node" 2>/dev/null | head -1)
        [ "$type" = "disk" ] && { echo "$node"; return 0; }
        parent=$(lsblk -ndo PKNAME "/dev/$node" 2>/dev/null | head -1)
        [ -z "$parent" ] && break
        node="$parent"
    done
    echo "$node"
}

# Refuse to operate on the disk backing the running system or live medium.
# Compare whole-disk kernel names, following parent links so normal partitions
# and common mapper/LVM/LUKS stacks resolve to the same backing disk.
target_disk=$(whole_disk_name "/dev/$device")
for _mnt in / /boot /boot/efi /live/boot-dev /live/to-ram /lib/live/mount/medium /run/live/medium; do
    _src=$(findmnt -fn -o SOURCE "$_mnt" 2>/dev/null) || continue
    _src=${_src%%[*}              # strip btrfs subvol suffix, e.g. /dev/sda2[/@]
    [ -b "$_src" ] || continue
    _pdisk=$(whole_disk_name "$_src")
    if [ "$target_disk" = "$_pdisk" ]; then
        echo "Refusing to format /dev/$device: it backs the running system or live medium" >&2
        exit 3
    fi
done

# Reject unknown filesystem types up front. Otherwise format_partitions()
# treats an unknown format as a no-op "success" (its *) branch returns 0 and
# checkerrorcode reports OK) -- but for a non-"part" mode the disk has already
# been wiped/repartitioned by then, so fail closed before any disk writes.
case "$format" in
    vfat|ext4|ntfs|exfat) ;;
    *) echo "Invalid format: $format" >&2; exit 2 ;;
esac

echo device is $device format is $format
echo label is $label

clearpartions(){
	live-usb-maker gui partition-clear --color=off -t "$device"
}

unmount_partitions()
{
	local umount_device=$device*
	if [ -n "$(df |grep $device)" ]; then
	    umount -q /dev/$umount_device 2>/dev/null
	    sleep 1
	    checkerrorcode "unmount partitions"

	fi
}

##clear_partitions from live-usb-maker by James Bowlin (BitJam) for antiX
clear_partitions()
{
    local dev=/dev/$device
    local bytes
    bytes=$(parted --script $dev unit B print 2>/dev/null | sed -rn "s/^Disk.*: ([0-9]+)B$/\1/ip")
    local block_size=512
    local pt_size=$((17 * 1024))
    local pt_cnt=$((pt_size / block_size))

    local total_blocks=$((bytes / $block_size))
    #echo -e "Total bytes: $bytes\nTotal blocks: $total_blocks"

    # Clear out previous primary partition table
    dd if=/dev/zero of=$dev bs=$block_size count=$pt_cnt

    checkerrorcode "primary partition table clear"

    sleep 1

    # Clear out sneaky iso-hybrid partition table
    dd if=/dev/zero of=$dev bs=$block_size count=$pt_cnt seek=64

    checkerrorcode "iso-hybrid partition table clear"

	sleep 1

    [ -n "$bytes" ] || return
    local offset=$((total_blocks - $pt_cnt))

    # Clear out secondary gpt partition table
    dd conv=notrunc if=/dev/zero of=$dev bs=$block_size count=$pt_cnt seek=$offset


    checkerrorcode "secondary partition table clear"

    sleep 1
    # Tell kernel the partition table has changed
    echo "refresh partitions info $dev"
    partprobe -s $dev

    checkerrorcode "refresh partitions info"
    sleep 1
}

create_partition()
{
	local dev=/dev/$device
	local option
	unmount_partitions
	sleep 1

	local bytes
	bytes=$(lsblk --bytes --nodeps --noheadings --output SIZE $dev 2>/dev/null)
	bytes=$((bytes / 1))
	local mbrlimit=$((2000 * 1024 * 1024 * 1024)) # about 2TB - devices over this size are made to gpt
	local parttabletype

    echo "bytes $bytes limit $mbrlimit"

	if [ "$part" = "defaults" ]; then
		if (($bytes <= $mbrlimit)) ; then # the <= prevents a race condition
			parttabletype="msdos"
			option="primary"
			echo "making new dos partition table"
			sleep 1
		fi

		if (($bytes > $mbrlimit)) ; then
			parttabletype="gpt"
			echo "making new gpt partition table"
			sleep 1
		fi
	fi

	if [ "$part" = "gpt" ]; then
		parttabletype="gpt"
		echo "making new gpt partition table"
		sleep 1
	fi

	if [ "$part" = "msdos" ]; then
		parttabletype="msdos"
		option="primary"
		echo "making new msdos partition table"
		sleep 1
	fi

	#set default part name if file system label is empty
	if [ "$parttabletype" = "gpt" ]; then
		if [ -z "$label" ]; then
			option="$label"
		else
			option="primary"
		fi
	fi

	parted -s "$dev" mklabel "$parttabletype"

	checkerrorcode "making new partition table"
	sleep 1
	partprobe "$dev"
	checkerrorcode "refresh partitions info"
	sleep 1
	#create partition. $option is "primary" (msdos) or empty (gpt, no name);
	#quote it when present so it can't word-split, but keep the original
	#argument vector when empty so parted's gpt "no name" case is unchanged.
	if [ -n "$option" ]; then
		parted -s -a optimal "$dev" mkpart "$option" 1 100%
	else
		parted -s -a optimal "$dev" mkpart 1 100%
	fi
	checkerrorcode "create new partition"
	sleep 1

	partprobe $dev
	checkerrorcode "refresh partitions info"
	sleep 1
}


makeusb(){
live-usb-maker gui --format="$format" --color=off -t "$device"
}

labelusb(){

if [ -z "$label" ]; then
	return
fi

#ensure device is unmounted
if [ -n "$(df |grep $device)" ]; then
    umount -q /dev/"$device$partnum" 2>/dev/null
    checkerrorcode "unmount partitions"
fi



case $format in

	vfat) fatlabel /dev/"$device$partnum" "$label"  ;;

	ext4)  e2label /dev/"$device$partnum" "$label"  ;;

	ntfs)  ntfslabel /dev/"$device$partnum" "$label"  ;;

	exfat) exfatlabel /dev/"$device$partnum" "$label"  ;;

	*)	echo "unknown format, exiting"  ;;

esac


checkerrorcode "label partition"
}

partitionrefresh(){

	##need device not partition

	local refreshdevice partition_to_mark filter partition_dev mark

	partition_dev="$device$partnum"


	#if mmc device
	if [[ "$device" == *"mmc"* ]]; then
		refreshdevice=${device%p*}
		filter=$refreshdevice
		filter+="p"
		partition_to_mark=${partition_dev#$filter}
	fi

	#if device is nvme, then partnum=p1
	if [[ "$device" == *"nvme"* ]]; then
		refreshdevice=${device%p*}
		filter=$refreshdevice
		filter+="p"
		partition_to_mark=${partition_dev#$filter}
	fi
	#if device is sdXY,
	if [[ "$device" == *"sd"* ]]; then
		refreshdevice=${device//[0-9]/}
		partition_to_mark=${partition_dev#$refreshdevice}
	fi

echo "Device $device"
echo "Refresh Partitions $refreshdevice"
echo "partition to mark $partition_to_mark"

#if gpt partition, use GUID for part type
if [ -n "$(blkid /dev/$refreshdevice | grep PTTYPE=\"gpt\")" ]; then
	mark="EBD0A0A2-B9E5-4433-87C0-68B6B72699C7"
	echo "mark $mark"
fi

#if msdos partition, use hexadecimal for part type, per format

if [ -n "$(blkid /dev/$refreshdevice | grep PTTYPE=\"dos\")" ]; then
	case $format in

	      vfat) mark="c"  ;;

	exfat | ntfs) mark="7"  ;;

	         *)  ;; #don't mark

	esac
	echo "mark $mark"
fi


#mark partition based on format if formating ntfs(7), exfat(7), or fat32(b)
case $format in

				vfat | exfat | ntfs) sfdisk /dev/$refreshdevice $partition_to_mark --part-type $mark
									 sleep 1

									 checkerrorcode "Setting Partition Type"	;;



							  *)  ;; #don't mark

esac

partprobe "/dev/$refreshdevice"

checkerrorcode "refresh partition info"

}

cleanuplog(){
	local dest="/var/log/formatusb.log" old="/var/log/formatusb.log.old" uid dir src
	# Where the GUI wrote its log: a user session logs to its private runtime
	# dir (/run/user/<uid>); a root run logs to /run. Both are non-world-writable.
	# World-writable /tmp is deliberately NOT searched as root: even with symlinks
	# rejected, its content and size would be attacker-controllable. The caller's
	# uid comes from pkexec/sudo/login.
	if [[ "$PKEXEC_UID" =~ ^[0-9]+$ ]]; then
		uid="$PKEXEC_UID"
	elif [[ "$SUDO_UID" =~ ^[0-9]+$ ]]; then
		uid="$SUDO_UID"
	else
		uid=$(id -u "$(logname)" 2>/dev/null)
	fi
	local -a logdirs=()
	[[ "$uid" =~ ^[0-9]+$ && "$uid" -ne 0 ]] && logdirs+=("/run/user/${uid}")
	logdirs+=("/run")
	for dir in "${logdirs[@]}"; do
		[[ -d "$dir" ]] || continue
		src="${dir}/formatusb.log"
		# Only act on a real regular file, never a symlink, so the source
		# cannot be redirected at a root-owned file.
		[[ -f "$src" && ! -L "$src" ]] || continue
		# Rotate the destination, never following a symlink left there.
		if [[ -L "$dest" ]]; then
			rm -f -- "$dest"
		elif [[ -f "$dest" ]]; then
			mv -- "$dest" "$old"
		fi
		# -P: if the source is swapped for a symlink in a race, copy the link
		# itself rather than dereferencing it to another file.
		cp -P -- "$src" "$dest"
		return 0
	done
}

partnumber()
{
	#ensure some partnums when working on devices

if [ "$part" != "part" ]; then

	partnum="1"

	#if device is mmc, then partnum=p1

	if [[ "$device" == *"mmc"* ]]; then
		partnum="p1"
	fi

	#if device is nvme, then partnum=p1
	if [[ "$device" == *"nvme"* ]]; then
		partnum="p1"
	fi
fi
}

format_partitions(){

	echo "formatting partitions $device$partnum"

	#ensure device is unmounted
	unmount_partitions

	case $format in

	vfat) mkfs.fat -F 32 /dev/"$device$partnum"  ;;

	ext4)  mkfs.ext4 -F /dev/"$device$partnum"
		   change_ownership;;

	ntfs)  mkfs.ntfs -Q /dev/"$device$partnum"  ;;

	exfat) mkfs.exfat /dev/"$device$partnum" ;;

	*)	echo "unknown format, exiting" ;;



esac

checkerrorcode "format partition"

}

disable_automount()
{
mkdir -p /run/udev/rules.d

checkerrorcode "hide disk from udev"
echo 'SUBSYSTEM=="block", ENV{UDISKS_IGNORE}="1"' > /run/udev/rules.d/91-mx-udisks-inhibit.rules
AUTOMOUNT_DISABLED=1   # tells the cleanup trap to re-enable automount on any exit
udevadm control --reload
udevadm trigger --subsystem-match=block

}

enable_automount()
{
rm -f /run/udev/rules.d/91-mx-udisks-inhibit.rules

checkerrorcode "make disk visible to udev"
AUTOMOUNT_DISABLED=""   # disarm the cleanup trap's re-enable once done normally
udevadm control --reload
udevadm trigger --subsystem-match=block
}

change_ownership()
{
	local USERNAME mountpoint
	# Resolve the calling user from the polkit-provided uid (fall back to
	# logname) rather than trusting the controlling terminal alone.
	if [[ "$PKEXEC_UID" =~ ^[0-9]+$ ]]; then
		USERNAME=$(getent passwd "$PKEXEC_UID" | cut -d: -f1)
	fi
	[[ -z "$USERNAME" ]] && USERNAME=$(/usr/bin/logname)
	[[ -z "$USERNAME" ]] && return

	# Use a fresh, root-owned, random mount point (mode 0700) instead of a
	# predictable world-writable /tmp path an attacker could pre-create/symlink.
	mountpoint=$(mktemp -d /run/formatusb-mnt.XXXXXX) || return
	# Register the temp mount for the global cleanup trap so it is always torn
	# down, even if a step below fails and exits (checkerrorcode exits on error).
	TEMP_MOUNT="$mountpoint"

	# Check the mount immediately: if it fails, chown/chmod would silently act
	# on the empty temp dir and the new filesystem would never be owned by the
	# user. checkerrorcode aborts here on failure (and the trap cleans up).
	mount -t ext4 /dev/"$device$partnum" "$mountpoint"
	checkerrorcode "Mounting new partition to set ownership"

	chown "$USERNAME":users "$mountpoint"
	checkerrorcode "Changing owner of partition to $USERNAME:users"

	chmod -R 775 "$mountpoint"
	checkerrorcode "Changing permissions of partition to $USERNAME:users"

	# Verify the unmount succeeded before reporting success, so a failed umount
	# is not logged OK while leaving the filesystem mounted.
	umount -q "$mountpoint"
	checkerrorcode "Unmounting temporary mount point"
	rmdir "$mountpoint"
	checkerrorcode "Removing temporary mount point"
	TEMP_MOUNT=""
}

checkerrorcode()
{
	retval=$?
	local msg="$1"
	#echo "retval is $retval"
	if [ ! $retval = 0 ]; then
		echo "$msg" "ERRROR"
		exit "$retval"
	else
		echo "$msg" "OK"
	fi
}

# Always-run cleanup: re-enable automount and tear down the temp mount point on
# ANY exit, including checkerrorcode's mid-run exit. Preserves the exit status.
_cleanup()
{
	local rc=$?
	if [ -n "$TEMP_MOUNT" ]; then
		umount -q "$TEMP_MOUNT" 2>/dev/null
		rmdir "$TEMP_MOUNT" 2>/dev/null
		TEMP_MOUNT=""
	fi
	if [ -n "$AUTOMOUNT_DISABLED" ]; then
		rm -f /run/udev/rules.d/91-mx-udisks-inhibit.rules
		udevadm control --reload 2>/dev/null
		udevadm trigger --subsystem-match=block 2>/dev/null
		AUTOMOUNT_DISABLED=""
	fi
	exit "$rc"
}

main(){

trap _cleanup EXIT

#echo main launched
unmount_partitions
disable_automount

if [ "$part" = "part" ]; then
	format_partitions
else
	clear_partitions
	sleep 1
	create_partition
	sleep 1
	partnumber
	format_partitions
	sleep 1
fi
	labelusb
	sleep 1
	partitionrefresh
	cleanuplog
	enable_automount
}

main
