ping_service/manager/apikeys.go

package main

import (
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"os"
	"sync"
	"time"
)

// APIKey represents an API key for external workers
type APIKey struct {
	Key         string    `json:"key"`          // The actual API key (hashed in storage)
	Name        string    `json:"name"`         // Human-readable name
	WorkerType  string    `json:"worker_type"`  // "ping" for now, could expand
	CreatedAt   time.Time `json:"created_at"`
	LastUsedAt  time.Time `json:"last_used_at,omitempty"`
	RequestCount int64    `json:"request_count"`
	Enabled     bool      `json:"enabled"`
}

// APIKeyStore manages API keys with encrypted storage
type APIKeyStore struct {
	keys   map[string]*APIKey // key -> APIKey (key is the actual API key)
	mu     sync.RWMutex
	file   string
	crypto *Crypto
}

func NewAPIKeyStore(filename string, crypto *Crypto) *APIKeyStore {
	ks := &APIKeyStore{
		keys:   make(map[string]*APIKey),
		file:   filename,
		crypto: crypto,
	}
	ks.load()
	return ks
}

// GenerateAPIKey creates a new API key (32 bytes = 256 bits)
func GenerateAPIKey() (string, error) {
	bytes := make([]byte, 32)
	if _, err := rand.Read(bytes); err != nil {
		return "", err
	}
	// Use base64 URL encoding (filesystem/URL safe)
	return base64.URLEncoding.EncodeToString(bytes), nil
}

// Add creates and stores a new API key
func (ks *APIKeyStore) Add(name, workerType string) (string, error) {
	ks.mu.Lock()
	defer ks.mu.Unlock()

	key, err := GenerateAPIKey()
	if err != nil {
		return "", err
	}

	apiKey := &APIKey{
		Key:        key,
		Name:       name,
		WorkerType: workerType,
		CreatedAt:  time.Now(),
		Enabled:    true,
	}

	ks.keys[key] = apiKey

	if err := ks.save(); err != nil {
		delete(ks.keys, key)
		return "", err
	}

	return key, nil
}

// Validate checks if an API key is valid and enabled
func (ks *APIKeyStore) Validate(key string) (*APIKey, bool) {
	ks.mu.RLock()
	defer ks.mu.RUnlock()

	apiKey, exists := ks.keys[key]
	if !exists || !apiKey.Enabled {
		return nil, false
	}

	return apiKey, true
}

// RecordUsage updates the last used timestamp and request count
func (ks *APIKeyStore) RecordUsage(key string) {
	ks.mu.Lock()
	defer ks.mu.Unlock()

	if apiKey, exists := ks.keys[key]; exists {
		apiKey.LastUsedAt = time.Now()
		apiKey.RequestCount++
		// Save async to avoid blocking requests
		go ks.save()
	}
}

// List returns all API keys (for admin UI)
func (ks *APIKeyStore) List() []*APIKey {
	ks.mu.RLock()
	defer ks.mu.RUnlock()

	list := make([]*APIKey, 0, len(ks.keys))
	for _, apiKey := range ks.keys {
		// Create a copy to avoid race conditions
		keyCopy := *apiKey
		list = append(list, &keyCopy)
	}
	return list
}

// Revoke disables an API key
func (ks *APIKeyStore) Revoke(key string) error {
	ks.mu.Lock()
	defer ks.mu.Unlock()

	apiKey, exists := ks.keys[key]
	if !exists {
		return fmt.Errorf("API key not found")
	}

	apiKey.Enabled = false
	return ks.save()
}

// Delete permanently removes an API key
func (ks *APIKeyStore) Delete(key string) error {
	ks.mu.Lock()
	defer ks.mu.Unlock()

	delete(ks.keys, key)
	return ks.save()
}

// save encrypts and writes keys to disk
func (ks *APIKeyStore) save() error {
	data, err := json.MarshalIndent(ks.keys, "", "  ")
	if err != nil {
		return err
	}

	// Encrypt the entire key store with server key
	encrypted, err := ks.crypto.EncryptWithServerKey(data)
	if err != nil {
		return err
	}

	return os.WriteFile(ks.file, encrypted, 0600)
}

// load decrypts and reads keys from disk
func (ks *APIKeyStore) load() error {
	data, err := os.ReadFile(ks.file)
	if err != nil {
		if os.IsNotExist(err) {
			return nil // File doesn't exist yet, that's okay
		}
		return err
	}

	// Decrypt with server key
	decrypted, err := ks.crypto.DecryptWithServerKey(data)
	if err != nil {
		return err
	}

	return json.Unmarshal(decrypted, &ks.keys)
}