pxmon/internal/cluster/store.go

package cluster

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"
)

const (
	envConfigPath    = "PXMON_CONFIG"
	envMasterKeyPath = "PXMON_MASTER_KEY"
	filePrefix       = "OBSCTL1:"
	masterKeyBytes   = 32
)

// Store persists encrypted cluster registry on disk.
type Store struct {
	path    string
	keyPath string
}

func NewStore(path string) (*Store, error) {
	if path == "" {
		var err error
		path, err = DefaultConfigPath()
		if err != nil {
			return nil, err
		}
	}

	keyPath, err := defaultMasterKeyPath(path)
	if err != nil {
		return nil, err
	}

	return &Store{path: path, keyPath: keyPath}, nil
}

func DefaultConfigPath() (string, error) {
	if p := os.Getenv(envConfigPath); p != "" {
		return p, nil
	}

	dir, err := os.UserConfigDir()
	if err != nil {
		return "", fmt.Errorf("resolve user config dir: %w", err)
	}

	return filepath.Join(dir, "pxmon", "clusters.enc"), nil
}

func defaultMasterKeyPath(configPath string) (string, error) {
	if p := os.Getenv(envMasterKeyPath); p != "" {
		return p, nil
	}
	if configPath == "" {
		return "", errors.New("empty config path")
	}
	return filepath.Join(filepath.Dir(configPath), "master.key"), nil
}

func (s *Store) Path() string {
	return s.path
}

func (s *Store) KeyPath() string {
	return s.keyPath
}

func (s *Store) LockerSessionPath() string {
	return filepath.Join(filepath.Dir(s.path), "locker.session")
}

func (s *Store) LockerAuditPath() string {
	return filepath.Join(filepath.Dir(s.path), "locker.audit.log")
}

func (s *Store) Load() (Registry, error) {
	f, err := os.Open(s.path)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			return newRegistry(), nil
		}
		return Registry{}, fmt.Errorf("open registry file: %w", err)
	}
	defer f.Close()

	data, err := io.ReadAll(f)
	if err != nil {
		return Registry{}, fmt.Errorf("read registry file: %w", err)
	}
	if len(strings.TrimSpace(string(data))) == 0 {
		return newRegistry(), nil
	}

	payload, err := s.decodePayload(data)
	if err != nil {
		return Registry{}, err
	}

	var reg Registry
	if err := json.Unmarshal(payload, &reg); err != nil {
		return Registry{}, fmt.Errorf("decode registry JSON: %w", err)
	}

	if reg.Version == 0 {
		reg.Version = currentVersion
	}
	if reg.Clusters == nil {
		reg.Clusters = []Cluster{}
	}
	for i := range reg.Clusters {
		reg.Clusters[i].Alerts = ensureAlertPolicy(reg.Clusters[i].Alerts)
		reg.Clusters[i].VMAlerts = ensureVMAlertPolicy(reg.Clusters[i].VMAlerts)
		reg.Clusters[i].AlertRouting = ensureAlertRoutingPolicy(reg.Clusters[i].AlertRouting)
		reg.Clusters[i].RunbookTrigger = ensureRunbookTrigger(reg.Clusters[i].RunbookTrigger)
		reg.Clusters[i].Drift = normalizeDriftControl(reg.Clusters[i].Drift)
		reg.Clusters[i].Tags = normalizeTagList(reg.Clusters[i].Tags)
		reg.Clusters[i].KVMTags = normalizeVMTagMap(reg.Clusters[i].KVMTags)
		reg.Clusters[i].Transport = normalizeTransport(reg.Clusters[i].Transport)
	}
	reg.Telegram = normalizeTelegram(reg.Telegram)
	reg.Locker = normalizeLocker(reg.Locker)
	reg.Backups = normalizeBackupConfig(reg.Backups)

	return reg, nil
}

func (s *Store) decodePayload(data []byte) ([]byte, error) {
	text := strings.TrimSpace(string(data))

	if strings.HasPrefix(text, "{") {
		// Backward compatibility with legacy unencrypted format.
		return []byte(text), nil
	}

	if !strings.HasPrefix(text, filePrefix) {
		return nil, errors.New("unsupported registry format")
	}

	blob := strings.TrimPrefix(text, filePrefix)
	raw, err := base64.StdEncoding.DecodeString(blob)
	if err != nil {
		return nil, fmt.Errorf("decode encrypted payload: %w", err)
	}

	key, err := s.loadOrCreateMasterKey()
	if err != nil {
		return nil, err
	}

	payload, err := decrypt(raw, key)
	if err != nil {
		return nil, fmt.Errorf("decrypt registry: %w", err)
	}
	return payload, nil
}

func (s *Store) Save(reg Registry) error {
	reg.Version = currentVersion
	if reg.Clusters == nil {
		reg.Clusters = []Cluster{}
	}
	for i := range reg.Clusters {
		reg.Clusters[i].Alerts = ensureAlertPolicy(reg.Clusters[i].Alerts)
		reg.Clusters[i].VMAlerts = ensureVMAlertPolicy(reg.Clusters[i].VMAlerts)
		reg.Clusters[i].AlertRouting = ensureAlertRoutingPolicy(reg.Clusters[i].AlertRouting)
		reg.Clusters[i].RunbookTrigger = ensureRunbookTrigger(reg.Clusters[i].RunbookTrigger)
		reg.Clusters[i].Drift = normalizeDriftControl(reg.Clusters[i].Drift)
		reg.Clusters[i].Tags = normalizeTagList(reg.Clusters[i].Tags)
		reg.Clusters[i].KVMTags = normalizeVMTagMap(reg.Clusters[i].KVMTags)
		reg.Clusters[i].Transport = normalizeTransport(reg.Clusters[i].Transport)
	}
	reg.Telegram = normalizeTelegram(reg.Telegram)
	reg.Locker = normalizeLocker(reg.Locker)
	reg.Backups = normalizeBackupConfig(reg.Backups)

	if err := os.MkdirAll(filepath.Dir(s.path), 0o700); err != nil {
		return fmt.Errorf("create config dir: %w", err)
	}

	payload, err := json.MarshalIndent(reg, "", "  ")
	if err != nil {
		return fmt.Errorf("encode registry JSON: %w", err)
	}

	key, err := s.loadOrCreateMasterKey()
	if err != nil {
		return err
	}

	encrypted, err := encrypt(payload, key)
	if err != nil {
		return fmt.Errorf("encrypt registry: %w", err)
	}

	content := filePrefix + base64.StdEncoding.EncodeToString(encrypted) + "\n"

	tmp := s.path + ".tmp"
	if err := os.WriteFile(tmp, []byte(content), 0o600); err != nil {
		return fmt.Errorf("write temp registry file: %w", err)
	}

	if err := os.Rename(tmp, s.path); err != nil {
		return fmt.Errorf("replace registry file: %w", err)
	}

	return nil
}

func (s *Store) loadOrCreateMasterKey() ([]byte, error) {
	if err := os.MkdirAll(filepath.Dir(s.keyPath), 0o700); err != nil {
		return nil, fmt.Errorf("create key dir: %w", err)
	}

	key, err := os.ReadFile(s.keyPath)
	if err == nil {
		if len(key) != masterKeyBytes {
			return nil, fmt.Errorf("invalid master key length: got %d", len(key))
		}
		return key, nil
	}
	if !errors.Is(err, os.ErrNotExist) {
		return nil, fmt.Errorf("read master key: %w", err)
	}

	key = make([]byte, masterKeyBytes)
	if _, err := rand.Read(key); err != nil {
		return nil, fmt.Errorf("generate master key: %w", err)
	}

	tmp := s.keyPath + ".tmp"
	if err := os.WriteFile(tmp, key, 0o600); err != nil {
		return nil, fmt.Errorf("write temp master key: %w", err)
	}
	if err := os.Rename(tmp, s.keyPath); err != nil {
		return nil, fmt.Errorf("replace master key: %w", err)
	}

	return key, nil
}

func encrypt(payload, key []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}

	nonce := make([]byte, gcm.NonceSize())
	if _, err := rand.Read(nonce); err != nil {
		return nil, err
	}

	sealed := gcm.Seal(nil, nonce, payload, nil)
	out := make([]byte, 0, len(nonce)+len(sealed))
	out = append(out, nonce...)
	out = append(out, sealed...)
	return out, nil
}

func decrypt(raw, key []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}

	nonceSize := gcm.NonceSize()
	if len(raw) <= nonceSize {
		return nil, errors.New("ciphertext too short")
	}

	nonce := raw[:nonceSize]
	ciphertext := raw[nonceSize:]

	payload, err := gcm.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		return nil, err
	}
	return payload, nil
}

type lockerSessionState struct {
	ExpiresAtUnix int64  `json:"expires_at_unix"`
	SigHex        string `json:"sig_hex"`
}

func (s *Store) SaveLockerSession(passwordHash string, ttl time.Duration) error {
	passwordHash = strings.TrimSpace(passwordHash)
	if passwordHash == "" {
		return errors.New("empty locker password hash")
	}
	if ttl <= 0 {
		ttl = 6 * time.Hour
	}

	key, err := s.loadOrCreateMasterKey()
	if err != nil {
		return err
	}

	expires := time.Now().UTC().Add(ttl).Unix()
	payload := strconv.FormatInt(expires, 10) + "|" + passwordHash
	mac := hmac.New(sha256.New, key)
	_, _ = mac.Write([]byte(payload))
	sig := hex.EncodeToString(mac.Sum(nil))

	state := lockerSessionState{
		ExpiresAtUnix: expires,
		SigHex:        sig,
	}
	data, err := json.Marshal(state)
	if err != nil {
		return err
	}
	data = append(data, '\n')

	path := s.LockerSessionPath()
	if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
		return err
	}
	tmp := path + ".tmp"
	if err := os.WriteFile(tmp, data, 0o600); err != nil {
		return err
	}
	return os.Rename(tmp, path)
}

func (s *Store) ValidateLockerSession(passwordHash string) (bool, time.Time, error) {
	passwordHash = strings.TrimSpace(passwordHash)
	if passwordHash == "" {
		return false, time.Time{}, nil
	}

	path := s.LockerSessionPath()
	raw, err := os.ReadFile(path)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			return false, time.Time{}, nil
		}
		return false, time.Time{}, err
	}

	var state lockerSessionState
	if err := json.Unmarshal(raw, &state); err != nil {
		return false, time.Time{}, nil
	}
	if state.ExpiresAtUnix <= 0 || strings.TrimSpace(state.SigHex) == "" {
		return false, time.Time{}, nil
	}
	expiresAt := time.Unix(state.ExpiresAtUnix, 0).UTC()
	if time.Now().UTC().After(expiresAt) {
		return false, expiresAt, nil
	}

	key, err := s.loadOrCreateMasterKey()
	if err != nil {
		return false, time.Time{}, err
	}

	payload := strconv.FormatInt(state.ExpiresAtUnix, 10) + "|" + passwordHash
	mac := hmac.New(sha256.New, key)
	_, _ = mac.Write([]byte(payload))
	expected := mac.Sum(nil)
	got, err := hex.DecodeString(strings.TrimSpace(state.SigHex))
	if err != nil {
		return false, expiresAt, nil
	}
	if !hmac.Equal(expected, got) {
		return false, expiresAt, nil
	}
	return true, expiresAt, nil
}

func (s *Store) ClearLockerSession() error {
	path := s.LockerSessionPath()
	if err := os.Remove(path); err != nil && !errors.Is(err, os.ErrNotExist) {
		return err
	}
	return nil
}

func (s *Store) AppendLockerAudit(event, detail string) error {
	path := s.LockerAuditPath()
	if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
		return err
	}
	f, err := os.OpenFile(path, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o600)
	if err != nil {
		return err
	}
	defer f.Close()

	ts := time.Now().UTC().Format(time.RFC3339)
	event = strings.TrimSpace(event)
	detail = strings.TrimSpace(detail)
	if event == "" {
		event = "event"
	}
	if detail == "" {
		detail = "-"
	}
	_, err = fmt.Fprintf(f, "%s event=%s detail=%s\n", ts, event, strconv.Quote(detail))
	return err
}