diff --git a/internal/pdcp/writer.go b/internal/pdcp/writer.go index 6e0541035..5dfda78ca 100644 --- a/internal/pdcp/writer.go +++ b/internal/pdcp/writer.go @@ -18,6 +18,7 @@ import ( "github.com/projectdiscovery/nuclei/v3/pkg/output" "github.com/projectdiscovery/retryablehttp-go" pdcpauth "github.com/projectdiscovery/utils/auth/pdcp" + "github.com/projectdiscovery/utils/env" errorutil "github.com/projectdiscovery/utils/errors" unitutils "github.com/projectdiscovery/utils/unit" updateutils "github.com/projectdiscovery/utils/update" @@ -30,11 +31,14 @@ const ( flushTimer = time.Minute MaxChunkSize = 4 * unitutils.Mega // 4 MB xidRe = `^[a-z0-9]{20}$` + teamIDHeader = "X-Team-Id" ) var ( xidRegex = regexp.MustCompile(xidRe) _ output.Writer = &UploadWriter{} + // teamID if given + teamID = env.GetEnvOrDefault("PDCP_TEAM_ID", "") ) // UploadWriter is a writer that uploads its output to pdcp @@ -244,6 +248,9 @@ 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) + } req.Header.Set("Content-Type", "application/octet-stream") req.Header.Set("Accept", "application/json") return req, nil