feat: add flag --distribution

Author: iseki0

api/create_sub_task.go

@@ -1,9 +1,10 @@
 package api
 
 import (
+	"time"
+
 	"github.com/murphysecurity/murphysec/model"
 	"github.com/murphysecurity/murphysec/version"
-	"time"
 )
 
 type CreateSubTaskRequest struct {
@@ -29,6 +30,7 @@ type CreateSubTaskRequest struct {
 	WebhookMode        *string          `json:"webhook_mode,omitempty"`
 	ExtraData          *string          `json:"extra_data,omitempty"`
 	IsAutonomous       bool             `json:"is_autonomous"`
+	Distribution       string           `json:"distribution"`
 }
 
 type CreateSubTaskResponse struct {

cmd/murphy/internal/common/distribution.go

@@ -0,0 +1,68 @@
+package common
+
+import (
+	"encoding/json"
+	"fmt"
+	"strconv"
+	"strings"
+
+	"github.com/spf13/pflag"
+)
+
+type DistributionFlag struct {
+	name string
+}
+
+func (d *DistributionFlag) String() string {
+	if d == nil {
+		return ""
+	}
+	return d.name
+}
+
+func (d *DistributionFlag) Set(s string) error {
+	s = strings.ToLower(strings.TrimSpace(s))
+	switch s {
+	case "external", "internal", "saas", "open_source", "":
+		d.name = s
+	default:
+		return fmt.Errorf("invalid distribution: %s", strconv.Quote(s))
+	}
+	return nil
+}
+
+func (d *DistributionFlag) Type() string {
+	return "DistributionFlag"
+}
+
+func (d *DistributionFlag) IsExternal() bool {
+	return d.name == "external"
+}
+
+func (d *DistributionFlag) IsInternal() bool {
+	return d.name == "internal"
+}
+
+func (d *DistributionFlag) IsSaaS() bool {
+	return d.name == "saas"
+}
+
+func (d *DistributionFlag) IsOpenSource() bool {
+	return d.name == "open_source"
+}
+
+func (d DistributionFlag) MarshalJSON() ([]byte, error) {
+	return json.Marshal(d.name)
+}
+
+func (d *DistributionFlag) UnmarshalJSON(b []byte) error {
+	var s string
+	if err := json.Unmarshal(b, &s); err != nil {
+		return err
+	}
+	return d.Set(s)
+}
+
+var _ pflag.Value = (*DistributionFlag)(nil)
+var _ json.Marshaler = (*DistributionFlag)(nil)
+var _ json.Unmarshaler = (*DistributionFlag)(nil)

cmd/murphy/internal/scan/cmd.go

@@ -6,6 +6,10 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
+	"os"
+	"path/filepath"
+
 	"github.com/murphysecurity/murphysec/api"
 	"github.com/murphysecurity/murphysec/cmd/murphy/internal/common"
 	"github.com/murphysecurity/murphysec/cmd/murphy/internal/cv"
@@ -22,9 +26,6 @@ import (
 	"github.com/repeale/fp-go"
 	"github.com/samber/lo"
 	"github.com/spf13/cobra"
-	"io"
-	"os"
-	"path/filepath"
 )
 
 var jsonOutput bool
@@ -49,6 +50,7 @@ var branch string
 var mavenModuleName []string
 var binaryOnly bool
 var scanProcess bool
+var distribution common.DistributionFlag
 
 func Cmd() *cobra.Command {
 	var c cobra.Command
@@ -69,6 +71,7 @@ func Cmd() *cobra.Command {
 	c.Flags().IntVarP(&concurrentNumber, "max-concurrent-uploads", "j", 1, "Set the maximum number of parallel uploads.")
 	c.Flags().StringVar(&webhookAddr, "webhook-addr", "", "specify the webhook address")
 	c.Flags().Var(&webhookMode, "webhook-mode", "specify the webhook mode, currently supports: simple, full")
+	c.Flags().Var(&distribution, "distribution", "specify the distribution, currently supports: external, internal, saas, open_source")
 	c.Flags().StringVar(&extraData, "extra-data", "", "specify the extra data")
 	c.Flags().BoolVar(&scanCodeHash, "scan-snippets", false, "Enable scanning of code snippets to detect SBOM and  vulnerabilities. Disabled by default")
 	c.Flags().BoolVar(&binaryOnly, "binary-only", false, "only scan binary files, skip source code scanning")
@@ -95,6 +98,7 @@ func DfCmd() *cobra.Command {
 	c.Flags().Var(&sbomOutputType, "sbom-format", "(Required) specify the SBOM format, currently supports: msdx1.1+json")
 	c.Flags().StringVar(&webhookAddr, "webhook-addr", "", "specify the webhook address")
 	c.Flags().Var(&webhookMode, "webhook-mode", "specify the webhook mode, currently supports: simple, full(default)")
+	c.Flags().Var(&distribution, "distribution", "specify the distribution, currently supports: external, internal, saas, open_source")
 	c.Flags().StringVar(&extraData, "extra-data", "", "specify the extra data")
 	c.Flags().BoolVar(&scanCodeHash, "scan-snippets", false, "Enable scanning of code snippets to detect SBOM and  vulnerabilities. Disabled by default")
 	c.Flags().StringArrayVar(&gradleProjectFilter.ProjectNames, "gradle-project-name", make([]string, 0), "specify the name of the Gradle project")

cmd/murphy/internal/scan/scan.go

@@ -4,10 +4,11 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"github.com/murphysecurity/murphysec/codehash"
 	"os"
 	"path/filepath"
 
+	"github.com/murphysecurity/murphysec/codehash"
+
 	"github.com/murphysecurity/murphysec/api"
 	"github.com/murphysecurity/murphysec/chunkupload"
 	"github.com/murphysecurity/murphysec/cmd/murphy/internal/common"
@@ -156,6 +157,8 @@ func scan(ctx context.Context, dir string, accessType model.AccessType, mode mod
 	createSubtask.PackagePrivateName = privateSourceName
 	createSubtask.ProjectTagNames = projectTagNames
 	createSubtask.IsAutonomous = scanCodeHash
+	createSubtask.Distribution = distribution.String()
+
 	if createSubtask.ProjectTagNames == nil {
 		createSubtask.ProjectTagNames = make([]string, 0)
 	}