Merge pull request #52 from yashviagrawal/develop

Appknox-Go: Region and Host functionality

README.md

@@ -33,6 +33,7 @@ Flags:
   -k, --insecure              Disable Security Checks
       --pac string            pac file path or url
       --proxy string          proxy url
+      --region string         Region names, e.g., global, saudi, uae. By default, global is used
       --version               version for appknox
 
 Use "appknox [command] --help" for more information about a command.

cmd/root.go

@@ -5,7 +5,7 @@ import (
 	"os"
 	"path/filepath"
 
-	"github.com/appknox/appknox-go/appknox"
+	// "github.com/appknox/appknox-go/appknox"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
 )
@@ -35,10 +35,15 @@ func init() {
 	viper.BindEnv("access-token", "APPKNOX_ACCESS_TOKEN")
 	viper.SetDefault("access-token", "")
 
-	RootCmd.PersistentFlags().String("host", appknox.DefaultAPIHost, "Appknox Server")
-	viper.BindPFlag("host", RootCmd.PersistentFlags().Lookup("host"))
-	viper.BindEnv("host", "APPKNOX_API_HOST")
-	viper.SetDefault("host", appknox.DefaultAPIHost)
+	RootCmd.PersistentFlags().String("host", "", "Appknox Server") // No default value here
+    viper.BindPFlag("host", RootCmd.PersistentFlags().Lookup("host"))
+    viper.BindEnv("host", "APPKNOX_API_HOST")
+
+
+	RootCmd.PersistentFlags().String("region", "", "Region names, e.g., global, saudi, uae. By default, global is used")
+    viper.BindPFlag("region", RootCmd.PersistentFlags().Lookup("region"))
+    viper.BindEnv("region", "APPKNOX_API_REGION")
+    viper.SetDefault("region", "global")
 
 	RootCmd.PersistentFlags().String("proxy", "", "proxy url")
 	viper.BindPFlag("proxy", RootCmd.PersistentFlags().Lookup("proxy"))

helper/clientInitialize.go

@@ -24,30 +24,89 @@ func getAppknoxAccessToken() string {
 	return accessToken
 }
 
+// GetHostMappings returns a map of host names to URLs.
+ func GetHostMappings() map[string]string {
+     return map[string]string{
+         "global": "https://api.appknox.com/",
+         "saudi":   "https://sa.secure.appknox.com/",
+         // Add more mappings as needed
+     }
+ }
+
+// ResolveHostAndRegion checks the host and region and returns the resolved base URL
+func ResolveHostAndRegion(host, region string, hostMappings map[string]string) (string, error) {
+    // If both region and host are provided, prioritize host and ignore region
+    if host != "" {
+        fmt.Printf("Both region and host provided. Using host URL: %s, ignoring region\n", host)
+        // Validate the host is a proper URL
+        _, err := url.ParseRequestURI(host)
+        if err != nil {
+            return "", fmt.Errorf("invalid host URL: %s", host)
+        }
+        return host, nil
+    }
+
+    // If region is provided, map it to the host URL
+    if region != "" {
+        if mappedHost, exists := hostMappings[region]; exists {
+            return mappedHost, nil
+        }
+        // Invalid region, return error and show available regions
+        availableRegions := make([]string, 0, len(hostMappings))
+        for key := range hostMappings {
+            availableRegions = append(availableRegions, key)
+        }
+        return "", fmt.Errorf("Invalid region name: %s. Available regions: %s", region, strings.Join(availableRegions, ", "))
+    }
+
+    // If neither host nor region are provided, default to the global host
+    return hostMappings["global"], nil
+}
+
+
 func getClient() *appknox.Client {
-	token := getAppknoxAccessToken()
-	host := viper.GetString("host")
-	client, err := appknox.NewClient(token)
-	if err != nil {
-		fmt.Println(err)
-		os.Exit(1)
-	}
-	proxyURL, err := GetProxy()
-	if err != nil {
-		fmt.Println(err)
-		os.Exit(1)
-	}
-	insecure := viper.GetBool("insecure")
-	client = client.SetHTTPTransportParams(proxyURL, insecure)
-	baseHost, err := url.Parse(host)
+    token := getAppknoxAccessToken()
+
+    // Check for region and host first
+    region := viper.GetString("region")
+    host := viper.GetString("host")
+
+    // Get the host mappings
+    hostMappings := GetHostMappings()
+
+    // Use the new function to resolve the host and region
+	resolvedHost, err := ResolveHostAndRegion(host, region, hostMappings)
 	if err != nil {
 		fmt.Println(err)
 		os.Exit(1)
 	}
-	client.BaseURL = baseHost
-	return client
+
+    client, err := appknox.NewClient(token)
+    if err != nil {
+        fmt.Println(err)
+        os.Exit(1)
+    }
+
+    proxyURL, err := GetProxy()
+    if err != nil {
+        fmt.Println(err)
+        os.Exit(1)
+    }
+
+    insecure := viper.GetBool("insecure")
+    client = client.SetHTTPTransportParams(proxyURL, insecure)
+
+    baseHost, err := url.Parse(resolvedHost)
+    if err != nil {
+        fmt.Println(err)
+        os.Exit(1)
+    }
+    client.BaseURL = baseHost
+    return client
 }
 
+
+
 // CheckToken checks if access token is valid.
 func CheckToken() (*appknox.Me, error) {
 	return GetMe()

helper/resolve_host_test.go

@@ -0,0 +1,57 @@
+package helper
+
+import (
+	"strings"
+	"testing"
+)
+
+func TestResolveHostAndRegion(t *testing.T) {
+	hostMappings := GetHostMappings()
+
+	tests := []struct {
+		name         string
+		host         string
+		region       string
+		expectedHost string
+		expectError  bool
+		errorMessage string
+	}{
+		// Case: Both host and region are empty, should default to global
+		{"Empty host and region", "", "", hostMappings["global"], false, ""},
+
+		// Case: Empty host, valid region
+		{"Empty host, valid region", "", "global", hostMappings["global"], false, ""},
+
+		// Case: Empty host, invalid region
+		{"Empty host, invalid region", "", "invalid-region", "", true, "Invalid region name: invalid-region. Available regions: global, saudi"},
+
+		// Case: Valid host, ignore region
+		{"Valid host, ignore region", "http://custom-host.com", "global", "http://custom-host.com", false, ""},
+
+		// Case: Empty host, valid region 'saudi'
+		{"Empty host, valid saudi region", "", "saudi", hostMappings["saudi"], false, ""},
+
+		// Case: Invalid host URL format
+		{"Invalid host URL format", "invalid_url", "", "", true, "invalid host URL: invalid_url"},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			result, err := ResolveHostAndRegion(test.host, test.region, hostMappings)
+
+			if (err != nil) != test.expectError {
+				t.Errorf("Expected error: %v, got: %v", test.expectError, err)
+			}
+
+			if test.expectError && err != nil {
+				if !strings.Contains(err.Error(), test.errorMessage) {
+					t.Errorf("Expected error message: %v, got: %v", test.errorMessage, err.Error())
+				}
+			}
+
+			if result != test.expectedHost && !test.expectError {
+				t.Errorf("Expected host: %s, got: %s", test.expectedHost, result)
+			}
+		})
+	}
+}