From eb97beadca47cee3aff6c0d3a72c54be6156956b Mon Sep 17 00:00:00 2001
From: Vladislav Yarmak <vladislav-ex-src@vm-0.com>
Date: Fri, 26 Mar 2021 22:09:37 +0200
Subject: [PATCH] seclient: add discover method

---
 main.go              |  7 +++++++
 seclient/messages.go | 19 ++++++++++++++++---
 seclient/seclient.go | 44 ++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 67 insertions(+), 3 deletions(-)

diff --git a/main.go b/main.go
index cb0029c..5442946 100644
--- a/main.go
+++ b/main.go
@@ -41,4 +41,11 @@ func main() {
 		log.Fatalln(err)
 	}
 	log.Printf("Geo List: %#v", geos)
+
+	log.Println("------------ GETTING IP LIST ------------")
+	ips, err := seclient.Discover(context.TODO(), "\"EU\",,")
+	if err != nil {
+		log.Fatalln(err)
+	}
+	log.Printf("IP List: %#v", ips)
 }
diff --git a/seclient/messages.go b/seclient/messages.go
index 2150718..9c1d59a 100644
--- a/seclient/messages.go
+++ b/seclient/messages.go
@@ -61,13 +61,26 @@ type SERegisterDeviceResponse struct {
 }
 
 type SEGeoEntry struct {
-	Country     string `json:"country"`
+	Country     string `json:"country,omitempty"`
 	CountryCode string `json:"country_code"`
 }
 
 type SEGeoListResponse struct {
 	Data struct {
 		Geos []SEGeoEntry `json:"geos"`
-	}
-	Status SEStatusPair
+	} `json:"data"`
+	Status SEStatusPair `json:"return_code"`
+}
+
+type SEIPEntry struct {
+	Geo   SEGeoEntry `json:"geo"`
+	IP    string     `json:"ip"`
+	Ports []uint16   `json:"ports"`
+}
+
+type SEDiscoverResponse struct {
+	Data struct {
+		IPs []SEIPEntry `json:"ips"`
+	} `json:"data"`
+	Status SEStatusPair `json:"return_code"`
 }
diff --git a/seclient/seclient.go b/seclient/seclient.go
index b0747f5..24ac92e 100644
--- a/seclient/seclient.go
+++ b/seclient/seclient.go
@@ -247,6 +247,50 @@ func (c *SEClient) GeoList(ctx context.Context) ([]SEGeoEntry, error) {
 	return geoListRes.Data.Geos, nil
 }
 
+func (c *SEClient) Discover(ctx context.Context, requestedGeo string) ([]SEIPEntry, error) {
+	geoListInput := url.Values{
+		"serial_no": {c.AssignedDeviceIDHash},
+		"requested_geo": {requestedGeo},
+	}
+	req, err := http.NewRequestWithContext(
+		ctx,
+		"POST",
+		c.Settings.Endpoints.Discover,
+		strings.NewReader(geoListInput.Encode()),
+	)
+	if err != nil {
+		return nil, err
+	}
+	c.populateRequest(req)
+	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
+	req.Header.Set("Accept", "application/json")
+
+	resp, err := c.HttpClient.Do(req)
+	if err != nil {
+		return nil, err
+	}
+
+	if resp.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("bad http status: %s", resp.Status)
+	}
+
+	decoder := json.NewDecoder(resp.Body)
+	var discoverRes SEDiscoverResponse
+	err = decoder.Decode(&discoverRes)
+	cleanupBody(resp.Body)
+
+	if err != nil {
+		return nil, err
+	}
+
+	if discoverRes.Status.Code != SE_STATUS_OK {
+		return nil, fmt.Errorf("API responded with error message: code=%d, msg=\"%s\"",
+			discoverRes.Status.Code, discoverRes.Status.Message)
+	}
+
+	return discoverRes.Data.IPs, nil
+}
+
 func (c *SEClient) populateRequest(req *http.Request) {
 	req.Header["SE-Client-Version"] = []string{c.Settings.ClientVersion}
 	req.Header["SE-Operating-System"] = []string{c.Settings.OperatingSystem}