common functions for grpc

cmd/ecosystems/npm.go

@@ -1,22 +1,16 @@
 package ecosystems
 
 import (
-	"context"
 	_ "embed"
 	"fmt"
-	"net/http"
 	"os"
 	"strings"
 
-	"buf.build/gen/go/safedep/api/grpc/go/safedep/services/malysis/v1/malysisv1grpc"
-	malysisv1pb "buf.build/gen/go/safedep/api/protocolbuffers/go/safedep/messages/malysis/v1"
 	packagev1 "buf.build/gen/go/safedep/api/protocolbuffers/go/safedep/messages/package/v1"
-	malysisv1 "buf.build/gen/go/safedep/api/protocolbuffers/go/safedep/services/malysis/v1"
-	drygrpc "github.com/safedep/dry/adapters/grpc"
 	"github.com/safedep/dry/log"
+	"github.com/safedep/pmg/pkg/analyser"
 	"github.com/safedep/pmg/pkg/common"
 	"github.com/spf13/cobra"
-	"google.golang.org/grpc"
 )
 
 var (
@@ -80,47 +74,25 @@ func wrapNpm() error {
 		name := parts[0]
 		version := parts[1]
 
-		tok := os.Getenv("SAFEDEP_API_KEY")
+		client, err := analyser.GetMalwareAnalysisClient()
-		tenantId := os.Getenv("SAFEDEP_TENANT_ID")
+		if err != nil {
-
+			return fmt.Errorf("Error while creating a malware analysis client: %s", err.Error())
-		if tok == "" || tenantId == "" {
-			return fmt.Errorf("SAFEDEP_API_KEY and SAFEDEP_TENANT_ID must be set")
 		}
 
-		headers := http.Header{}
+		resp, err := analyser.SubmitPackageForAnalysis(client, packagev1.Ecosystem_ECOSYSTEM_NPM, name, version)
-		headers.Set("x-tenant-id", tenantId)
-		cc, err := drygrpc.GrpcClient("pmg-pkg-scan", "api.safedep.io", "443",
-			tok, headers, []grpc.DialOption{})
-		client := malysisv1grpc.NewMalwareAnalysisServiceClient(cc)
-
-		req := &malysisv1.AnalyzePackageRequest{
-			Target: &malysisv1pb.PackageAnalysisTarget{
-				PackageVersion: &packagev1.PackageVersion{
-					Package: &packagev1.Package{
-						Ecosystem: packagev1.Ecosystem_ECOSYSTEM_NPM,
-						Name:      name,
-					},
-					Version: version,
-				},
-			},
-		}
-
-		resp, err := client.AnalyzePackage(context.Background(), req)
 		if err != nil {
 			log.Debugf("Failed to analyze %s@%s: %v\n", name, version, err)
 			continue
 		}
 		log.Debugf("Submitted %s@%s | Analysis ID: %s\n", name, version, resp.GetAnalysisId())
 
-		analysisReportReq := &malysisv1.GetAnalysisReportRequest{
+		reportResp, err := analyser.GetAnalysisReport(client, resp.GetAnalysisId())
-			AnalysisId: resp.GetAnalysisId(),
+		if err != nil {
+			log.Debugf("Failed to get analysis report for %s:%s %v\n", name, resp.GetAnalysisId(), err)
+			continue
 		}
 
-		reportResp, err := client.GetAnalysisReport(context.Background(), analysisReportReq)
+		_ = reportResp.GetReport()
-		report := reportResp.GetReport()
-
-		log.Debugf("Report: ", report.GetWarnings())
-
 	}
 
 	// TODO:

pkg/analyser/analysis.go

@@ -12,7 +12,7 @@ import (
 
 func SubmitPackageForAnalysis(client malysisv1grpc.MalwareAnalysisServiceClient,
 	ecosystem packagev1.Ecosystem, name string,
-	version string) (string, error) {
+	version string) (*malysisv1.AnalyzePackageResponse, error) {
 	req := &malysisv1.AnalyzePackageRequest{
 		Target: &malysisv1pb.PackageAnalysisTarget{
 			PackageVersion: &packagev1.PackageVersion{
@@ -26,12 +26,13 @@ func SubmitPackageForAnalysis(client malysisv1grpc.MalwareAnalysisServiceClient,
 	}
 	resp, err := client.AnalyzePackage(context.Background(), req)
 	if err != nil {
-		return "", fmt.Errorf("failed to analyze %s@%s: %v", name, version, err)
+		return nil, fmt.Errorf("failed to analyze %s@%s: %v", name, version, err)
 	}
-	return resp.GetAnalysisId(), nil
+	return resp, nil
 }
+
 func GetAnalysisReport(client malysisv1grpc.MalwareAnalysisServiceClient,
-	analysisId string) (*malysisv1pb.Report, error) {
+	analysisId string) (*malysisv1.GetAnalysisReportResponse, error) {
 	analysisReportReq := &malysisv1.GetAnalysisReportRequest{
 		AnalysisId: analysisId,
 	}
@@ -39,5 +40,5 @@ func GetAnalysisReport(client malysisv1grpc.MalwareAnalysisServiceClient,
 	if err != nil {
 		return nil, fmt.Errorf("failed to get analysis report: %v", err)
 	}
-	return reportResp.GetReport(), nil
+	return reportResp, nil
 }

pkg/analyser/client.go

@@ -2,24 +2,13 @@ package analyser
 
 import (
 	"fmt"
-	"net/http"
-	"os"
 
 	"buf.build/gen/go/safedep/api/grpc/go/safedep/services/malysis/v1/malysisv1grpc"
-	drygrpc "github.com/safedep/dry/adapters/grpc"
+	"github.com/safedep/pmg/pkg/common"
-	"google.golang.org/grpc"
 )
 
 func GetMalwareAnalysisClient() (malysisv1grpc.MalwareAnalysisServiceClient, error) {
-	tok := os.Getenv("SAFEDEP_API_KEY")
+	cc, err := common.NewCloudClientConnection()
-	tenantId := os.Getenv("SAFEDEP_TENANT_ID")
-	if tok == "" || tenantId == "" {
-		return nil, fmt.Errorf("SAFEDEP_API_KEY and SAFEDEP_TENANT_ID must be set")
-	}
-	headers := http.Header{}
-	headers.Set("x-tenant-id", tenantId)
-	cc, err := drygrpc.GrpcClient("pmg-pkg-scan", "api.safedep.io", "443",
-		tok, headers, []grpc.DialOption{})
 	if err != nil {
 		return nil, fmt.Errorf("failed to create gRPC client: %v", err)
 	}

pkg/common/grpc.go

@@ -0,0 +1,34 @@
+package common
+
+import (
+	"fmt"
+	"net/http"
+	"os"
+
+	drygrpc "github.com/safedep/dry/adapters/grpc"
+	"google.golang.org/grpc"
+)
+
+func NewCloudClientConnection() (*grpc.ClientConn, error) {
+	tok := os.Getenv("SAFEDEP_API_KEY")
+	tenantId := os.Getenv("SAFEDEP_TENANT_ID")
+	if tok == "" || tenantId == "" {
+		return nil, fmt.Errorf("SAFEDEP_API_KEY and SAFEDEP_TENANT_ID must be set")
+	}
+	headers := http.Header{}
+	headers.Set("x-tenant-id", tenantId)
+
+	cc, err := newGrpcClient(headers, tok, "pmg-pkg-scan", "api.safedep.io", "443")
+	if err != nil {
+		return nil, fmt.Errorf("failed to create gRPC client: %v", err)
+	}
+	return cc, nil
+}
+
+func newGrpcClient(headers http.Header, token, clientName, host, port string) (*grpc.ClientConn, error) {
+	cc, err := drygrpc.GrpcClient(clientName, host, port, token, headers, []grpc.DialOption{})
+	if err != nil {
+		return nil, err
+	}
+	return cc, nil
+}