From aed3ec4b73ab52e94a10c2523cb2a1902325cd6d Mon Sep 17 00:00:00 2001 From: Karlo Kegljevic Date: Mon, 2 Feb 2026 18:26:00 +0100 Subject: [PATCH] Add custom HTTP header support --- docs/man/trufflehog.1 | 3 +++ main.go | 31 +++++++++++++++++++++++++++++++ pkg/common/http.go | 12 ++++++++++++ pkg/detectors/http.go | 1 + pkg/detectors/readme/readme.go | 2 +- pkg/feature/feature.go | 20 ++++++++++++++++++++ pkg/sources/huggingface/client.go | 4 +++- 7 files changed, 71 insertions(+), 2 deletions(-) diff --git a/docs/man/trufflehog.1 b/docs/man/trufflehog.1 index 666b21e51a0d..a27293f2a0ec 100644 --- a/docs/man/trufflehog.1 +++ b/docs/man/trufflehog.1 @@ -50,6 +50,9 @@ Specifies which type(s) of results to output: verified (confirmed valid by API), \fB--no-color\fR Disable colorized output .TP +\fB--header=HEADER\fR +Custom header to add to all outbound HTTP requests; repeatable. Format 'Name: value' or 'Name=value'. +.TP \fB--allow-verification-overlap\fR Allow verification of similar credentials across detectors .TP diff --git a/main.go b/main.go index a73df6c0bc73..3895b96ad605 100644 --- a/main.go +++ b/main.go @@ -61,6 +61,7 @@ var ( results = cli.Flag("results", "Specifies which type(s) of results to output: verified (confirmed valid by API), unknown (verification failed due to error), unverified (detected but not verified), filtered_unverified (unverified but would have been filtered out). Defaults to verified,unverified,unknown.").String() noColor = cli.Flag("no-color", "Disable colorized output").Bool() noColour = cli.Flag("no-colour", "Alias for --no-color").Hidden().Bool() + customHeaders = cli.Flag("header", "Custom header to add to all outbound HTTP requests; repeatable. Format 'Name: value' or 'Name=value'.").Strings() allowVerificationOverlap = cli.Flag("allow-verification-overlap", "Allow verification of similar credentials across detectors").Bool() filterUnverified = cli.Flag("filter-unverified", "Only output first unverified result per chunk per detector if there are more than one results.").Bool() @@ -512,6 +513,36 @@ func run(state overseer.State, logSync func() error) { feature.UserAgentSuffix.Store(*userAgentSuffix) } + // Parse and set any custom headers specified via --header + if customHeaders != nil && len(*customHeaders) > 0 { + hdr := http.Header{} + for _, h := range *customHeaders { + var key, val string + colonIdx := strings.Index(h, ":") + equalsIdx := strings.Index(h, "=") + switch { + case colonIdx == -1 && equalsIdx == -1: + logger.V(1).Info("skipping invalid --header format; expected 'Name: value' or 'Name=value'", "header", h) + continue + case colonIdx == -1 || (equalsIdx != -1 && equalsIdx < colonIdx): + key = strings.TrimSpace(h[:equalsIdx]) + val = strings.TrimSpace(h[equalsIdx+1:]) + default: + key = strings.TrimSpace(h[:colonIdx]) + val = strings.TrimSpace(h[colonIdx+1:]) + } + if key == "" || val == "" { + logger.V(1).Info("skipping invalid --header with empty key or value", "header", h) + continue + } + hdr.Add(key, val) + } + if len(hdr) > 0 { + feature.CustomHeaders.Store(hdr) + logger.V(2).Info("custom headers configured", "count", len(hdr)) + } + } + // OSS Default APK handling on feature.EnableAPKHandler.Store(true) diff --git a/pkg/common/http.go b/pkg/common/http.go index 330e8a86509b..c872a9cc0262 100644 --- a/pkg/common/http.go +++ b/pkg/common/http.go @@ -96,8 +96,20 @@ func UserAgent() string { return "TruffleHog" } +// ApplyCustomHeaders adds any globally configured custom headers to req. +func ApplyCustomHeaders(req *http.Request) { + if hdr := feature.CustomHeaders.Load(); hdr != nil { + for k, vals := range hdr { + for _, v := range vals { + req.Header.Add(k, v) + } + } + } +} + func (t *CustomTransport) RoundTrip(req *http.Request) (*http.Response, error) { req.Header.Add("User-Agent", UserAgent()) + ApplyCustomHeaders(req) return t.T.RoundTrip(req) } diff --git a/pkg/detectors/http.go b/pkg/detectors/http.go index 8624d8346b42..e89cb0922615 100644 --- a/pkg/detectors/http.go +++ b/pkg/detectors/http.go @@ -75,6 +75,7 @@ type detectorTransport struct { func (t *detectorTransport) RoundTrip(req *http.Request) (*http.Response, error) { req.Header.Add("User-Agent", userAgent()) + common.ApplyCustomHeaders(req) return t.T.RoundTrip(req) } diff --git a/pkg/detectors/readme/readme.go b/pkg/detectors/readme/readme.go index 12a286017c26..b85a10465faf 100644 --- a/pkg/detectors/readme/readme.go +++ b/pkg/detectors/readme/readme.go @@ -48,7 +48,7 @@ func (s Scanner) FromData(ctx context.Context, verify bool, data []byte) (result } req.SetBasicAuth(resMatch, "") req.Header.Add("accept", "application/json") - res, err := http.DefaultClient.Do(req) + res, err := detectors.DetectorHttpClientWithLocalAddresses.Do(req) if err == nil { defer res.Body.Close() if res.StatusCode >= 200 && res.StatusCode < 300 { diff --git a/pkg/feature/feature.go b/pkg/feature/feature.go index 3aa92a3759a6..c664f3f18d03 100644 --- a/pkg/feature/feature.go +++ b/pkg/feature/feature.go @@ -1,6 +1,7 @@ package feature import ( + "net/http" "sync/atomic" ) @@ -16,6 +17,7 @@ var ( GitlabProjectsPerPage atomic.Int64 UseGithubGraphQLAPI atomic.Bool // use github graphql api to fetch issues, pr's and comments HTMLDecoderEnabled atomic.Bool + CustomHeaders AtomicHeader ) type AtomicString struct { @@ -41,3 +43,21 @@ func (as *AtomicString) Swap(newValue string) string { as.Store(newValue) return oldValue } + +// AtomicHeader stores an http.Header atomically for global access. +type AtomicHeader struct { + value atomic.Value +} + +// Load returns the current http.Header value; may be nil if unset. +func (ah *AtomicHeader) Load() http.Header { + if v := ah.value.Load(); v != nil { + return v.(http.Header) + } + return nil +} + +// Store sets the http.Header value. +func (ah *AtomicHeader) Store(h http.Header) { + ah.value.Store(h) +} diff --git a/pkg/sources/huggingface/client.go b/pkg/sources/huggingface/client.go index d236265803b4..ecfdf093ac3f 100644 --- a/pkg/sources/huggingface/client.go +++ b/pkg/sources/huggingface/client.go @@ -8,6 +8,7 @@ import ( "time" "github.com/trufflesecurity/trufflehog/v3/pkg/context" + "github.com/trufflesecurity/trufflehog/v3/pkg/common" ) // Maps for API and HTML paths @@ -125,7 +126,8 @@ func NewHFClient(baseURL, apiKey string, timeout time.Duration) *HFClient { BaseURL: baseURL, APIKey: apiKey, HTTPClient: &http.Client{ - Timeout: timeout, + Timeout: timeout, + Transport: common.NewInstrumentedTransport(common.NewCustomTransport(nil)), }, } }