From ebac39dc332ec0fe17de3a17df6cf94a7cbdc9a0 Mon Sep 17 00:00:00 2001
From: Pavel Muhortov <pavel.muhortov@hmp.today>
Date: Sun, 7 Jan 2024 15:49:39 +0300
Subject: [PATCH] improved execution logic

---
 wg-client-management.sh | 163 +++++++++++++++++++++++++---------------
 1 file changed, 104 insertions(+), 59 deletions(-)

diff --git a/wg-client-management.sh b/wg-client-management.sh
index 4db859c..5237ff7 100644
--- a/wg-client-management.sh
+++ b/wg-client-management.sh
@@ -163,7 +163,7 @@ startsendtlgm() {
 }
 
 #######################################
-# Create wireguard client certificates.
+# Create wireguard client peer.
 # Globals:
 #   clientname
 #   clientaddr
@@ -171,22 +171,38 @@ startsendtlgm() {
 # Arguments:
 #   None
 #######################################
-createcert() {
-  wg genkey | tee "/etc/wireguard/pki/${clientname}-private.key" | wg pubkey > "/etc/wireguard/pki/${clientname}-public.key"
+createpeer() {
+  if [ -f "/etc/wireguard/pki/${clientname}-private.key" ]; then
+    addtologs "${clientname} private key exists, create skipped"
+  else
+    wg genkey | tee "/etc/wireguard/pki/${clientname}-private.key" > /dev/null 2>&1
+    addtologs "created ${clientname} new private key"
+  fi
+  wg pubkey > "/etc/wireguard/pki/${clientname}-public.key" < "/etc/wireguard/pki/${clientname}-private.key"
+  addtologs "created ${clientname} wireguard certificates"
+
   clientpublkey=$(cat "/etc/wireguard/pki/${clientname}-public.key")
-  clientprivkey=$(cat "/etc/wireguard/pki/${clientname}-private.key")
-  wg set "${iface_name}" peer "${clientpublkey}" \
-                         allowed-ips "${clientaddr}/32" \
-                         persistent-keepalive 5
-  {
-    printf "%s\n" "[Peer]"
-    printf "%s\n" "  PublicKey = ${clientpublkey}"
-    printf "%s\n" "  AllowedIPs = ${clientaddr}/32"
-    printf "%s\n" "  PersistentKeepalive = 5"
-  } >> "${servercfgname}"
-  addtologs "created ${clientname} wireguard certificate"
-  ip -4 route add "${clientaddr}/32" dev "${iface_name}"
-  addtologs "created route to ${clientname} peer"
+  if grep -q -w "${clientpublkey}" "${servercfgname}"; then
+    addtologs "${clientname} peer configuration exists, create skipped"
+  else
+    wg set "${iface_name}" peer "${clientpublkey}" \
+                           allowed-ips "${clientaddr}/32" \
+                           persistent-keepalive 5
+    {
+      printf "%s\n" "[Peer]"
+      printf "%s\n" "  PublicKey = ${clientpublkey}"
+      printf "%s\n" "  AllowedIPs = ${clientaddr}/32"
+      printf "%s\n" "  PersistentKeepalive = 5"
+    } >> "${servercfgname}"
+    addtologs "created ${clientname} wireguard peer configuration"
+  fi
+
+  if ip ro | grep -q -w "${clientaddr}"; then
+    addtologs "${clientname} peer address exists in routes, create skipped"
+  else
+    ip -4 route add "${clientaddr}/32" dev "${iface_name}"
+    addtologs "created ${clientname} peer route"
+  fi
 }
 
 #######################################
@@ -201,6 +217,7 @@ createcert() {
 #   None
 #######################################
 createconf() {
+  clientprivkey=$(cat "/etc/wireguard/pki/${clientname}-private.key")
   clientconf=$(cat "${clientconfdef}")
   clientconf=${clientconf//clientaddr/${clientaddr}}
   clientconf=${clientconf//clientprivkey/${clientprivkey}}
@@ -224,14 +241,14 @@ createinfo() {
     "" \
     "${additional}" \
     "" \
-    "Peer Address: \`${clientaddr}\`" \
+    "Peer Address: \`$(grep 'Address' "/etc/wireguard/${clientname}.conf" | awk '{print $3}')\`" \
   )
   qrencode -o "/etc/wireguard/${clientname}.png" -t png -s 6 < "/etc/wireguard/${clientname}.conf"
   addtologs "created ${clientname} qr code"
 }
 
 #######################################
-# Delete wireguard client certificates.
+# Delete wireguard client peer.
 # Globals:
 #   clientname
 #   clientpublkey
@@ -239,24 +256,43 @@ createinfo() {
 # Arguments:
 #   None
 #######################################
-deletecert() {
-  clientpublkey=$(cat "/etc/wireguard/pki/${clientname}-public.key")
-  clientprivkey=$(cat "/etc/wireguard/pki/${clientname}-private.key")
-  wg set "${iface_name}" peer "${clientpublkey}" remove
-  rm -f "/etc/wireguard/pki/${clientname}-public.key"
-  rm -f "/etc/wireguard/pki/${clientname}-private.key"
-  addtologs "deleted ${clientname} wireguard certificate"
-  # PublicKey =
-  s2=$(grep -n "${clientpublkey}" "${servercfgname}" | cut -d":" -f1)
-  # [Peer]
-  s1=$(( s2 - 1 ))
-  # AllowedIPs =
-  s3=$(( s2 + 1 ))
-  # PersistentKeepalive =
-  s4=$(( s2 + 2 ))
-  sed -i "${s1}d;${s2}d;${s3}d;${s4}d" "${servercfgname}"
-  ip -4 route del "${clientaddr}/32" dev "${iface_name}"
-  addtologs "deleted route to ${clientname} peer"
+deletepeer() {
+  if [ -f "/etc/wireguard/pki/${clientname}-private.key" ]; then
+    wg pubkey > "/etc/wireguard/pki/${clientname}-public.key" < "/etc/wireguard/pki/${clientname}-private.key"
+  fi
+
+  if [ -f "/etc/wireguard/pki/${clientname}-public.key" ]; then
+    clientpublkey=$(cat "/etc/wireguard/pki/${clientname}-public.key")
+    if grep -q -w "${clientpublkey}" "${servercfgname}"; then
+      wg set "${iface_name}" peer "${clientpublkey}" remove
+      # PublicKey =
+      s2=$(grep -n "${clientpublkey}" "${servercfgname}" | cut -d":" -f1)
+      # [Peer]
+      s1=$(( s2 - 1 ))
+      # AllowedIPs =
+      s3=$(( s2 + 1 ))
+      # PersistentKeepalive =
+      s4=$(( s2 + 2 ))
+      sed -i "${s1}d;${s2}d;${s3}d;${s4}d" "${servercfgname}"
+      addtologs "deleted ${clientname} wireguard peer configuration"
+    else
+      addtologs "${clientname} peer configuration does not exist, delete skipped"
+    fi
+
+    rm -f "/etc/wireguard/pki/${clientname}-public.key" > /dev/null 2>&1
+    rm -f "/etc/wireguard/pki/${clientname}-private.key" > /dev/null 2>&1
+    addtologs "deleted ${clientname} wireguard certificates"
+  else
+    addtologs "${clientname} certificates do not exist, delete skipped"
+  fi
+
+  if ip ro | grep -q -w "${clientaddr} dev ${iface_name}"; then
+    ip -4 route del "${clientaddr}/32" dev "${iface_name}"
+    addtologs "deleted ${clientname} peer route"
+  else
+    addtologs "${clientname} peer route does not exist, delete skipped"
+  fi
+
 }
 
 #######################################
@@ -267,8 +303,12 @@ deletecert() {
 #   None
 #######################################
 deleteconf() {
-  rm -f "/etc/wireguard/${clientname}.conf"
-  addtologs "deleted ${clientname} wireguard config file"
+  if [ -f "/etc/wireguard/${clientname}.conf" ]; then
+    rm -f "/etc/wireguard/${clientname}.conf"
+    addtologs "deleted ${clientname} wireguard config file"
+  else
+    addtologs "${clientname} wireguard config does not exist, delete skipped"
+  fi
 }
 
 #######################################
@@ -279,8 +319,12 @@ deleteconf() {
 #   None
 #######################################
 deleteinfo() {
-  rm -f "/etc/wireguard/${clientname}.png"
-  addtologs "deleted ${clientname} qr code"
+  if [ -f "/etc/wireguard/${clientname}.png" ]; then
+    rm -f "/etc/wireguard/${clientname}.png"
+    addtologs "deleted ${clientname} qr code"
+  else
+    addtologs "${clientname} qr code does not exist, delete skipped"
+  fi
 }
 
 #
@@ -307,27 +351,33 @@ done
 time=$(date +%s)
 logs=/dev/null
 
-if ! command -v qrencode &> /dev/null || \
-  ! command -v /usr/local/bin/sendmail.py &> /dev/null || \
-  ! command -v python3 &> /dev/null; then
-  execerror "Not found dependencies"
-fi
-
 #
 # MAIN:
 #
 
 if checkroot; then
+  if ! command -v qrencode &> /dev/null || \
+     ! command -v grepcidr &> /dev/null || \
+     ! command -v /usr/local/bin/sendmail.py &> /dev/null || \
+     ! command -v python3 &> /dev/null || \
+     ! command -v curl &> /dev/null; then
+    execerror "Not found dependencies"
+  fi
   getconfig
   if [ "${clienttodo}" == "add" ] && \
+     [ -n "${iface_name}" ] && \
      [ -n "${clientname}" ] && \
-     grepcidr "0.0.0.0/0" <(echo "${clientaddr}") >/dev/null; then
-    if [ -f "/etc/wireguard/${clientname}.conf" ] || \
+     grepcidr "0.0.0.0/0" <(echo "${clientaddr}") > /dev/null; then
+    if ip ro | grep -q -w "${clientaddr}" || \
       grep -q -w "${clientaddr}/32" "${servercfgname}"; then
-      execerror "wireguard config exist or address used, exit"
+      addtologs "${clientaddr} address used, create peer, conf skipped"
     else
-      createcert
+      createpeer
       createconf
+    fi
+    if [ ! -f "/etc/wireguard/${clientname}.conf" ]; then
+      addtologs "${clientname}.conf not found, create info skipped"
+    else
       createinfo
       startsendmail
       startsendtlgm
@@ -337,17 +387,12 @@ if checkroot; then
       systemctl restart "wg-quick@${iface_name}"
     fi
   elif [ "${clienttodo}" == "del" ] && \
+     [ -n "${iface_name}" ] && \
      [ -n "${clientname}" ] && \
      grepcidr "0.0.0.0/0" <(echo "${clientaddr}") >/dev/null; then
-    if [ -f "/etc/wireguard/${clientname}.conf" ]; then
-      deleteconf
-    fi
-    if [ -f "/etc/wireguard/${clientname}.png" ]; then
-      deleteinfo
-    fi
-    if grep -q -w "${clientaddr}/32" "${servercfgname}"; then
-      deletecert
-    fi
+    deleteconf
+    deleteinfo
+    deletepeer
     if [ "${resetforce}" -eq 1 ];then
       addtologs "restarting wg-quick@${iface_name}..."
       systemctl "restart wg-quick@${iface_name}"