tyto/backend/internal/server/registry.go

// Package server implements the central Tyto server for multi-device monitoring.
package server

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sync"
	"time"
)

// AgentStatus represents the current state of an agent.
type AgentStatus string

const (
	AgentStatusPending   AgentStatus = "pending"
	AgentStatusApproved  AgentStatus = "approved"
	AgentStatusConnected AgentStatus = "connected"
	AgentStatusOffline   AgentStatus = "offline"
	AgentStatusRevoked   AgentStatus = "revoked"
)

// AgentRecord stores metadata about a registered agent.
type AgentRecord struct {
	ID           string      `json:"id"`
	Name         string      `json:"name,omitempty"`
	Hostname     string      `json:"hostname"`
	OS           string      `json:"os"`
	Architecture string      `json:"architecture"`
	Version      string      `json:"version"`
	Capabilities []string    `json:"capabilities,omitempty"`
	Status       AgentStatus `json:"status"`
	CertSerial   string      `json:"certSerial,omitempty"`
	CertExpiry   time.Time   `json:"certExpiry,omitempty"`
	LastSeen     time.Time   `json:"lastSeen,omitempty"`
	RegisteredAt time.Time   `json:"registeredAt"`
	Tags         []string    `json:"tags,omitempty"`
}

// Registry manages agent registrations.
type Registry struct {
	mu       sync.RWMutex
	agents   map[string]*AgentRecord
	filePath string
}

// NewRegistry creates a new agent registry.
func NewRegistry(filePath string) *Registry {
	r := &Registry{
		agents:   make(map[string]*AgentRecord),
		filePath: filePath,
	}

	// Load existing registrations
	if filePath != "" {
		r.load()
	}

	return r
}

// Register adds or updates an agent registration.
func (r *Registry) Register(agent *AgentRecord) error {
	r.mu.Lock()
	defer r.mu.Unlock()

	existing, exists := r.agents[agent.ID]
	if exists {
		// Update existing record
		existing.Hostname = agent.Hostname
		existing.OS = agent.OS
		existing.Architecture = agent.Architecture
		existing.Version = agent.Version
		existing.Capabilities = agent.Capabilities
		existing.LastSeen = time.Now()

		// Don't change status if already approved/connected
		if existing.Status == AgentStatusRevoked {
			return fmt.Errorf("agent %s is revoked", agent.ID)
		}
	} else {
		// New registration
		agent.RegisteredAt = time.Now()
		agent.LastSeen = time.Now()
		if agent.Status == "" {
			agent.Status = AgentStatusPending
		}
		r.agents[agent.ID] = agent
	}

	return r.save()
}

// Get returns an agent record by ID.
func (r *Registry) Get(id string) (*AgentRecord, bool) {
	r.mu.RLock()
	defer r.mu.RUnlock()

	agent, exists := r.agents[id]
	if !exists {
		return nil, false
	}

	// Return a copy
	copy := *agent
	return &copy, true
}

// List returns all registered agents.
func (r *Registry) List() []*AgentRecord {
	r.mu.RLock()
	defer r.mu.RUnlock()

	result := make([]*AgentRecord, 0, len(r.agents))
	for _, agent := range r.agents {
		copy := *agent
		result = append(result, &copy)
	}
	return result
}

// ListPending returns agents awaiting approval.
func (r *Registry) ListPending() []*AgentRecord {
	r.mu.RLock()
	defer r.mu.RUnlock()

	var result []*AgentRecord
	for _, agent := range r.agents {
		if agent.Status == AgentStatusPending {
			copy := *agent
			result = append(result, &copy)
		}
	}
	return result
}

// Approve marks an agent as approved.
func (r *Registry) Approve(id string) error {
	r.mu.Lock()
	defer r.mu.Unlock()

	agent, exists := r.agents[id]
	if !exists {
		return fmt.Errorf("agent %s not found", id)
	}

	if agent.Status == AgentStatusRevoked {
		return fmt.Errorf("agent %s is revoked", id)
	}

	agent.Status = AgentStatusApproved
	return r.save()
}

// Reject removes a pending agent registration.
func (r *Registry) Reject(id string) error {
	r.mu.Lock()
	defer r.mu.Unlock()

	agent, exists := r.agents[id]
	if !exists {
		return fmt.Errorf("agent %s not found", id)
	}

	if agent.Status != AgentStatusPending {
		return fmt.Errorf("agent %s is not pending", id)
	}

	delete(r.agents, id)
	return r.save()
}

// Revoke marks an agent as revoked.
func (r *Registry) Revoke(id string) error {
	r.mu.Lock()
	defer r.mu.Unlock()

	agent, exists := r.agents[id]
	if !exists {
		return fmt.Errorf("agent %s not found", id)
	}

	agent.Status = AgentStatusRevoked
	return r.save()
}

// Remove deletes an agent registration.
func (r *Registry) Remove(id string) error {
	r.mu.Lock()
	defer r.mu.Unlock()

	if _, exists := r.agents[id]; !exists {
		return fmt.Errorf("agent %s not found", id)
	}

	delete(r.agents, id)
	return r.save()
}

// IsApproved checks if an agent is approved to connect.
func (r *Registry) IsApproved(id string) bool {
	r.mu.RLock()
	defer r.mu.RUnlock()

	agent, exists := r.agents[id]
	if !exists {
		return false
	}

	return agent.Status == AgentStatusApproved || agent.Status == AgentStatusConnected
}

// UpdateStatus updates an agent's connection status.
func (r *Registry) UpdateStatus(id string, status AgentStatus) error {
	r.mu.Lock()
	defer r.mu.Unlock()

	agent, exists := r.agents[id]
	if !exists {
		return fmt.Errorf("agent %s not found", id)
	}

	agent.Status = status
	agent.LastSeen = time.Now()
	return r.save()
}

// UpdateLastSeen updates the last seen timestamp.
func (r *Registry) UpdateLastSeen(id string) {
	r.mu.Lock()
	defer r.mu.Unlock()

	if agent, exists := r.agents[id]; exists {
		agent.LastSeen = time.Now()
	}
}

// load reads the registry from disk.
func (r *Registry) load() error {
	data, err := os.ReadFile(r.filePath)
	if err != nil {
		if os.IsNotExist(err) {
			return nil
		}
		return err
	}

	var agents []*AgentRecord
	if err := json.Unmarshal(data, &agents); err != nil {
		return err
	}

	for _, agent := range agents {
		r.agents[agent.ID] = agent
	}

	return nil
}

// save writes the registry to disk.
func (r *Registry) save() error {
	if r.filePath == "" {
		return nil
	}

	agents := make([]*AgentRecord, 0, len(r.agents))
	for _, agent := range r.agents {
		agents = append(agents, agent)
	}

	data, err := json.MarshalIndent(agents, "", "  ")
	if err != nil {
		return err
	}

	// Ensure directory exists
	dir := filepath.Dir(r.filePath)
	if err := os.MkdirAll(dir, 0755); err != nil {
		return err
	}

	return os.WriteFile(r.filePath, data, 0644)
}