#!/bin/bash
echo "---Ensuring UID: ${UID} matches user---"
usermod -u ${UID} ${USER}
echo "---Ensuring GID: ${GID} matches user---"
groupmod -g ${GID} ${USER} > /dev/null 2>&1 ||:
usermod -g ${GID} ${USER}
echo "---Setting umask to ${UMASK}---"
umask ${UMASK}

echo "---Checking for optional scripts---"
cp -f /opt/custom/user.sh /opt/scripts/start-user.sh > /dev/null 2>&1 ||:
cp -f /opt/scripts/user.sh /opt/scripts/start-user.sh > /dev/null 2>&1 ||:

if [ -f /opt/scripts/start-user.sh ]; then
    echo "---Found optional script, executing---"
    chmod -f +x /opt/scripts/start-user.sh ||:
    /opt/scripts/start-user.sh || echo "---Optional Script has thrown an Error---"
else
    echo "---No optional script found, continuing---"
fi

if [ ! -f /etc/sudoers.d/${USER} ]; then
  echo "---Creating sudoers file for user: ${USER}---"
  echo "${USER} ALL=(ALL) NOPASSWD: /usr/bin/wg-quick up *
${USER} ALL=(ALL) NOPASSWD: /usr/bin/wg-quick down *
${USER} ALL=(ALL) NOPASSWD: /usr/bin/wg show wg*" > /etc/sudoers.d/${USER}
else
  echo "---Found sudoers file for user: ${USER}---"
fi

if [ "${SETUP_IPTABLES}" == "true" ]; then
  echo "---Setting up iptables---"
  iptables -t nat -A PREROUTING -i wg${WG_CONFIG_ID} -p ${NAT_PROTOCOL} --dport ${NAT_TUNNEL_PORT} -j DNAT --to-destination ${NAT_DESTINATION}
  iptables -A FORWARD -i wg${WG_CONFIG_ID} -o eth0 -p ${NAT_PROTOCOL} --dport ${NAT_DESTINATION#*:} -d ${NAT_DESTINATION%%:*} -j ACCEPT
  iptables -A FORWARD -i eth0 -o wg${WG_CONFIG_ID} -p ${NAT_PROTOCOL} --sport ${NAT_DESTINATION#*:} -s ${NAT_DESTINATION%%:*} -j ACCEPT
  iptables -t nat -A POSTROUTING -o eth0 -j MASQUERADE
  iptables -A FORWARD -i wg${WG_CONFIG_ID} -j DROP
elif [ "${SETUP_IPTABLES}" == "custom" ]; then
  if [ ! -f ${DATA_DIR}/wg${WG_CONFIG_ID}/iptables ]; then
    echo "---No custom iptables found, make sure that the file iptables exists---"
  else
    echo "---Applying custom iptables---"
    while IFS= read -r line; do
      if [ -z "$line" ]; then
        continue
      fi
      echo "Applying: $line"
      eval "$line"
    done < "${DATA_DIR}/wg${WG_CONFIG_ID}/iptables"
  fi
fi

create_servercfg() {
  if [ ! -d ${DATA_DIR}/wg${WG_CONFIG_ID}/server ]; then
    mkdir -p ${DATA_DIR}/wg${WG_CONFIG_ID}/server
  fi
  cat <<EOF > ${DATA_DIR}/wg${WG_CONFIG_ID}/server/wg${WG_CONFIG_ID}.conf
[Interface]
Address = ${WG_NET_IP}/${WG_NET_SUBNET}
PrivateKey = ${1}
ListenPort = ${SERVER_LISTEN_PORT}

EOF
}

create_peer() {
  ADDRESS_NO=${1}
  cat <<EOF >> ${DATA_DIR}/wg${WG_CONFIG_ID}/server/wg${WG_CONFIG_ID}.conf
[Peer] # Client ${1}
PublicKey = ${2}
$(if [ ! -z "${3}" ]; then echo "PresharedKey = ${3}"; fi)
AllowedIPs = ${WG_NET_IP%.*}.$(( ${WG_NET_IP##*.} + ADDRESS_NO ))/32

EOF
}

create_config() {
  if [ ! -d ${DATA_DIR}/wg${WG_CONFIG_ID}/client ]; then
    mkdir -p ${DATA_DIR}/wg${WG_CONFIG_ID}/client
  fi
  ADDRESS_NO=${1}
  cat <<EOF > ${DATA_DIR}/wg${WG_CONFIG_ID}/client/peer${1}.conf
[Interface]
Address = ${WG_NET_IP%.*}.$(( ${WG_NET_IP##*.} + ADDRESS_NO ))/32
PrivateKey = ${2}

[Peer]
PublicKey = ${3}
$(if [ ! -z "${4}" ]; then echo "PresharedKey = ${4}"; fi)
Endpoint = ${SERVER_IP}:${SERVER_LISTEN_PORT}
AllowedIPs = ${WG_NET_IP}/32

EOF
if [ "${GENERATE_QR}" == "true" ]; then
  qrencode -t png -o ${DATA_DIR}/wg${WG_CONFIG_ID}/client/peer${1}.png -r ${DATA_DIR}/wg${WG_CONFIG_ID}/client/peer${1}.conf
fi
}

if [ ! -f ${DATA_DIR}/wg${WG_CONFIG_ID}/server/wg${WG_CONFIG_ID}.conf ]; then
  echo "---Generating wg${WG_CONFIG_ID}.conf---"
  PRIV_KEY=$(wg genkey)
  create_servercfg "${PRIV_KEY}"
  unset PRIV_KEY PUB_KEY
else
  echo "---Configuration wg${WG_CONFIG_ID}.conf found!---"
fi

for ((i=1; i<=PEERS; i++))
do
  if ! grep -q "^\[Peer\] # Client ${i}" ${DATA_DIR}/wg${WG_CONFIG_ID}/server/wg${WG_CONFIG_ID}.conf ; then
    echo "---Generating configuration for Peer ${i}---"
    PRIV_KEY=$(wg genkey)
    PUBLIC_KEY=$(echo ${PRIV_KEY} | wg pubkey)
    if [ "${GENERATE_PSK}" == "true" ]; then
      WG_PSK=$(wg genpsk)
    fi
    create_peer "${i}" "${PUBLIC_KEY}" "${WG_PSK}"
    SRV_PRIV_KEY=$(grep -A 5 "\[Interface\]" ${DATA_DIR}/wg${WG_CONFIG_ID}/server/wg${WG_CONFIG_ID}.conf | awk '/PrivateKey = /{print $NF; exit}')
    SRV_PUBLIC_KEY=$(echo $SRV_PRIV_KEY | wg pubkey)
    create_config "${i}" "${PRIV_KEY}" "${SRV_PUBLIC_KEY}" "${WG_PSK}"
    unset SRV_PRIV_KEY SRV_PUBLIC_KEY PRIV_KEY PUBLIC_KEY WG_PSK
  else
    echo "---Client ${i} already existing---"
  fi
done

cp ${DATA_DIR}/wg${WG_CONFIG_ID}/server/wg${WG_CONFIG_ID}.conf /etc/wireguard/wg${WG_CONFIG_ID}.conf

echo "---Taking ownership of data...---"
chown -R root:${GID} /opt/scripts
chown root:root /etc/wireguard/wg${WG_CONFIG_ID}.conf
chmod 440 /etc/sudoers.d/${USER}
chmod -R 750 /opt/scripts
chown ${UID}:${GID} ${DATA_DIR}
chown -R ${UID}:${GID} ${DATA_DIR}/wg${WG_CONFIG_ID}

echo "---Starting...---"
term_handler() {
  echo "---Stopping WireGuard tunnel wg${WG_CONFIG_ID}---"
  wg-quick down wg${WG_CONFIG_ID} > /dev/null 2>&1
  kill $(pidof sleep)
  exit 143;
}

trap 'kill ${!}; term_handler' SIGTERM
su ${USER} -c "/opt/scripts/start-server.sh" &
killpid="$!"
while true
do
	wait $killpid
	exit 0;
done