Skip to content

Commit f2cd191

Browse files
McPatateshahzadhaider1MuneebUllahKhan222
authored
feat: add OpenRouter detector (#4500)
* feat: add OpenRouter detector closes #4499 * refactor: correct doc comment referencing OpenAI to OpenRouter * refactor: remove unnecessary `AnalysisInfo` field * fix: data field types Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com> * fix: handle float fmt correctly Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com> * fix detector type add feature flag for openrouter detector * fix integration tests --------- Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com> Co-authored-by: Shahzad Haider <shahzadhaider.se@gmail.com> Co-authored-by: Shahzad Haider <76992801+shahzadhaider1@users.noreply.github.com> Co-authored-by: Muneeb Ullah Khan <100969065+MuneebUllahKhan222@users.noreply.github.com>
1 parent 00155c9 commit f2cd191

9 files changed

Lines changed: 360 additions & 7 deletions

File tree

‎main.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -551,6 +551,7 @@ func run(state overseer.State, logSync func() error) {
551551
feature.PgAnalyzeReadKeyDetectorEnabled.Store(true)
552552
feature.RedHatPyxisDetectorEnabled.Store(true)
553553
feature.OctopusDeployDetectorEnabled.Store(true)
554+
feature.OpenRouterDetectorEnabled.Store(true)
554555

555556
conf := &config.Config{}
556557
if *configFilename != "" {

‎pkg/detectors/openrouter/openrouter.go‎

Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,138 @@
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/detector_typepb"
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-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 OpenRouter 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: s.Type(),
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+
SecretParts: map[string]string{"key": token},
53+
}
54+
55+
if verify {
56+
client := s.client
57+
if client == nil {
58+
client = defaultClient
59+
}
60+
61+
verified, extraData, verificationErr := verifyToken(ctx, client, token)
62+
s1.Verified = verified
63+
s1.ExtraData = extraData
64+
s1.SetVerificationError(verificationErr)
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": fmtFloatPtr(key.Limit),
100+
"usage": fmt.Sprintf("%f", key.Usage),
101+
"is_free_tier": strconv.FormatBool(key.IsFreeTier),
102+
"limit_remaining": fmtFloatPtr(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() detector_typepb.DetectorType {
114+
return detector_typepb.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+
func fmtFloatPtr(f *float64) string {
122+
if f == nil {
123+
return "<nil>"
124+
}
125+
return fmt.Sprintf("%f", *f)
126+
}
127+
128+
type keyResponse struct {
129+
Data key `json:"data"`
130+
}
131+
132+
type key struct {
133+
Label string `json:"label"`
134+
Limit *float64 `json:"limit"`
135+
Usage float64 `json:"usage"`
136+
IsFreeTier bool `json:"is_free_tier"`
137+
LimitRemaining *float64 `json:"limit_remaining"`
138+
}

‎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/detector_typepb"
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", "detectors6")
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: detector_typepb.DetectorType_OpenRouter,
53+
Redacted: inactiveSecret[:12] + "..." + inactiveSecret[70:],
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: detector_typepb.DetectorType_OpenRouter,
70+
Verified: true,
71+
Redacted: secret[:12] + "..." + secret[70:],
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].SecretParts = 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+
}

‎pkg/engine/defaults/defaults.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -535,6 +535,7 @@ import (
535535
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/openai"
536536
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/openaiadmin"
537537
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/opencagedata"
538+
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/openrouter"
538539
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/openuv"
539540
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/openvpn"
540541
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/openweather"
@@ -1436,6 +1437,7 @@ func buildDetectorList() []detectors.Detector {
14361437
&openai.Scanner{},
14371438
&openaiadmin.Scanner{},
14381439
&opencagedata.Scanner{},
1440+
&openrouter.Scanner{},
14391441
&openuv.Scanner{},
14401442
&openvpn.Scanner{},
14411443
&openweather.Scanner{},
@@ -1816,6 +1818,8 @@ func buildDetectorList() []detectors.Detector {
18161818
return !feature.RedHatPyxisDetectorEnabled.Load()
18171819
case *octopusdeploy.Scanner:
18181820
return !feature.OctopusDeployDetectorEnabled.Load()
1821+
case *openrouter.Scanner:
1822+
return !feature.OpenRouterDetectorEnabled.Load()
18191823
default:
18201824
return false
18211825
}

‎pkg/engine/defaults/defaults_test.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,7 @@ var excludedFromDefaultList = map[detector_typepb.DetectorType]struct{}{
138138
detector_typepb.DetectorType_PgAnalyzeReadKey: {},
139139
detector_typepb.DetectorType_RedHatPyxis: {},
140140
detector_typepb.DetectorType_OctopusDeploy: {},
141+
detector_typepb.DetectorType_OpenRouter: {},
141142

142143
// Reserved / special types.
143144
detector_typepb.DetectorType_CustomRegex: {}, // added dynamically via engine config, not via buildDetectorList()

‎pkg/feature/feature.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ var (
3030
RedHatPyxisDetectorEnabled atomic.Bool
3131
OctopusDeployDetectorEnabled atomic.Bool
3232
DropUnverifiedJWTResults atomic.Bool
33+
OpenRouterDetectorEnabled atomic.Bool
3334
)
3435

3536
type AtomicString struct {

0 commit comments

Comments
 (0)