feat: add enrollment token validation for agent enrollment

internal/models/agent.go

@@ -169,48 +169,99 @@ func (m *Model) SaveAgentInfo(data *nats.AgentReport, servers string, autoAdmitA
 		// This is a new agent, we must set the nickname to the agent's hostname initially
 		query.SetNickname(data.Hostname)
 
-		// Set the associated site
-		if data.Site == "" {
-			s, err := m.GetDefaultSite()
+		// Check if agent has an enrollment token
+		usedEnrollmentToken := false
+		if data.EnrollmentToken != "" {
+			// Try to validate and use the enrollment token
+			tokenTenantID, tokenSiteID, err := m.ValidateEnrollmentToken(data.EnrollmentToken)
 			if err != nil {
-				log.Printf("[ERROR]: could not get default site, reason: %v", err)
-				return err
-			}
-			query.AddSite(s)
-		} else {
-			siteID, err := strconv.Atoi(data.Site)
-			if err != nil {
-				log.Printf("[ERROR]: could not convert site ID to int, reason: %v", err)
-				return err
+				log.Printf("[WARN]: could not validate enrollment token, reason: %v", err)
+			} else if tokenTenantID > 0 {
+				// Token is valid, use its tenant/site
+				if tokenSiteID > 0 {
+					// Token has a specific site
+					valid, err := m.ValidateTenantAndSite(tokenTenantID, tokenSiteID)
+					if err != nil {
+						log.Printf("[ERROR]: could not validate token's tenant and site, reason: %v", err)
+						return err
+					}
+					if valid {
+						query.AddSiteIDs(tokenSiteID)
+						usedEnrollmentToken = true
+						log.Printf("[INFO]: agent enrolled using token to tenant %d, site %d", tokenTenantID, tokenSiteID)
+					}
+				} else {
+					// Token only has tenant, use default site of that tenant
+					defaultSite, err := m.GetDefaultSiteForTenant(tokenTenantID)
+					if err != nil {
+						log.Printf("[WARN]: could not get default site for tenant %d, reason: %v", tokenTenantID, err)
+					} else if defaultSite != nil {
+						query.AddSite(defaultSite)
+						usedEnrollmentToken = true
+						log.Printf("[INFO]: agent enrolled using token to tenant %d, default site %d", tokenTenantID, defaultSite.ID)
+					}
+				}
 			}
+		}
 
-			tenantID, err := strconv.Atoi(data.Tenant)
-			if err != nil {
-				log.Printf("[ERROR]: could not convert tenant ID to int, reason: %v", err)
-				return err
-			}
+		// If enrollment token was not used, fall back to original logic
+		if !usedEnrollmentToken {
+			// Set the associated site
+			if data.Site == "" {
+				s, err := m.GetDefaultSite()
+				if err != nil {
+					log.Printf("[ERROR]: could not get default site, reason: %v", err)
+					return err
+				}
+				query.AddSite(s)
+			} else {
+				siteID, err := strconv.Atoi(data.Site)
+				if err != nil {
+					log.Printf("[ERROR]: could not convert site ID to int, reason: %v", err)
+					return err
+				}
 
-			// Check if tenantID is right and associated with the site
-			valid, err := m.ValidateTenantAndSite(tenantID, siteID)
-			if err != nil {
-				log.Printf("[ERROR]: could not check if tenant and site are valid, reason: %v", err)
-				return err
-			}
+				tenantID, err := strconv.Atoi(data.Tenant)
+				if err != nil {
+					log.Printf("[ERROR]: could not convert tenant ID to int, reason: %v", err)
+					return err
+				}
 
-			if valid {
-				query.AddSiteIDs(siteID)
-			} else {
-				log.Printf("[ERROR]: tenant and site are not valid")
-				return errors.New("tenant and site are not valid")
+				// Check if tenantID is right and associated with the site
+				valid, err := m.ValidateTenantAndSite(tenantID, siteID)
+				if err != nil {
+					log.Printf("[ERROR]: could not check if tenant and site are valid, reason: %v", err)
+					return err
+				}
+
+				if valid {
+					query.AddSiteIDs(siteID)
+				} else {
+					log.Printf("[ERROR]: tenant and site are not valid")
+					return errors.New("tenant and site are not valid")
+				}
 			}
 		}
 
-		return query.
+		err = query.
 			SetFirstContact(time.Now()).
 			SetLastContact(time.Now()).
 			OnConflictColumns(agent.FieldID).
 			UpdateNewValues().
 			Exec(context.Background())
+
+		if err != nil {
+			return err
+		}
+
+		// If enrollment token was used successfully, increment its usage counter
+		if usedEnrollmentToken {
+			if err := m.UseEnrollmentToken(data.EnrollmentToken); err != nil {
+				log.Printf("[WARN]: could not increment enrollment token usage, reason: %v", err)
+			}
+		}
+
+		return nil
 	}
 }
 
@@ -846,6 +897,10 @@ func (m *Model) GetDefaultSite() (*ent.Site, error) {
 	return m.Client.Site.Query().Where(site.IsDefault(true), site.HasTenantWith(tenant.ID(t.ID))).Only(context.Background())
 }
 
+func (m *Model) GetDefaultSiteForTenant(tenantID int) (*ent.Site, error) {
+	return m.Client.Site.Query().Where(site.IsDefault(true), site.HasTenantWith(tenant.ID(tenantID))).Only(context.Background())
+}
+
 func (m *Model) ValidateTenantAndSite(tenantID, siteID int) (bool, error) {
 	return m.Client.Site.Query().Where(site.ID(siteID), site.HasTenantWith(tenant.ID(tenantID))).Exist(context.Background())
 }

internal/models/enrollmenttoken.go

@@ -0,0 +1,100 @@
+package models
+
+import (
+	"context"
+	"time"
+
+	"github.com/open-uem/ent"
+	"github.com/open-uem/ent/enrollmenttoken"
+)
+
+// ValidateEnrollmentToken validates an enrollment token and returns the associated tenant/site IDs
+// Returns (tenantID, siteID, error) where siteID may be 0 if no site is associated
+func (m *Model) ValidateEnrollmentToken(token string) (int, int, error) {
+	if token == "" {
+		return 0, 0, nil
+	}
+
+	ctx := context.Background()
+
+	et, err := m.Client.EnrollmentToken.Query().
+		Where(enrollmenttoken.Token(token)).
+		WithTenant().
+		WithSite().
+		Only(ctx)
+	if err != nil {
+		if ent.IsNotFound(err) {
+			return 0, 0, nil // Token not found, not an error - just return 0,0
+		}
+		return 0, 0, err
+	}
+
+	// Check if token is active
+	if !et.Active {
+		return 0, 0, nil
+	}
+
+	// Check if token is expired
+	if et.ExpiresAt != nil && et.ExpiresAt.Before(time.Now()) {
+		return 0, 0, nil
+	}
+
+	// Check if max uses reached (0 = unlimited)
+	if et.MaxUses > 0 && et.CurrentUses >= et.MaxUses {
+		return 0, 0, nil
+	}
+
+	// Get tenant ID
+	tenantID := 0
+	if et.Edges.Tenant != nil {
+		tenantID = et.Edges.Tenant.ID
+	}
+
+	// Get site ID (may be nil)
+	siteID := 0
+	if et.Edges.Site != nil {
+		siteID = et.Edges.Site.ID
+	}
+
+	return tenantID, siteID, nil
+}
+
+// UseEnrollmentToken increments the current_uses counter for an enrollment token
+func (m *Model) UseEnrollmentToken(token string) error {
+	if token == "" {
+		return nil
+	}
+
+	ctx := context.Background()
+
+	_, err := m.Client.EnrollmentToken.Update().
+		Where(enrollmenttoken.Token(token)).
+		AddCurrentUses(1).
+		Save(ctx)
+
+	return err
+}
+
+// GetEnrollmentTokenInfo returns tenant and site IDs for a valid token, or falls back to provided IDs
+func (m *Model) GetEnrollmentTokenInfo(token string, fallbackTenantID, fallbackSiteID int) (int, int, bool, error) {
+	if token == "" {
+		return fallbackTenantID, fallbackSiteID, false, nil
+	}
+
+	tenantID, siteID, err := m.ValidateEnrollmentToken(token)
+	if err != nil {
+		return fallbackTenantID, fallbackSiteID, false, err
+	}
+
+	// If token validation returned 0,0, use fallback
+	if tenantID == 0 {
+		return fallbackTenantID, fallbackSiteID, false, nil
+	}
+
+	// If site is not set in token, use fallback site if tenant matches
+	if siteID == 0 {
+		return tenantID, fallbackSiteID, true, nil
+	}
+
+	return tenantID, siteID, true, nil
+}