✅ add tests for newly moved queries

vulncheck-oss · Nov 1, 2024 · a44650d · a44650d
1 parent fce1851
commit a44650d
Show file tree

Hide file tree

Showing 4 changed files with 163 additions and 66 deletions.
diff --git a/pkg/cmd/offline/ipintel/ipintel.go b/pkg/cmd/offline/ipintel/ipintel.go
@@ -2,7 +2,6 @@ package ipintel
 
 import (
 	"fmt"
-	"github.com/package-url/packageurl-go"
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
 	"github.com/vulncheck-oss/cli/pkg/cache"
@@ -11,7 +10,6 @@ import (
 	"github.com/vulncheck-oss/cli/pkg/ui"
 	"github.com/vulncheck-oss/cli/pkg/utils"
 	"slices"
-	"strings"
 )
 
 func Command() *cobra.Command {
@@ -53,7 +51,7 @@ func Command() *cobra.Command {
 				return fmt.Errorf("index ipintel-%s is required for this command, and is not cached", args[0])
 			}
 
-			query := buildQuery(country, asn, cidr, countryCode, hostname, id)
+			query := search.QueryIPIntel(country, asn, cidr, countryCode, hostname, id)
 
 			if !jsonOutput && !config.IsCI() {
 				ui.Info(fmt.Sprintf("Searching index %s, last updated on %s", index.Name, utils.ParseDate(index.LastUpdated)))
@@ -140,64 +138,3 @@ func AliasCommands() []*cobra.Command {
 
 	return commands
 }
-
-func buildQuery(country, asn, cidr, countryCode, hostname, id string) string {
-	var conditions []string
-
-	if country != "" {
-		conditions = append(conditions, fmt.Sprintf(".country == %q", country))
-	}
-	if asn != "" {
-		conditions = append(conditions, fmt.Sprintf(".asn == %q", asn))
-	}
-	if cidr != "" {
-		// Note: CIDR matching would require additional logic
-		conditions = append(conditions, fmt.Sprintf(".ip == %q", cidr))
-		// conditions = append(conditions, fmt.Sprintf(".ip | startswith(%q)", strings.Split(cidr, "/")[0]))
-	}
-	if countryCode != "" {
-		conditions = append(conditions, fmt.Sprintf(".country_code == %q", countryCode))
-	}
-	if hostname != "" {
-		conditions = append(conditions, fmt.Sprintf(".hostnames | any(. == %q)", hostname))
-	}
-	if id != "" {
-		conditions = append(conditions, fmt.Sprintf(".type.id == %q", id))
-	}
-
-	if len(conditions) == 0 {
-		return "true"
-	}
-	return strings.Join(conditions, " and ")
-}
-
-func BuildPurlQuery(instance packageurl.PackageURL) string {
-	seperator := "/"
-	var conditions []string
-
-	if instance.Type == "maven" {
-		seperator = ":"
-	}
-
-	if instance.Namespace == "alpine" {
-		conditions = append(conditions, fmt.Sprintf(".package_name == %q", instance.Name))
-	} else {
-		if instance.Namespace != "" {
-			conditions = append(conditions, fmt.Sprintf(".name == %q", fmt.Sprintf("%s%s%s", instance.Namespace, seperator, instance.Name)))
-		} else {
-			conditions = append(conditions, fmt.Sprintf(".name == %q", instance.Name))
-		}
-	}
-
-	if instance.Version != "" {
-		if instance.Type == "golang" {
-			// For golang, match the version at the end of the string
-			// conditions = append(conditions, fmt.Sprintf(".version | contains(%q)", instance.Version))
-			conditions = append(conditions, fmt.Sprintf("(.version | index(%q)) != null and (.version | rindex(%q)) == (.version | length - %d)", instance.Version, instance.Version, len(instance.Version)))
-		} else {
-			// For other types, keep the contains check
-			conditions = append(conditions, fmt.Sprintf(".version == %q", instance.Version))
-		}
-	}
-	return strings.Join(conditions, " and ")
-}
diff --git a/pkg/cmd/offline/purl/purl.go b/pkg/cmd/offline/purl/purl.go
@@ -5,7 +5,6 @@ import (
 	"github.com/package-url/packageurl-go"
 	"github.com/spf13/cobra"
 	"github.com/vulncheck-oss/cli/pkg/cache"
-	"github.com/vulncheck-oss/cli/pkg/cmd/offline/ipintel"
 	"github.com/vulncheck-oss/cli/pkg/cmd/offline/packages"
 	"github.com/vulncheck-oss/cli/pkg/cmd/offline/sync"
 	"github.com/vulncheck-oss/cli/pkg/config"
@@ -55,7 +54,7 @@ func Command() *cobra.Command {
 
 			index := indices.GetIndex(indexName)
 
-			query := ipintel.BuildPurlQuery(instance)
+			query := search.QueryPURL(instance)
 
 			if !jsonOutput && !config.IsCI() {
 				if err := ui.PurlInstance(instance); err != nil {

diff --git a/pkg/search/search.go b/pkg/search/search.go
@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"github.com/itchyny/gojq"
+	"github.com/package-url/packageurl-go"
 	"github.com/tidwall/gjson"
 	"github.com/vulncheck-oss/cli/pkg/config"
 	"github.com/vulncheck-oss/cli/pkg/ui"
@@ -66,6 +67,67 @@ type Stats struct {
 	Query        string
 }
 
+func QueryIPIntel(country, asn, cidr, countryCode, hostname, id string) string {
+	var conditions []string
+
+	if country != "" {
+		conditions = append(conditions, fmt.Sprintf(".country == %q", country))
+	}
+	if asn != "" {
+		conditions = append(conditions, fmt.Sprintf(".asn == %q", asn))
+	}
+	if cidr != "" {
+		// Note: CIDR matching would require additional logic
+		conditions = append(conditions, fmt.Sprintf(".ip == %q", cidr))
+		// conditions = append(conditions, fmt.Sprintf(".ip | startswith(%q)", strings.Split(cidr, "/")[0]))
+	}
+	if countryCode != "" {
+		conditions = append(conditions, fmt.Sprintf(".country_code == %q", countryCode))
+	}
+	if hostname != "" {
+		conditions = append(conditions, fmt.Sprintf(".hostnames | any(. == %q)", hostname))
+	}
+	if id != "" {
+		conditions = append(conditions, fmt.Sprintf(".type.id == %q", id))
+	}
+
+	if len(conditions) == 0 {
+		return "true"
+	}
+	return strings.Join(conditions, " and ")
+}
+
+func QueryPURL(instance packageurl.PackageURL) string {
+	seperator := "/"
+	var conditions []string
+
+	if instance.Type == "maven" {
+		seperator = ":"
+	}
+
+	if instance.Namespace == "alpine" {
+		conditions = append(conditions, fmt.Sprintf(".package_name == %q", instance.Name))
+	} else {
+		if instance.Namespace != "" {
+			conditions = append(conditions, fmt.Sprintf(".name == %q", fmt.Sprintf("%s%s%s", instance.Namespace, seperator, instance.Name)))
+		} else {
+			conditions = append(conditions, fmt.Sprintf(".name == %q", instance.Name))
+		}
+	}
+
+	if instance.Version != "" {
+		if instance.Type == "golang" {
+			// For golang, match the version at the end of the string
+			// conditions = append(conditions, fmt.Sprintf(".version | contains(%q)", instance.Version))
+			conditions = append(conditions, fmt.Sprintf("(.version | index(%q)) != null and (.version | rindex(%q)) == (.version | length - %d)", instance.Version, instance.Version, len(instance.Version)))
+		} else {
+			// For other types, keep the contains check
+			conditions = append(conditions, fmt.Sprintf(".version == %q", instance.Version))
+		}
+	}
+	return strings.Join(conditions, " and ")
+}
+
 func Index(indexName, query string) ([]Entry, *Stats, error) {
 	startTime := time.Now()
 	var stats Stats

diff --git a/pkg/search/search_test.go b/pkg/search/search_test.go
@@ -3,6 +3,7 @@ package search
 import (
 	"fmt"
 	"github.com/itchyny/gojq"
+	"github.com/package-url/packageurl-go"
 	"os"
 	"path/filepath"
 	"sort"
@@ -12,6 +13,104 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
+// ... existing code ...
+
+func TestQueryIPIntel(t *testing.T) {
+	testCases := []struct {
+		name          string
+		country       string
+		asn           string
+		cidr          string
+		countryCode   string
+		hostname      string
+		id            string
+		expectedQuery string
+	}{
+		{
+			name:          "All fields",
+			country:       "United States",
+			asn:           "AS15169",
+			cidr:          "172.217.0.0/16",
+			countryCode:   "US",
+			hostname:      "google.com",
+			id:            "initial-access",
+			expectedQuery: `.country == "United States" and .asn == "AS15169" and .ip == "172.217.0.0/16" and .country_code == "US" and .hostnames | any(. == "google.com") and .type.id == "initial-access"`,
+		},
+		{
+			name:          "Partial fields",
+			country:       "Canada",
+			asn:           "AS577",
+			expectedQuery: `.country == "Canada" and .asn == "AS577"`,
+		},
+		{
+			name:          "No fields",
+			expectedQuery: "true",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			result := QueryIPIntel(tc.country, tc.asn, tc.cidr, tc.countryCode, tc.hostname, tc.id)
+			assert.Equal(t, tc.expectedQuery, result, "Unexpected query for IP Intel")
+		})
+	}
+}
+
+func TestQueryPURL(t *testing.T) {
+	testCases := []struct {
+		name          string
+		purl          packageurl.PackageURL
+		expectedQuery string
+	}{
+		{
+			name: "Maven package",
+			purl: packageurl.PackageURL{
+				Type:      "maven",
+				Namespace: "org.apache",
+				Name:      "log4j",
+				Version:   "2.14.1",
+			},
+			expectedQuery: `.name == "org.apache:log4j" and .version == "2.14.1"`,
+		},
+		{
+			name: "NPM package",
+			purl: packageurl.PackageURL{
+				Type:    "npm",
+				Name:    "lodash",
+				Version: "4.17.21",
+			},
+			expectedQuery: `.name == "lodash" and .version == "4.17.21"`,
+		},
+		{
+			name: "Alpine package",
+			purl: packageurl.PackageURL{
+				Type:      "alpine",
+				Namespace: "alpine",
+				Name:      "busybox",
+				Version:   "1.33.1-r3",
+			},
+			expectedQuery: `.package_name == "busybox" and .version == "1.33.1-r3"`,
+		},
+		{
+			name: "Golang package",
+			purl: packageurl.PackageURL{
+				Type:      "golang",
+				Namespace: "github.com/gin-gonic",
+				Name:      "gin",
+				Version:   "v1.7.2",
+			},
+			expectedQuery: `.name == "github.com/gin-gonic/gin" and (.version | index("v1.7.2")) != null and (.version | rindex("v1.7.2")) == (.version | length - 6)`,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			result := QueryPURL(tc.purl)
+			assert.Equal(t, tc.expectedQuery, result, "Unexpected query for PURL")
+		})
+	}
+}
+
 func TestQuickFilter(t *testing.T) {
 	testCases := []struct {
 		name     string