#!/usr/bin/env bash
export $(cat .env | xargs)

if [ "$#" -ne 2 ]; then
  echo "Usage: $0 <team_count> <peer_count>"
  exit 1
else
  TEAM_COUNT=$1
  PEER_COUNT=$2
fi

# If WG_DOMAIN or WG_PORT are not set, exit
if [ -z "$WG_HOST" ] || [ -z "$WG_PORT" ]; then
  echo "Error: WG_DOMAIN and WG_PORT must be set in the .env file."
  exit 1
fi


WG_DIR="wireguard"
WG_CONF="$WG_DIR/server.conf"
SERVER_PRIV_FILE="$WG_DIR/config/server_private.key"
SERVER_PUB_FILE="$WG_DIR/config/server_public.key"
SERVER_ENDPOINT="$WG_HOST:$WG_PORT"

# If wireguard is already set, exit

if [ -f "$WG_DIR" ]; then
  read -p "$WG_DIR already exists. Do you want to overwrite it? (y/n): " choice
  rm -rf $WG_DIR
fi

umask 077 # Set restrictive permissions

mkdir -p "$WG_DIR/config"
echo "Generating WireGuard server keys..."
wg genkey | tee $SERVER_PRIV_FILE | wg pubkey > $SERVER_PUB_FILE

cat >> $WG_CONF <<EOL
[Interface]
Address = 10.201.1.1/15
ListenPort = $WG_PORT
PrivateKey = $(cat $SERVER_PRIV_FILE)
PostUp = sysctl -w net.ipv4.ip_forward=1
PostUp = iptables -t nat -A POSTROUTING -o %i -j MASQUERADE
PostUp = iptables -A FORWARD -i %i -j ACCEPT
PostUp = iptables -A FORWARD -o %i -m state --state RELATED,ESTABLISHED -j ACCEPT
PostUp = iptables -t raw -I PREROUTING 1 -s 10.201.1.1 -d 10.202.0.0/15 -j ACCEPT
PostUp = iptables -t filter -I FORWARD 1 -d 10.202.0.0/15 -j ACCEPT

PostDown = iptables -t nat -A POSTROUTING -o %i -j MASQUERADE
PostDown = iptables -D FORWARD -i %i -j ACCEPT
PostDown = iptables -D FORWARD -o %i -m state --state RELATED,ESTABLISHED -j ACCEPT
PostDown = iptables -t raw -D PREROUTING -s 10.201.1.1 -d 10.202.0.0/15 -j ACCEPT
PostDown = iptables -t filter -D FORWARD -d 10.202.0.0/15 -j ACCEPT

EOL

function generate_peer() {
  local peer_conf=$1
  local peer_ip=$2
  local name=$3
  local peer_private_key=$(wg genkey)
  local peer_public_key=$(echo $peer_private_key | wg pubkey)

  cat >> $WG_CONF <<EOL
[Peer]
# $name
PublicKey = $peer_public_key
AllowedIPs = $peer_ip/32
EOL

  mkdir -p $(dirname $peer_conf)
  cat > $peer_conf <<EOL
[Interface]
PrivateKey = $peer_private_key
Address = $peer_ip/32

[Peer]
PublicKey = $(cat $SERVER_PUB_FILE)
Endpoint = $SERVER_ENDPOINT
AllowedIPs = $peer_ip/15, 10.202.0.0/24
PersistentKeepalive = 25
EOL

  echo "Generated configuration for $name at $peer_conf"
}

# NOP
VULNBOX_IP=$(printf "$CTF_IPPATTERN" 0)
PEER_NAME="nop_vulnbox"
PEER_CONF="$WG_DIR/nop_vulnbox.conf"

generate_peer $PEER_CONF $VULNBOX_IP $PEER_NAME

# Team and Peer generation loop
for team_id in $(seq 1 $TEAM_COUNT); do
  VULNBOX_IP=$(printf "$CTF_IPPATTERN" $team_id)
  PEER_NAME="team${team_id}_vulnbox"
  PEER_CONF="$WG_DIR/${team_id}/vulnbox.conf"
  generate_peer $PEER_CONF $VULNBOX_IP $PEER_NAME
  for peer_id in $(seq 1 $PEER_COUNT); do
    PEER_IP=$(echo $VULNBOX_IP | awk -F. -v last=$((100 + $peer_id)) '{print $1"."$2"."$3"."last}') # Replace last octet with 100 + peer_id
    PEER_NAME="team${team_id}_peer${peer_id}"
    PEER_CONF="$WG_DIR/${team_id}/${PEER_NAME}.conf"
    generate_peer $PEER_CONF $PEER_IP $PEER_NAME
  done
done

echo "WireGuard configuration generation complete."
echo "To start the WireGuard server, run: wg-quick up $WG_CONF"
