#!/bin/bash
#
# Copyright (c) 2020 Bubble, Inc.  All rights reserved. For personal (non-commercial) use, see license: https://getbubblenow.com/bubble-license/
#
# Connect an Ubuntu 20.04 system to a Bubble
# You must run this as a user who has "sudo" privileges
#
# Usage:
#
#     ubuntu_connect_bubble [hostname]
#
#   hostname : bubble to connect this Ubuntu system to
#              If not specified, an env var will be used (see below)
#              If neither BUBBLE_HOST (hostname) or BUBBLE_API (URL) are specified,
#                then you'll be asked to enter the Bubble hostname
#
# Environment Variables
#
#   BUBBLE_HOST      : hostname of the Bubble to use. The API base URI will be https://<hostname>/api
#   BUBBLE_API       : full base URL of the Bubble API to use. Typically this ends with /api
#   BUBBLE_USER      : account to use
#   BUBBLE_PASS      : password for account
#   BUBBLE_CERT_WIZ  : Set this to 'true' to run the Ubuntu certificate wizard, even if it looks
#                      like the certificate is already installed
#
# This command will:
#  * Disconnect from any current Bubble
#  * Check if the device already exists for the Bubble. If not, creates the device
#  * Check if we have a local copy of the certificate and vpn.conf files. If not, download them
#  * Ensure routing configuration is correct
#  * Install the certificate (this step requires your interaction)
#  * Install the VPN config
#  * Start the VPN
#
function die() {
  echo 1>&2 "$0: fatal error: ${1}"
  exit 1
}

# from https://stackoverflow.com/a/34407620
function uriencode { jq -nr --arg v "$1" '$v|@uri'; }

SCRIPT_DIR="$(cd "$(dirname "${0}")" && pwd)"

INSTALL_PACKAGES=""
if [[ -z "$(which wg)" ]] ; then
  echo "Adding WireGuard to install list ..."
  INSTALL_PACKAGES="${INSTALL_PACKAGES} wireguard"
fi

if [[ -z "$(which bget)" ]] ; then
  export PATH=${PATH}:${SCRIPT_DIR}
  if [[ -z "$(which bget)" ]] ; then
    die "bget command not found, even after adding ${SCRIPT_DIR} to PATH"
  fi
fi

if [[ -z "$(which java)" ]] ; then
  echo "Adding Java 11 to install list ..."
  INSTALL_PACKAGES="${INSTALL_PACKAGES} openjdk-11-jre"
else
  JAVA_VERSION="$(java -version 2>&1 | grep "openjdk version" | awk -F '"' '{print $2}' | awk -F '.' '{print $1}')"
  if [[ -z "${JAVA_VERSION}" ]] ; then
    echo "Adding Java 11 to install list (unknown java version found) ..."
    INSTALL_PACKAGES="${INSTALL_PACKAGES} openjdk-11-jre"
  elif [[ ${JAVA_VERSION} -lt 11 ]] ; then
    echo "Adding Java 11 to install list (unexpected java version: ${JAVA_VERSION}) ..."
    INSTALL_PACKAGES="${INSTALL_PACKAGES} openjdk-11-jre"
  fi
fi

if [[ -z "$(which jq)" ]] ; then
  echo "Adding jq to install list ..."
  INSTALL_PACKAGES="${INSTALL_PACKAGES} jq"
fi
if [[ -z "$(which route)" ]] ; then
  echo "Adding net-tools to install list ..."
  INSTALL_PACKAGES="${INSTALL_PACKAGES} net-tools"
fi
if [[ -z "$(which resolvconf)" ]] ; then
  echo "Adding resolvconf to install list ..."
  INSTALL_PACKAGES="${INSTALL_PACKAGES} resolvconf"
fi

if [[ -n "${INSTALL_PACKAGES}" ]] ; then
  echo "
The following apt packages will be installed:

  ${INSTALL_PACKAGES}

Press Enter to continue with installation, or Control-C to exit"
  read -r DUMMY || die "Bubble connection canceled"

  # Do not quote INSTALL_PACKAGES, we want to install all packages...
  sudo apt update && sudo apt install -y ${INSTALL_PACKAGES} || die "Error installing packages: ${INSTALL_PACKAGES}"
fi

ARG_HOST="${1}"
if [[ -z "${ARG_HOST}" ]] ; then
  if [[ -z "${BUBBLE_API}" ]] ; then
    if [[ -z "${BUBBLE_HOST}" ]] ; then
      echo -n "No hostname argument provided and neither BUBBLE_API nor BUBBLE_HOST env vars were defined.
Enter Bubble hostname: "
      read -r ARG_HOST
      echo
      if [[ -z "${ARG_HOST}" ]] ; then
        die "Empty Bubble hostname"
      fi
      BUBBLE_API="https://${ARG_HOST}/api"
    else
      BUBBLE_API="https://${BUBBLE_HOST}/api"
    fi
  fi
else
  BUBBLE_API="https://${ARG_HOST}/api"
fi

if [[ -z "${BUBBLE_USER}" ]] ; then
  echo -n "BUBBLE_USER env var not defined.
Enter Bubble username: "
  read -r BUBBLE_USER
  echo
fi
if [[ -z "${BUBBLE_PASS}" ]] ; then
  echo -n "BUBBLE_PASS env var not defined.
Enter Bubble password: "
  read -rs BUBBLE_PASS
  echo
fi

BUBBLE_HOST="$(echo -n "${BUBBLE_API}" | awk -F '/' '{print $3}' | awk -F ':' '{print $1}')"
if [[ -z "${BUBBLE_HOST}" ]] ; then
  die "No hostname could be determined from BUBBLE_API: ${BUBBLE_API}"
fi

echo "Logging in to Bubble ${BUBBLE_API} ..."
export BUBBLE_API
export BUBBLE_USER
export BUBBLE_PASS
export BUBBLE_SKIP_PATH_WARNING=true
bget me | jq .email > /dev/null || die "Error logging into Bubble with API: ${BUBBLE_API}"

# Ensure WireGuard is not running
if [[ -n "$(sudo wg show)" ]] ; then
  echo "Stopping WireGuard VPN ..."
  sudo wg-quick down wg0
fi

BUBBLE_DEVICE_BASE="${HOME}/bubble_devices"
BUBBLE_DIR="${BUBBLE_DEVICE_BASE}/${BUBBLE_HOST}"
mkdir -p "${BUBBLE_DIR}" || die "Error creating directory: ${BUBBLE_DIR}"

CERT_FILE="${BUBBLE_DIR}"/bubble-${BUBBLE_HOST}.crt
VPN_FILE="${BUBBLE_DIR}"/vpn.conf

DEVICE_NAME_FILE=${BUBBLE_DIR}/bubble_device_name
if [[ ! -f "${DEVICE_NAME_FILE}" ]] ; then
  DISTRO="$(cat /etc/os-release  | grep ^NAME= | awk -F '=' '{print $2}' | tr -d '"')"
  echo -n "Linux_${DISTRO}_$(hostname -s)" >"${DEVICE_NAME_FILE}"
  echo "Initialized device: ${DEVICE_NAME_FILE}"
fi
DEVICE_NAME="$(cat "${DEVICE_NAME_FILE}")"

# Check API to see if device is already registered
if [[ $(bgetn me/devices | grep -c "${DEVICE_NAME}") -eq 0 ]] ; then
  echo "Device not found, registering now: ${DEVICE_NAME}"
  echo "{
  \"name\": \"${DEVICE_NAME}\",
  \"deviceType\": \"Linux\"
}" | bput me/devices - > "${BUBBLE_DIR}/device.json" || die "Error creating device"
  rm -f "${VPN_FILE}" "${CERT_FILE}" || die "Error removing obsolete vpn.conf and cert files"
else
  echo "Device already registered: ${DEVICE_NAME}"
  if [[ ! -f "${BUBBLE_DIR}/device.json" ]] ; then
    echo "Downloading device.json file ..."
    bget "me/devices/$(uriencode "${DEVICE_NAME}")" > "${BUBBLE_DIR}/device.json" || die "Error downloading device JSON file"
  fi
fi

# Do we have both a cert file and a vpn.conf? If so, we are done
if [[ ! -s "${CERT_FILE}" ]] ; then
  # Download cert file
  echo "Downloading certificate ..."
  bget auth/cacert?deviceType=Linux --raw >"${CERT_FILE}" || die "Error downloading certificate file"
fi
if [[ ! -s "${VPN_FILE}" ]] ; then
  # Download vpn.conf file
  echo "Downloading vpn.conf ..."
  bget "me/devices/${DEVICE_NAME}/vpn/vpn.conf" --raw >"${VPN_FILE}" || die "Error downloading vpn.conf file"
fi

CURRENT_DEVICE_DIR="${BUBBLE_DEVICE_BASE}/current"
CREATE_CURRENT_SYMLINK=0
if [[ -e "${CURRENT_DEVICE_DIR}" ]] ; then
  if [[ "$(basename "$(readlink -f "${CURRENT_DEVICE_DIR}")")" != "${BUBBLE_HOST}" ]] ; then
    CREATE_CURRENT_SYMLINK=1
  fi
else
  CREATE_CURRENT_SYMLINK=1
fi
if [[ ${CREATE_CURRENT_SYMLINK} -eq 1 ]] ; then
  echo "Marking ${BUBBLE_HOST} as current Bubble (cd ${BUBBLE_DEVICE_BASE} && ln -sf ${BUBBLE_HOST} current)..."
  cd "${BUBBLE_DEVICE_BASE}" && ln -sf "${BUBBLE_HOST}" current || die "Error creating symlink ${BUBBLE_DEVICE_BASE}/current -> ${BUBBLE_DEVICE_BASE}/${BUBBLE_HOST}"
fi
CURRENT_DEVICE_JSON="$(cat "${CURRENT_DEVICE_DIR}/device.json")"
if [[ -z "${CURRENT_DEVICE_JSON}" ]] ; then
  die "File was empty: ${CURRENT_DEVICE_JSON}"
fi
DEVICE_UUID="$(echo "${CURRENT_DEVICE_JSON}" | jq -r .uuid)"
if [[ -z "${DEVICE_UUID}" ]] ; then
  die "No device UUID could be read from JSON: ${CURRENT_DEVICE_JSON}"
fi

CERTS_DIR=/usr/share/ca-certificates/extra
if [[ ! -d "${CERTS_DIR}" ]] ; then
  echo "Creating ${CERTS_DIR}"
  sudo mkdir -p ${CERTS_DIR}
fi
CERT_DEST="${CERTS_DIR}/$(basename "${CERT_FILE}")"
RUN_CERT_WIZ=0
if cmp "${CERT_FILE}" "${CERT_DEST}" ; then
  if [[ -z "${BUBBLE_CERT_WIZ}" ]] ; then
    echo "Certificate already installed: ${CERT_DEST}"
  elif [[ "${BUBBLE_CERT_WIZ}" == "true" ]] ; then
    RUN_CERT_WIZ=1
  else
    die "Invalid value for BUBBLE_CERT_WIZ env var: ${BUBBLE_CERT_WIZ} (expected 'true' or nothing)"
  fi
else
  RUN_CERT_WIZ=1
  echo "Copying certificate to ${CERT_DEST} ..."
  sudo cp "${CERT_FILE}" "${CERT_DEST}" || die "Error copying certificate: ${CERT_FILE} -> ${CERT_DEST}"
fi

if [[ ${RUN_CERT_WIZ} -eq 1 ]] ; then
  echo "
### Finishing Certificate Installation for Ubuntu

We're going to run the Ubuntu certificate wizard via:

   sudo dpkg-reconfigure ca-certificates

When the wizard opens:
 - Press Enter with 'Yes' selected for 'Trust new certificates from certificate authorities?'
 - Press the space bar to check the box for your Bubble certificate
 - Press Enter to commit the changes

To continue and run the Ubuntu certificate wizard, press Enter now
"
  read -r DUMMY
  sudo dpkg-reconfigure ca-certificates || die "Error reconfiguring system CA certificates"
fi

# Ensure ssh route table exists
if [[ $(sudo cat /etc/iproute2/rt_tables | grep -c ssh) -eq 0 ]] ; then
  sudo bash -c 'echo "2 ssh" >> /etc/iproute2/rt_tables'
else
  echo "ssh table already exists in rt_tables"
fi

REBOOT_FILE=/tmp/reboot-required.bubble

# Set sysctl vars
SYSCTL="/etc/sysctl.conf"
SYS_RP_FILTER="net.ipv4.conf.all.rp_filter"
if [[ $(sudo cat ${SYSCTL} | grep -v '^#' | grep -c "${SYS_RP_FILTER}") -eq 0 ]] ; then
  sudo bash -c 'echo "'"${SYS_RP_FILTER}"' = 2" >> '"${SYSCTL}"'' || die "Error setting ${SYS_RP_FILTER} = 2 in ${SYSCTL}"
  touch ${REBOOT_FILE}
else
  echo "${SYS_RP_FILTER} already defined in ${SYSCTL}"
fi

SYS_IP_FORWARD="net.ipv4.ip_forward"
if [[ $(cat ${SYSCTL} | grep -v '^#' | grep -c "${SYS_IP_FORWARD}") -eq 0 ]] ; then
  sudo bash -c 'echo "'"${SYS_IP_FORWARD}"' = 1" >> '"${SYSCTL}"'' || die "Error setting ${SYS_IP_FORWARD} = 1 in ${SYSCTL}"
  touch ${REBOOT_FILE}
else
  echo "${SYS_IP_FORWARD} already defined in ${SYSCTL}"
fi

if [[ -f ${REBOOT_FILE} ]] ; then
  echo "
sysctl settings changed, reboot required.

Press Enter to reboot, or Control-C to exit and reboot later"
  read -r DUMMY || die "Please reboot later to ensure sysctl changes take effect"
  sudo bash -c 'sync && shutdown -r now' && exit 1 || die "Error rebooting"
fi

WG_CONF=/etc/wireguard/wg0.conf
echo "Copying vpn.conf to ${WG_CONF} ..."
sudo cp "${VPN_FILE}" ${WG_CONF} || die "Error copying vpn.conf: ${VPN_FILE} -> ${WG_CONF}"

if [[ $(sudo cat "${WG_CONF}" | grep -c PostUp) -ne 0 ]] ; then
  echo "${WG_CONF} already contains PostUp directives"
else
  GATEWAY=$(route -n | grep "^0.0.0.0" | awk '{print $2}')
  if [[ -z "${GATEWAY}" ]] ; then
    die "Error determining gateway IP using 'route -n'"
  fi

  IFACE=$(route -n | grep "^0.0.0.0" | awk '{print $8}')
  if [[ -z "${IFACE}" ]] ; then
    die "Error determining gateway interface using 'route -n'"
  fi

  WG_TEMP=$(mktemp /tmp/wg.conf.XXXXXXX)
  sudo cat ${WG_CONF} | grep -A 3 '\[Interface\]' >> "${WG_TEMP}"
  echo "PostUp = ip route add default via ${GATEWAY} dev ${IFACE} table ssh
  PostUp = ip rule add fwmark 0x2 table ssh
  PostUp = /sbin/iptables -A OUTPUT -t mangle -o wg0 -p tcp --sport 22 -j MARK --set-mark 2
  PreDown = /sbin/iptables -D OUTPUT -t mangle -o wg0 -p tcp --sport 22 -j MARK --set-mark 2
  PreDown = ip rule del fwmark 0x2 table ssh
  PreDown = ip route del default via ${GATEWAY} dev ${IFACE} table ssh
  " >> "${WG_TEMP}"
  sudo cat ${WG_CONF} | grep -A 10 '\[Peer\]' >> "${WG_TEMP}"

  sudo mv "${WG_TEMP}" ${WG_CONF} || die "Error installing updated ${WG_CONF} file"
fi

echo "Starting WireGuard VPN ..."
sudo wg-quick up wg0 || die "Error starting WireGuard"

echo "
===============================================================
======= Ubuntu device successfully connected to Bubble! =======
===============================================================
Device UUID : ${DEVICE_UUID}
Device Name : ${DEVICE_NAME}
Bubble Host : ${BUBBLE_HOST}
Device Dir  : ${BUBBLE_DIR}
Certificate : ${CERT_DEST}
VPN Config  : ${WG_CONF}
Public IP   : $(curl -s http://checkip.amazonaws.com/)
"

if [[ -n "$(which firefox)" ]] ; then
  echo "
===========================================================
===================== Firefox Warning =====================
===========================================================
Firefox does not use the Ubuntu certificate store. It has its
own certificate store.

Firefox will not properly connect to most websites until you
install the Bubble certificate in Firefox.

To install the Bubble Certificate in Firefox:
 * Open Firefox
 * In the search box (top right of page), enter \"certificate\"
 * Scroll down, click \"View Certificates\" button
 * Click \"Import\" button
 * Select your certificate file: ${CERT_FILE}
 * Check the box \"Trust this CA to identify websites\" and click OK.

More detailed instructions (with screenshots) can be found here:
  https://git.bubblev.org/bubblev/bubble-docs/src/branch/master/cert_instructions/firefox_cert.md
"
fi