#!/bin/sh

# SPDX-FileCopyrightText: 2022 unman <unman@thirdeyesecurity.org>
# SPDX-FileCopyrightText: 2023 - 2024 Benjamin Grande M. S. <ben.grande.b@gmail.com>
#
# SPDX-License-Identifier: AGPL-3.0-or-later

set -eu

test "$(id -u)" = "0" || exec sudo "$0" "$@"

usage(){
  echo "Usage: ${0##*/} [QUBE]"
  exit "${1:-1}"
}

case "${1-}" in
  -h|--help) usage 0;;
  -*) usage 1;;
  "") qube="sys-wireguard";;
  *) qube="${1}";;
esac

if ! qvm-check -q -- "$qube" >/dev/null 2>&1; then
  echo "Qube '$qube' doesn't exist" >&2
  usage 1
fi

user_conf="/home/user/wireguard.conf"
system_conf="/etc/wireguard/wireguard.conf"

qvm-run "$qube" -- "test -f ${user_conf}" || {
  echo "File '${user_conf}' was not found" >&2
  if qvm-check -q --running -- "$qube" >/dev/null 2>&1; then
    qvm-pause --verbose -- "$qube"
  fi
  echo "Firewalling $qube to drop all connections"
  qvm-firewall --verbose -- "$qube" reset
  qvm-firewall --verbose -- "$qube" del --rule-no 0
  qvm-firewall --verbose -- "$qube" add drop
  if qvm-check -q --paused -- "$qube" >/dev/null 2>&1; then
    qvm-unpause --verbose -- "$qube"
  fi
  exit 1
}

qvm-run -u root "$qube" -- "cp ${user_conf} ${system_conf}"

## TOFU
# shellcheck disable=SC2016
endpoint="$(qvm-run -p -u root "$qube" -- awk '/Endpoint/{print $3}' "${system_conf}")"
if echo "${endpoint}" | grep -qF "["; then
  ip="${ip##[\[]}"
  ip="${ip%%\]*}"
  port="${endpoint##*:}"
else
  ip="${endpoint%%:*}"
  port="${endpoint##*:}"
fi

if test -z "$ip" || test -z "$port";then
  echo "Endpoint (IP:Port) not found: ${system_conf}" >&2
  exit 1
fi

if qvm-check -q --running -- "$qube" >/dev/null 2>&1; then
  qvm-pause --verbose -- "$qube"
fi

echo "Firewalling $qube to reach only '$ip:$port'"
qvm-firewall --verbose -- "$qube" reset
qvm-firewall --verbose -- "$qube" del --rule-no 0
qvm-firewall --verbose -- "$qube" add accept dsthost="$ip" dstports="$port" proto=udp
qvm-firewall --verbose -- "$qube" add accept dsthost="$ip" dstports="$port" proto=tcp
qvm-firewall --verbose -- "$qube" add drop

if qvm-check -q --paused -- "$qube" >/dev/null 2>&1; then
  qvm-unpause --verbose -- "$qube"
fi

qvm-run -u root "$qube" -- "systemctl restart wg-quick@wireguard"
qvm-run -u root "$qube" -- "/rw/config/network-hooks.d/50-sys-wireguard"