Skip to content

Commit f436f10

Browse files
McPatateMcPatate
committed
feat: add OpenRouter detector
closes #4499
1 parent 1ef44e7 commit f436f10

File tree

4 files changed

+355
-14
lines changed

4 files changed

+355
-14
lines changed

‎pkg/detectors/openrouter/openrouter.go‎

Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,131 @@
1+
package openrouter
2+
3+
import (
4+
"context"
5+
"encoding/json"
6+
"fmt"
7+
"io"
8+
"net/http"
9+
"strconv"
10+
11+
regexp "github.com/wasilibs/go-re2"
12+
13+
"github.com/trufflesecurity/trufflehog/v3/pkg/common"
14+
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors"
15+
"github.com/trufflesecurity/trufflehog/v3/pkg/pb/detectorspb"
16+
)
17+
18+
type Scanner struct {
19+
client *http.Client
20+
}
21+
22+
// Ensure the Scanner satisfies the interface at compile time.
23+
var _ detectors.Detector = (*Scanner)(nil)
24+
25+
var (
26+
defaultClient = common.SaneHttpClient()
27+
28+
keyPat = regexp.MustCompile(`\b(sk-or-v1-[0-9a-fA-F]{64})\b`)
29+
)
30+
31+
// Keywords are used for efficiently pre-filtering chunks.
32+
// Use identifiers in the secret preferably, or the provider name.
33+
func (s Scanner) Keywords() []string {
34+
return []string{"sk-or-v1-"}
35+
}
36+
37+
// FromData will find and optionally verify OpenAI secrets in a given set of bytes.
38+
func (s Scanner) FromData(ctx context.Context, verify bool, data []byte) (results []detectors.Result, err error) {
39+
dataStr := string(data)
40+
41+
uniqueMatches := make(map[string]struct{})
42+
for _, match := range keyPat.FindAllStringSubmatch(dataStr, -1) {
43+
uniqueMatches[match[1]] = struct{}{}
44+
}
45+
46+
for token := range uniqueMatches {
47+
s1 := detectors.Result{
48+
DetectorType: detectorspb.DetectorType_OpenRouter,
49+
// NOTE: we redact the same way it is done in the `Label` field
50+
Redacted: token[:12] + "..." + token[70:],
51+
Raw: []byte(token),
52+
}
53+
54+
if verify {
55+
client := s.client
56+
if client == nil {
57+
client = defaultClient
58+
}
59+
60+
verified, extraData, verificationErr := verifyToken(ctx, client, token)
61+
s1.Verified = verified
62+
s1.ExtraData = extraData
63+
s1.SetVerificationError(verificationErr)
64+
s1.AnalysisInfo = map[string]string{"key": token}
65+
}
66+
67+
results = append(results, s1)
68+
}
69+
70+
return results, err
71+
}
72+
73+
func verifyToken(ctx context.Context, client *http.Client, token string) (bool, map[string]string, error) {
74+
req, err := http.NewRequestWithContext(ctx, "GET", "https://openrouter.ai/api/v1/key", nil)
75+
if err != nil {
76+
return false, nil, err
77+
}
78+
79+
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
80+
res, err := client.Do(req)
81+
if err != nil {
82+
return false, nil, err
83+
}
84+
defer func() {
85+
_, _ = io.Copy(io.Discard, res.Body)
86+
_ = res.Body.Close()
87+
}()
88+
89+
switch res.StatusCode {
90+
case http.StatusOK:
91+
var keyResponse keyResponse
92+
if err = json.NewDecoder(res.Body).Decode(&keyResponse); err != nil {
93+
return false, nil, err
94+
}
95+
96+
key := keyResponse.Data
97+
extraData := map[string]string{
98+
"label": key.Label,
99+
"limit": fmt.Sprintf("%d", key.Limit),
100+
"usage": fmt.Sprintf("%d", key.Usage),
101+
"is_free_tier": strconv.FormatBool(key.IsFreeTier),
102+
"limit_remaining": fmt.Sprintf("%d", key.LimitRemaining),
103+
}
104+
return true, extraData, nil
105+
case http.StatusUnauthorized:
106+
// Invalid
107+
return false, nil, nil
108+
default:
109+
return false, nil, fmt.Errorf("unexpected HTTP response status %d", res.StatusCode)
110+
}
111+
}
112+
113+
func (s Scanner) Type() detectorspb.DetectorType {
114+
return detectorspb.DetectorType_OpenRouter
115+
}
116+
117+
func (s Scanner) Description() string {
118+
return "OpenRouter provides a unified API that gives you access to hundreds of AI models through a single endpoint, while automatically handling fallbacks and selecting the most cost-effective options."
119+
}
120+
121+
type keyResponse struct {
122+
Data key `json:"data"`
123+
}
124+
125+
type key struct {
126+
Label string `json:"label"`
127+
Limit int32 `json:"limit"`
128+
Usage int32 `json:"usage"`
129+
IsFreeTier bool `json:"is_free_tier"`
130+
LimitRemaining int32 `json:"limit_remaining"`
131+
}

‎pkg/detectors/openrouter/openrouter_integration_test.go‎

Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,125 @@
1+
//go:build detectors
2+
// +build detectors
3+
4+
package openrouter
5+
6+
import (
7+
"context"
8+
"fmt"
9+
"testing"
10+
"time"
11+
12+
"github.com/kylelemons/godebug/pretty"
13+
"github.com/trufflesecurity/trufflehog/v3/pkg/common"
14+
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors"
15+
"github.com/trufflesecurity/trufflehog/v3/pkg/pb/detectorspb"
16+
)
17+
18+
func TestOpenRouter_FromChunk(t *testing.T) {
19+
ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
20+
defer cancel()
21+
testSecrets, err := common.GetSecret(ctx, "trufflehog-testing", "detectors5")
22+
if err != nil {
23+
t.Fatalf("could not get test secrets from GCP: %s", err)
24+
}
25+
26+
secret := testSecrets.MustGetField("OPENROUTER")
27+
inactiveSecret := testSecrets.MustGetField("OPENROUTER_INACTIVE")
28+
29+
type args struct {
30+
ctx context.Context
31+
data []byte
32+
verify bool
33+
}
34+
35+
tests := []struct {
36+
name string
37+
s Scanner
38+
args args
39+
want []detectors.Result
40+
wantErr bool
41+
}{
42+
{
43+
name: "Found, unverified OpenRouter token sk-or-v1-",
44+
s: Scanner{},
45+
args: args{
46+
ctx: context.Background(),
47+
data: []byte(fmt.Sprintf("You can find an OpenRouter secret %s within", inactiveSecret)),
48+
verify: true,
49+
},
50+
want: []detectors.Result{
51+
{
52+
DetectorType: detectorspb.DetectorType_OpenRouter,
53+
Redacted: "sk-or-v1-3dd...aa5",
54+
Verified: false,
55+
},
56+
},
57+
wantErr: false,
58+
},
59+
{
60+
name: "Found, verified OpenRouter token sk-or-v1-",
61+
s: Scanner{},
62+
args: args{
63+
ctx: context.Background(),
64+
data: []byte(fmt.Sprintf("You can find an OpenRouter secret %s within", secret)),
65+
verify: true,
66+
},
67+
want: []detectors.Result{
68+
{
69+
DetectorType: detectorspb.DetectorType_OpenRouter,
70+
Verified: true,
71+
Redacted: "sk-or-v1-753...1a5",
72+
},
73+
},
74+
wantErr: false,
75+
},
76+
{
77+
name: "not found",
78+
s: Scanner{},
79+
args: args{
80+
ctx: context.Background(),
81+
data: []byte("You cannot find the secret within"),
82+
verify: true,
83+
},
84+
want: nil,
85+
wantErr: false,
86+
},
87+
}
88+
for _, tt := range tests {
89+
t.Run(tt.name, func(t *testing.T) {
90+
s := Scanner{}
91+
got, err := s.FromData(tt.args.ctx, tt.args.verify, tt.args.data)
92+
if (err != nil) != tt.wantErr {
93+
t.Errorf("OpenRouter.FromData() error = %v, wantErr %v", err, tt.wantErr)
94+
return
95+
}
96+
for i := range got {
97+
if len(got[i].Raw) == 0 {
98+
t.Fatal("no raw secret present")
99+
}
100+
got[i].Raw = nil
101+
got[i].ExtraData = nil
102+
got[i].AnalysisInfo = nil
103+
}
104+
if diff := pretty.Compare(got, tt.want); diff != "" {
105+
t.Errorf("OpenRouter.FromData() %s diff: (-got +want)\n%s", tt.name, diff)
106+
}
107+
})
108+
}
109+
}
110+
111+
func BenchmarkFromData(benchmark *testing.B) {
112+
ctx := context.Background()
113+
s := Scanner{}
114+
for name, data := range detectors.MustGetBenchmarkData() {
115+
benchmark.Run(name, func(b *testing.B) {
116+
b.ResetTimer()
117+
for n := 0; n < b.N; n++ {
118+
_, err := s.FromData(ctx, false, data)
119+
if err != nil {
120+
b.Fatal(err)
121+
}
122+
}
123+
})
124+
}
125+
}

‎pkg/detectors/openrouter/openrouter_test.go‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
package openrouter
2+
3+
import (
4+
"context"
5+
"testing"
6+
7+
"github.com/google/go-cmp/cmp"
8+
9+
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors"
10+
"github.com/trufflesecurity/trufflehog/v3/pkg/engine/ahocorasick"
11+
)
12+
13+
func TestOpenRouter_Pattern(t *testing.T) {
14+
d := Scanner{}
15+
ahoCorasickCore := ahocorasick.NewAhoCorasickCore([]detectors.Detector{d})
16+
tests := []struct {
17+
name string
18+
input string
19+
want []string
20+
}{
21+
{
22+
name: "API key",
23+
input: `OPENROUTER_API_KEY = "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"`,
24+
want: []string{"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"},
25+
},
26+
{
27+
name: "invalid pattern",
28+
input: `
29+
[INFO] Sending request to the openrouter API
30+
[DEBUG] Using Key=sk-or-v1-a2Cy8xCLyvrAf7lZKfhQhyCr4RAID9D
31+
[ERROR] Response received: 401 UnAuthorized
32+
`,
33+
want: []string{},
34+
},
35+
}
36+
37+
for _, test := range tests {
38+
t.Run(test.name, func(t *testing.T) {
39+
detectorMatches := ahoCorasickCore.FindDetectorMatches([]byte(test.input))
40+
if len(detectorMatches) == 0 {
41+
t.Errorf("keywords '%v' not matched by: %s", d.Keywords(), test.input)
42+
return
43+
}
44+
45+
results, err := d.FromData(context.Background(), false, []byte(test.input))
46+
if err != nil {
47+
t.Errorf("error = %v", err)
48+
return
49+
}
50+
51+
if len(results) != len(test.want) {
52+
if len(results) == 0 {
53+
t.Errorf("did not receive result")
54+
} else {
55+
t.Errorf("expected %d results, only received %d", len(test.want), len(results))
56+
}
57+
return
58+
}
59+
60+
actual := make(map[string]struct{}, len(results))
61+
for _, r := range results {
62+
if len(r.RawV2) > 0 {
63+
actual[string(r.RawV2)] = struct{}{}
64+
} else {
65+
actual[string(r.Raw)] = struct{}{}
66+
}
67+
}
68+
expected := make(map[string]struct{}, len(test.want))
69+
for _, v := range test.want {
70+
expected[v] = struct{}{}
71+
}
72+
73+
if diff := cmp.Diff(expected, actual); diff != "" {
74+
t.Errorf("%s diff: (-want +got)\n%s", test.name, diff)
75+
}
76+
})
77+
}
78+
}

0 commit comments

Comments
 (0)