Netpulse_SasS/server/internal/store/enrollment.go

package store

import (
	"context"
	"crypto/sha256"
	"errors"
	"fmt"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/netpulse/netpulse/server/internal/auth"
)

// ErrEnrollmentInvalid — токен не існує, згорів або вже використаний.
//
// Одна помилка на всі три випадки навмисно: розрізняти їх у відповіді
// означало б підказувати тому, хто підбирає токени, наскільки він
// близько.
var ErrEnrollmentInvalid = errors.New("запрошення недійсне")

// EnrollmentTTL — скільки живе запрошення.
//
// Доба: людина копіює команду, йде до стійки й ставить агента. Довший
// строк перетворює запрошення на постійний ключ, який лежить у чиємусь
// чаті; коротший не переживає навіть обіду.
const EnrollmentTTL = 24 * time.Hour

// Enrollment — запрошення на реєстрацію зонда.
type Enrollment struct {
	ID        string     `json:"id"`
	Prefix    string     `json:"prefix"`
	NameHint  string     `json:"name_hint,omitempty"`
	Modules   []string   `json:"modules"`
	ExpiresAt time.Time  `json:"expires_at"`
	UsedAt    *time.Time `json:"used_at,omitempty"`
	UsedBy    string     `json:"used_by,omitempty"`
	AgentName string     `json:"agent_name,omitempty"`
	CreatedAt time.Time  `json:"created_at"`
	// Token заповнюється рівно один раз — у відповіді на створення.
	Token string `json:"token,omitempty"`
}

// CreateEnrollment видає одноразовий токен.
func (s *Store) CreateEnrollment(ctx context.Context, tenantID, userID, nameHint string,
	modules []string, ttl time.Duration) (*Enrollment, error) {

	token, err := auth.NewToken("np_enr_")
	if err != nil {
		return nil, err
	}
	sum := sha256.Sum256([]byte(token))
	if ttl <= 0 {
		ttl = EnrollmentTTL
	}
	if len(modules) == 0 {
		modules = []string{"icmp"}
	}

	e := &Enrollment{
		Token:     token,
		Prefix:    token[:min(len(token), 14)],
		NameHint:  nameHint,
		Modules:   modules,
		ExpiresAt: time.Now().Add(ttl),
	}

	err = s.InTenantTx(ctx, tenantID, func(tx pgx.Tx) error {
		return tx.QueryRow(ctx, `
			INSERT INTO core.agent_enrollments
				(tenant_id, token_hash, prefix, name_hint, modules, expires_at, created_by)
			VALUES ($1, $2, $3, NULLIF($4,''), $5::core.slug[], $6, $7)
			RETURNING id::text, created_at
		`, tenantID, sum[:], e.Prefix, nameHint, modules, e.ExpiresAt, nullUUID(userID)).
			Scan(&e.ID, &e.CreatedAt)
	})
	if err != nil {
		return nil, err
	}
	return e, nil
}

// ListEnrollments — запрошення тенанта.
func (s *Store) ListEnrollments(ctx context.Context, tenantID string) ([]Enrollment, error) {
	out := []Enrollment{}
	err := s.InTenantTx(ctx, tenantID, func(tx pgx.Tx) error {
		rows, err := tx.Query(ctx, `
			SELECT e.id::text, e.prefix, COALESCE(e.name_hint,''), e.modules::text[],
			       e.expires_at, e.used_at, COALESCE(e.used_by::text,''),
			       COALESCE(a.name,''), e.created_at
			FROM core.agent_enrollments e
			LEFT JOIN core.agents a ON a.id = e.used_by
			WHERE e.tenant_id = $1
			ORDER BY e.created_at DESC
			LIMIT 100
		`, tenantID)
		if err != nil {
			return err
		}
		defer rows.Close()
		for rows.Next() {
			var e Enrollment
			if err := rows.Scan(&e.ID, &e.Prefix, &e.NameHint, &e.Modules,
				&e.ExpiresAt, &e.UsedAt, &e.UsedBy, &e.AgentName, &e.CreatedAt); err != nil {
				return err
			}
			out = append(out, e)
		}
		return rows.Err()
	})
	return out, err
}

// DeleteEnrollment відкликає невикористане запрошення.
func (s *Store) DeleteEnrollment(ctx context.Context, tenantID, id string) error {
	return s.InTenantTx(ctx, tenantID, func(tx pgx.Tx) error {
		ct, err := tx.Exec(ctx,
			`DELETE FROM core.agent_enrollments WHERE id = $1 AND tenant_id = $2`, id, tenantID)
		if err != nil {
			return err
		}
		if ct.RowsAffected() == 0 {
			return ErrNotFound
		}
		return nil
	})
}

// EnrolledAgent — те, що зонд отримує після реєстрації.
type EnrolledAgent struct {
	ID    string
	Name  string
	Token string
}

// RedeemEnrollment обмінює одноразовий токен на постійний зонд.
//
// Уся операція в одній транзакції під `FOR UPDATE`: два агенти, що
// стартували з однієї команди встановлення (скопійований скрипт,
// повторний запуск), інакше створили б два зонди з одного запрошення.
//
// Токен зонда повертається один раз і в базі лежить лише хешем — так
// само, як refresh-токен сесії.
func (s *Store) RedeemEnrollment(ctx context.Context, token, hostname, requestedName,
	version, osName, arch string) (*EnrolledAgent, error) {

	sum := sha256.Sum256([]byte(token))

	var (
		enrollID, tenantID string
		nameHint           string
		modules            []string
		siteID             *string
		createdBy          *string
		expires            time.Time
		usedAt             *time.Time
	)

	// Без InTenantTx: тенант тут ще невідомий — саме його ми й
	// з'ясовуємо з токена. RLS натомість не потрібна: вибірка йде за
	// унікальним хешем, який знає лише той, кому його видали.
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return nil, err
	}
	defer func() { _ = tx.Rollback(ctx) }()

	err = tx.QueryRow(ctx, `
		SELECT id::text, tenant_id::text, COALESCE(name_hint,''), modules::text[],
		       site_id::text, created_by::text, expires_at, used_at
		FROM core.agent_enrollments
		WHERE token_hash = $1
		FOR UPDATE
	`, sum[:]).Scan(&enrollID, &tenantID, &nameHint, &modules, &siteID,
		&createdBy, &expires, &usedAt)

	if errors.Is(err, pgx.ErrNoRows) {
		return nil, ErrEnrollmentInvalid
	}
	if err != nil {
		return nil, err
	}
	if usedAt != nil || time.Now().After(expires) {
		return nil, ErrEnrollmentInvalid
	}

	name := firstNonEmpty(requestedName, nameHint, hostname, "зонд")
	name, err = uniqueAgentName(ctx, tx, tenantID, name)
	if err != nil {
		return nil, err
	}

	agentToken, err := auth.NewToken("np_agt_")
	if err != nil {
		return nil, err
	}
	agentSum := sha256.Sum256([]byte(agentToken))

	var agentID string
	if err := tx.QueryRow(ctx, `
		INSERT INTO core.agents
			(tenant_id, site_id, name, token_hash, version, os, arch, hostname,
			 status, enabled_modules, enrolled_at, enrolled_by)
		VALUES ($1, $2, $3, $4, NULLIF($5,''), NULLIF($6,''), NULLIF($7,''), NULLIF($8,''),
		        'pending', $9::core.slug[], now(), $10)
		RETURNING id::text
	`, tenantID, siteID, name, agentSum[:], version, osName, arch, hostname,
		modules, createdBy).Scan(&agentID); err != nil {
		return nil, err
	}

	if _, err := tx.Exec(ctx, `
		UPDATE core.agent_enrollments SET used_at = now(), used_by = $2 WHERE id = $1
	`, enrollID, agentID); err != nil {
		return nil, err
	}

	if err := tx.Commit(ctx); err != nil {
		return nil, err
	}
	return &EnrolledAgent{ID: agentID, Name: name, Token: agentToken}, nil
}

// uniqueAgentName підбирає вільне ім'я.
//
// Суфікс числом, а не відмова: людина, яка ставить п'ятий агент на
// однакових машинах, не має вигадувати імена — вона хоче, щоб він
// просто запрацював.
func uniqueAgentName(ctx context.Context, tx pgx.Tx, tenantID, base string) (string, error) {
	name := base
	for i := 2; i < 1000; i++ {
		var taken bool
		if err := tx.QueryRow(ctx, `
			SELECT EXISTS (SELECT 1 FROM core.agents WHERE tenant_id = $1 AND name = $2)
		`, tenantID, name).Scan(&taken); err != nil {
			return "", err
		}
		if !taken {
			return name, nil
		}
		name = fmt.Sprintf("%s-%d", base, i)
	}
	return "", errors.New("не вдалося підібрати ім'я зонда")
}

func firstNonEmpty(vals ...string) string {
	for _, v := range vals {
		if v != "" {
			return v
		}
	}
	return ""
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}