add team-id option

Author: RamanaReddy0M

cmd/httpx/httpx.go

@@ -138,5 +138,8 @@ func setupOptionalAssetUpload(opts *runner.Options) *pdcp.UploadWriter {
 		// silently ignore
 		writer.SetAssetGroupName(opts.AssetName)
 	}
+	if opts.TeamID != "" {
+		writer.SetTeamID(opts.TeamID)
+	}
 	return writer
 }

internal/pdcp/writer.go

@@ -35,8 +35,6 @@ const (
 
 var (
 	xidRegex = regexp.MustCompile(xidRe)
-	// teamID if given
-	teamID = env.GetEnvOrDefault("PDCP_TEAM_ID", "")
 	// EnableeUpload if set to true enables the upload feature
 	HideAutoSaveMsg   = env.GetEnvOrDefault("DISABLE_CLOUD_UPLOAD_WRN", false)
 	EnableCloudUpload = env.GetEnvOrDefault("ENABLE_CLOUD_UPLOAD", false)
@@ -54,6 +52,7 @@ type UploadWriter struct {
 	assetGroupName string
 	counter        atomic.Int32
 	closed         atomic.Bool
+	TeamID         string
 }
 
 // NewUploadWriterCallback creates a new upload writer callback
@@ -63,9 +62,10 @@ func NewUploadWriterCallback(ctx context.Context, creds *pdcpauth.PDCPCredential
 		return nil, fmt.Errorf("no credentials provided")
 	}
 	u := &UploadWriter{
-		creds: creds,
-		done:  make(chan struct{}, 1),
-		data:  make(chan runner.Result, 8), // default buffer size
+		creds:  creds,
+		done:   make(chan struct{}, 1),
+		data:   make(chan runner.Result, 8), // default buffer size
+		TeamID: "",
 	}
 	var err error
 	tmp, err := urlutil.Parse(creds.Server)
@@ -111,6 +111,11 @@ func (u *UploadWriter) SetAssetGroupName(name string) {
 	u.assetGroupName = name
 }
 
+// SetTeamID sets the team id for the upload writer
+func (u *UploadWriter) SetTeamID(id string) {
+	u.TeamID = id
+}
+
 func (u *UploadWriter) autoCommit(ctx context.Context) {
 	// wait for context to be done
 	defer func() {
@@ -244,8 +249,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 != "" {
+		req.Header.Set(teamIDHeader, u.TeamID)
 	}
 	req.Header.Set("Content-Type", "application/octet-stream")
 	req.Header.Set("Accept", "application/json")

runner/options.go

@@ -25,7 +25,7 @@ import (
 	"github.com/projectdiscovery/httpx/common/httpx"
 	"github.com/projectdiscovery/httpx/common/stringz"
 	"github.com/projectdiscovery/networkpolicy"
-	"github.com/projectdiscovery/utils/auth/pdcp"
+	pdcpauth "github.com/projectdiscovery/utils/auth/pdcp"
 	"github.com/projectdiscovery/utils/env"
 	fileutil "github.com/projectdiscovery/utils/file"
 	sliceutil "github.com/projectdiscovery/utils/slice"
@@ -41,7 +41,10 @@ const (
 	DefaultOutputDirectory = "output"
 )
 
-var PDCPApiKey = ""
+var (
+	PDCPApiKey = ""
+	TeamIDEnv  = env.GetEnvOrDefault("PDCP_TEAM_ID", "")
+)
 
 // OnResultCallback (hostResult)
 type OnResultCallback func(Result)
@@ -318,6 +321,7 @@ type Options struct {
 	AssetID string
 	// AssetFileUpload
 	AssetFileUpload string
+	TeamID          string
 	// OnClose adds a callback function that is invoked when httpx is closed
 	// to be exact at end of existing closures
 	OnClose func()
@@ -509,6 +513,7 @@ func ParseOptions() *Options {
 	flagSet.CreateGroup("cloud", "Cloud",
 		flagSet.DynamicVar(&options.PdcpAuth, "auth", "true", "configure projectdiscovery cloud (pdcp) api key"),
 		flagSet.BoolVarP(&options.AssetUpload, "dashboard", "pd", false, "upload / view output in projectdiscovery cloud (pdcp) UI dashboard"),
+		flagSet.StringVarP(&options.TeamID, "team-id", "tid", TeamIDEnv, "upload asset results to given team id (optional)"),
 		flagSet.StringVarP(&options.AssetID, "asset-id", "aid", "", "upload new assets to existing asset id (optional)"),
 		flagSet.StringVarP(&options.AssetName, "asset-name", "aname", "", "assets group name to set (optional)"),
 		flagSet.StringVarP(&options.AssetFileUpload, "dashboard-upload", "pdu", "", "upload httpx output file (jsonl) in projectdiscovery cloud (pdcp) UI dashboard"),
@@ -540,9 +545,9 @@ func ParseOptions() *Options {
 		AuthWithPDCP()
 	} else if len(options.PdcpAuth) == 36 {
 		PDCPApiKey = options.PdcpAuth
-		ph := pdcp.PDCPCredHandler{}
-		if _, err := ph.GetCreds(); err == pdcp.ErrNoCreds {
-			apiServer := env.GetEnvOrDefault("PDCP_API_SERVER", pdcp.DefaultApiServer)
+		ph := pdcpauth.PDCPCredHandler{}
+		if _, err := ph.GetCreds(); err == pdcpauth.ErrNoCreds {
+			apiServer := env.GetEnvOrDefault("PDCP_API_SERVER", pdcpauth.DefaultApiServer)
 			if validatedCreds, err := ph.ValidateAPIKey(PDCPApiKey, apiServer, "httpx"); err == nil {
 				_ = ph.SaveCreds(validatedCreds)
 			}