add key system with RWPX prefix

internal/handler/keys.go

@@ -0,0 +1,108 @@
+package handler
+
+import (
+"encoding/json"
+"fmt"
+"math/rand"
+"net/http"
+"os"
+"sync"
+"time"
+)
+
+type KeyData struct {
+Name     string    \`json:"name"\`
+Key      string    \`json:"key"\`
+Requests int       \`json:"requests"\`
+Tokens   int       \`json:"tokens"\`
+Created  time.Time \`json:"created"\`
+LastUsed time.Time \`json:"last_used"\`
+}
+
+type KeyStore struct {
+mu   sync.Mutex
+Keys map[string]*KeyData \`json:"keys"\`
+}
+
+var store = &KeyStore{Keys: make(map[string]*KeyData)}
+var storePath = "/data/keys.json"
+
+func init() {
+os.MkdirAll("/data", 0755)
+data, err := os.ReadFile(storePath)
+if err == nil {
+json.Unmarshal(data, store)
+}
+if store.Keys == nil {
+store.Keys = make(map[string]*KeyData)
+}
+}
+
+func saveStore() {
+data, _ := json.MarshalIndent(store, "", "  ")
+os.WriteFile(storePath, data, 0644)
+}
+
+func GenerateKey(name string) string {
+const chars = "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
+b := make([]byte, 8)
+for i := range b {
+b[i] = chars[rand.Intn(len(chars))]
+}
+key := "RWPX-" + string(b)
+store.mu.Lock()
+store.Keys[key] = &KeyData{
+Name:    name,
+Key:     key,
+Created: time.Now(),
+}
+saveStore()
+store.mu.Unlock()
+return key
+}
+
+func ValidateKey(key string) (bool, *KeyData) {
+if key == "free" {
+return true, &KeyData{Name: "guest"}
+}
+store.mu.Lock()
+defer store.mu.Unlock()
+kd, ok := store.Keys[key]
+return ok, kd
+}
+
+func TrackUsage(key string, tokens int) {
+if key == "free" {
+return
+}
+store.mu.Lock()
+defer store.mu.Unlock()
+if kd, ok := store.Keys[key]; ok {
+kd.Requests++
+kd.Tokens += tokens
+kd.LastUsed = time.Now()
+saveStore()
+}
+}
+
+func HandleGenKey(w http.ResponseWriter, r *http.Request) {
+if r.Method != "POST" {
+http.Error(w, "Method not allowed", 405)
+return
+}
+name := r.FormValue("name")
+if name == "" {
+name = "usuario"
+}
+key := GenerateKey(name)
+w.Header().Set("Content-Type", "application/json")
+fmt.Fprintf(w, "{"key":"%s","name":"%s"}", key, name)
+}
+
+func HandleStats(w http.ResponseWriter, r *http.Request) {
+store.mu.Lock()
+defer store.mu.Unlock()
+w.Header().Set("Content-Type", "application/json")
+data, _ := json.Marshal(store)
+w.Write(data)
+}

main.go

@@ -17,6 +17,8 @@ func main() {
 	version.StartVersionUpdater()
 
 	http.HandleFunc("/", handler.HandleIndex)
+http.HandleFunc("/genkey", handler.HandleGenKey)
+http.HandleFunc("/stats", handler.HandleStats)
 http.HandleFunc("/v1/models", handler.HandleModels)
 	http.HandleFunc("/v1/chat/completions", handler.HandleChatCompletions)
 	http.HandleFunc("/v1/messages", handler.HandleMessages)