From 2f7eea410df292ab18c657e8aa34240a07475c6a Mon Sep 17 00:00:00 2001 From: Ramana Reddy <90540245+RamanaReddy0M@users.noreply.github.com> Date: Fri, 16 Aug 2024 13:27:26 +0530 Subject: [PATCH] Add `team-id` option (#5523) * add team-id option * fix dashboard url when uploading to team --------- Co-authored-by: Tarun Koyalwar --- cmd/nuclei/main.go | 2 ++ internal/pdcp/utils.go | 10 +++++++++- internal/pdcp/writer.go | 25 ++++++++++++++++++------- internal/runner/runner.go | 3 +++ pkg/types/types.go | 2 ++ 5 files changed, 34 insertions(+), 8 deletions(-) diff --git a/cmd/nuclei/main.go b/cmd/nuclei/main.go index 945af907ea..de29d6581e 100644 --- a/cmd/nuclei/main.go +++ b/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)"), diff --git a/internal/pdcp/utils.go b/internal/pdcp/utils.go index 5d4fa4e1f0..3385d5cad5 100644 --- a/internal/pdcp/utils.go +++ b/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() } diff --git a/internal/pdcp/writer.go b/internal/pdcp/writer.go index 2197041137..aa2d5134d0 100644 --- a/internal/pdcp/writer.go +++ b/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, - done: make(chan struct{}, 1), + creds: creds, + 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 != "" { - req.Header.Set(teamIDHeader, teamID) + if u.TeamID != NoneTeamID && u.TeamID != "" { + req.Header.Set(teamIDHeader, u.TeamID) } req.Header.Set("Content-Type", "application/octet-stream") req.Header.Set("Accept", "application/json") diff --git a/internal/runner/runner.go b/internal/runner/runner.go index bc436500af..bfb2bc64be 100644 --- a/internal/runner/runner.go +++ b/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) } diff --git a/pkg/types/types.go b/pkg/types/types.go index 32c7f6f6a2..cab1aacf54 100644 --- a/pkg/types/types.go +++ b/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