diff --git a/init.sh b/init.sh
index 73b9af1..b1f54d2 100755
--- a/init.sh
+++ b/init.sh
@@ -7,19 +7,26 @@ WG_IF=${WG_INTERFACE:-wg0}
 PHY_IF=${WG_HOST_INTERFACE:-eth0}
 ADDRESS=${WG_ADDRESS:-192.168.2.1/24}
 
+function shutdown() {
+ wg-quick down $WG_IF
+}
+
 if [ ! -f "/etc/wireguard/$WG_IF.conf" ]; then
  mkdir -p /etc/wireguard/keys
  wg genkey | tee /etc/wireguard/keys/$WG_IF | wg pubkey > /etc/wireguard/keys/$WG_IF.pub
  
  echo "[Interface]" > /etc/wireguard/$WG_IF.conf
  echo "PrivateKey = $(cat /etc/wireguard/keys/$WG_IF)" >> /etc/wireguard/$WG_IF.conf
- echo "Address = $ADDRESS" >> /etc/wireguard/$WG_IF.conf
  echo "ListenPort = 51820" >> /etc/wireguard/$WG_IF.conf
  echo "PostUp = iptables -A FORWARD -i $WG_IF -j ACCEPT; iptables -t nat -A POSTROUTING -o $PHY_IF -j MASQUERADE" >> /etc/wireguard/$WG_IF.conf
  echo "PostDown = iptables -D FORWARD -i $WG_IF -j ACCEPT; iptables -t nat -D POSTROUTING -o $PHY_IF -j MASQUERADE" >> /etc/wireguard/$WG_IF.conf
  echo "SaveConfig = true" >> /etc/wireguard/$WG_IF.conf
+ echo "[Peer]"
+ echo "Endpoint = $ADDRESS" >> /etc/wireguard/$WG_IF.conf
  
  echo "Generated Config for $WG_IF, Publickey: $(cat /etc/wireguard/keys/$WG_IF.pub)"
 fi
 
+trap shutdown EXIT
+
 /usr/bin/wireguard-go -f $WG_IF