From 2ff72c0c5cf25dddb97372d41ce58e84e7680813 Mon Sep 17 00:00:00 2001
From: Sheldon Lee <sheldon@sheldonlee.com>
Date: Fri, 17 May 2024 15:52:58 +0800
Subject: [PATCH] Remove polkit privilege escalation and move old script to
 attach-usb-polkit

---
 .config/scripts/libvirt/attach-usb        | 24 +------
 .config/scripts/libvirt/attach-usb-polkit | 76 +++++++++++++++++++++++
 2 files changed, 77 insertions(+), 23 deletions(-)
 create mode 100755 .config/scripts/libvirt/attach-usb-polkit

diff --git a/.config/scripts/libvirt/attach-usb b/.config/scripts/libvirt/attach-usb
index c77395e..67207d7 100755
--- a/.config/scripts/libvirt/attach-usb
+++ b/.config/scripts/libvirt/attach-usb
@@ -1,34 +1,12 @@
 #!/bin/bash
 script_location="$(readlink -f "$0")"
 cd "$(dirname $script_location)"
-
-if [ "$(id -u)" -ne "0" ]; then
-	pkttyagent --process $$ &
-	pkexec --keep-cwd sh -c "\
-		WAYLAND_DISPLAY=\"$WAYLAND_DISPLAY\"\
-		DISPLAY=\"$DISPLAY\"\
-		XDG_RUNTIME_DIR=\"$XDG_RUNTIME_DIR\"\
-		\"$0\""
-	exit 0
-fi
-
-run_as_user() {
-	if [ -z "$PKEXEC_UID" ]; then
-		USER_ID=1000
-	else
-		USER_ID="$PKEXEC_UID"
-	fi
-	username="$(id -nu "$USER_ID")"
-	cmd="$1"
-	echo "$(su "$username" sh -c "source \"/home/$username/.bash_profile\"; $cmd")"
-}
-
 menu() {
 	if [ -z "$DISPLAY" ]; then
 		echo -e "$1" | fzf --prompt="$2"
 		return 0 
 	fi
-	run_as_user "echo -e '$1' | menucmd -p '$2'"
+	echo -e "$1" | menucmd -p "$2"
 }
 
 csv="$(virsh list | tail --line=+3 | tr ' ' '\n' | uniq | grep -v -e "^[[:space:]]*$" | tr '\n' ',')"
diff --git a/.config/scripts/libvirt/attach-usb-polkit b/.config/scripts/libvirt/attach-usb-polkit
new file mode 100755
index 0000000..c77395e
--- /dev/null
+++ b/.config/scripts/libvirt/attach-usb-polkit
@@ -0,0 +1,76 @@
+#!/bin/bash
+script_location="$(readlink -f "$0")"
+cd "$(dirname $script_location)"
+
+if [ "$(id -u)" -ne "0" ]; then
+	pkttyagent --process $$ &
+	pkexec --keep-cwd sh -c "\
+		WAYLAND_DISPLAY=\"$WAYLAND_DISPLAY\"\
+		DISPLAY=\"$DISPLAY\"\
+		XDG_RUNTIME_DIR=\"$XDG_RUNTIME_DIR\"\
+		\"$0\""
+	exit 0
+fi
+
+run_as_user() {
+	if [ -z "$PKEXEC_UID" ]; then
+		USER_ID=1000
+	else
+		USER_ID="$PKEXEC_UID"
+	fi
+	username="$(id -nu "$USER_ID")"
+	cmd="$1"
+	echo "$(su "$username" sh -c "source \"/home/$username/.bash_profile\"; $cmd")"
+}
+
+menu() {
+	if [ -z "$DISPLAY" ]; then
+		echo -e "$1" | fzf --prompt="$2"
+		return 0 
+	fi
+	run_as_user "echo -e '$1' | menucmd -p '$2'"
+}
+
+csv="$(virsh list | tail --line=+3 | tr ' ' '\n' | uniq | grep -v -e "^[[:space:]]*$" | tr '\n' ',')"
+
+vms="$(echo "$csv" | awk -F ',' -v n='3' '{
+for (i = 1; i+1 <= NF; i++) {
+	printf "%s ", $i
+	if (i % n == 0) {
+		printf "\n"
+	}
+}
+}')"
+
+if [ -z "$vms" ]; then
+	menu "No VMs running" "Error"
+	exit 1
+fi
+
+vm="$(menu "$vms" "VM")"
+vm_name=$(echo "$vm" | awk '{print $2}')
+
+directory="./usb-devices"
+
+item_list=""
+
+for file in "$directory"/*; do
+	item_list="${item_list}$(basename $file)\n"
+done
+
+# strip last newline
+end=${#item_list}-2
+item_list=${item_list:0:end}
+
+device="${directory}/$(menu "$item_list" "Device")"
+[ -z "$device" ] && exit 1
+
+output=$(virsh attach-device "$vm_name" "$device" 2>&1) && echo "Attached $device to $vm_name"
+if [ $? -ne 0 ]; then
+	virsh detach-device "$vm_name" "$device" &> /dev/null && echo "Detached $device from $vm_name"
+	if [ $? -ne 0 ]; then
+		attached_vm_name="$(echo "$output" | grep domain | rev | awk '{print $1}' | rev)"
+		virsh detach-device "$attached_vm_name" "$device" &> /dev/null && echo "Detached $device from $attached_vm_name"
+		virsh attach-device "$vm_name" "$device" &> /dev/null && echo "Attached $device to $vm_name"
+	fi
+fi