Search for Vulns via Artifact

cmd/guacone/cmd/vulnerability.go

@@ -22,24 +22,25 @@ import (
 	"os"
 	"strings"
 
-	"github.com/guacsec/guac/pkg/guacanalytics"
-	"go.uber.org/zap"
-
 	"github.com/guacsec/guac/internal/testing/ptrfrom"
-
-	"github.com/Khan/genqlient/graphql"
 	model "github.com/guacsec/guac/pkg/assembler/clients/generated"
 	"github.com/guacsec/guac/pkg/assembler/helpers"
 	"github.com/guacsec/guac/pkg/cli"
+	"github.com/guacsec/guac/pkg/guacanalytics"
 	"github.com/guacsec/guac/pkg/logging"
+
+	"github.com/Khan/genqlient/graphql"
 	"github.com/jedib0t/go-pretty/v6/table"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
 )
 
 const (
-	guacType   string = "guac"
-	noVulnType string = "novuln"
+	guacType     string = "guac"
+	noVulnType   string = "novuln"
+	purlType     string = "purl"
+	uriType      string = "uri"
+	artifactType string = "artifact"
 )
 
 type queryOptions struct {
@@ -50,14 +51,28 @@ type queryOptions struct {
 	vulnerabilityID string
 	depth           int
 	pathsToReturn   int
+	inputType       string
 }
 
 var queryVulnCmd = &cobra.Command{
-	Use:   "vuln [flags] <purl/sbomURI>",
+	Use:   "vuln [flags] <type> <input>",
 	Short: "query if a package is affected by the specified vulnerability",
+	Long: `The vuln command allows you to query whether a specific package, SBOM URI, or artifact is affected by a given vulnerability.
+
+Positional Arguments:
+  <type>    Specify the input type: 'artifact', 'uri', or 'purl'
+  <input>   The corresponding input based on the specified type`,
 	Run: func(cmd *cobra.Command, args []string) {
 		ctx := logging.WithLogger(context.Background())
 
+		// Ensure exactly two arguments are provided: <type> and <input>
+		if len(args) != 2 {
+			fmt.Println("error: Exactly two arguments must be provided: <type> and <input>.")
+			_ = cmd.Help()
+			os.Exit(1)
+		}
+
+		// Extract and validate other flags
 		opts, err := validateQueryVulnFlags(
 			viper.GetString("gql-addr"),
 			viper.GetString("header-file"),
@@ -67,7 +82,7 @@ var queryVulnCmd = &cobra.Command{
 			args,
 		)
 		if err != nil {
-			fmt.Printf("unable to validate flags: %v\n", err)
+			fmt.Printf("Unable to validate flags: %v\n", err)
 			_ = cmd.Help()
 			os.Exit(1)
 		}
@@ -80,8 +95,9 @@ var queryVulnCmd = &cobra.Command{
 		tTemp.Render()
 		t.AppendHeader(rowHeader)
 
+		// Process based on the specified input type
 		if opts.vulnerabilityID != "" {
-			printVulnInfoByVulnId(ctx, gqlclient, t, opts)
+			printVulnInfoByVulnID(ctx, gqlclient, t, opts)
 		} else {
 			printVulnInfo(ctx, gqlclient, t, opts)
 		}
@@ -124,52 +140,117 @@ func getPkgResponseFromPurl(ctx context.Context, gqlclient graphql.Client, purl
 
 func printVulnInfo(ctx context.Context, gqlclient graphql.Client, t table.Writer, opts queryOptions) {
 	logger := logging.FromContext(ctx)
-	var path []string
+	var paths []string
 	var tableRows []table.Row
 
-	// The primaryCall parameter in searchForSBOMViaPkg is there for us to know whether
-	// the searchString is expected to be a PURL, and we are searching via a purl.
-	depVulnPath, depVulnTableRows, err := guacanalytics.SearchForSBOMViaPkg(ctx, gqlclient, opts.searchString, opts.depth, opts.isPurl)
-	if err != nil {
-		logger.Fatalf("error searching via hasSBOM: %v", err)
-	}
-
-	path = append(path, depVulnPath...)
-	tableRows = append(tableRows, depVulnTableRows...)
+	var depVulnPaths []string
+	var depVulnTableRows []table.Row
 
-	if len(depVulnPath) == 0 {
-		occur := searchArtToPkg(ctx, gqlclient, opts.searchString, logger)
+	switch opts.inputType {
+	case artifactType:
+		// If it's an artifact, search for SBOMs via artifact
+		var err error
+		depVulnPaths, depVulnTableRows, err = guacanalytics.SearchForSBOMViaArtifact(ctx, gqlclient, opts.searchString, opts.depth)
+		if err != nil {
+			logger.Fatalf("error searching via hasSBOM for artifact: %v", err)
+		}
 
-		if occur != nil && len(occur.IsOccurrence) > 0 {
-			subjectPackage, ok := occur.IsOccurrence[0].Subject.(*model.AllIsOccurrencesTreeSubjectPackage)
-			if ok {
-				// The primaryCall parameter in searchForSBOMViaPkg is there for us to know that
-				// the searchString is expected to be an artifact, but isn't, so we have to check via PURLs instead of artifacts.
-				depVulnPath, depVulnTableRows, err = guacanalytics.SearchForSBOMViaPkg(ctx, gqlclient, subjectPackage.Namespaces[0].Names[0].Versions[0].Id, opts.depth, false)
-				if err != nil {
-					logger.Fatalf("error searching via hasSBOM: %v", err)
-				}
+		if len(depVulnPaths) == 0 {
+			depVulnPaths, depVulnTableRows, err = findConnectedPkgAndSearchViaPkg(ctx, gqlclient, opts)
+			if err != nil {
+				logger.Fatalf("error finding purl connected to artifact and searching via package: %v", err)
+			}
+		}
+	default:
+		// Otherwise, search for SBOMs via package
+		var err error
+		depVulnPaths, depVulnTableRows, err = guacanalytics.SearchForSBOMViaPkg(ctx, gqlclient, opts.searchString, opts.depth, opts.isPurl)
+		if err != nil {
+			logger.Fatalf("error searching via hasSBOM for package: %v", err)
+		}
 
-				path = append(path, depVulnPath...)
-				tableRows = append(tableRows, depVulnTableRows...)
+		if len(depVulnPaths) == 0 && opts.inputType == purlType {
+			depVulnPaths, depVulnTableRows, err = findConnectedArtAndSearchViaArt(ctx, gqlclient, opts)
+			if err != nil {
+				logger.Fatalf("error finding artifact connected to package and searching via artifact: %v", err)
 			}
 		}
 	}
 
-	if len(path) > 0 {
-		t.AppendRows(tableRows)
+	paths = append(paths, depVulnPaths...)
+	tableRows = append(tableRows, depVulnTableRows...)
 
+	if len(paths) > 0 {
+		t.AppendRows(tableRows)
 		fmt.Println(t.Render())
-		fmt.Printf("Visualizer url: http://localhost:3000/?path=%v\n", strings.Join(removeDuplicateValuesFromPath(path), `,`))
+		fmt.Printf("Visualizer URL: http://localhost:3000/?path=%v\n", strings.Join(removeDuplicateValuesFromPath(paths), `,`))
 	} else {
 		fmt.Printf("No path to vulnerabilities found!\n")
 	}
 }
 
-func searchArtToPkg(ctx context.Context, gqlclient graphql.Client, searchString string, logger *zap.SugaredLogger) *model.OccurrencesResponse {
+// findConnectedArtAndSearchViaArt finds the artifact attached to the packages with the given purl.
+// After finding the artifact, the graph is searched via that artifact.
+func findConnectedArtAndSearchViaArt(ctx context.Context, gqlclient graphql.Client, opts queryOptions) ([]string, []table.Row, error) {
+	var depVulnPaths []string
+	var depVulnTableRows []table.Row
+
+	// convert the package to an artifact via an isOccurrence.
+	pkgSpec, err := helpers.PurlToPkgFilter(opts.searchString)
+	if err != nil {
+		return nil, nil, fmt.Errorf("error converting purl to pkg %v", err)
+	}
+
+	occ, err := model.Occurrences(ctx, gqlclient, model.IsOccurrenceSpec{
+		Subject: &model.PackageOrSourceSpec{
+			Package: &pkgSpec,
+		},
+	})
+	if err != nil {
+		return nil, nil, fmt.Errorf("error getting occurrences for package: %v", err)
+	}
+
+	if len(occ.IsOccurrence) > 0 {
+		art := occ.IsOccurrence[0].Artifact
+
+		newSearchString := art.Algorithm + ":" + art.Digest
+
+		depVulnPaths, depVulnTableRows, err = guacanalytics.SearchForSBOMViaArtifact(ctx, gqlclient, newSearchString, opts.depth)
+		if err != nil {
+			return nil, nil, fmt.Errorf("error searching via hasSBOM for artifact: %v", err)
+		}
+	}
+
+	return depVulnPaths, depVulnTableRows, nil
+}
+
+// findConnectedPkgAndSearchViaPkg finds the pkg attached to the artifact.
+// After finding the pkg, the graph is searched via that package.
+func findConnectedPkgAndSearchViaPkg(ctx context.Context, gqlclient graphql.Client, opts queryOptions) ([]string, []table.Row, error) {
+	var depVulnPaths []string
+	var depVulnTableRows []table.Row
+
+	occurrence, err := searchArtToPkg(ctx, gqlclient, opts.searchString)
+	if err != nil {
+		return nil, nil, fmt.Errorf("error searching for package via artifact: %v", err)
+	}
+
+	pkg, ok := occurrence.IsOccurrence[0].Subject.(*model.AllIsOccurrencesTreeSubjectPackage)
+	if !ok {
+		return nil, nil, fmt.Errorf("error converting isOccurrence to package subject")
+	}
+
+	depVulnPaths, depVulnTableRows, err = guacanalytics.SearchForSBOMViaPkg(ctx, gqlclient, pkg.Namespaces[0].Names[0].Versions[0].Purl, opts.depth, true)
+	if err != nil {
+		return nil, nil, fmt.Errorf("error searching via hasSBOM for artifact: %v", err)
+	}
+	return depVulnPaths, depVulnTableRows, nil
+}
+
+func searchArtToPkg(ctx context.Context, gqlclient graphql.Client, searchString string) (*model.OccurrencesResponse, error) {
 	split := strings.Split(searchString, ":")
 	if len(split) != 2 {
-		logger.Fatalf("failed to parse artifact. Needs to be in algorithm:digest form")
+		return nil, fmt.Errorf("failed to parse artifact. Needs to be in algorithm:digest form")
 	}
 	artifactFilter := model.ArtifactSpec{
 		Algorithm: ptrfrom.String(strings.ToLower(split[0])),
@@ -180,13 +261,13 @@ func searchArtToPkg(ctx context.Context, gqlclient graphql.Client, searchString
 		Artifact: &artifactFilter,
 	})
 	if err != nil {
-		logger.Fatalf("error querying for occurrences: %v", err)
+		return nil, fmt.Errorf("error querying for occurrences: %v", err)
 	}
 
-	return o
+	return o, nil
 }
 
-func printVulnInfoByVulnId(ctx context.Context, gqlclient graphql.Client, t table.Writer, opts queryOptions) {
+func printVulnInfoByVulnID(ctx context.Context, gqlclient graphql.Client, t table.Writer, opts queryOptions) {
 	logger := logging.FromContext(ctx)
 	var tableRows []table.Row
 
@@ -201,7 +282,8 @@ func printVulnInfoByVulnId(ctx context.Context, gqlclient graphql.Client, t tabl
 	}
 	var path []string
 
-	if opts.isPurl {
+	switch opts.inputType {
+	case purlType:
 		pkgResponse, err := getPkgResponseFromPurl(ctx, gqlclient, opts.searchString)
 		if err != nil {
 			logger.Fatalf("getPkgResponseFromPurl - error: %v", err)
@@ -211,7 +293,22 @@ func printVulnInfoByVulnId(ctx context.Context, gqlclient graphql.Client, t tabl
 		if vulnNeighborError != nil {
 			logger.Fatalf("error querying neighbor: %v", err)
 		}
-	} else {
+	case artifactType:
+		split := strings.Split(opts.searchString, ":")
+
+		occur, err := searchArtToPkg(ctx, gqlclient, split[0]+":"+split[1])
+		if err != nil {
+			logger.Fatalf("error searching for package via artifact: %v", err)
+		}
+		subjectPackage, ok := occur.IsOccurrence[0].Subject.(*model.AllIsOccurrencesTreeSubjectPackage)
+		if ok {
+			var vulnNeighborError error
+			path, tableRows, vulnNeighborError = queryVulnsViaVulnNodeNeighbors(ctx, gqlclient, subjectPackage.Namespaces[0].Names[0].Versions[0].Id, vulnResponse.Vulnerabilities, opts.depth, opts.pathsToReturn)
+			if vulnNeighborError != nil {
+				logger.Fatalf("error querying neighbor: %v", err)
+			}
+		}
+	case uriType:
 		foundHasSBOM, err := model.HasSBOMs(ctx, gqlclient, model.HasSBOMSpec{Uri: &opts.searchString})
 		if err != nil {
 			logger.Fatalf("failed getting hasSBOM via URI: %s with error: %w", opts.searchString, err)
@@ -223,7 +320,10 @@ func printVulnInfoByVulnId(ctx context.Context, gqlclient graphql.Client, t tabl
 				logger.Fatalf("error querying neighbor: %v", err)
 			}
 		} else if artResponse, ok := foundHasSBOM.HasSBOM[0].Subject.(*model.AllHasSBOMTreeSubjectArtifact); ok {
-			occur := searchArtToPkg(ctx, gqlclient, artResponse.Algorithm+":"+artResponse.Digest, logger)
+			occur, err := searchArtToPkg(ctx, gqlclient, artResponse.Algorithm+":"+artResponse.Digest)
+			if err != nil {
+				logger.Fatalf("error searching for package via artifact: %v", err)
+			}
 			subjectPackage, ok := occur.IsOccurrence[0].Subject.(*model.AllIsOccurrencesTreeSubjectPackage)
 			if ok {
 				var vulnNeighborError error
@@ -236,6 +336,7 @@ func printVulnInfoByVulnId(ctx context.Context, gqlclient graphql.Client, t tabl
 			logger.Fatalf("located hasSBOM does not have a subject that is a package or artifact")
 		}
 	}
+
 	if len(path) > 0 {
 		t.AppendRows(tableRows)
 		fmt.Println(t.Render())
@@ -450,19 +551,58 @@ func validateQueryVulnFlags(graphqlEndpoint, headerFile, vulnID string, depth, p
 	opts.pathsToReturn = path
 
 	if len(args) > 0 {
-		_, err := helpers.PurlToPkg(args[0])
+		validTypes := []string{artifactType, uriType, purlType}
+
+		// Initialize variables to hold type and input
+		var typeArg, inputArg string
+
+		// Iterate through arguments to identify type and input, because they might not be in order
+		for _, arg := range args {
+			lowered := strings.ToLower(arg)
+			if contains(validTypes, lowered) {
+				if typeArg != "" {
+					fmt.Println("error: Multiple types provided. Please specify only one type.")
+					os.Exit(1)
+				}
+				typeArg = lowered
+			} else {
+				if inputArg != "" {
+					fmt.Println("error: Multiple inputs provided. Please specify only one input.")
+					os.Exit(1)
+				}
+				inputArg = arg
+			}
+		}
+
+		// Validate that typeArg has been set
+		if typeArg == "" {
+			fmt.Printf("error: Input type not specified or invalid. Valid types are: %v\n", validTypes)
+			os.Exit(1)
+		}
+
+		_, err := helpers.PurlToPkg(inputArg)
 		if err != nil {
 			opts.isPurl = false
 		} else {
 			opts.isPurl = true
 		}
-		opts.searchString = args[0]
+		opts.searchString = inputArg
+		opts.inputType = typeArg
 	} else {
 		return opts, fmt.Errorf("expected subject input to be purl or SBOM URI")
 	}
 	return opts, nil
 }
 
+func contains(slice []string, item string) bool {
+	for _, a := range slice {
+		if a == item {
+			return true
+		}
+	}
+	return false
+}
+
 func init() {
 	set, err := cli.BuildFlags([]string{"vuln-id", "search-depth", "num-path"})
 	if err != nil {