From 8b3379aa4c18b3b841dc13ca91a4860aa16b10c5 Mon Sep 17 00:00:00 2001
From: Tarun Koyalwar <45962551+tarunKoyalwar@users.noreply.github.com>
Date: Sun, 16 Jun 2024 03:18:01 +0530
Subject: [PATCH] add team-id env input support (#5295)

---
 internal/pdcp/writer.go | 7 +++++++
 1 file changed, 7 insertions(+)

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