Skip to content

Commit

Permalink
Cleaned up code and fix some functionality
Browse files Browse the repository at this point in the history 
Signed-off-by: nathannaveen <[email protected]>
  • Loading branch information
@nathannaveen
nathannaveen committed Sep 30, 2024
1 parent 29fdcd7 commit f292dbc
Show file tree
Hide file tree
Showing 2 changed files with 301 additions and 155 deletions.
253 changes: 190 additions & 63 deletions cmd/guacone/cmd/vulnerability.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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"),
Expand All @@ -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)
}
Expand All @@ -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)
}
Expand Down Expand Up @@ -124,48 +140,114 @@ 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

var depVulnPath []string
var depVulnPaths []string
var depVulnTableRows []table.Row
//check if artifact
split := strings.Split(opts.searchString, ":")
if len(split) == 2 {

switch opts.inputType {
case artifactType:
// If it's an artifact, search for SBOMs via artifact
var err error
// 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.SearchForSBOMViaArtifact(ctx, gqlclient, opts.searchString, opts.depth)
depVulnPaths, depVulnTableRows, err = guacanalytics.SearchForSBOMViaArtifact(ctx, gqlclient, opts.searchString, opts.depth)
if err != nil {
logger.Fatalf("error searching via hasSBOM: %v", err)
logger.Fatalf("error searching via hasSBOM for artifact: %v", err)
}
} else {

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
// 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)
depVulnPaths, depVulnTableRows, err = guacanalytics.SearchForSBOMViaPkg(ctx, gqlclient, opts.searchString, opts.depth, opts.isPurl)
if err != nil {
logger.Fatalf("error searching via hasSBOM: %v", err)
logger.Fatalf("error searching via hasSBOM for package: %v", err)
}

if len(depVulnPaths) == 0 {
depVulnPaths, depVulnTableRows, err = findConnectedArtAndSearchViaArt(ctx, gqlclient, opts)
if err != nil {
logger.Fatalf("error finding artifact connected to package and searching via artifact: %v", err)
}
}
}

path = append(path, depVulnPath...)
paths = append(paths, depVulnPaths...)
tableRows = append(tableRows, depVulnTableRows...)

if len(path) > 0 {
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)
}

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])),
Expand All @@ -176,13 +258,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

Expand All @@ -197,7 +279,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)
Expand All @@ -207,11 +290,37 @@ 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, ":")
if len(split) == 2 {
occur := searchArtToPkg(ctx, gqlclient, split[0]+":"+split[1], logger)

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)
}
if pkgResponse, ok := foundHasSBOM.HasSBOM[0].Subject.(*model.AllHasSBOMTreeSubjectPackage); ok {
var vulnNeighborError error
path, tableRows, vulnNeighborError = queryVulnsViaVulnNodeNeighbors(ctx, gqlclient, pkgResponse.Namespaces[0].Names[0].Versions[0].Id, vulnResponse.Vulnerabilities, opts.depth, opts.pathsToReturn)
if vulnNeighborError != nil {
logger.Fatalf("error querying neighbor: %v", err)
}
} else if artResponse, ok := foundHasSBOM.HasSBOM[0].Subject.(*model.AllHasSBOMTreeSubjectArtifact); ok {
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
Expand All @@ -221,31 +330,10 @@ func printVulnInfoByVulnId(ctx context.Context, gqlclient graphql.Client, t tabl
}
}
} else {
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)
}
if pkgResponse, ok := foundHasSBOM.HasSBOM[0].Subject.(*model.AllHasSBOMTreeSubjectPackage); ok {
var vulnNeighborError error
path, tableRows, vulnNeighborError = queryVulnsViaVulnNodeNeighbors(ctx, gqlclient, pkgResponse.Namespaces[0].Names[0].Versions[0].Id, vulnResponse.Vulnerabilities, opts.depth, opts.pathsToReturn)
if vulnNeighborError != nil {
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)
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)
}
}
} else {
logger.Fatalf("located hasSBOM does not have a subject that is a package or artifact")
}
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())
Expand Down Expand Up @@ -460,19 +548,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 {
Expand Down
Loading

0 comments on commit f292dbc

Please sign in to comment.