common functions for grpc

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"
-	"google.golang.org/grpc"
+	"github.com/safedep/pmg/pkg/common"
 )
 
 func GetMalwareAnalysisClient() (malysisv1grpc.MalwareAnalysisServiceClient, 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 := drygrpc.GrpcClient("pmg-pkg-scan", "api.safedep.io", "443",
-		tok, headers, []grpc.DialOption{})
+	cc, err := common.NewCloudClientConnection()
 	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
+}