#!/usr/bin/env bash

##
## This script will install and configure WireGuard installation for tunnelling entire traffic, 
## it is intended to be run from cloud-init during creation of a VM
##
## put this line into 'user data' or 'custom script' when you create a machine:
## wget -O - https://hmage.net/wireguard.sh | bash

function main() {
set -eE
set -o pipefail
set -x

unset INTERACTIVE
[ -t 1 ] && INTERACTIVE=yes

PACKAGES=(
wireguard
wireguard-tools
netfilter-persistent # for NAT
iptables-persistent  # for NAT
curl                 # for getting our ip
jq                   # for getting our interface name
)

export DEBIAN_FRONTEND=noninteractive
apt-get update
apt-get upgrade -y --no-install-recommends
apt-get install -y "${PACKAGES[@]}" --no-install-recommends

# Generate WireGuard keys
umask 077
wg genkey | tee /etc/wireguard/server_private.key | wg pubkey > /etc/wireguard/server_public.key
wg genkey | tee /etc/wireguard/client_private.key | wg pubkey > /etc/wireguard/client_public.key
wg genpsk > /etc/wireguard/preshared.key

PORT=1637
IPADDR=$(curl -s whatismyip.akamai.com)
IFACE=$(ip -j route | jq -r '.[] | select(.dst == "default") | .dev')

echo Interface is $IFACE

## Take down existing wg0 if it exists
wg-quick down wg0 || true

##
## WireGuard configs
##
cat << EOF > /etc/wireguard/wg0.conf
[Interface]
PrivateKey = $(cat /etc/wireguard/server_private.key)
Address = 10.104.104.1/24
ListenPort = $PORT
MTU = 1500

[Peer]
PublicKey = $(cat /etc/wireguard/client_public.key)
PresharedKey = $(cat /etc/wireguard/preshared.key)
AllowedIPs = 10.104.104.2/32
EOF

cat << EOF > /root/"$IPADDR-client.conf"
[Interface]
PrivateKey = $(cat /etc/wireguard/client_private.key)
Address = 10.104.104.2/24
DNS = 8.8.8.8, 8.8.4.4
MTU = 1500

[Peer]
PublicKey = $(cat /etc/wireguard/server_public.key)
PresharedKey = $(cat /etc/wireguard/preshared.key)
AllowedIPs = 0.0.0.0/0
Endpoint = $IPADDR:$PORT
PersistentKeepalive = 15
EOF

# Enable and start WireGuard
systemctl enable wg-quick@wg0
systemctl restart wg-quick@wg0
wg-quick up wg0 || true

## Set up iptables
if ! iptables -C INPUT -t filter -p udp -m udp --dport $PORT -j ACCEPT; then
    iptables -A INPUT -t filter -p udp -m udp --dport $PORT -j ACCEPT
fi
if ! iptables -C FORWARD -t filter -i wg0 -j ACCEPT; then
    iptables -A FORWARD -t filter -i wg0 -j ACCEPT
fi
if ! iptables -C FORWARD -t filter -i $IFACE -o wg0 -j ACCEPT; then
    iptables -A FORWARD -t filter -i $IFACE -o wg0 -j ACCEPT
fi
if ! iptables -C POSTROUTING -t nat -o $IFACE -j MASQUERADE; then
    iptables -A POSTROUTING -t nat -o $IFACE -j MASQUERADE
fi
if ! iptables -t mangle -C FORWARD -p tcp --tcp-flags SYN,RST SYN -j TCPMSS --clamp-mss-to-pmtu; then
    iptables -t mangle -A FORWARD -p tcp --tcp-flags SYN,RST SYN -j TCPMSS --clamp-mss-to-pmtu
fi

service netfilter-persistent save

# Enable IP forwarding
echo 'net.ipv4.ip_forward=1' >> /etc/sysctl.d/nat.conf
service procps restart

echo 'Successfully configured WireGuard. Now copy the config:'
echo "$ scp root@$IPADDR:*.conf ."
}

main