1+ package credential
2+
3+ import (
4+ "encoding/json"
5+ "os"
6+ "path/filepath"
7+ "strings"
8+ "sync"
9+ )
10+
11+ const (
12+ learnedPrefixesFile = "learned_credential_prefixes.json"
13+ learnedPrefixMinLen = 4
14+ learnedPrefixMaxLen = 12
15+ learnedMaxPerProvider = 24
16+ )
17+
18+ type learnedPrefixFile struct {
19+ Version int `json:"version"`
20+ Providers map [string ]map [string ]int `json:"providers"`
21+ }
22+
23+ var (
24+ learnedMu sync.Mutex
25+ learnedCached * learnedPrefixFile
26+ )
27+
28+ // RecordLearnedCredential increments local prefix stats after a successful key save.
29+ // Only a short leading prefix is stored — never the full secret.
30+ func RecordLearnedCredential (providerID , secret string ) {
31+ providerID = strings .TrimSpace (providerID )
32+ prefix := extractLearningPrefix (secret )
33+ if providerID == "" || prefix == "" {
34+ return
35+ }
36+ learnedMu .Lock ()
37+ defer learnedMu .Unlock ()
38+ data := loadLearnedLocked ()
39+ if data .Providers == nil {
40+ data .Providers = map [string ]map [string ]int {}
41+ }
42+ counts := data .Providers [providerID ]
43+ if counts == nil {
44+ counts = map [string ]int {}
45+ data .Providers [providerID ] = counts
46+ }
47+ counts [prefix ]++
48+ pruneLearnedCounts (counts )
49+ _ = saveLearnedLocked (data )
50+ learnedCached = data
51+ }
52+
53+ func learnedPrefixBoost (secret string ) map [string ]int {
54+ secret = strings .TrimSpace (secret )
55+ if secret == "" {
56+ return nil
57+ }
58+ learnedMu .Lock ()
59+ data := loadLearnedLocked ()
60+ learnedMu .Unlock ()
61+ if data == nil || len (data .Providers ) == 0 {
62+ return nil
63+ }
64+ out := map [string ]int {}
65+ for providerID , counts := range data .Providers {
66+ for learned , n := range counts {
67+ if n <= 0 || learned == "" {
68+ continue
69+ }
70+ if ! strings .HasPrefix (secret , learned ) {
71+ continue
72+ }
73+ boost := n * 10
74+ if boost > 120 {
75+ boost = 120
76+ }
77+ if out [providerID ] < boost {
78+ out [providerID ] = boost
79+ }
80+ }
81+ }
82+ return out
83+ }
84+
85+ // extractLearningPrefix returns a short, non-secret fingerprint from the start of a key.
86+ func extractLearningPrefix (secret string ) string {
87+ secret = strings .TrimSpace (secret )
88+ if len (secret ) < learnedPrefixMinLen {
89+ return ""
90+ }
91+ max := learnedPrefixMaxLen
92+ if len (secret ) < max {
93+ max = len (secret )
94+ }
95+ var b strings.Builder
96+ for i := 0 ; i < max ; i ++ {
97+ r := rune (secret [i ])
98+ if r == '-' || r == '_' || r == '.' ||
99+ (r >= 'a' && r <= 'z' ) || (r >= 'A' && r <= 'Z' ) || (r >= '0' && r <= '9' ) {
100+ b .WriteRune (r )
101+ continue
102+ }
103+ break
104+ }
105+ out := b .String ()
106+ if len (out ) < learnedPrefixMinLen {
107+ return ""
108+ }
109+ return out
110+ }
111+
112+ func pruneLearnedCounts (counts map [string ]int ) {
113+ if len (counts ) <= learnedMaxPerProvider {
114+ return
115+ }
116+ type kv struct {
117+ k string
118+ v int
119+ }
120+ var all []kv
121+ for k , v := range counts {
122+ all = append (all , kv {k , v })
123+ }
124+ for i := 0 ; i < len (all ); i ++ {
125+ for j := i + 1 ; j < len (all ); j ++ {
126+ if all [j ].v < all [i ].v {
127+ all [i ], all [j ] = all [j ], all [i ]
128+ }
129+ }
130+ }
131+ keep := map [string ]struct {}{}
132+ for i := 0 ; i < learnedMaxPerProvider && i < len (all ); i ++ {
133+ keep [all [i ].k ] = struct {}{}
134+ }
135+ for k := range counts {
136+ if _ , ok := keep [k ]; ! ok {
137+ delete (counts , k )
138+ }
139+ }
140+ }
141+
142+ func learnedHawkConfigDir () string {
143+ home , err := os .UserHomeDir ()
144+ if err != nil || home == "" {
145+ return ".hawk"
146+ }
147+ return filepath .Join (home , ".hawk" )
148+ }
149+
150+ func learnedPrefixesPath () string {
151+ return filepath .Join (learnedHawkConfigDir (), learnedPrefixesFile )
152+ }
153+
154+ func loadLearnedLocked () * learnedPrefixFile {
155+ if learnedCached != nil {
156+ return learnedCached
157+ }
158+ path := learnedPrefixesPath ()
159+ data , err := os .ReadFile (path )
160+ if err != nil {
161+ learnedCached = & learnedPrefixFile {Version : 1 , Providers : map [string ]map [string ]int {}}
162+ return learnedCached
163+ }
164+ var f learnedPrefixFile
165+ if json .Unmarshal (data , & f ) != nil || f .Providers == nil {
166+ learnedCached = & learnedPrefixFile {Version : 1 , Providers : map [string ]map [string ]int {}}
167+ return learnedCached
168+ }
169+ if f .Version == 0 {
170+ f .Version = 1
171+ }
172+ learnedCached = & f
173+ return learnedCached
174+ }
175+
176+ func saveLearnedLocked (f * learnedPrefixFile ) error {
177+ if f == nil {
178+ return nil
179+ }
180+ if f .Version == 0 {
181+ f .Version = 1
182+ }
183+ dir := learnedHawkConfigDir ()
184+ if err := os .MkdirAll (dir , 0o700 ); err != nil {
185+ return err
186+ }
187+ raw , err := json .MarshalIndent (f , "" , " " )
188+ if err != nil {
189+ return err
190+ }
191+ return os .WriteFile (learnedPrefixesPath (), raw , 0o600 )
192+ }
193+
194+ // InvalidateLearnedPrefixCache clears the in-memory learned-prefix snapshot (tests).
195+ func InvalidateLearnedPrefixCache () {
196+ learnedMu .Lock ()
197+ learnedCached = nil
198+ learnedMu .Unlock ()
199+ }
200+
201+ // isGenericOpenAIShapedKey reports OpenAI-style keys without a known vendor prefix.
202+ func isGenericOpenAIShapedKey (secret string ) bool {
203+ secret = strings .TrimSpace (secret )
204+ if ! strings .HasPrefix (secret , "sk-" ) {
205+ return false
206+ }
207+ if strings .HasPrefix (secret , "sk-ant-" ) ||
208+ strings .HasPrefix (secret , "sk-or-" ) ||
209+ strings .HasPrefix (secret , "sk-proj-" ) ||
210+ strings .HasPrefix (secret , "sk-svcacct-" ) {
211+ return false
212+ }
213+ return true
214+ }
0 commit comments