#!/bin/dash
# CLI find samba shares and create rox app
# jamesbond 2024 - GPLv3
# uses: nmblookup smbclient nbtscan mpscan mount.cifs

### configurations
APPTITLE="Find SMB Shares"
ROXAPPDIR=/usr/share/Shares
CONFIG=.config
MNTPOINT=mnt-point
DRIVE_ICON=/usr/share/pixmaps/midi-icons/drive48.png
DRIVE_ICON_MOUNTED=/usr/share/pixmaps/midi-icons/drivemounter48.png
DEFAULT_ROXOPT=
MYSELF="$(readlink -f "$0")"

### run-time
UIMODE=   # xdialog,dialog,plain, blank=autoconfigure
#SERVERS= # separate by spaces
SCANONLY= # if true, don't bother to create roxapp
DEV=${DEV:-wlan0}
#SUBNET=  # 192.168.0.4/24
NETMASK=
PREFIX=
IPRANGE=
ROXOPT=



########## UI helpers #########

# $1-info
msgbox() {
	case "$UIMODE" in
		xdialog) Xdialog --title "$APPTITLE" --msgbox "$1" 0 0 10000 ;; 
		dialog)  dialog --backtitle "$APPTITLE" --msgbox "$1" 0 0 > /dev/stderr ;;
		plain)   echo "$1" > /dev/stderr ;;
	esac
}

# $1-info $2-timeout (in ms)
infobox() {
	local timeout=${2:-10000}
	case "$UIMODE" in
		xdialog) Xdialog --title "$APPTITLE" --infobox "$1" 0 0 $timeout ;;
		dialog)  dialog --backtitle "$APPTITLE" --no-cancel --pause "$1" 12 60 $(( timeout/1000 )) > /dev/stderr ;;
		plain)   echo "$1" > /dev/stderr ;;
	esac
}


# $1-text
yesno() {
	local yesno
	case "$UIMODE" in
		xdialog) Xdialog --title "$APPTITLE" --yesno "$1" 0 0 ;;
		dialog)  dialog --title "$APPTITLE" --yesno "$1" 0 0 > /dev/stderr ;;
		plain)   read -p "$1 [yes/no]: " yesno > /dev/stderr < /dev/tty
		         case "$yesno" in
					yes|Yes|YES) return 0 ;;
					no|No|NO|"") return 1 ;;
		         esac
		         ;;
	esac
}

# $1-prompt, $2-current value
inputbox() {
	local input
	case "$UIMODE" in
		xdialog) Xdialog --title "$APPTITLE" --stdout --inputbox "$1" 0 0 "$2" ;;
		dialog)  dialog --title "$APPTITLE" --stdout --inputbox "$1" 0 0 "$2"  ;;	
		plain)   read -p "$1: " input > /dev/stderr < /dev/tty
				 echo "$input"
		         ;;
	esac
}

# $1-text
splash() {
	if [ "$XPID" ]; then
		kill $XPID
		XPID=
	fi

	[ "$1" ] &&
	case "$UIMODE" in
		xdialog) Xdialog --title "$APPTITLE" --no-buttons --infobox "$1" 0 0 1000000 & XPID=$! ;;
		dialog)  dialog --backtitle "$APPTITLE" --infobox "$1" 0 0 > /dev/stderr ;;
		plain)   echo "$1" > /dev/stderr ;;
	esac
}

# $1-text
die() {
	infobox "$1"
	exit
}

# true if not root, false if root
not_root() {
	if [ $(id -u) -ne 0 ]; then
		infobox "You must be root do to this."
		return 0
	fi
	return 1
}




############### IPv4 utilities ###############

# in: $1-decimal
printhex() {
	dc -O 16 -e "$1 p"
}

# slow unoptimised power function
# $1-base $-exponent
# out: stdout
pow() {
	local base=$1 exp=$2 accum=1
	while [ $exp -gt 0 ]; do
		: $(( accum *= base ))
		: $(( exp = exp - 1 ))
	done
	echo $accum
}

# in: $1-dotted IPv4 address (optionally cidr)
# out: stdout
ipv4_to_decimal() {
	local p ipval=0
	OIFS="$IFS"; IFS=.
	for p in ${1%/*}; do
		: $(( ipval = ipval*256 + p ))
		#echo $p, $ipval
	done
	IFS="$OIFS"
	#printhex $ipval
	echo $ipval
}

# in: $1-decimal
# out: stdout dotted IPv4 address
decimal_to_ipv4() {
	local count=4 divisor rem=$1 out=""
	while [ $count -gt 0 ]; do
		divisor=$(pow 2 $(( 8 * (count-1) )) )
		#echo $divisor
		out="${out}.$(( rem / divisor ))"
		: $(( rem = rem % divisor ))
		: $(( count = count - 1 ))
	done
	echo "${out#.}"
}

# in: $1-prefix
# out: stdout decimal
prefix_to_netmask() {
	local a=$(pow 2 32)
	local b=$(pow 2 $(( 32 - $1 )))
	local netmask=$(( a - b ))
	#printhex $netmask
	echo $netmask
}

# in: $1-netmask decimal
# out: prefix
netmask_to_prefix() {
	local a=$(pow 2 32)
	local bcast=$(( a - $1 - 1))
	local prefix=32 val=$bcast
	while [ $val -gt 0 ]; do
		: $(( val = val / 2 ))
		: $(( prefix = prefix - 1 ))
	done
	echo $prefix
}

# in: $1-IP CIDR (dotted/prefix)
# out: stdout subnet address and netmask (dotted)
get_subnet_netmask() {
	local ip=${1%/*} prefix=${1#*/}
	local ipval=$(ipv4_to_decimal $ip)
	local netmask=$(prefix_to_netmask $prefix)
	local net=$(( ipval & netmask ))
	echo $(decimal_to_ipv4 $net) $(decimal_to_ipv4 $netmask) 
}

# in: $1-subnet $2-netmask (dotted)
# out: stdout first IP and broadcast address (dotted)
get_iprange() {
	local net=$(ipv4_to_decimal $1)
	local mask=$(ipv4_to_decimal $2)
	
	local a=$(pow 2 32)
	local bcast=$(( a - mask - 1))
	local last=$(( net + bcast ))
	local first=$(( net + 1 ))
	[ "$bcast" -eq 0 ] && first=$net # special case

	echo $(decimal_to_ipv4 $first)-$(decimal_to_ipv4 $last)
}

# tests
#ipcidr=192.168.0.4/24
#ipv4_to_decimal ${ipcidr%/*}
#decimal_to_ipv4 3232235524
#prefix_to_netmask ${ipcidr#*/}
#netmask_to_prefix $(prefix_to_netmask ${ipcidr#*/})
#get_subnet_netmask $ipcidr
#get_iprange $(get_subnet_netmask $ipcidr)
#exit





############### auto-configure ###############

# in: DEV SUBNET
# out: SUBNET NETMASK PREFIX IPRANGE
autoconfigure_SUBNET() {
	local ipcidr
	if [ -z "$SUBNET" ]; then
		ipcidr=$(ip address show dev $DEV | sed '/inet /!d; s/.*inet //; s/ .*//')
		SUBNET="$(get_subnet_netmask $ipcidr)"
	else
		SUBNET="$(get_subnet_netmask $SUBNET)"
	fi
	NETMASK=${SUBNET#* }
	PREFIX=$(netmask_to_prefix $(ipv4_to_decimal $NETMASK))
	SUBNET=${SUBNET% *}
	IPRANGE=$(get_iprange $SUBNET $NETMASK)
	#echo $SUBNET/$PREFIX $NETMASK $IPRANGE
}

autoconfigure_UIMODE() {
	[ "$UIMODE" ] && return
	[ "$DISPLAY" ] && UIMODE=xdialog || UIMODE=dialog
}



############### SMB utilities ###############

# in: $1-subnet $2-prefix
scan_nbtscan() {
	nbtscan -er $1/$2 2> /dev/null
}

# in: $1-range $2-port #3-timeout
scan_mpscan() {
	local first=${1%-*} last=${1#*-}
	local timeout=
	[ "$3" ] && timeout="-t $3"
	mpscan $timeout -p $2 $first - $last
}

# in: $1-IP
# out: stdout machine name (if any)
get_server_name() {
	local name
	name="$(nmblookup -A $1 | awk '/<20>/ { sub(/^[ \t]+/, "", $1); print $1; exit }' 2> /dev/null)"
	if [ "$name" ]; then
		echo "$name"
	else
		smbclient -g -N -L $1 2> /dev/null |
		awk -F "|" '/^Server/ { if ($2) print $2 }; /^Workgroup/ { if ($3) print $3 }'
	fi
}

# in: $1-IP $2-option (pass -N for open shares, pass USER/PASSWD env for protected shares)
# out: stdout open shares if any
find_shares() {
	smbclient $2 -g -L $1 2> /dev/null |
	awk -F "|" '/^Disk/ { if ($2) print $2 };'
}

# in: $1-IP
find_open_shares() {
	find_shares "$1" "-N"
}

# in: $1-IP $2-user $3-passwd
find_protected_shares() {
	USER="$2" PASSWD="$3" find_shares "$1"
}

# $1-name
# out: stdout
name_to_ip() {
	{
		getent hosts "$1"
		nmblookup "$1" | grep -v failed
	} 2> /dev/null | awk '{print $1}' | sort | uniq
}


############### Find SMB functions ###############

# in: SUBNET PREFIX IPRANGE $1-timeout
# out: stdout
scan_all() {
	local timeout=
	[ $# -gt 0 ] && timeout="$1" && shift
	while [ $# -gt 0 ]; do
		case $1 in
			nbt)   scan_nbtscan $SUBNET $PREFIX         | awk '{print $1}' ;;
			mp139) scan_mpscan  $IPRANGE 139 "$timeout" | awk '{print $2}' ;;
			mp445) scan_mpscan  $IPRANGE 445 "$timeout" | awk '{print $2}' ;;
		esac 
		shift
	done 2> /dev/null | sort | uniq
}

# $1-servers
# out: stdout 
get_shares() {
	local shares xxshares ip uu pp ss
	shares="IP|Name|Share|User|Passwd"
	for ip in $1; do
		splash "Getting server name for $ip ..."
		name="$(get_server_name $ip)"

		unset uu pp
		if yesno "$ip ($name): Scan for protected shares?"; then
			uu=$(inputbox "$ip ($name): Username?")
			if [ "$uu" ]; then
				pp=$(inputbox "$ip ($name): Password?")
			fi
		fi

		splash "Finding shares from $ip ($name)..."
		xxshares="$(find_open_shares $ip "-N" | 
		while read -r ss; do
			echo "$ip|$name|$ss"
		done)"
		[ "$xxshares" ] &&shares="$shares
$xxshares"

		if [ "$uu" ]; then
		xxshares="$(find_protected_shares $ip "$uu" "$pp" | while read -r ss; do
			echo "$ip|$name|$ss|$uu|$pp"
		done)"
			[ "$xxshares" ] && shares="$shares
$xxshares"	
		fi
	done
	splash ""
	echo "$shares"
}

# $1-shares
create_rox_app() {
	local ss folder
	echo "$1" | while read -r ss; do
		OIFS="$IFS" IFS="|"
		set -- $ss # $1-ip $2-name $3-share $4-user $5-pass
		IFS="$OIFS"	

		if yesno "Create ROX App for $1 ($2) - $3?"; then
			folder="$2++$3"
			[ -z "$2" ] && folder="$1$folder"
			folder="$ROXAPPDIR/$folder"
			mkdir -p "$folder"
			
			# handler, icon, config
			ln -sf "$MYSELF" "$folder"/AppRun
			ln -sf $DRIVE_ICON "$folder"/.DirIcon
			cat > "$folder"/$CONFIG << EOF
IP='$1'
NAME='$2'
SHARE='$3'
USER='$4'
PASS='$(echo "$5" | base32)'
ROXOPT='$ROXOPT'
EOF
			# and right-click
			cat > "$folder"/AppInfo.xml << EOF
<?xml version="1.0"?>
<AppInfo>
	<Summary>SMB share</Summary>
	<About>
		<Purpose>Samba share mount</Purpose>
		<Authors>jamesbond</Authors>
	</About>
	<AppMenu>
		<Item option="unmount" icon="gtk-harddisk"><Label>Unmount</Label></Item>
		<Item option="edit-config" icon="gtk-edit"><Label>Edit Mount Settings...</Label></Item>
	</AppMenu>
</AppInfo>
EOF

		fi		
	done
}


############### ROX APP personality ###############

# $1-mount/unmount
roxapp() {
	autoconfigure_UIMODE
	
	local APPDIR="$(dirname "$0")"
	! [ -r "$APPDIR/$CONFIG" ] && die "$APPDIR improperly configured".
	. "$APPDIR/$CONFIG" # IP, NAME, SHARE, USER, PASS, ROXOPT
	case "$1" in
		""|mount)       do_mount ;;
		unmount|umount) do_unmount ;;
		edit-config)    $EDITOR "$APPDIR/$CONFIG" ;;
	esac
	exit 0
}

do_mount() {
	local ipaddr options
	# 1. make mount point if it does not exist
	! [ -d "$APPDIR/$MNTPOINT" ] && mkdir -p "$APPDIR/$MNTPOINT"

	# 2. see if it's already mounted
	if mountpoint -q "$APPDIR/$MNTPOINT"; then
		rox -d "$APPDIR/$MNTPOINT"
		exit 0
	fi

	# 3. else, mount it
	# 3.1 map name to IP address
	ipaddr=$(name_to_ip "$NAME")
	[ -z "$ipaddr" ] && ipaddr=$IP # if fail uses cache

	# 3.2 construct options
	options="ip=$ipaddr"
	[ "$USER" ] && options="$options,user=$USER"
	[ "$PASS" ] && options="$options,pass=$(echo "$PASS"|base32 -d)"
	[ "$DEFAULT_ROXOPT" ] && options="$options,$DEFAULT_ROXOPT"
	[ "$ROXOPT" ] && options="$options,$ROXOPT"

	# 3.3 mount it
	#echo $options
	if mount.cifs "//$NAME/$SHARE" "$APPDIR/$MNTPOINT" -o "$options" ||
	   mount.cifs "//$NAME/$SHARE" "$APPDIR/$MNTPOINT" -o "$options,vers=1.0"; then
		rox -d "$APPDIR/$MNTPOINT"
		ln -sf "$DRIVE_ICON_MOUNTED" "$APPDIR"/.DirIcon
		rox -x "$ROXAPPDIR"
	else
		infobox "Failed to mount $APPDIR."
	fi
}

do_unmount() {
	# 1. ignore if mount point does not exist
	! [ -d "$APPDIR/$MNTPOINT" ] && exit 1
	rox -D "$APPDIR/$MNTPOINT"
	ln -sf "$DRIVE_ICON" "$APPDIR"/.DirIcon
	rox -x "$ROXAPPDIR"
	
	# 2. ignore if not mounted
	! mountpoint -q "$APPDIR/$MNTPOINT" && exit 2

	# 3. unmount
	umount -l "$APPDIR/$MNTPOINT"
	infobox "$APPDIR un-mounted."
}


##################################

usage() {
	cat << EOF
${0##*/} [flags]

Flags:
--dev dev         : use subnet from device (default: $DEV)
--subnet subnet   : subnet to search (CIDR) (default: determined from $DEV)
--roxopt options  : additional options when creating roxapp
--uimode mode     : xdialog, dialog, or plain
--servers servers : bypass searching and use the given servers (separate by spaces)

EOF
	exit 0
}



############### MAIN ###############

# Must be root
if not_root; then
	exec gtksu "$APPTITLE" "$0" "$@"
fi

# 0. Dual-personality - check if we're invoked as rox-app, or as normal
case "$0" in
	*/AppRun) roxapp "$@" ;; # does not return
esac

# 1. params
while [ $# -gt 0 ]; do
	case "$1" in
		-h|--help) usage ;;
		--subnet) shift; SUBNET="$1" ;;
		--dev)    shift; DEV="$1"    ;;
		--roxopt) shift; ROXOPT="$1" ;;
		--uimode) shift; UIMODE="$1" ;;
		--scanonly) SCANONLY=yes     ;;
		--servers)  shift; SERVERS="$1" ;;
	esac
	shift
done

# 2. autoconfigure
autoconfigure_UIMODE
autoconfigure_SUBNET

# 3. scan
if [ -z "$SERVERS" ]; then
	splash "Scanning Network for SMB Shares, please wait ..."
	servers=$(scan_all "" nbt mp139 mp445)
	splash ""
else
	servers="$SERVERS"
fi

# 4. get shares
shares="$(get_shares "$servers" | sed '1d')"
if [ "$shares" ]; then
	msgbox "
$(echo "$shares"|wc -l) shares found:
=====================================
$shares
"
	# 5. create rox app
	[ -z "$SCANONLY" ] &&
	create_rox_app "$shares"
fi
