vet/cmd/inspect/malware.go

package inspect

import (
	"context"
	"fmt"
	"os"
	"strings"
	"time"

	"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"
	malysisv1 "buf.build/gen/go/safedep/api/protocolbuffers/go/safedep/services/malysis/v1"
	"github.com/jedib0t/go-pretty/v6/table"
	"github.com/jedib0t/go-pretty/v6/text"
	"github.com/safedep/dry/api/pb"
	"github.com/safedep/dry/utils"
	"github.com/safedep/vet/internal/auth"
	"github.com/safedep/vet/internal/ui"
	"github.com/safedep/vet/pkg/malysis"
	"github.com/spf13/cobra"
)

var (
	malwareAnalysisPackageUrl string
	malwareAnalysisTimeout    time.Duration
	malwareAnalysisReportJSON string
)

func newPackageMalwareInspectCommand() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "malware",
		Short: "Inspect an OSS package for malware",
		Long:  `Inspect an OSS package for malware using SafeDep Malware Analysis API`,
		RunE: func(cmd *cobra.Command, args []string) error {
			err := executeMalwareAnalysis()
			if err != nil {
				ui.PrintError("Failed: %v", err)
			}

			return nil
		},
	}

	cmd.Flags().StringVar(&malwareAnalysisPackageUrl, "purl", "",
		"Package URL to inspect for malware")
	cmd.Flags().DurationVar(&malwareAnalysisTimeout, "timeout", 5*time.Minute,
		"Timeout for malware analysis")
	cmd.Flags().StringVar(&malwareAnalysisReportJSON, "report-json", "",
		"Path to save malware analysis report in JSON format")

	_ = cmd.MarkFlagRequired("purl")

	return cmd
}

func executeMalwareAnalysis() error {
	cc, err := auth.MalwareAnalysisClientConnection("malware-analysis")
	if err != nil {
		return err
	}

	service := malysisv1grpc.NewMalwareAnalysisServiceClient(cc)

	purl, err := pb.NewPurlPackageVersion(malwareAnalysisPackageUrl)
	if err != nil {
		return err
	}

	ctx := context.Background()
	ctx, cancelFun := context.WithTimeout(ctx, malwareAnalysisTimeout)

	defer cancelFun()

	analyzePackageResponse, err := service.AnalyzePackage(ctx, &malysisv1.AnalyzePackageRequest{
		Target: &malysisv1pb.PackageAnalysisTarget{
			PackageVersion: purl.PackageVersion(),
		},
	})

	if err != nil {
		return fmt.Errorf("failed to submit package for malware analysis: %v", err)
	}

	ui.PrintMsg("Submitted package for malware analysis with ID: %s",
		analyzePackageResponse.GetAnalysisId())

	ui.StartSpinner("Waiting for malware analysis to complete")
	var report *malysisv1pb.Report

	for {
		reportResponse, err := service.GetAnalysisReport(ctx, &malysisv1.GetAnalysisReportRequest{
			AnalysisId: analyzePackageResponse.GetAnalysisId(),
		})

		if err != nil {
			return fmt.Errorf("failed to get malware analysis report: %v", err)
		}

		if reportResponse.GetStatus() == malysisv1.AnalysisStatus_ANALYSIS_STATUS_FAILED {
			return fmt.Errorf("malware analysis failed: %s", reportResponse.GetErrorMessage())
		}

		if reportResponse.GetStatus() == malysisv1.AnalysisStatus_ANALYSIS_STATUS_COMPLETED {
			report = reportResponse.GetReport()
			break
		}

		time.Sleep(5 * time.Second)
	}

	ui.StopSpinner()

	if report == nil {
		return fmt.Errorf("malware analysis report is empty")
	}

	ui.PrintSuccess("Malware analysis completed successfully")

	err = renderToJSON(report)
	if err != nil {
		ui.PrintError("Failed to render malware analysis report in JSON format: %v", err)
	}

	return renderMalwareAnalysisReport(malwareAnalysisPackageUrl,
		analyzePackageResponse.GetAnalysisId(), report)
}

func renderToJSON(report *malysisv1pb.Report) error {
	if malwareAnalysisReportJSON == "" {
		return nil
	}

	data, err := utils.ToPbJson(report, "  ")
	if err != nil {
		return err
	}

	return os.WriteFile(malwareAnalysisReportJSON, []byte(data), 0644)
}

func renderMalwareAnalysisReport(purl string, analysisId string, report *malysisv1pb.Report) error {
	ui.PrintMsg("Malware analysis report for package: %s", purl)

	tbl := table.NewWriter()
	tbl.SetOutputMirror(os.Stdout)
	tbl.SetStyle(table.StyleLight)

	tbl.AppendHeader(table.Row{"Package URL", "Status", "Confidence"})

	status := text.FgHiGreen.Sprint("SAFE")
	if report.GetInference().GetIsMalware() {
		status = text.FgHiRed.Sprint("MALWARE")
	}

	confidence := report.GetInference().GetConfidence().String()
	confidence = strings.TrimPrefix(confidence, "CONFIDENCE_")

	tbl.AppendRow(table.Row{purl, status, confidence})
	tbl.Render()

	fmt.Println()
	fmt.Println(text.FgHiYellow.Sprintf("** The full report is available at: %s",
		reportVisualizationUrl(analysisId)))
	fmt.Println()

	return nil
}

func reportVisualizationUrl(analysisId string) string {
	return malysis.ReportURL(analysisId)
}