From 98430961ba16f32055c1b1d3fdc6e5748280beca Mon Sep 17 00:00:00 2001 From: Lakshman Patel Date: Wed, 15 Jul 2026 11:16:28 +0530 Subject: [PATCH] Add opt-in live secret verification and exact token counting MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit DetectSecretsVerified adds TruffleHog-style live-credential verification (GitHub, Slack, Stripe, AWS key pairs) behind an explicit opt-in call — DetectSecrets itself is unchanged and still makes zero network calls. CountTokensExact adds an opt-in path to get real Claude/Gemini token counts via their count-tokens APIs, since neither publishes an offline tokenizer; the existing cl100k_base-approximated EstimateTokens family is unchanged, with doc comments now stating plainly that it's an approximation for those two providers. --- internal/core/exact_counter.go | 130 +++++++++++++++++ internal/core/exact_counter_test.go | 106 ++++++++++++++ internal/secrets/secrets.go | 68 ++++++++- internal/secrets/verify.go | 189 +++++++++++++++++++++++++ internal/secrets/verify_test.go | 207 ++++++++++++++++++++++++++++ secrets.go | 13 ++ tok.go | 14 ++ tokenizer.go | 13 +- 8 files changed, 737 insertions(+), 3 deletions(-) create mode 100644 internal/core/exact_counter.go create mode 100644 internal/core/exact_counter_test.go create mode 100644 internal/secrets/verify.go create mode 100644 internal/secrets/verify_test.go diff --git a/internal/core/exact_counter.go b/internal/core/exact_counter.go new file mode 100644 index 000000000..f588dd7f9 --- /dev/null +++ b/internal/core/exact_counter.go @@ -0,0 +1,130 @@ +package core + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/url" + "strings" + "time" +) + +// exactCountClient is used for opt-in exact-count API calls. It is a var +// (not embedded per-call) so tests can override it, and is only ever +// touched by CountTokensExact — never by EstimateTokens/EstimateTokensPrecise, +// which remain fully offline and unchanged. +var exactCountClient = &http.Client{Timeout: 10 * time.Second} + +// anthropicCountTokensURL and googleCountTokensURL are vars (not consts) so +// tests can point them at a local mock server. +var ( + anthropicCountTokensURL = "https://api.anthropic.com/v1/messages/count_tokens" + googleCountTokensURL = "https://generativelanguage.googleapis.com/v1beta/models/%s:countTokens" +) + +// CountTokensExact returns the exact token count for text under the real +// tokenizer of the model's provider (Anthropic or Google), by calling that +// provider's live count-tokens API. This exists because neither Anthropic +// nor Google publish an offline BPE tokenizer library — EstimateTokens and +// friends only approximate Claude/Gemini counts via OpenAI's cl100k_base +// encoding (see ClaudeBase/GeminiBase in the tok package). +// +// apiKey is required and is never read from the environment implicitly: +// this function makes a real network call using the given text and key, and +// is never invoked by any other function in this package — callers must opt +// in explicitly. model must have a "claude" or "gemini" prefix (case +// insensitive); any other prefix returns an error rather than silently +// falling back to an approximation. +func CountTokensExact(ctx context.Context, apiKey, model, text string) (int, error) { + if apiKey == "" { + return 0, fmt.Errorf("tok: CountTokensExact: apiKey is required") + } + lower := strings.ToLower(strings.TrimSpace(model)) + switch { + case strings.HasPrefix(lower, "claude"): + return countTokensAnthropic(ctx, apiKey, model, text) + case strings.HasPrefix(lower, "gemini"): + return countTokensGoogle(ctx, apiKey, model, text) + default: + return 0, fmt.Errorf("tok: CountTokensExact: unsupported model %q (must be a claude* or gemini* model)", model) + } +} + +func countTokensAnthropic(ctx context.Context, apiKey, model, text string) (int, error) { + body, err := json.Marshal(map[string]any{ + "model": model, + "messages": []map[string]string{ + {"role": "user", "content": text}, + }, + }) + if err != nil { + return 0, fmt.Errorf("tok: CountTokensExact: marshal request: %w", err) + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, anthropicCountTokensURL, bytes.NewReader(body)) + if err != nil { + return 0, fmt.Errorf("tok: CountTokensExact: create request: %w", err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("x-api-key", apiKey) + req.Header.Set("anthropic-version", "2023-06-01") + + resp, err := exactCountClient.Do(req) + if err != nil { + return 0, fmt.Errorf("tok: CountTokensExact: request failed: %w", err) + } + defer func() { _ = resp.Body.Close() }() + + if resp.StatusCode != http.StatusOK { + data, _ := io.ReadAll(resp.Body) + return 0, fmt.Errorf("tok: CountTokensExact: anthropic API returned %d: %s", resp.StatusCode, string(data)) + } + + var out struct { + InputTokens int `json:"input_tokens"` + } + if err := json.NewDecoder(resp.Body).Decode(&out); err != nil { + return 0, fmt.Errorf("tok: CountTokensExact: decode response: %w", err) + } + return out.InputTokens, nil +} + +func countTokensGoogle(ctx context.Context, apiKey, model, text string) (int, error) { + body, err := json.Marshal(map[string]any{ + "contents": []map[string]any{ + {"parts": []map[string]string{{"text": text}}}, + }, + }) + if err != nil { + return 0, fmt.Errorf("tok: CountTokensExact: marshal request: %w", err) + } + + endpoint := fmt.Sprintf(googleCountTokensURL, url.PathEscape(model)) + "?key=" + url.QueryEscape(apiKey) + req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(body)) + if err != nil { + return 0, fmt.Errorf("tok: CountTokensExact: create request: %w", err) + } + req.Header.Set("Content-Type", "application/json") + + resp, err := exactCountClient.Do(req) + if err != nil { + return 0, fmt.Errorf("tok: CountTokensExact: request failed: %w", err) + } + defer func() { _ = resp.Body.Close() }() + + if resp.StatusCode != http.StatusOK { + data, _ := io.ReadAll(resp.Body) + return 0, fmt.Errorf("tok: CountTokensExact: google API returned %d: %s", resp.StatusCode, string(data)) + } + + var out struct { + TotalTokens int `json:"totalTokens"` + } + if err := json.NewDecoder(resp.Body).Decode(&out); err != nil { + return 0, fmt.Errorf("tok: CountTokensExact: decode response: %w", err) + } + return out.TotalTokens, nil +} diff --git a/internal/core/exact_counter_test.go b/internal/core/exact_counter_test.go new file mode 100644 index 000000000..b4c4896e7 --- /dev/null +++ b/internal/core/exact_counter_test.go @@ -0,0 +1,106 @@ +package core + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func TestCountTokensExact_RequiresAPIKey(t *testing.T) { + if _, err := CountTokensExact(context.Background(), "", "claude-3-5-sonnet", "hello"); err == nil { + t.Fatal("expected an error when apiKey is empty, got nil") + } +} + +func TestCountTokensExact_UnsupportedModel(t *testing.T) { + if _, err := CountTokensExact(context.Background(), "sk-fake", "gpt-4o", "hello"); err == nil { + t.Fatal("expected an error for a non-claude/gemini model, got nil") + } +} + +func TestCountTokensExact_Anthropic(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if got := r.Header.Get("x-api-key"); got != "sk-ant-fake" { + t.Errorf("x-api-key = %q, want %q", got, "sk-ant-fake") + } + if got := r.Header.Get("anthropic-version"); got == "" { + t.Error("anthropic-version header missing") + } + var body map[string]any + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + t.Fatalf("decode request body: %v", err) + } + if body["model"] != "claude-3-5-sonnet-20241022" { + t.Errorf("model = %v, want claude-3-5-sonnet-20241022", body["model"]) + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]int{"input_tokens": 7}) + })) + defer srv.Close() + + orig := anthropicCountTokensURL + anthropicCountTokensURL = srv.URL + defer func() { anthropicCountTokensURL = orig }() + + got, err := CountTokensExact(context.Background(), "sk-ant-fake", "claude-3-5-sonnet-20241022", "hello world") + if err != nil { + t.Fatalf("CountTokensExact() error: %v", err) + } + if got != 7 { + t.Errorf("CountTokensExact() = %d, want 7", got) + } +} + +func TestCountTokensExact_AnthropicError(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(`{"error": "invalid api key"}`)) + })) + defer srv.Close() + + orig := anthropicCountTokensURL + anthropicCountTokensURL = srv.URL + defer func() { anthropicCountTokensURL = orig }() + + if _, err := CountTokensExact(context.Background(), "bad-key", "claude-3-5-sonnet", "hi"); err == nil { + t.Fatal("expected an error for a 401 response, got nil") + } else if !strings.Contains(err.Error(), "401") { + t.Errorf("error = %v, want it to mention status 401", err) + } +} + +func TestCountTokensExact_Google(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if got := r.URL.Query().Get("key"); got != "goog-fake" { + t.Errorf("key query param = %q, want %q", got, "goog-fake") + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]int{"totalTokens": 5}) + })) + defer srv.Close() + + orig := googleCountTokensURL + googleCountTokensURL = srv.URL + "/%s" + defer func() { googleCountTokensURL = orig }() + + got, err := CountTokensExact(context.Background(), "goog-fake", "gemini-2.5-pro", "hello world") + if err != nil { + t.Fatalf("CountTokensExact() error: %v", err) + } + if got != 5 { + t.Errorf("CountTokensExact() = %d, want 5", got) + } +} + +func TestCountTokensExact_NeverCalledByEstimateTokens(t *testing.T) { + // EstimateTokens and friends must remain fully offline. This test + // doesn't hit the network by construction (no http.Client involved), + // which is the point: there is no code path from EstimateTokens into + // CountTokensExact or exactCountClient. + if got := EstimateTokens("hello world"); got <= 0 { + t.Errorf("EstimateTokens() = %d, want > 0", got) + } +} diff --git a/internal/secrets/secrets.go b/internal/secrets/secrets.go index af078e9af..80d18546a 100644 --- a/internal/secrets/secrets.go +++ b/internal/secrets/secrets.go @@ -1,7 +1,9 @@ package secrets import ( + "context" "math" + "net/http" "regexp" "sort" "strings" @@ -15,11 +17,22 @@ type SecretMatch struct { Masked string StartPos int EndPos int + + // Verified reports whether a live-credential check confirmed this + // secret is an active credential. nil means verification was not + // attempted — this is always the case for DetectSecrets, which never + // makes network calls. A non-nil value means DetectSecretsVerified + // attempted a live check for this match's type; verification errors + // (network failures, timeouts, etc.) also leave this nil rather than + // reporting a false negative. + Verified *bool } type secretPattern struct { Type string Pattern *regexp.Regexp + // Verify is optional; nil means no live verifier exists for this type. + Verify Verifier } // SecretDetector detects and redacts secrets from text using compiled regex patterns. @@ -93,7 +106,7 @@ func compilePatterns() []secretPattern { if err != nil { continue } - patterns = append(patterns, secretPattern{Type: r.Type, Pattern: compiled}) + patterns = append(patterns, secretPattern{Type: r.Type, Pattern: compiled, Verify: patternVerifiers[r.Type]}) } return patterns } @@ -138,6 +151,59 @@ func (sd *SecretDetector) DetectSecrets(text string) []SecretMatch { return matches } +// DetectSecretsVerified is like DetectSecrets, but additionally makes a live +// network call per detected secret (where a verifier exists for its type) +// to confirm whether the credential is actually an active one — following +// TruffleHog's live-verification model. This is the only function in this +// package that makes network calls; DetectSecrets and RedactSecrets never +// do, and remain unchanged whether or not this function is ever called. +// +// client may be nil, in which case a default 5-second-timeout client is +// used. AWS Access Keys are only verified when a paired AWS Secret Key was +// also matched in the same text — an access key ID alone can't be signed, +// so without a pair the match's Verified field stays nil, identical to +// DetectSecrets. +func (sd *SecretDetector) DetectSecretsVerified(ctx context.Context, client *http.Client, text string) []SecretMatch { + if client == nil { + client = defaultVerifyClient + } + matches := sd.DetectSecrets(text) + + var awsSecretKey string + for _, m := range matches { + if m.Type == "AWS Secret Key" { + awsSecretKey = m.Value + break + } + } + + for i := range matches { + m := &matches[i] + if m.Type == "AWS Access Key" { + if awsSecretKey == "" { + continue + } + ok, err := verifyAWSKeyPair(ctx, client, m.Value, awsSecretKey) + if err != nil { + continue + } + m.Verified = &ok + continue + } + + verify := patternVerifiers[m.Type] + if verify == nil { + continue + } + ok, err := verify(ctx, client, m.Value) + if err != nil { + continue + } + m.Verified = &ok + } + return matches +} + // RedactSecrets replaces detected secrets with [REDACTED] markers. func (sd *SecretDetector) RedactSecrets(text string) string { matches := sd.DetectSecrets(text) diff --git a/internal/secrets/verify.go b/internal/secrets/verify.go new file mode 100644 index 000000000..3146b1adb --- /dev/null +++ b/internal/secrets/verify.go @@ -0,0 +1,189 @@ +package secrets + +import ( + "context" + "crypto/hmac" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "net/http" + "net/url" + "strings" + "time" +) + +// Verifier makes a live network call to confirm whether a detected secret +// value is an active credential, following TruffleHog's live-verification +// model. Verifiers are never invoked by DetectSecrets — only by +// DetectSecretsVerified, which callers must opt into explicitly. +type Verifier func(ctx context.Context, client *http.Client, value string) (bool, error) + +// defaultVerifyClient is used when DetectSecretsVerified is called with a +// nil client. +var defaultVerifyClient = &http.Client{Timeout: 5 * time.Second} + +// patternVerifiers maps a secretPattern's Type to its live-verification +// function. Types with no entry here are never verified (Verified stays nil +// even under DetectSecretsVerified) — either because no reasonable API +// exists to check them (e.g. "Generic Password") or because verification +// needs more than the single matched value (AWS Access Key, handled +// separately in DetectSecretsVerified via pairing with AWS Secret Key). +var patternVerifiers = map[string]Verifier{ + "GitHub Token": verifyGitHubToken, + "GitHub Fine-grained Token": verifyGitHubToken, + "Slack Bot Token": verifySlackToken, + "Slack User Token": verifySlackToken, + "Stripe Secret Key": verifyStripeKey, +} + +func verifyGitHubToken(ctx context.Context, client *http.Client, token string) (bool, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, "https://api.github.com/user", nil) + if err != nil { + return false, err + } + req.Header.Set("Authorization", "token "+token) + req.Header.Set("Accept", "application/vnd.github+json") + + resp, err := client.Do(req) + if err != nil { + return false, err + } + defer func() { _ = resp.Body.Close() }() + return resp.StatusCode == http.StatusOK, nil +} + +func verifySlackToken(ctx context.Context, client *http.Client, token string) (bool, error) { + form := url.Values{"token": {token}} + req, err := http.NewRequestWithContext(ctx, http.MethodPost, "https://slack.com/api/auth.test", strings.NewReader(form.Encode())) + if err != nil { + return false, err + } + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + + resp, err := client.Do(req) + if err != nil { + return false, err + } + defer func() { _ = resp.Body.Close() }() + if resp.StatusCode != http.StatusOK { + return false, nil + } + + var body struct { + OK bool `json:"ok"` + } + if err := json.NewDecoder(resp.Body).Decode(&body); err != nil { + return false, err + } + return body.OK, nil +} + +func verifyStripeKey(ctx context.Context, client *http.Client, key string) (bool, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, "https://api.stripe.com/v1/balance", nil) + if err != nil { + return false, err + } + req.SetBasicAuth(key, "") + + resp, err := client.Do(req) + if err != nil { + return false, err + } + defer func() { _ = resp.Body.Close() }() + return resp.StatusCode == http.StatusOK, nil +} + +// awsSTSBaseURL is the AWS STS endpoint used for AWS key-pair verification. +// It is a var (not a const) so tests can point it at a local mock server. +var awsSTSBaseURL = "https://sts.amazonaws.com" + +// verifyAWSKeyPair verifies an AWS access key + secret key pair via STS +// GetCallerIdentity, signed with AWS Signature Version 4. Unlike the other +// verifiers, this needs two values (an access key ID alone can't be +// signed) — DetectSecretsVerified only calls this when both an "AWS Access +// Key" and an "AWS Secret Key" were matched in the same scan. +func verifyAWSKeyPair(ctx context.Context, client *http.Client, accessKey, secretKey string) (bool, error) { + const ( + region = "us-east-1" + service = "sts" + ) + + base, err := url.Parse(awsSTSBaseURL) + if err != nil { + return false, err + } + host := base.Host + + now := time.Now().UTC() + amzDate := now.Format("20060102T150405Z") + dateStamp := now.Format("20060102") + + canonicalQuery := url.Values{ + "Action": {"GetCallerIdentity"}, + "Version": {"2011-06-15"}, + }.Encode() + + canonicalHeaders := "host:" + host + "\n" + "x-amz-date:" + amzDate + "\n" + signedHeaders := "host;x-amz-date" + payloadHash := sha256Hex("") + + canonicalRequest := strings.Join([]string{ + http.MethodGet, + "/", + canonicalQuery, + canonicalHeaders, + signedHeaders, + payloadHash, + }, "\n") + + credentialScope := dateStamp + "/" + region + "/" + service + "/aws4_request" + stringToSign := strings.Join([]string{ + "AWS4-HMAC-SHA256", + amzDate, + credentialScope, + sha256Hex(canonicalRequest), + }, "\n") + + signingKey := deriveAWSSigningKey(secretKey, dateStamp, region, service) + signature := hex.EncodeToString(hmacSHA256(signingKey, stringToSign)) + + authHeader := fmt.Sprintf( + "AWS4-HMAC-SHA256 Credential=%s/%s, SignedHeaders=%s, Signature=%s", + accessKey, credentialScope, signedHeaders, signature, + ) + + reqURL := awsSTSBaseURL + "/?" + canonicalQuery + req, err := http.NewRequestWithContext(ctx, http.MethodGet, reqURL, nil) + if err != nil { + return false, err + } + req.Host = host + req.Header.Set("X-Amz-Date", amzDate) + req.Header.Set("Authorization", authHeader) + + resp, err := client.Do(req) + if err != nil { + return false, err + } + defer func() { _ = resp.Body.Close() }() + return resp.StatusCode == http.StatusOK, nil +} + +func sha256Hex(s string) string { + h := sha256.Sum256([]byte(s)) + return hex.EncodeToString(h[:]) +} + +func hmacSHA256(key []byte, data string) []byte { + h := hmac.New(sha256.New, key) + h.Write([]byte(data)) + return h.Sum(nil) +} + +func deriveAWSSigningKey(secretKey, dateStamp, region, service string) []byte { + kDate := hmacSHA256([]byte("AWS4"+secretKey), dateStamp) + kRegion := hmacSHA256(kDate, region) + kService := hmacSHA256(kRegion, service) + return hmacSHA256(kService, "aws4_request") +} diff --git a/internal/secrets/verify_test.go b/internal/secrets/verify_test.go new file mode 100644 index 000000000..71097a4e4 --- /dev/null +++ b/internal/secrets/verify_test.go @@ -0,0 +1,207 @@ +package secrets + +import ( + "context" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "testing" +) + +func TestDetectSecrets_NeverCallsNetwork(t *testing.T) { + // A client whose transport always fails — if DetectSecrets ever tried to + // use it, the test would need to observe a failure. Since DetectSecrets + // doesn't accept an http.Client at all, this test instead asserts the + // stronger property directly: DetectSecrets returns unverified matches + // (Verified == nil) with no way to have made a network call. + sd := NewDetector() + matches := sd.DetectSecrets("token ghp_0123456789abcdefghijklmnopqrstuvwxyz end") + if len(matches) == 0 { + t.Fatal("expected a match") + } + for _, m := range matches { + if m.Verified != nil { + t.Errorf("DetectSecrets produced a non-nil Verified field for %q; it must never attempt verification", m.Type) + } + } +} + +func TestVerifyGitHubToken(t *testing.T) { + cases := []struct { + name string + status int + want bool + }{ + {"live token", http.StatusOK, true}, + {"revoked token", http.StatusUnauthorized, false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if got := r.Header.Get("Authorization"); got != "token abc123" { + t.Errorf("Authorization header = %q, want %q", got, "token abc123") + } + w.WriteHeader(tc.status) + })) + defer srv.Close() + + // verifyGitHubToken hardcodes the GitHub API host, so exercise + // the HTTP-mechanics/status-interpretation logic directly + // against the mock server via a RoundTripper redirect instead + // of trying to override the URL. + client := &http.Client{Transport: redirectTransport{target: srv.URL}} + got, err := verifyGitHubToken(context.Background(), client, "abc123") + if err != nil { + t.Fatalf("verifyGitHubToken() error: %v", err) + } + if got != tc.want { + t.Errorf("verifyGitHubToken() = %v, want %v", got, tc.want) + } + }) + } +} + +func TestVerifySlackToken(t *testing.T) { + cases := []struct { + name string + body string + want bool + }{ + {"live token", `{"ok": true}`, true}, + {"revoked token", `{"ok": false, "error": "invalid_auth"}`, false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(tc.body)) + })) + defer srv.Close() + + client := &http.Client{Transport: redirectTransport{target: srv.URL}} + got, err := verifySlackToken(context.Background(), client, "xoxb-fake") + if err != nil { + t.Fatalf("verifySlackToken() error: %v", err) + } + if got != tc.want { + t.Errorf("verifySlackToken() = %v, want %v", got, tc.want) + } + }) + } +} + +func TestVerifyStripeKey(t *testing.T) { + cases := []struct { + name string + status int + want bool + }{ + {"live key", http.StatusOK, true}, + {"revoked key", http.StatusUnauthorized, false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + user, _, ok := r.BasicAuth() + if !ok || user != "sk_test_fake" { + t.Errorf("expected basic auth with key as username, got user=%q ok=%v", user, ok) + } + w.WriteHeader(tc.status) + })) + defer srv.Close() + + client := &http.Client{Transport: redirectTransport{target: srv.URL}} + got, err := verifyStripeKey(context.Background(), client, "sk_test_fake") + if err != nil { + t.Fatalf("verifyStripeKey() error: %v", err) + } + if got != tc.want { + t.Errorf("verifyStripeKey() = %v, want %v", got, tc.want) + } + }) + } +} + +func TestVerifyAWSKeyPair(t *testing.T) { + cases := []struct { + name string + status int + want bool + }{ + {"live pair", http.StatusOK, true}, + {"invalid pair", http.StatusForbidden, false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + var gotAuth string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotAuth = r.Header.Get("Authorization") + w.WriteHeader(tc.status) + })) + defer srv.Close() + + orig := awsSTSBaseURL + awsSTSBaseURL = srv.URL + defer func() { awsSTSBaseURL = orig }() + + got, err := verifyAWSKeyPair(context.Background(), srv.Client(), "AKIAIOSFODNN7EXAMPLE", "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY") + if err != nil { + t.Fatalf("verifyAWSKeyPair() error: %v", err) + } + if got != tc.want { + t.Errorf("verifyAWSKeyPair() = %v, want %v", got, tc.want) + } + if !strings.HasPrefix(gotAuth, "AWS4-HMAC-SHA256 Credential=AKIAIOSFODNN7EXAMPLE/") { + t.Errorf("Authorization header = %q, want AWS4-HMAC-SHA256 Credential=AKIAIOSFODNN7EXAMPLE/... prefix", gotAuth) + } + }) + } +} + +func TestDetectSecretsVerified_AWSRequiresPair(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + t.Error("verifier should not be called when no AWS Secret Key is present") + })) + defer srv.Close() + orig := awsSTSBaseURL + awsSTSBaseURL = srv.URL + defer func() { awsSTSBaseURL = orig }() + + sd := NewDetector() + matches := sd.DetectSecretsVerified(context.Background(), srv.Client(), "key=AKIAIOSFODNN7EXAMPLE here, no secret key present") + + for _, m := range matches { + if m.Type == "AWS Access Key" && m.Verified != nil { + t.Errorf("AWS Access Key without a paired secret should stay unverified, got Verified=%v", *m.Verified) + } + } +} + +func TestDetectSecretsVerified_UnknownTypeStaysNil(t *testing.T) { + sd := NewDetector() + matches := sd.DetectSecretsVerified(context.Background(), nil, "password=supersecret123") + for _, m := range matches { + if m.Type == "Generic Password" && m.Verified != nil { + t.Errorf("Generic Password has no verifier and should stay unverified, got Verified=%v", *m.Verified) + } + } +} + +// redirectTransport rewrites every request's scheme+host to target's, +// preserving path/query, so verifier functions that hardcode a real +// external URL can still be exercised against a local httptest.Server. +type redirectTransport struct { + target string +} + +func (rt redirectTransport) RoundTrip(req *http.Request) (*http.Response, error) { + target, err := url.Parse(rt.target + req.URL.RequestURI()) + if err != nil { + return nil, err + } + clone := req.Clone(req.Context()) + clone.URL = target + clone.Host = target.Host + return http.DefaultTransport.RoundTrip(clone) +} diff --git a/secrets.go b/secrets.go index 155378783..ab57120cd 100644 --- a/secrets.go +++ b/secrets.go @@ -1,6 +1,8 @@ package tok import ( + "context" + "net/http" "sync" "github.com/GrayCodeAI/tok/internal/secrets" @@ -38,6 +40,17 @@ func (sd *SecretDetector) DetectSecrets(text string) []SecretMatch { return sd.inner.DetectSecrets(text) } +// DetectSecretsVerified is like DetectSecrets, but additionally makes a live +// network call per detected secret (where a verifier exists for its type) +// to confirm whether the credential is actually active — following +// TruffleHog's live-verification model. This is the only method on +// SecretDetector that makes network calls; DetectSecrets never does, and +// callers must opt into this explicitly. client may be nil to use a default +// 5-second-timeout client. +func (sd *SecretDetector) DetectSecretsVerified(ctx context.Context, client *http.Client, text string) []SecretMatch { + return sd.inner.DetectSecretsVerified(ctx, client, text) +} + // RedactSecrets replaces detected secrets with [REDACTED] markers. func (sd *SecretDetector) RedactSecrets(text string) string { return sd.inner.RedactSecrets(text) diff --git a/tok.go b/tok.go index 78bcc1f50..6832951b0 100644 --- a/tok.go +++ b/tok.go @@ -106,6 +106,20 @@ func EstimateTokensForModel(text string, model string) int { return core.EstimateTokensForModel(text, model) } +// CountTokensExact returns the exact token count for text under the real +// tokenizer of the model's provider, by calling that provider's live +// count-tokens API. Unlike EstimateTokensForModel (which approximates +// Claude/Gemini counts via OpenAI's cl100k_base encoding, since neither +// provider publishes an offline tokenizer), this makes a real network call +// using apiKey and returns the provider's own exact count. +// +// model must have a "claude" or "gemini" prefix. apiKey is required and is +// never read from the environment implicitly — this is an explicit opt-in; +// no other function in this package makes network calls. +func CountTokensExact(ctx context.Context, apiKey, model, text string) (int, error) { + return core.CountTokensExact(ctx, apiKey, model, text) +} + // WarmupTokenizer pre-initializes the BPE tokenizer in the background. // Call at application startup to avoid latency on the first Compress call. func WarmupTokenizer() { diff --git a/tokenizer.go b/tokenizer.go index 532f28b0d..17da4fd53 100644 --- a/tokenizer.go +++ b/tokenizer.go @@ -10,9 +10,18 @@ const ( Cl100kBase Encoding = "cl100k_base" // O200kBase is used by GPT-4o, GPT-4.1, o1, o3, o4-mini. O200kBase Encoding = "o200k_base" - // ClaudeBase is used by Claude models (uses cl100k_base as closest approximation). + // ClaudeBase is an approximation only: Anthropic publishes no offline + // BPE tokenizer, so counts using this encoding are OpenAI's cl100k_base + // token boundaries, not Claude's real tokenizer, and will diverge from + // the real count. For an exact count, use CountTokensExact, which calls + // Anthropic's /v1/messages/count_tokens API (requires an API key). ClaudeBase Encoding = "cl100k_base" - // GeminiBase is used by Gemini models (uses cl100k_base as closest approximation). + // GeminiBase is an approximation only: Google publishes no offline BPE + // tokenizer for Gemini, so counts using this encoding are OpenAI's + // cl100k_base token boundaries, not Gemini's real tokenizer, and will + // diverge from the real count. For an exact count, use + // CountTokensExact, which calls Google's countTokens API (requires an + // API key). GeminiBase Encoding = "cl100k_base" // R50kBase is used by older GPT-3 models. R50kBase Encoding = "r50k_base"