#!/bin/sh
# wifi-connect - scan, pick a network and connect, using iwd and dhcpcd
#
#   wifi-connect                  scan and choose from a list
#   wifi-connect <ssid>           connect by name, asks for the passphrase if needed
#   wifi-connect <ssid> <psk>     without asking
#   wifi-connect --status         what the wireless device is doing
#   wifi-connect --forget <ssid>  drop a saved network
#
# iwd remembers every network it joins and reconnects to it after a reboot
set -u
PATH=/usr/bin

if [ -t 1 ]; then
	R=$(printf '\033[0m') G=$(printf '\033[32m') W=$(printf '\033[36m') E=$(printf '\033[31m')
else
	R= G= W= E=
fi
say() { printf '%s::%s %s\n' "$W" "$R" "$*"; }
ok()  { printf '%s::%s %s\n' "$G" "$R" "$*"; }
die() { printf '%s!!%s %s\n' "$E" "$R" "$*" >&2; exit 1; }

if [ "$(id -u)" != 0 ]; then
	command -v doas >/dev/null 2>&1 && exec doas "$0" "$@"
	die "run it as root"
fi
command -v iwctl >/dev/null 2>&1 || die "iwd is not installed: scraps add iwd"

# iwd runs the radio and dhcpcd the address, both start at boot from now on
for s in iwd dhcpcd; do
	[ -e "/etc/dinit.d/boot.d/$s" ] || dinitctl enable "$s" >/dev/null 2>&1 || \
		die "could not start $s, see: dinitctl status $s"
done

iface=""
for i in /sys/class/net/*; do
	[ -d "$i/wireless" ] || [ -L "$i/phy80211" ] || continue
	iface=${i##*/}; break
done
[ -n "$iface" ] || die "no wireless device found, it may need firmware: scraps add linux-firmware"

for k in /sys/class/rfkill/rfkill*; do
	[ -w "$k/soft" ] && [ "$(cat "$k/soft")" = 1 ] && echo 0 >"$k/soft"
done

plain() { sed 's/\x1b\[[0-9;]*m//g'; }

case "${1:-}" in
--status)
	iwctl station "$iface" show | plain
	ip -4 addr show dev "$iface" | sed -n 's/^ *inet \([^ ]*\).*/  address  \1/p'
	exit 0 ;;
--forget)
	[ -n "${2:-}" ] || die "usage: wifi-connect --forget <ssid>"
	iwctl known-networks "$2" forget && ok "forgot $2"
	exit $? ;;
-h|--help)
	sed -n '2,9p' "$0" | sed 's/^# \{0,1\}//'
	exit 0 ;;
esac

ssid=${1:-}
psk=${2:-}
if [ -z "$ssid" ]; then
	say "scanning on $iface"
	iwctl station "$iface" scan >/dev/null 2>&1
	sleep 3
	# name, security, signal, the name may have spaces so it is everything before the last two
	list=$(iwctl station "$iface" get-networks | plain | awk '
		/^-+$/ { d++; next }
		d >= 2 && NF >= 3 {
			s = ($1 == ">") ? 2 : 1; n = $s
			for (i = s + 1; i <= NF - 2; i++) n = n " " $i
			printf "%s\t%s\t%s\n", n, $(NF-1), $NF
		}')
	[ -n "$list" ] || die "no networks found"
	printf '%s\n' "$list" | awk -F'\t' '{ printf "  %2d  %-32s %-6s %s\n", NR, $1, $2, $3 }'
	printf 'network number: '
	read -r pick
	ssid=$(printf '%s\n' "$list" | awk -F'\t' -v n="$pick" 'NR == n { print $1 }')
	[ -n "$ssid" ] || die "no network number $pick"
	sec=$(printf '%s\n' "$list" | awk -F'\t' -v n="$pick" 'NR == n { print $2 }')
	if [ "$sec" != open ] && ! iwctl known-networks list | plain | grep -q "^ *$ssid "; then
		stty -echo 2>/dev/null; printf 'passphrase for %s: ' "$ssid"; read -r psk; stty echo 2>/dev/null; echo
	fi
fi

say "connecting to $ssid"
if [ -n "$psk" ]; then
	iwctl --passphrase "$psk" station "$iface" connect "$ssid" || die "could not connect to $ssid"
else
	iwctl station "$iface" connect "$ssid" || die "could not connect to $ssid"
fi

n=0
while [ "$n" -lt 20 ]; do
	addr=$(ip -4 addr show dev "$iface" | sed -n 's/^ *inet \([^ ]*\).*/\1/p' | head -1)
	[ -n "$addr" ] && { ok "connected to $ssid, address $addr"; exit 0; }
	sleep 1; n=$((n + 1))
done
die "connected to $ssid but got no address, see: dinitctl status dhcpcd"
