From d2d5f2bb4d9a5b7ae3dbc26a538facabc42009b2 Mon Sep 17 00:00:00 2001
From: bcmmbaga <bethuelmbaga12@gmail.com>
Date: Wed, 31 Jan 2024 19:40:49 +0300
Subject: [PATCH] add support for country wide checks

---
 .../server/http/posture_checks_handler.go     |  7 +-
 management/server/posture/geo_location.go     | 15 ++--
 .../server/posture/geo_location_test.go       | 88 +++++++++++++++++++
 3 files changed, 101 insertions(+), 9 deletions(-)

diff --git a/management/server/http/posture_checks_handler.go b/management/server/http/posture_checks_handler.go
index 8ca68fba31b..c0c86d156e5 100644
--- a/management/server/http/posture_checks_handler.go
+++ b/management/server/http/posture_checks_handler.go
@@ -3,6 +3,7 @@ package http
 import (
 	"encoding/json"
 	"net/http"
+	"regexp"
 
 	"github.com/gorilla/mux"
 	"github.com/rs/xid"
@@ -247,8 +248,10 @@ func validatePostureChecksUpdate(req api.PostureCheckUpdate) error {
 			if loc.CountryCode == "" {
 				return status.Errorf(status.InvalidArgument, "country code for geolocation check shouldn't be empty")
 			}
-			if loc.CityName == "" {
-				return status.Errorf(status.InvalidArgument, "city name for geolocation check shouldn't be empty")
+
+			countryCodeRegex := regexp.MustCompile("^[a-zA-Z]{2}$")
+			if !countryCodeRegex.MatchString(loc.CountryCode) {
+				return status.Errorf(status.InvalidArgument, "country code must be 2 letters (ISO 3166-1 alpha-2 format)")
 			}
 		}
 
diff --git a/management/server/posture/geo_location.go b/management/server/posture/geo_location.go
index c52813e5ae8..303c1f1fd25 100644
--- a/management/server/posture/geo_location.go
+++ b/management/server/posture/geo_location.go
@@ -31,16 +31,17 @@ type GeoLocationCheck struct {
 
 func (g *GeoLocationCheck) Check(peer nbpeer.Peer) (bool, error) {
 	for _, loc := range g.Locations {
-		if loc.CountryCode == peer.Meta.Location.CountryCode && loc.CityName == peer.Meta.Location.CityName {
-			switch g.Action {
-			case GeoLocationActionDeny:
-				return false, nil
-			case GeoLocationActionAllow:
-				return true, nil
+		if loc.CountryCode == peer.Meta.Location.CountryCode {
+			if loc.CityName == "" || loc.CityName == peer.Meta.Location.CityName {
+				switch g.Action {
+				case GeoLocationActionDeny:
+					return false, nil
+				case GeoLocationActionAllow:
+					return true, nil
+				}
 			}
 		}
 	}
-
 	// At this point, no location in the list matches the peer's location
 	// For action deny and no location match, allow the peer
 	if g.Action == GeoLocationActionDeny {
diff --git a/management/server/posture/geo_location_test.go b/management/server/posture/geo_location_test.go
index ddaa7ef72ae..a2345f09ef0 100644
--- a/management/server/posture/geo_location_test.go
+++ b/management/server/posture/geo_location_test.go
@@ -42,6 +42,27 @@ func TestGeoLocationCheck_Check(t *testing.T) {
 			wantErr: false,
 			isValid: true,
 		},
+		{
+			name: "Peer location matches the location in the allow country only",
+			input: peer.Peer{
+				Meta: peer.PeerSystemMeta{
+					Location: Location{
+						CountryCode: "DE",
+						CityName:    "Berlin",
+					},
+				},
+			},
+			check: GeoLocationCheck{
+				Locations: []Location{
+					{
+						CountryCode: "DE",
+					},
+				},
+				Action: GeoLocationActionAllow,
+			},
+			wantErr: false,
+			isValid: true,
+		},
 		{
 			name: "Peer location doesn't match the location in the allow sets",
 			input: peer.Peer{
@@ -68,6 +89,27 @@ func TestGeoLocationCheck_Check(t *testing.T) {
 			wantErr: false,
 			isValid: false,
 		},
+		{
+			name: "Peer location doesn't match the location in the allow country only",
+			input: peer.Peer{
+				Meta: peer.PeerSystemMeta{
+					Location: Location{
+						CountryCode: "DE",
+						CityName:    "Frankfurt am Main",
+					},
+				},
+			},
+			check: GeoLocationCheck{
+				Locations: []Location{
+					{
+						CountryCode: "US",
+					},
+				},
+				Action: GeoLocationActionAllow,
+			},
+			wantErr: false,
+			isValid: false,
+		},
 		{
 			name: "Peer location matches the location in the deny sets",
 			input: peer.Peer{
@@ -94,6 +136,30 @@ func TestGeoLocationCheck_Check(t *testing.T) {
 			wantErr: false,
 			isValid: false,
 		},
+		{
+			name: "Peer location matches the location in the deny country only",
+			input: peer.Peer{
+				Meta: peer.PeerSystemMeta{
+					Location: Location{
+						CountryCode: "DE",
+						CityName:    "Berlin",
+					},
+				},
+			},
+			check: GeoLocationCheck{
+				Locations: []Location{
+					{
+						CountryCode: "DE",
+					},
+					{
+						CountryCode: "US",
+					},
+				},
+				Action: GeoLocationActionDeny,
+			},
+			wantErr: false,
+			isValid: false,
+		},
 		{
 			name: "Peer location doesn't match the location in the deny sets",
 			input: peer.Peer{
@@ -120,6 +186,28 @@ func TestGeoLocationCheck_Check(t *testing.T) {
 			wantErr: false,
 			isValid: true,
 		},
+		{
+			name: "Peer location doesn't match the location in the deny country only",
+			input: peer.Peer{
+				Meta: peer.PeerSystemMeta{
+					Location: Location{
+						CountryCode: "DE",
+						CityName:    "Frankfurt am Main",
+					},
+				},
+			},
+			check: GeoLocationCheck{
+				Locations: []Location{
+					{
+						CountryCode: "US",
+						CityName:    "Los Angeles",
+					},
+				},
+				Action: GeoLocationActionDeny,
+			},
+			wantErr: false,
+			isValid: true,
+		},
 	}
 
 	for _, tt := range tests {