fix: look up auth, user info once during snyk enrich

Closes #46.
snyk · Feb 5, 2024 · a77584c · a77584c
1 parent bc887c3
commit a77584c
Show file tree

Hide file tree

Showing 6 changed files with 89 additions and 47 deletions.
diff --git a/internal/commands/snyk/packages.go b/internal/commands/snyk/packages.go
@@ -20,11 +20,33 @@ func NewPackageCommand(logger zerolog.Logger) *cobra.Command {
 			if err != nil {
 				logger.Fatal().Err(err).Msg("Not a valid purl")
 			}
-			logger.Debug().Str("purl", args[0]).Msg("Looking up package vulnerabilities from Snyk")
-			resp, err := snyk.GetPackageVulnerabilities(purl)
+
+			logger.
+				Debug().
+				Str("purl", args[0]).
+				Msg("Looking up package vulnerabilities from Snyk")
+
+			auth, err := snyk.AuthFromToken(snyk.APIToken())
+			if err != nil {
+				logger.
+					Fatal().
+					Err(err).
+					Msg("Failed to get API credentials.")
+			}
+
+			orgID, err := snyk.SnykOrgID(auth)
+			if err != nil {
+				logger.
+					Fatal().
+					Err(err).
+					Msg("Failed to look up user info.")
+			}
+
+			resp, err := snyk.GetPackageVulnerabilities(&purl, auth, orgID)
 			if err != nil {
 				logger.Fatal().Err(err).Msg("An error occurred")
 			}
+
 			fmt.Print(string(resp.Body))
 		},
 	}

diff --git a/lib/snyk/enrich_cyclonedx.go b/lib/snyk/enrich_cyclonedx.go
@@ -35,6 +35,20 @@ func enrichCycloneDX(bom *cdx.BOM, logger zerolog.Logger) *cdx.BOM {
 		return bom
 	}
 
+	auth, err := AuthFromToken(APIToken())
+	if err != nil {
+		// TODO: log error when logger instance available.
+		// See https://github.com/snyk/parlay/pull/49
+		return nil
+	}
+
+	orgID, err := SnykOrgID(auth)
+	if err != nil {
+		// TODO: log error when logger instance available.
+		// See https://github.com/snyk/parlay/pull/49
+		return nil
+	}
+
 	wg := sizedwaitgroup.New(20)
 	var mutex = &sync.Mutex{}
 	vulnerabilities := make(map[cdx.Component][]issues.CommonIssueModelVTwo)
@@ -53,7 +67,7 @@ func enrichCycloneDX(bom *cdx.BOM, logger zerolog.Logger) *cdx.BOM {
 				return
 			}
 
-			resp, err := GetPackageVulnerabilities(purl)
+			resp, err := GetPackageVulnerabilities(&purl, auth, orgID)
 			if err != nil {
 				logger.Err(err).
 					Str("purl", purl.ToString()).

diff --git a/lib/snyk/enrich_spdx.go b/lib/snyk/enrich_spdx.go
@@ -36,6 +36,20 @@ const (
 )
 
 func enrichSPDX(bom *spdx.Document, logger zerolog.Logger) *spdx.Document {
+	auth, err := AuthFromToken(APIToken())
+	if err != nil {
+		// TODO: log error when logger instance available.
+		// See https://github.com/snyk/parlay/pull/49
+		return nil
+	}
+
+	orgID, err := SnykOrgID(auth)
+	if err != nil {
+		// TODO: log error when logger instance available.
+		// See https://github.com/snyk/parlay/pull/49
+		return nil
+	}
+
 	mutex := &sync.Mutex{}
 	wg := sizedwaitgroup.New(20)
 	vulnerabilities := make(map[*spdx_2_3.Package][]issues.CommonIssueModelVTwo)
@@ -54,27 +68,18 @@ func enrichSPDX(bom *spdx.Document, logger zerolog.Logger) *spdx.Document {
 				return
 			}
 
-			resp, err := GetPackageVulnerabilities(*purl)
-			if err != nil {
-				logger.Err(err).
-					Str("purl", purl.String()).
-					Msg("Failed to fetch vulnerabilities for package.")
-				return
-			}
-
-			packageData := resp.Body
-			var packageDoc issues.IssuesWithPurlsResponse
-			if err := json.Unmarshal(packageData, &packageDoc); err != nil {
-				logger.Err(err).
-					Str("status", resp.Status()).
-					Msg("Failed to decode Snyk vulnerability response.")
-				return
-			}
-
-			if packageDoc.Data != nil {
-				mutex.Lock()
-				vulnerabilities[pkg] = *packageDoc.Data
-				mutex.Unlock()
+			resp, err := GetPackageVulnerabilities(purl, auth, orgID)
+
+			if err == nil {
+				packageData := resp.Body
+				var packageDoc issues.IssuesWithPurlsResponse
+				if err := json.Unmarshal(packageData, &packageDoc); err == nil {
+					if packageDoc.Data != nil {
+						mutex.Lock()
+						vulnerabilities[pkg] = *packageDoc.Data
+						mutex.Unlock()
+					}
+				}
 			}
 		}(pkg, i)
 	}

diff --git a/lib/snyk/package.go b/lib/snyk/package.go
@@ -18,11 +18,9 @@ package snyk
 
 import (
 	"context"
-	"errors"
-	"fmt"
-	"os"
 
 	"github.com/deepmap/oapi-codegen/pkg/securityprovider"
+	"github.com/google/uuid"
 	"github.com/package-url/packageurl-go"
 
 	"github.com/snyk/parlay/snyk/issues"
@@ -31,29 +29,14 @@ import (
 const snykServer = "https://api.snyk.io/rest"
 const version = "2023-04-28"
 
-func GetPackageVulnerabilities(purl packageurl.PackageURL) (*issues.FetchIssuesPerPurlResponse, error) {
-	token := os.Getenv("SNYK_TOKEN")
-	if token == "" {
-		return nil, errors.New("Must provide a SNYK_TOKEN environment variable")
-	}
-
-	auth, err := securityprovider.NewSecurityProviderApiKey("header", "Authorization", fmt.Sprintf("token %s", token))
-	if err != nil {
-		return nil, err
-	}
-
-	org, err := getSnykOrg(auth)
-	if err != nil {
-		return nil, err
-	}
-
+func GetPackageVulnerabilities(purl *packageurl.PackageURL, auth *securityprovider.SecurityProviderApiKey, orgID *uuid.UUID) (*issues.FetchIssuesPerPurlResponse, error) {
 	client, err := issues.NewClientWithResponses(snykServer, issues.WithRequestEditorFn(auth.Intercept))
 	if err != nil {
 		return nil, err
 	}
 
 	params := issues.FetchIssuesPerPurlParams{Version: version}
-	resp, err := client.FetchIssuesPerPurlWithResponse(context.Background(), *org, purl.ToString(), &params)
+	resp, err := client.FetchIssuesPerPurlWithResponse(context.Background(), *orgID, purl.ToString(), &params)
 	if err != nil {
 		return nil, err
 	}

diff --git a/lib/snyk/self.go b/lib/snyk/self.go
@@ -22,6 +22,7 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
+	"os"
 
 	"github.com/deepmap/oapi-codegen/pkg/securityprovider"
 	"github.com/google/uuid"
@@ -37,7 +38,7 @@ type selfDocument struct {
 	}
 }
 
-func getSnykOrg(auth *securityprovider.SecurityProviderApiKey) (*uuid.UUID, error) {
+func SnykOrgID(auth *securityprovider.SecurityProviderApiKey) (*uuid.UUID, error) {
 	experimental, err := users.NewClientWithResponses(snykServer, users.WithRequestEditorFn(auth.Intercept))
 	if err != nil {
 		return nil, err
@@ -64,3 +65,20 @@ func getSnykOrg(auth *securityprovider.SecurityProviderApiKey) (*uuid.UUID, erro
 
 	return nil, errors.New("Failed to get org ID.")
 }
+
+func AuthFromToken(token string) (*securityprovider.SecurityProviderApiKey, error) {
+	if token == "" {
+		return nil, errors.New("Must provide a SNYK_TOKEN environment variable")
+	}
+
+	auth, err := securityprovider.NewSecurityProviderApiKey("header", "Authorization", fmt.Sprintf("token %s", token))
+	if err != nil {
+		return nil, err
+	}
+
+	return auth, nil
+}
+
+func APIToken() string {
+	return os.Getenv("SNYK_TOKEN")
+}
diff --git a/lib/snyk/self_test.go b/lib/snyk/self_test.go
@@ -38,7 +38,7 @@ func TestGetSnykOrg_Success(t *testing.T) {
 		httpmock.NewJsonResponderOrPanic(http.StatusOK, httpmock.File("testdata/self.json")),
 	)
 
-	actualOrg, err := getSnykOrg(auth)
+	actualOrg, err := SnykOrgID(auth)
 	assert.NoError(t, err)
 	assert.Equal(t, expectedOrg, *actualOrg)
 }
@@ -53,7 +53,7 @@ func TestGetSnykOrg_Unauthorized(t *testing.T) {
 		httpmock.NewJsonResponderOrPanic(http.StatusUnauthorized, []byte(`{"msg":"unauthorized"}`)),
 	)
 
-	actualOrg, err := getSnykOrg(auth)
+	actualOrg, err := SnykOrgID(auth)
 	assert.ErrorContains(t, err, "Failed to get user info (401)")
 	assert.Nil(t, actualOrg)
 }