oikos/internal/checkdefaults/defaults.go

package checkdefaults

import (
	"context"
	"encoding/json"
	"fmt"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
)

type CheckDef struct {
	Kind       string
	Script     string
	Host       string
	User       string
	Port       int
	Thresholds map[string]any
	Extra      map[string]any
}

func ResolveHost(attrs map[string]any) string {
	if ip, ok := attrs["lan_ip"].(string); ok && ip != "" {
		return ip
	}
	if mesh, ok := attrs["mesh"].(map[string]any); ok {
		if nb, ok := mesh["netbird"].(map[string]any); ok {
			if ip, ok := nb["ip"].(string); ok && ip != "" {
				return ip
			}
		}
	}
	if ip, ok := attrs["mesh_ip"].(string); ok && ip != "" {
		return ip
	}
	return ""
}

func resolveSSHUser(attrs map[string]any) string {
	if ssh, ok := attrs["ssh"].(map[string]any); ok {
		if u, ok := ssh["user"].(string); ok && u != "" {
			return u
		}
	}
	return "root"
}

func resolveSSHPort(attrs map[string]any) int {
	if ssh, ok := attrs["ssh"].(map[string]any); ok {
		switch p := ssh["port"].(type) {
		case float64:
			return int(p)
		case int:
			return p
		}
	}
	return 22
}

func ForEntityType(entityType string, attrs map[string]any) []CheckDef {
	host := ResolveHost(attrs)
	user := resolveSSHUser(attrs)
	port := resolveSSHPort(attrs)

	ssh := func(script string) CheckDef {
		return CheckDef{Kind: "ssh-script", Script: script, Host: host, User: user, Port: port}
	}

	switch entityType {
	case "proxmox-host", "standalone-server":
		if host == "" {
			return nil
		}
		return []CheckDef{
			{Kind: "ping", Host: host},
			ssh("cpu_check.sh"),
			ssh("memory_check.sh"),
			ssh("load_check.sh"),
			ssh("disk_usage_check.sh"),
			ssh("updates_check.sh"),
		}
	case "workstation":
		if host == "" {
			return nil
		}
		return []CheckDef{
			{Kind: "ping", Host: host},
			ssh("cpu_check.sh"),
			ssh("memory_check.sh"),
			ssh("load_check.sh"),
		}
	case "lxc":
		if host == "" {
			return nil
		}
		return []CheckDef{
			ssh("cpu_check.sh"),
			ssh("memory_check.sh"),
			ssh("load_check.sh"),
			ssh("disk_usage_check.sh"),
		}
	case "vm":
		if host == "" {
			return nil
		}
		return []CheckDef{
			{Kind: "ping", Host: host},
		}
	case "service":
		if host == "" {
			return nil
		}
		n, _ := attrs["name"].(string)
		if n == "" {
			return nil
		}
		return []CheckDef{
			{Kind: "ssh-script", Script: "process_check.sh", Host: host, User: user, Port: port,
				Extra: map[string]any{"args": n}},
		}
	}
	return nil
}

func ShortSlug(slug string) string {
	const n = 8
	if len(slug) > n {
		return slug[len(slug)-n:]
	}
	return slug
}

func DefaultInterval(kind string) int32 {
	switch kind {
	case "ping":
		return 30
	case "ssh-script":
		return 60
	default:
		return 300
	}
}

func Ensure(ctx context.Context, tx pgx.Tx, entityID uuid.UUID, slug, entityType string, attrsJSON []byte) {
	_, _ = tx.Exec(ctx,
		`INSERT INTO entity_status (entity_id, health, updated_at)
		 VALUES ($1, 'unknown', now())
		 ON CONFLICT (entity_id) DO NOTHING`,
		entityID)

	var attrs map[string]any
	if len(attrsJSON) > 0 {
		json.Unmarshal(attrsJSON, &attrs)
	}
	if attrs == nil {
		attrs = map[string]any{}
	}

	defs := ForEntityType(entityType, attrs)
	if len(defs) == 0 {
		return
	}

	for i, def := range defs {
		checkID, err := uuid.NewV7()
		if err != nil {
			checkID = uuid.New()
		}
		checkSlug := fmt.Sprintf("check:%s:%s:%d", def.Kind, ShortSlug(slug), i)

		_, _ = tx.Exec(ctx,
			`INSERT INTO entities (id, slug, type, name, state, attributes, version, created_at, updated_at)
			 VALUES ($1, $2, 'check', $2, 'active', '{}', 1, now(), now())
			 ON CONFLICT (slug) DO NOTHING`,
			checkID, checkSlug)

		configMap := map[string]any{}
		if def.Script != "" {
			configMap["script"] = def.Script
		}
		if def.Host != "" {
			configMap["host"] = def.Host
		}
		if def.User != "" && def.User != "root" {
			configMap["user"] = def.User
		}
		if def.Port != 0 && def.Port != 22 {
			configMap["port"] = def.Port
		}
		if def.Thresholds != nil {
			configMap["thresholds"] = def.Thresholds
		}
		for k, v := range def.Extra {
			configMap[k] = v
		}
		configJSON, _ := json.Marshal(configMap)

		_, _ = tx.Exec(ctx,
			`INSERT INTO check_defs (entity_id, kind, config, interval_s, timeout_s, enabled)
			 VALUES ($1, $2, $3, $4, 30, true)
			 ON CONFLICT (entity_id) DO NOTHING`,
			checkID, def.Kind, configJSON, DefaultInterval(def.Kind))
	}
}