#!/bin/bash
#
# Enables a WiFi network connection.

# Exit on first failure.
set -e

print_help() {
  cat <<EOF
Usage: ${0##*/} [--help] --country COUNTRY --ssid SSID [--psk]
Enables a WiFi network connection.
  --help              Optional. Display this help and exit.
  --country COUNTRY   A two-digit country code, as of ISO 3166-1 alpha-2.
  --ssid SSID         The name (SSID) of the WiFi network.
  --psk               Optional. Authenticate with the WiFi network using a
                      password read from stdin. The password must be 8-63
                      characters in length.

Note: this script only stores the WiFi configuration and triggers its
activation, but it doesn't (and cannot) verify whether the device can actually
connect to the wireless network successfully. In that sense, enabling WiFi
means that the device will *attempt* to connect to the wireless network, and
otherwise fall back to the wired Ethernet connection.
EOF
}

# Parse command-line arguments.
while (( "$#" > 0 )); do
  case "$1" in
    --help)
      print_help
      exit
      ;;
    --country)
      if (( "$#" < 2 )); then
        shift
        break
      fi
      WIFI_COUNTRY="$2"
      shift # For flag name.
      shift # For flag value.
      ;;
    --ssid)
      if (( "$#" < 2 )); then
        shift
        break
      fi
      WIFI_SSID="$2"
      shift # For flag name.
      shift # For flag value.
      ;;
    --psk)
      # Read the password from stdin. Otherwise, prompt the user.
      if [[ -p /dev/stdin ]]; then
        WIFI_PSK="$(cat -)"
      else
        read -r -s -p 'Provide the WiFi password: ' WIFI_PSK
        echo
      fi
      shift # For flag name.
      ;;
    *)
      >&2 echo "Unknown flag/argument: $1"
      >&2 echo "Use the '--help' flag for more information"
      exit 1
      ;;
  esac
done
readonly WIFI_COUNTRY
readonly WIFI_SSID
readonly WIFI_PSK="${WIFI_PSK:-}"

if [[ -z "${WIFI_COUNTRY}" ]]; then
  >&2 echo 'Missing argument: COUNTRY'
  >&2 echo "Use the '--help' flag for more information"
  exit 1
fi

# According to ISO 3166-1 alpha-2, the country code has to contain 2 letters.
COUNTRY_LENGTH="$(echo -n "${WIFI_COUNTRY}" | wc --bytes)"
readonly COUNTRY_LENGTH
if (( "${COUNTRY_LENGTH}" != 2 )); then
  >&2 echo 'Invalid argument: COUNTRY'
  >&2 echo "Use the '--help' flag for more information"
  exit 1
fi

if [[ -z "${WIFI_SSID}" ]]; then
  >&2 echo 'Missing argument: SSID'
  >&2 echo "Use the '--help' flag for more information"
  exit 1
fi

# If a password is specified, it has to be 8-63 characters in length.
if [[ -n "${WIFI_PSK}" ]]; then
  PSK_LENGTH="$(echo -n "${WIFI_PSK}" | wc --bytes)"
  readonly PSK_LENGTH
  if (( "${PSK_LENGTH}" < 8 || "${PSK_LENGTH}" > 63 )); then
    >&2 echo 'Invalid PSK.'
    >&2 echo "Use the '--help' flag for more information"
    exit 1
  fi
fi

# Echo commands before executing them, by default to stderr.
set -x

# Exit on unset variable.
set -u

# Set country code (persists to /boot/firmware/cmdline.txt and applies via
# iw reg set).
raspi-config nonint do_wifi_country "${WIFI_COUNTRY}"

# Enable the wireless radio in case it was previously disabled.
nmcli radio wifi on

readonly CONNECTION_NAME='tinypilot-wlan0'

# Ensure exactly one tinypilot-wlan0 connection exists (singleton).
CONN_COUNT="$(nmcli --get-values connection.uuid \
  connection show "${CONNECTION_NAME}" 2>/dev/null \
  | wc -l | tr -d ' ')"
readonly CONN_COUNT

if (( CONN_COUNT == 0 )); then
  # No connection exists — create one. An SSID is required when creating a
  # wifi connection.
  nmcli connection add \
    type wifi \
    con-name "${CONNECTION_NAME}" \
    ifname wlan0 \
    wifi.ssid "${WIFI_SSID}" \
    connection.autoconnect yes \
    connection.autoconnect-priority 0
elif (( CONN_COUNT > 1 )); then
  # Duplicates — delete all and recreate.
  nmcli connection delete "${CONNECTION_NAME}" 2>/dev/null || true
  nmcli connection add \
    type wifi \
    con-name "${CONNECTION_NAME}" \
    ifname wlan0 \
    wifi.ssid "${WIFI_SSID}" \
    connection.autoconnect yes \
    connection.autoconnect-priority 0
fi

# Modify to WiFi client mode. Clear AP-specific fields (band, channel) and
# any static IP settings that may have been set previously.
nmcli connection modify "${CONNECTION_NAME}" \
  wifi.mode infrastructure \
  wifi.ssid "${WIFI_SSID}" \
  wifi.band "" \
  wifi.channel "" \
  ipv4.method auto \
  ipv4.addresses "" \
  ipv4.gateway "" \
  ipv4.dns ""

# Set security. For open networks (no password), remove the security setting
# entirely — NM rejects empty wifi-sec.key-mgmt values.
# Disable command tracing to avoid leaking the password.
set +x
if [[ -n "${WIFI_PSK}" ]]; then
  nmcli connection modify "${CONNECTION_NAME}" \
    wifi-sec.key-mgmt wpa-psk \
    wifi-sec.psk "${WIFI_PSK}" \
    wifi-sec.psk-flags 0
  set -x
else
  set -x
  nmcli connection modify "${CONNECTION_NAME}" \
    remove 802-11-wireless-security 2>/dev/null || true
fi

# Return as soon as NM accepts the activation request. With the default 30s
# wait, a wrong PSK or out-of-range AP would block the script and surface as a
# script-level error even though the configuration was saved successfully.
nmcli --wait 0 connection up "${CONNECTION_NAME}"
