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/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