Add url parsing for both cases

Signed-off-by: Christian Zunker <[email protected]>
mondoohq · Dec 5, 2023 · 80b5674 · 80b5674
1 parent 015e1cb
commit 80b5674
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 9 deletions.
diff --git a/gen/gen.go b/gen/gen.go
@@ -75,7 +75,6 @@ func generateSchema(token string, basePath string) error {
 
 // loadSchema loads the GraphQL schema from the Mondoo API.
 func loadSchema(token string) (schema interface{}, err error) {
-	apiHost := "us.api.mondoo.com"
 	introspection := `
 {
   __schema {
@@ -177,14 +176,7 @@ fragment TypeRef on __Type {
   }
 }
 `
-	apiEndpoint := "https://" + apiHost + "/query"
-	endpoint, ok := os.LookupEnv("MONDOO_API_ENDPOINT")
-	if ok {
-		apiEndpoint, err = url.JoinPath(endpoint, "/query")
-		if err != nil {
-			log.Fatalf("invalid MONDOO_API_ENDPOINT: %v", err)
-		}
-	}
+	apiEndpoint, apiHost := getAPIEndpoint()
 	fmt.Printf("using endpoint %s\n", apiEndpoint)
 	// do introspection query
 	req, err := http.NewRequest(
@@ -348,3 +340,27 @@ func parseTemplate(text string) *template.Template {
 		},
 	}).Parse(text))
 }
+
+func getAPIEndpoint() (string, string) {
+	apiHost := "us.api.mondoo.com"
+	apiEndpoint, err := url.JoinPath("https://", apiHost, "/query")
+	if err != nil {
+		log.Fatalf("invalid MONDOO_API_ENDPOINT: %v", err)
+	}
+	endpoint, ok := os.LookupEnv("MONDOO_API_ENDPOINT")
+	if ok {
+		if !strings.HasPrefix(endpoint, "http://") && !strings.HasPrefix(endpoint, "https://") {
+			endpoint = "https://" + endpoint
+		}
+		apiEndpoint, err = url.JoinPath(endpoint, "/query")
+		if err != nil {
+			log.Fatalf("invalid MONDOO_API_ENDPOINT: %v", err)
+		}
+	}
+	parsedUrl, err := url.Parse(apiEndpoint)
+	if err != nil {
+		log.Fatalf("invalid API url: %v", err)
+	}
+
+	return apiEndpoint, parsedUrl.Host
+}
diff --git a/gen/gen_test.go b/gen/gen_test.go
@@ -18,3 +18,26 @@ func TestRun(t *testing.T) {
 	err := generateSchema(token, "..")
 	require.NoError(t, err)
 }
+func TestGetAPIEndpoint(t *testing.T) {
+	// Test case 1: MONDOO_API_ENDPOINT not set
+	expectedEndpoint := "https://us.api.mondoo.com/query"
+	expectedHost := "us.api.mondoo.com"
+	actualEndpoint, actualHost := getAPIEndpoint()
+	require.Equal(t, expectedEndpoint, actualEndpoint, "Unexpected API endpoint")
+	require.Equal(t, expectedHost, actualHost, "Unexpected host")
+
+	// Test case 2: MONDOO_API_ENDPOINT set
+	os.Setenv("MONDOO_API_ENDPOINT", "custom.api.endpoint")
+	expectedEndpoint = "https://custom.api.endpoint/query"
+	expectedHost = "custom.api.endpoint"
+	actualEndpoint, actualHost = getAPIEndpoint()
+	require.Equal(t, expectedEndpoint, actualEndpoint, "Unexpected API endpoint")
+	require.Equal(t, expectedHost, actualHost, "Unexpected host")
+
+	os.Setenv("MONDOO_API_ENDPOINT", "http://custom.api.endpoint:1234")
+	expectedEndpoint = "http://custom.api.endpoint:1234/query"
+	expectedHost = "custom.api.endpoint:1234"
+	actualEndpoint, actualHost = getAPIEndpoint()
+	require.Equal(t, expectedEndpoint, actualEndpoint, "Unexpected API endpoint")
+	require.Equal(t, expectedHost, actualHost, "Unexpected host")
+}