From ed775ab0700f8e2620f4017477b8ae68c09dcaab Mon Sep 17 00:00:00 2001
From: Vladislav Yarmak <vladislav-ex-src@vm-0.com>
Date: Fri, 25 Jun 2021 04:53:53 +0300
Subject: [PATCH] main: add location picking

---
 main.go | 29 +++++++++++++++++++++++++++--
 1 file changed, 27 insertions(+), 2 deletions(-)

diff --git a/main.go b/main.go
index 659c1d5..dfd237e 100644
--- a/main.go
+++ b/main.go
@@ -8,6 +8,7 @@ import (
 	"encoding/json"
 	"errors"
 	"flag"
+	"math/rand"
 	"fmt"
 	"io/ioutil"
 	"log"
@@ -219,7 +220,11 @@ func run() int {
 		}
 		proxyHostname = bestLocation.Hostname
 	} else {
-		return 0
+		proxyHostname = pickServer(serverList, args.location)
+		if proxyHostname == "" {
+			mainLogger.Critical("Server for location \"%s\" not found.", args.location)
+			return 13
+		}
 	}
 
 	//runTicker(context.Background(), args.refresh, args.refreshRetry, func(ctx context.Context) error {
@@ -263,7 +268,6 @@ func run() int {
 		}
 	}
 
-	// TODO: set servername
 	proxyNetAddr := net.JoinHostPort(proxyHostname, strconv.FormatUint(uint64(ASSUMED_PROXY_PORT), 10))
 	handlerDialer := NewProxyDialer(proxyNetAddr, proxyHostname, auth, caPool, dialer)
 	mainLogger.Info("Endpoint: %s", proxyNetAddr)
@@ -336,6 +340,27 @@ func printProxies(username, password string, serverList wndclient.ServerList) in
 	return 0
 }
 
+func pickServer(serverList wndclient.ServerList, location string) string {
+	var candidates []string
+	for _, country := range serverList {
+		for _, group := range country.Groups {
+			for _, host := range group.Hosts {
+				if country.Name + "/" + group.City == location {
+					candidates = append(candidates, host.Hostname)
+				}
+			}
+		}
+	}
+
+	if len(candidates) == 0 {
+		return ""
+	}
+
+	rnd := rand.New(rand.NewSource(time.Now().UnixNano()))
+
+	return candidates[rnd.Intn(len(candidates))]
+}
+
 func loadState(filename string) (*wndclient.WndClientState, error) {
 	file, err := os.Open(filename)
 	if err != nil {