Add team-id option (#5523)

* add team-id option

* fix dashboard url when uploading to team

---------

Co-authored-by: Tarun Koyalwar <tarun@projectdiscovery.io>

cmd/nuclei/main.go

@@ -12,6 +12,7 @@ import (
 	"strings"
 	"time"
 
+	_pdcp "github.com/projectdiscovery/nuclei/v3/internal/pdcp"
 	"github.com/projectdiscovery/utils/auth/pdcp"
 	"github.com/projectdiscovery/utils/env"
 	_ "github.com/projectdiscovery/utils/pprof"
@@ -418,6 +419,7 @@ on extensive configurability, massive extensibility and ease of use.`)
 
 	flagSet.CreateGroup("cloud", "Cloud",
 		flagSet.DynamicVar(&pdcpauth, "auth", "true", "configure projectdiscovery cloud (pdcp) api key"),
+		flagSet.StringVarP(&options.TeamID, "team-id", "tid", _pdcp.TeamIDEnv, "upload scan results to given team id (optional)"),
 		flagSet.BoolVarP(&options.EnableCloudUpload, "cloud-upload", "cup", false, "upload scan results to pdcp dashboard"),
 		flagSet.StringVarP(&options.ScanID, "scan-id", "sid", "", "upload scan results to existing scan id (optional)"),
 		flagSet.StringVarP(&options.ScanName, "scan-name", "sname", "", "scan name to set (optional)"),

internal/pdcp/utils.go

@@ -5,9 +5,17 @@ import (
 	urlutil "github.com/projectdiscovery/utils/url"
 )
 
-func getScanDashBoardURL(id string) string {
+func getScanDashBoardURL(id string, teamID string) string {
 	ux, _ := urlutil.Parse(pdcpauth.DashBoardURL)
 	ux.Path = "/scans/" + id
+	if ux.Params == nil {
+		ux.Params = urlutil.NewOrderedParams()
+	}
+	if teamID != "" {
+		ux.Params.Add("team_id", teamID)
+	} else {
+		ux.Params.Add("team_id", NoneTeamID)
+	}
 	ux.Update()
 	return ux.String()
 }

internal/pdcp/writer.go

@@ -32,13 +32,14 @@ const (
 	MaxChunkSize   = 4 * unitutils.Mega // 4 MB
 	xidRe          = `^[a-z0-9]{20}$`
 	teamIDHeader   = "X-Team-Id"
+	NoneTeamID     = "none"
 )
 
 var (
 	xidRegex               = regexp.MustCompile(xidRe)
 	_        output.Writer = &UploadWriter{}
 	// teamID if given
-	teamID = env.GetEnvOrDefault("PDCP_TEAM_ID", "")
+	TeamIDEnv = env.GetEnvOrDefault("PDCP_TEAM_ID", NoneTeamID)
 )
 
 // UploadWriter is a writer that uploads its output to pdcp
@@ -53,6 +54,7 @@ type UploadWriter struct {
 	scanID    string
 	scanName  string
 	counter   atomic.Int32
+	TeamID    string
 }
 
 // NewUploadWriter creates a new upload writer
@@ -61,8 +63,9 @@ func NewUploadWriter(ctx context.Context, creds *pdcpauth.PDCPCredentials) (*Upl
 		return nil, fmt.Errorf("no credentials provided")
 	}
 	u := &UploadWriter{
-		creds: creds,
+		creds:  creds,
-		done:  make(chan struct{}, 1),
+		done:   make(chan struct{}, 1),
+		TeamID: NoneTeamID,
 	}
 	var err error
 	reader, writer := io.Pipe()
@@ -110,6 +113,14 @@ func (u *UploadWriter) SetScanName(name string) {
 	u.scanName = name
 }
 
+func (u *UploadWriter) SetTeamID(id string) {
+	if id == "" {
+		u.TeamID = NoneTeamID
+	} else {
+		u.TeamID = id
+	}
+}
+
 func (u *UploadWriter) autoCommit(ctx context.Context, r *io.PipeReader) {
 	reader := bufio.NewReader(r)
 	ch := make(chan string, 4)
@@ -136,7 +147,7 @@ func (u *UploadWriter) autoCommit(ctx context.Context, r *io.PipeReader) {
 		if u.scanID == "" {
 			gologger.Verbose().Msgf("Scan results upload to cloud skipped, no results found to upload")
 		} else {
-			gologger.Info().Msgf("%v Scan results uploaded to cloud, you can view scan results at %v", u.counter.Load(), getScanDashBoardURL(u.scanID))
+			gologger.Info().Msgf("%v Scan results uploaded to cloud, you can view scan results at %v", u.counter.Load(), getScanDashBoardURL(u.scanID, u.TeamID))
 		}
 	}()
 	// temporary buffer to store the results
@@ -189,7 +200,7 @@ func (u *UploadWriter) uploadChunk(buff *bytes.Buffer) error {
 	// if successful, reset the buffer
 	buff.Reset()
 	// log in verbose mode
-	gologger.Warning().Msgf("Uploaded results chunk, you can view scan results at %v", getScanDashBoardURL(u.scanID))
+	gologger.Warning().Msgf("Uploaded results chunk, you can view scan results at %v", getScanDashBoardURL(u.scanID, u.TeamID))
 	return nil
 }
 
@@ -248,8 +259,8 @@ func (u *UploadWriter) getRequest(bin []byte) (*retryablehttp.Request, error) {
 	req.URL.Update()
 
 	req.Header.Set(pdcpauth.ApiKeyHeaderName, u.creds.APIKey)
-	if teamID != "" {
+	if u.TeamID != NoneTeamID && u.TeamID != "" {
-		req.Header.Set(teamIDHeader, teamID)
+		req.Header.Set(teamIDHeader, u.TeamID)
 	}
 	req.Header.Set("Content-Type", "application/octet-stream")
 	req.Header.Set("Accept", "application/json")

internal/runner/runner.go

@@ -426,6 +426,9 @@ func (r *Runner) setupPDCPUpload(writer output.Writer) output.Writer {
 	if r.options.ScanName != "" {
 		uploadWriter.SetScanName(r.options.ScanName)
 	}
+	if r.options.TeamID != "" {
+		uploadWriter.SetTeamID(r.options.TeamID)
+	}
 	return output.NewMultiWriter(writer, uploadWriter)
 }
 

pkg/types/types.go

@@ -384,6 +384,8 @@ type Options struct {
 	ScanID string
 	// ScanName is the name of the scan to be uploaded
 	ScanName string
+	// TeamID is the team ID to use for cloud upload
+	TeamID string
 	// JsConcurrency is the number of concurrent js routines to run
 	JsConcurrency int
 	// SecretsFile is file containing secrets for nuclei