#!/bin/bash
#
# Enable TinyPilot WiFi access point.

# Exit on first failure.
set -e

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

# Exit on unset variable.
set -u

if (( "${EUID}" != 0 )); then
  >&2 echo 'This script requires root privileges.'
  >&2 echo 'Please re-run with sudo:'
  >&2 echo "  sudo $0 $*"
  >&2 exit 1
fi

NETWORK_NAME='TinyPilotWiFi'
NETWORK_COUNTRY='US'

print_help() {
  cat <<EOF
Usage: ${0##*/} [--help] [--password NETWORK_PASSWORD] [--ssid NETWORK_NAME] [--country NETWORK_COUNTRY]
Enable TinyPilot WiFi access point.
  --help                      Display this help and exit.
  --password NETWORK_PASSWORD The network password. Must be 8–63 characters in
                              length. If not specified, you will be prompted for
                              a password.
  --ssid NETWORK_NAME         Optional. The network name. Defaults to: ${NETWORK_NAME}
  --country NETWORK_COUNTRY   Optional. The network country in ISO 3166-1 alpha-2 format.
                              See https://en.wikipedia.org/wiki/ISO_3166-1_alpha-2
                              Defaults to: ${NETWORK_COUNTRY}

Note: Running this script will clear any potentially existing custom static IP
      configuration and also any potentially existing WiFi settings that might
      have been set up previously.

EOF
}

# Parse command-line arguments.
NETWORK_PASSWORD=''
while (( "$#" > 0 )); do
  case "$1" in
    --help)
      print_help
      exit
      ;;
    --password)
      if (( "$#" < 2 )); then
        shift;
        break;
      fi
      NETWORK_PASSWORD="$2"
      shift # For flag name.
      shift # For flag value.
      ;;
    --ssid)
      if (( "$#" < 2 )); then
        shift;
        break;
      fi
      NETWORK_NAME="$2"
      shift # For flag name.
      shift # For flag value.
      ;;
    --country)
      if (( "$#" < 2 )); then
        shift;
        break;
      fi
      NETWORK_COUNTRY="$2"
      shift # For flag name.
      shift # For flag value.
      ;;
    *)
      >&2 echo "Unknown agrument: $1"
      >&2 print_help
      exit 1
      ;;
  esac
done
readonly NETWORK_NAME
readonly NETWORK_COUNTRY

# Validate command-line arguments.
if [[ -z "${NETWORK_PASSWORD}" ]]; then
  read -r -s -p 'Set a WiFi password: ' NETWORK_PASSWORD
fi
readonly NETWORK_PASSWORD

if [[ -z "${NETWORK_PASSWORD}" ]] || \
  (( "${#NETWORK_PASSWORD}" < 8 )) || \
  (( "${#NETWORK_PASSWORD}" > 63 )); then
  >&2 echo 'The WiFi password must be 8–63 characters in length'
  >&2 print_help
  exit 1
fi

if [[ -z "${NETWORK_NAME}" ]]; then
  >&2 echo "The network name can't be empty"
  >&2 print_help
  exit 1
fi

if (( "${#NETWORK_COUNTRY}" != 2 )); then
  >&2 echo "Invalid network country: ${NETWORK_COUNTRY}"
  >&2 echo 'The country must be a 2 character country code (e.g., US)'
  >&2 print_help
  exit 1
fi

# Set country code (persists to /boot/firmware/cmdline.txt and applies via
# iw reg set).
raspi-config nonint do_wifi_country "${NETWORK_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 "${NETWORK_NAME}" \
    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 "${NETWORK_NAME}" \
    connection.autoconnect yes \
    connection.autoconnect-priority 0
fi

# AP's static IP address in CIDR notation. NM uses this as the gateway
# address for wlan0 and as the subnet for its internal dnsmasq DHCP server.
readonly NETWORK_AP_IP_ADDRESS='192.168.50.1/24'

# Modify to AP mode. NM's shared mode automatically runs an internal dnsmasq
# instance for DHCP serving to AP clients.
# Avoid leaking the password in command tracing.
set +x
nmcli connection modify "${CONNECTION_NAME}" \
  wifi.mode ap \
  wifi.band bg \
  wifi.channel 7 \
  wifi.ssid "${NETWORK_NAME}" \
  wifi-sec.key-mgmt wpa-psk \
  wifi-sec.psk "${NETWORK_PASSWORD}" \
  wifi-sec.psk-flags 0 \
  ipv4.method shared \
  ipv4.addresses "${NETWORK_AP_IP_ADDRESS}" \
  ipv4.gateway "" \
  ipv4.dns ""
set -x

# Return as soon as NM accepts the activation request, so the script doesn't
# block on activation-layer problems when the configuration has already been
# saved successfully.
nmcli --wait 0 connection up "${CONNECTION_NAME}"

cat <<EOF
TinyPilot is now running a WiFi access point with the following details:
SSID: ${NETWORK_NAME}
Country: ${NETWORK_COUNTRY}
IP Range: ${NETWORK_AP_IP_ADDRESS}
EOF
