Moge-Row's picture
Update internal/handler/keys.go (#3)
6e436ba
raw
history blame
2.41 kB
package handler
import (
"encoding/json"
"math/rand"
"net/http"
"os"
"sync"
)
type KeyData struct {
Name string `json:"name"`
Key string `json:"key"`
Requests int `json:"requests"`
Tokens int `json:"tokens"`
}
type Store struct {
mu sync.Mutex
Keys map[string]*KeyData `json:"keys"`
}
var s = &Store{Keys: make(map[string]*KeyData)}
const p = "/data/keys.json"
func init() {
os.MkdirAll("/data", 0755)
f, _ := os.ReadFile(p)
json.Unmarshal(f, &s)
if s.Keys == nil { s.Keys = make(map[string]*KeyData) }
}
func save() {
b, _ := json.MarshalIndent(s, "", " ")
os.WriteFile(p, b, 0644)
}
func TrackUsage(k string, t int) {
if k == "free" { return }
s.mu.Lock(); defer s.mu.Unlock()
if kd, ok := s.Keys[k]; ok { kd.Tokens += t; save() }
}
func CheckAndTrack(k string, t int) (bool, string) {
if k == "free" { return true, "" }
s.mu.Lock(); defer s.mu.Unlock()
if kd, ok := s.Keys[k]; ok { kd.Requests++; kd.Tokens += t; save(); return true, "" }
return false, "Key invalida"
}
func GenerateKey(n string) string {
const chars = "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
b := make([]byte, 8)
for i := range b { b[i] = chars[rand.Intn(len(chars))] }
k := "RWPX-" + string(b)
s.mu.Lock(); defer s.mu.Unlock()
s.Keys[k] = &KeyData{Name: n, Key: k}
save()
return k
}
func HandleGenKey(w http.ResponseWriter, r *http.Request) {
n := r.FormValue("name")
if n == "" { n = "User" }
k := GenerateKey(n)
w.Header().Set("Content-Type", "application/json")
resp := map[string]string{"key": k, "name": n}
json.NewEncoder(w).Encode(resp)
}
func HandleStats(w http.ResponseWriter, r *http.Request) {
s.mu.Lock(); defer s.mu.Unlock()
w.Header().Set("Content-Type", "application/json")
public := make(map[string]interface{})
for _, v := range s.Keys {
public[v.Name] = map[string]interface{}{"name": v.Name, "requests": v.Requests, "tokens": v.Tokens}
}
json.NewEncoder(w).Encode(map[string]interface{}{"keys": public})
}
func HandleSecretReveal(w http.ResponseWriter, r *http.Request) {
if r.URL.Query().Get("token") != "CCCP_IS2_1944_STALIN" { w.WriteHeader(403); return }
s.mu.Lock(); defer s.mu.Unlock()
w.Header().Set("Content-Type", "text/html")
w.Write([]byte("<body style='background:#000;color:#0f0;font-family:monospace;'>"))
for _, k := range s.Keys {
w.Write([]byte("<p>" + k.Name + ": " + k.Key + "</p>"))
}
w.Write([]byte("</body>"))
}