diff --git a/internal/models/agent.go b/internal/models/agent.go index 0b06c1f..0bdd74c 100644 --- a/internal/models/agent.go +++ b/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()) } diff --git a/internal/models/enrollmenttoken.go b/internal/models/enrollmenttoken.go new file mode 100644 index 0000000..3122c42 --- /dev/null +++ b/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 +}