add team-id env input support (#5295)

This commit is contained in:
Tarun Koyalwar 2024-06-16 03:18:01 +05:30 committed by GitHub
parent baf7e36658
commit 8b3379aa4c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

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