From fbbf94055d48b1fef9e563a5151bd6b695e6b9fd Mon Sep 17 00:00:00 2001
From: Kavin <20838718+FireMasterK@users.noreply.github.com>
Date: Mon, 4 Sep 2023 20:52:06 +0100
Subject: [PATCH] Add wireguard servers to list from api.

---
 .../privateinternetaccess/updater/api.go      |  1 +
 .../updater/hosttoserver.go                   | 97 +++++++++++++------
 .../privateinternetaccess/updater/servers.go  |  6 ++
 3 files changed, 76 insertions(+), 28 deletions(-)

diff --git a/internal/provider/privateinternetaccess/updater/api.go b/internal/provider/privateinternetaccess/updater/api.go
index fc17d2aa6..5c249d0f0 100644
--- a/internal/provider/privateinternetaccess/updater/api.go
+++ b/internal/provider/privateinternetaccess/updater/api.go
@@ -27,6 +27,7 @@ type regionData struct {
 	Servers     struct {
 		UDP []serverData `json:"ovpnudp"`
 		TCP []serverData `json:"ovpntcp"`
+		WG  []serverData `json:"wg"`
 	} `json:"servers"`
 }
 
diff --git a/internal/provider/privateinternetaccess/updater/hosttoserver.go b/internal/provider/privateinternetaccess/updater/hosttoserver.go
index 76cd6efc8..0927b72f8 100644
--- a/internal/provider/privateinternetaccess/updater/hosttoserver.go
+++ b/internal/provider/privateinternetaccess/updater/hosttoserver.go
@@ -7,51 +7,92 @@ import (
 	"github.com/qdm12/gluetun/internal/models"
 )
 
-type nameToServer map[string]models.Server
+type nameToServer map[string][]models.Server
 
 func (nts nameToServer) add(name, hostname, region string,
 	tcp, udp, portForward bool, ip netip.Addr) (change bool) {
-	server, ok := nts[name]
-	if !ok {
-		change = true
-		server.VPN = vpn.OpenVPN
-		server.ServerName = name
-		server.Hostname = hostname
-		server.Region = region
-		server.PortForward = portForward
-	}
 
-	if !server.TCP && tcp {
-		change = true
-		server.TCP = tcp
-	}
-	if !server.UDP && udp {
-		change = true
-		server.UDP = udp
-	}
+	var server models.Server
 
-	ipFound := false
-	for _, existingIP := range server.IPs {
-		if ip == existingIP {
-			ipFound = true
+	// Check for existing server for this name.
+	servers := nts[name]
+	for i, existingServer := range servers {
+		if existingServer.ServerName != name {
+			continue
+		}
+
+		server = existingServer
+
+		switch existingServer.VPN {
+		case vpn.OpenVPN:
+			// Update OpenVPN supported protocols and return
+			if !existingServer.TCP {
+				servers[i].TCP = tcp
+				change = true
+			}
+			if !existingServer.UDP {
+				servers[i].UDP = udp
+				change = true
+			}
+			ipFound := false
+			for _, existingIP := range server.IPs {
+				if ip == existingIP {
+					ipFound = true
+					break
+				}
+			}
+			if !ipFound {
+				server.IPs = append(server.IPs, ip)
+				change = true
+			}
+			break
+		case vpn.Wireguard:
+			// Update IPs and return
+			ipFound := false
+			for _, existingIP := range server.IPs {
+				if ip == existingIP {
+					ipFound = true
+					break
+				}
+			}
+			if !ipFound {
+				server.IPs = append(server.IPs, ip)
+				change = true
+			}
 			break
 		}
+
+		break
 	}
 
-	if !ipFound {
+	if server.ServerName == "" {
 		change = true
-		server.IPs = append(server.IPs, ip)
+		nts[name] = append(servers, models.Server{
+			VPN:         vpn.OpenVPN,
+			Region:      region,
+			Hostname:    hostname,
+			PortForward: portForward,
+			ServerName:  name,
+			TCP:         tcp,
+			UDP:         udp,
+		})
+		nts[name] = append(servers, models.Server{
+			VPN:         vpn.Wireguard,
+			Region:      region,
+			Hostname:    hostname,
+			PortForward: portForward,
+			ServerName:  name,
+			IPs:         []netip.Addr{ip},
+		})
 	}
 
-	nts[name] = server
-
 	return change
 }
 
 func (nts nameToServer) toServersSlice() (servers []models.Server) {
 	servers = make([]models.Server, 0, len(nts))
-	for _, server := range nts {
-		servers = append(servers, server)
+	for _, hostServers := range nts {
+		servers = append(servers, hostServers...)
 	}
 	return servers
 }
diff --git a/internal/provider/privateinternetaccess/updater/servers.go b/internal/provider/privateinternetaccess/updater/servers.go
index 96d75187c..fa7e99024 100644
--- a/internal/provider/privateinternetaccess/updater/servers.go
+++ b/internal/provider/privateinternetaccess/updater/servers.go
@@ -89,6 +89,12 @@ func addData(regions []regionData, nts nameToServer) (change bool) {
 				change = true
 			}
 		}
+
+		for _, server := range region.Servers.WG {
+			if nts.add(server.CN, region.DNS, region.Name, false, false, region.PortForward, server.IP) {
+				change = true
+			}
+		}
 	}
 
 	return change