#!/usr/bin/with-contenv bash

mkdir -p /config/{templates,coredns}

echo "Uname info: $(uname -a)"
# check for wireguard module
ip link del dev test 2>/dev/null
if ip link add dev test type wireguard; then
  echo "**** It seems the wireguard module is already active. Skipping kernel header install and module compilation. ****"
  SKIP_COMPILE="true"
  ip link del dev test
else
  echo "**** The wireguard module is not active, will attempt kernel header install and module compilation. ****"
fi  

# install headers if necessary
if [ "$SKIP_COMPILE" != "true" ] && [ ! -e /lib/modules/$(uname -r)/build ]; then
  echo "**** Attempting kernel header install ****"
  apt-get update
  if apt-cache show linux-headers-$(uname -r) 2&>1 >/dev/null; then
    apt-get install -y \
      linux-headers-$(uname -r)
  elif (uname -r | grep -q 'v7+') || (uname -r | grep -q 'v7l+') || (uname -r | grep -q 'v8+'); then
    echo "**** Raspbian kernel naming convention detected, attempting to install raspbian kernel headers ****"
    curl -s http://archive.raspberrypi.org/debian/raspberrypi.gpg.key | apt-key add -
    echo -e \
      "deb http://archive.raspberrypi.org/debian/ buster main\ndeb-src http://archive.raspberrypi.org/debian/ buster main" \
      > /etc/apt/sources.list.d/raspbian.list
    apt-get update
    apt-get install -y \
      raspberrypi-kernel-headers
  elif uname -v | grep -q 'Ubuntu'; then
    echo "**** Ubuntu kernel detected, but likely not Bionic. ****"
    echo "**** Attempting to install kernel headers from Ubuntu Xenial repo ****"
    if uname -m | grep -q 'x86_64'; then
      echo -e \
        "deb http://archive.ubuntu.com/ubuntu/ xenial main restricted\ndeb-src http://archive.ubuntu.com/ubuntu/ xenial main restricted\n\ndeb http://archive.ubuntu.com/ubuntu/ xenial-updates main restricted\ndeb-src http://archive.ubuntu.com/ubuntu/ xenial-updates main restricted" \
        > /etc/apt/sources.list.d/xenial-focal.list
    else
      echo -e \
      "deb http://ports.ubuntu.com/ubuntu-ports/ xenial main restricted\ndeb-src http://ports.ubuntu.com/ubuntu-ports/ xenial main restricted\n\ndeb http://ports.ubuntu.com/ubuntu-ports/ xenial-updates main restricted\ndeb-src http://ports.ubuntu.com/ubuntu-ports/ xenial-updates main restricted" \
        > /etc/apt/sources.list.d/xenial-focal.list
    fi
    apt-get update
    if apt-cache show linux-headers-$(uname -r) 2&>1 >/dev/null; then
      apt-get install -y \
        linux-headers-$(uname -r)
    else
      echo "**** No kernel headers found in the Ubuntu Xenial repo!! Trying Ubuntu Focal. ****"
      sed -i 's/xenial/focal/g' /etc/apt/sources.list.d/xenial-focal.list
      apt-get update
      if apt-cache show linux-headers-$(uname -r) 2&>1 >/dev/null; then
        apt-get install -y \
          linux-headers-$(uname -r)
      else
        echo "**** No kernel headers found in the Ubuntu repos!! Will try the headers from host (if mapped), may or may not work ****"
        rm -rf /etc/apt/sources.list.d/xenial-focal.list
      fi
    fi
  elif uname -v | grep -q 'Debian'; then
    echo "**** Debian host detected, attempting to install kernel headers from Debian Buster repo ****"
    curl -s https://ftp-master.debian.org/keys/archive-key-10.asc | apt-key add -
    curl -s https://ftp-master.debian.org/keys/archive-key-10-security.asc | apt-key add -
    cat <<DUDE > /etc/apt/sources.list.d/debian.list
deb http://deb.debian.org/debian buster main contrib non-free
deb-src http://deb.debian.org/debian buster main contrib non-free
deb http://deb.debian.org/debian-security/ buster/updates main contrib non-free
deb-src http://deb.debian.org/debian-security/ buster/updates main contrib non-free
deb http://deb.debian.org/debian buster-updates main contrib non-free
deb-src http://deb.debian.org/debian buster-updates main contrib non-free
deb http://deb.debian.org/debian buster-backports main contrib non-free
deb-src http://deb.debian.org/debian buster-backports main contrib non-free
DUDE
    apt-get update
    if apt-cache show linux-headers-$(uname -r) 2&>1 >/dev/null; then
      if uname -r | grep -qs "bpo"; then
        echo "**** Backported kernel detected ****"
        apt-get install -y -t buster-backports \
          linux-headers-$(uname -r)
      else
        apt-get install -y \
          linux-headers-$(uname -r)
      fi
    else
      echo "**** Attempting to install kernel headers from the Debian Stretch repo ****"
      curl -s https://ftp-master.debian.org/keys/archive-key-9.asc | apt-key add -
      curl -s https://ftp-master.debian.org/keys/archive-key-9-security.asc | apt-key add -
      sed -i 's/buster/stretch/g' /etc/apt/sources.list.d/debian.list
      apt-get update
      if apt-cache show linux-headers-$(uname -r) 2&>1 >/dev/null; then
        if uname -r | grep -qs "bpo"; then
          echo "**** Backported kernel detected ****"
          apt-get install -y -t stretch-backports \
            linux-headers-$(uname -r)
        else
          apt-get install -y \
            linux-headers-$(uname -r)
        fi
      else
        echo "**** No kernel headers found in Debian repos!! Will try the headers from host (if mapped), may or may not work ****"
        rm -rf /etc/apt/sources.list.d/debian.list
      fi
    fi
  else
    echo "**** No kernel headers found in the Ubuntu or Debian repos!! Will try the headers from host (if mapped), may or may not work ****"
  fi
fi

if [ "$SKIP_COMPILE" != "true" ]; then
  if [ -e /lib/modules/$(uname -r)/build ]; then 
    echo "**** Kernel headers seem to be present, attempting to build the wireguard module. . . ****"
    if [ ! -f /lib/modules/$(uname -r)/build/certs/signing_key.pem ]; then
      mkdir -p /lib/modules/$(uname -r)/build/certs
      cd /lib/modules/$(uname -r)/build/certs
      cat <<DUDE >> x509.genkey
[ req ]
default_bits = 4096
distinguished_name = req_distinguished_name
prompt = no
string_mask = utf8only
x509_extensions = myexts

[ req_distinguished_name ]
CN = Modules

[ myexts ]
basicConstraints=critical,CA:FALSE
keyUsage=digitalSignature
subjectKeyIdentifier=hash
authorityKeyIdentifier=keyid
DUDE
      echo "**** Generating signing key ****"
      openssl req -new -nodes -utf8 -sha512 -days 36500 -batch -x509 -config x509.genkey -outform DER -out signing_key.x509 -keyout signing_key.pem
    fi
    cd /app
    echo "**** Building the module ****"
    make -C wireguard-linux-compat/src -j$(nproc)
    make -C wireguard-linux-compat/src install
    echo "**** Let's test our new module. ****"
    ip link del dev test 2>/dev/null
    if ip link add dev test type wireguard; then
      echo "**** The module is active, moving forward with setup. ****"
      ip link del dev test
    else
      echo "**** The module is not active, review the logs. Sleeping now. . . ****"
      sleep infinity
    fi
  else
    echo "**** Kernel headers don't seem to be available, can't compile the module. Sleeping now. . . ****"
    sleep infinity
  fi
fi

# prepare symlinks
rm -rf /etc/wireguard
mkdir -p /etc/wireguard
ln -s /config/wg0.conf /etc/wireguard/wg0.conf
# prepare templates
[[ ! -f /config/templates/server.conf ]] && \
  cp /defaults/server.conf /config/templates/server.conf
[[ ! -f /config/templates/peer.conf ]] && \
  cp /defaults/peer.conf /config/templates/peer.conf

generate_confs () {
  mkdir -p /config/server
  if [ ! -f /config/server/privatekey-server ]; then
    umask 077
    wg genkey | tee /config/server/privatekey-server | wg pubkey > /config/server/publickey-server
  fi
  eval "`printf %s`
  cat <<DUDE > /config/wg0.conf
`cat /config/templates/server.conf`

DUDE"
  for i in ${PEERS_ARRAY[@]}; do
    if [[ "${i}" =~ ^[0-9]+$ ]]; then
      PEER_ID="peer${i}"
    else
      PEER_ID="peer_${i//[^[:alnum:]_-]/}"
    fi
    mkdir -p /config/${PEER_ID}
    if [ ! -f "/config/${PEER_ID}/privatekey-${PEER_ID}" ]; then
      umask 077
      wg genkey | tee /config/${PEER_ID}/privatekey-${PEER_ID} | wg pubkey > /config/${PEER_ID}/publickey-${PEER_ID}
    fi
    if [ -f "/config/${PEER_ID}/${PEER_ID}.conf" ]; then
      CLIENT_IP=$(cat /config/${PEER_ID}/${PEER_ID}.conf | grep "Address" | awk '{print $NF}')
      if [ -n "${ORIG_INTERFACE}" ] && [ "${INTERFACE}" != "${ORIG_INTERFACE}" ]; then
        CLIENT_IP=$(echo "${CLIENT_IP}" | sed "s|${ORIG_INTERFACE}|${INTERFACE}|")
      fi
    else
      for idx in {2..254}; do
        PROPOSED_IP="${INTERFACE}.${idx}"
        if ! grep -q -R "${PROPOSED_IP}" /config/peer*/*.conf && ([ -z "${ORIG_INTERFACE}" ] || ! grep -q -R "${ORIG_INTERFACE}.${idx}" /config/peer*/*.conf); then
          CLIENT_IP="${PROPOSED_IP}"
          break
        fi
      done
    fi
    eval "`printf %s`
    cat <<DUDE > /config/${PEER_ID}/${PEER_ID}.conf
`cat /config/templates/peer.conf`
DUDE"
    cat <<DUDE >> /config/wg0.conf
[Peer]
# ${PEER_ID}
PublicKey = $(cat /config/${PEER_ID}/publickey-${PEER_ID})
AllowedIPs = ${CLIENT_IP}/32

DUDE
    echo "PEER ${i} QR code:"
    qrencode -t ansiutf8 < /config/${PEER_ID}/${PEER_ID}.conf
    qrencode -o /config/${PEER_ID}/${PEER_ID}.png < /config/${PEER_ID}/${PEER_ID}.conf
  done
}

save_vars () {
  cat <<DUDE > /config/.donoteditthisfile
ORIG_SERVERURL="$SERVERURL"
ORIG_SERVERPORT="$SERVERPORT"
ORIG_PEERDNS="$PEERDNS"
ORIG_PEERS="$PEERS"
ORIG_INTERFACE="$INTERFACE"
ORIG_ALLOWEDIPS="$ALLOWEDIPS"
DUDE
}

if [ -n "$PEERS" ]; then
  echo "**** Server mode is selected ****"
  if [[ "$PEERS" =~ ^[0-9]+$ ]] && ! [[ "$PEERS" =~ *,* ]]; then
    PEERS_ARRAY=($(seq 1 $PEERS))
  else
    PEERS_ARRAY=($(echo "$PEERS" | tr ',' ' '))
  fi
  PEERS_COUNT=$(echo "${#PEERS_ARRAY[@]}")
  if [ -z "$SERVERURL" ] || [ "$SERVERURL" = "auto" ]; then
    SERVERURL=$(curl -s icanhazip.com)
    echo "**** SERVERURL var is either not set or is set to \"auto\", setting external IP to auto detected value of $SERVERURL ****"
  else
    echo "**** External server address is set to $SERVERURL ****"
  fi
  SERVERPORT=${SERVERPORT:-51820}
  echo "**** External server port is set to ${SERVERPORT}. Make sure that port is properly forwarded to port 51820 inside this container ****"
  INTERNAL_SUBNET=${INTERNAL_SUBNET:-10.13.13.0}
  echo "**** Internal subnet is set to $INTERNAL_SUBNET ****"
  INTERFACE=$(echo "$INTERNAL_SUBNET" | awk 'BEGIN{FS=OFS="."} NF--')
  ALLOWEDIPS=${ALLOWEDIPS:-0.0.0.0/0, ::/0}
  echo "**** AllowedIPs for peers $ALLOWEDIPS ****"
  if [ -z "$PEERDNS" ] || [ "$PEERDNS" = "auto" ]; then
    PEERDNS="${INTERFACE}.1"
    echo "**** PEERDNS var is either not set or is set to \"auto\", setting peer DNS to ${INTERFACE}.1 to use wireguard docker host's DNS. ****"
  else
    echo "**** Peer DNS servers will be set to $PEERDNS ****"
  fi
  if [ ! -f /config/wg0.conf ]; then
    echo "**** No wg0.conf found (maybe an initial install), generating 1 server and ${PEERS} peer/client confs ****"
    generate_confs
    save_vars
  else
    echo "**** Server mode is selected ****"
    [[ -f /config/.donoteditthisfile ]] && \
      . /config/.donoteditthisfile
    if [ "$SERVERURL" != "$ORIG_SERVERURL" ] || [ "$SERVERPORT" != "$ORIG_SERVERPORT" ] || [ "$PEERDNS" != "$ORIG_PEERDNS" ] || [ "$PEERS" != "$ORIG_PEERS" ] || [ "$INTERFACE" != "$ORIG_INTERFACE" ] || [ "$ALLOWEDIPS" != "$ORIG_ALLOWEDIPS" ]; then
      echo "**** Server related environment variables changed, regenerating 1 server and ${PEERS} peer/client confs ****"
      generate_confs
      save_vars
    else
      echo "**** No changes to parameters. Existing configs are used. ****"
    fi
  fi
else
  echo "**** Client mode selected. ****"
  if [ ! -f /config/wg0.conf ]; then
    echo "**** No client conf found. Provide your own client conf as \"/config/wg0.conf\" and restart the container. ****"
    sleep infinity
  fi
  echo "**** Disabling CoreDNS ****"
  rm -rf /etc/services.d/coredns
fi

# set up CoreDNS
[[ ! -f /config/coredns/Corefile ]] && \
  cp /defaults/Corefile /config/coredns/Corefile

# permissions
chown -R abc:abc \
	/config
