#!/bin/sh

set -e

usage () {
  if [ "$#" -eq "0" ]; then
    usage set_hostname enable_ssh set_wlan set_keymap set_timezone
    return 0
  fi
  echo "Usage: "
  for arg in "$@"; do
    case "$arg" in
      set_hostname)
        echo "  $0 set_hostname HOSTNAME"
        ;;
      import_ssh_id)
        echo "  $0 import_ssh_id USERID1 [USERID2]..."
        ;;
      enable_ssh)
        echo "  $0 enable_ssh [-k|--key-only]|[-p|--pass-auth] [-d|--disabled] [KEY_LINE1 [KEY_LINE2]...]"
        ;;
      set_wlan)
        echo "  $0 set_wlan [-h|--hidden] [-p|--plain] SSID [PASS [COUNTRY]]"
        ;;
      set_wlan_country)
        echo "  $0 set_wlan_country COUNTRY]"
        ;;
      set_keymap)
        echo "  $0 set_keymap KEYMAP"
        ;;
      set_timezone)
        echo "  $0 set_timezone TIMEZONE"
        ;;
    esac
  done
  }

set_hostname () (
  if [ "$#" -ne 1 ]; then
    usage set_hostname
    exit 1
  fi
  HOSTNAME="$1"
  raspi-config nonint do_hostname "$HOSTNAME"
  echo "$HOSTNAME" > /etc/hostname
)

FIRSTUSER=$(getent passwd 1000 | cut -d: -f1)
FIRSTUSERHOME=$(getent passwd 1000 | cut -d: -f6)
SSH_DIR="$FIRSTUSERHOME/.ssh"
AUTHORISED_KEYS_FILE="$SSH_DIR/authorized_keys"

add_ssh_keys () (
  if ! [ -d "$SSH_DIR" ]; then
    install -o "$FIRSTUSER" -g "$FIRSTUSER" -m 700 -d "$SSH_DIR"
  fi
  for key in "$@"; do
    echo "$key" >> "$AUTHORISED_KEYS_FILE"
  done
  if [ -f "$AUTHORISED_KEYS_FILE" ]; then
    chmod 600 "$AUTHORISED_KEYS_FILE"
    chown "$FIRSTUSER:$FIRSTUSER" "$AUTHORISED_KEYS_FILE"
  fi
)

enable_ssh () (
  ENABLE=1
  KEY_ONLY_SED_STR='s/^[#\s]*PasswordAuthentication\s\+\S\+$/PasswordAuthentication no/'
  PASSAUTH_SED_STR='s/^[#\s]*PasswordAuthentication\s\+\S\+$/PasswordAuthentication yes/'
  for arg in "$@"; do
    if [ "$arg" = "-k" ] || [ "$arg" = "--key-only" ]; then
      sed -i "$KEY_ONLY_SED_STR" /etc/ssh/sshd_config
    elif [ "$arg" = "-p" ] || [ "$arg" = "--pass-auth" ]; then
      sed -i "$PASSAUTH_SED_STR" /etc/ssh/sshd_config
    elif [ "$arg" = "-d" ] || [ "$arg" = "--disabled" ]; then
      ENABLE=0
    else
      add_ssh_keys "$arg"
    fi
  done
  if [ "$ENABLE" = 1 ]; then
    systemctl -q enable ssh
  fi
)

set_wlan_country () (
  if [ "$#" -ne 1 ]; then
    usage set_wlan_country
    exit 1
  fi
  # shellcheck disable=SC2030
  COUNTRY="$1"
  raspi-config nonint do_wifi_country "$COUNTRY"
)


set_wlan () (
  HIDDEN=0
  PLAIN=0
  SCRIPT='/var/lib/raspberrypi-sys-mods/set-wlan'
  for arg in "$@"; do
    # shellcheck disable=SC2031
    if [ "$arg" = "-h" ] || [ "$arg" = "--hidden" ]; then
      HIDDEN=1
    elif [ "$arg" = "-p" ] || [ "$arg" = "--plain" ]; then
      PLAIN=1
    elif [ -z "${SSID+set}" ]; then
      SSID="$arg"
    elif [ -z "${PASS+set}" ]; then
      PASS="$arg"
    elif [ -z "${COUNTRY+set}" ]; then
      COUNTRY="$arg"
    else
      usage set_wlan
      exit 1
    fi
  done
  if [ -z "${SSID+set}" ]; then
    usage set_wlan
    exit 1
  fi
  if [ -e /boot/wpa_supplicant.conf ]; then
    echo "Ignoring network configuration: /boot/wpa_supplicant.conf exists"
    exit 0
  fi

  if [ -n "$COUNTRY" ]; then
    set_wlan_country "$COUNTRY"
  fi

  # Replace ' with '\'' in SSID and PASS to support single-quote characters
  SSID=$(printf '%s' "$SSID" | sed "s/'/'\\\\''/g")
  PASS=$(printf '%s' "$PASS" | sed "s/'/'\\\\''/g")

  mkdir -p "$(dirname "$SCRIPT")"
  # shellcheck disable=SC2094
  cat <<- EOF > "$SCRIPT"
	#!/bin/sh
	COUNTER=0
	while [ "\$COUNTER" -lt 10 ]; do
	  COUNTER=\$((COUNTER + 1))
	  if raspi-config nonint do_wifi_ssid_passphrase  '$SSID' '$PASS' '$HIDDEN' '$PLAIN'; then
	    break
	  fi
	  sleep 5
	done
	systemctl stop set-wlan.timer
	systemctl disable set-wlan.timer
	rm -f "\$0"
	rm -f /etc/systemd/system/set-wlan.timer
	rm -f /etc/systemd/system/set-wlan.service
	rmdir --ignore-fail-on-non-empty '$(dirname "$SCRIPT")'
	EOF
  chmod 700 "$SCRIPT"

  cat <<- EOF > /etc/systemd/system/set-wlan.timer
		[Unit]
		Description=Configure WLAN for rpi-imager
		[Timer]
		OnBootSec=1
		OnUnitActiveSec=10
		[Install]
		WantedBy=timers.target
	EOF

  cat <<- EOF > /etc/systemd/system/set-wlan.service
		[Unit]
		Description=Configure WLAN for rpi-imager
		After=NetworkManager.service dhcpcd.service
		ConditionPathIsDirectory=|/run/wpa_supplicant
		ConditionPathExists=|/run/dhcpcd.pid
		[Service]
		Type=oneshot
		ExecStart=$SCRIPT
	EOF
  ln -f -s /etc/systemd/system/set-wlan.timer \
  /etc/systemd/system/timers.target.wants/set-wlan.timer
)

import_ssh_id () (
  SCRIPT='/var/lib/raspberrypi-sys-mods/import-ssh'
  if [ "$#" -eq 0 ]; then
    usage import_ssh_id
    exit 1
  fi
  if ! command -v ssh-import-id > /dev/null; then
    echo "ssh-import-id not available"
    exit 1
  fi
  mkdir -p "$(dirname "$SCRIPT")"
  # shellcheck disable=SC2094
  cat <<- EOF > "$SCRIPT"
	#!/bin/sh
	COUNTER=0
	while [ "\$COUNTER" -lt 10 ]; do
	  COUNTER=\$((COUNTER + 1))
	  if runuser -u \$(getent passwd 1000 | cut -d: -f1) -- ssh-import-id $@; then
	    break
	  fi
	  sleep 5
	done
	systemctl stop import-ssh.timer
	systemctl disable import-ssh.timer
	rm -f "\$0"
	rm -f /etc/systemd/system/import-ssh.timer
	rm -f /etc/systemd/system/import-ssh.service
	rmdir --ignore-fail-on-non-empty "$(dirname "$SCRIPT")"
	EOF
  chmod 700 "$SCRIPT"

  cat <<- EOF > /etc/systemd/system/import-ssh.timer
		[Unit]
		Description=Import SSH keys using ssh-import-id
		[Timer]
		OnBootSec=1
		OnUnitActiveSec=10
		[Install]
		WantedBy=timers.target
	EOF

  cat <<- EOF > /etc/systemd/system/import-ssh.service
		[Unit]
		Description=Import SSH keys using ssh-import-id
		After=network-online.target userconfig.service
		[Service]
		Type=oneshot
		ExecStart=$SCRIPT
	EOF
  ln -f -s /etc/systemd/system/import-ssh.timer \
  /etc/systemd/system/timers.target.wants/import-ssh.timer
)

set_keymap () (
  if [ "$#" -ne 1 ]; then
    usage set_keymap
    exit 1
  fi
  raspi-config nonint do_configure_keyboard "$1"
)

set_timezone () (
  if [ "$#" -ne 1 ]; then
    usage set_timezone
    exit 1
  fi
  raspi-config nonint do_change_timezone "$1"
)

if [ "$#" -eq 0 ]; then
  echo "No command specified"
  usage
  exit 1
fi

command="$1"; shift
case "$command" in
  set_hostname|import_ssh_id|enable_ssh|set_wlan_country|set_wlan|set_keymap|set_timezone)
    "$command" "$@"
    ;;
  *)
    echo "Unsupported command: $command"
    usage
    exit 1
    ;;
esac