Add support for bulk transfer of hosts for team (#761)

- Relevant datastore method.
- Expose via API endpoint /teams/:id/hosts.
This commit is contained in:
Zach Wasserman 2021-05-12 19:05:45 -07:00 committed by GitHub
parent 61c5eceb97
commit 1a29a408ce
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
9 changed files with 136 additions and 1 deletions

View file

@ -98,5 +98,6 @@ var TestFunctions = [...]func(*testing.T, kolide.Datastore){
testTeamUsers,
testUserTeams,
testUserCreateWithTeams,
testTeamAddHostsToTeam,
testSaveHostSoftware,
}

View file

@ -1,11 +1,15 @@
package datastore
import (
"fmt"
"testing"
"time"
"github.com/fleetdm/fleet/server/kolide"
"github.com/fleetdm/fleet/server/test"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gopkg.in/guregu/null.v3"
)
func testTeamGetSetDelete(t *testing.T, ds kolide.Datastore) {
@ -53,7 +57,6 @@ func testTeamUsers(t *testing.T, ds kolide.Datastore) {
require.NoError(t, err)
team2, err := ds.NewTeam(&kolide.Team{Name: "team2"})
require.NoError(t, err)
_ = team2
team1, err = ds.Team(team1.ID)
require.NoError(t, err)
@ -101,3 +104,50 @@ func testTeamUsers(t *testing.T, ds kolide.Datastore) {
assert.ElementsMatch(t, team2Users, team2.Users)
}
func testTeamAddHostsToTeam(t *testing.T, ds kolide.Datastore) {
team1, err := ds.NewTeam(&kolide.Team{Name: "team1"})
require.NoError(t, err)
team2, err := ds.NewTeam(&kolide.Team{Name: "team2"})
require.NoError(t, err)
for i := 0; i < 10; i++ {
test.NewHost(t, ds, fmt.Sprint(i), "", "key"+fmt.Sprint(i), "uuid"+fmt.Sprint(i), time.Now())
}
for i := 1; i <= 10; i++ {
host, err := ds.Host(uint(i))
require.NoError(t, err)
assert.Equal(t, null.Int{}, host.TeamID)
}
require.NoError(t, ds.AddHostsToTeam(&team1.ID, []uint{1, 2, 3}))
require.NoError(t, ds.AddHostsToTeam(&team2.ID, []uint{3, 4, 5}))
for i := 1; i <= 10; i++ {
host, err := ds.Host(uint(i))
require.NoError(t, err)
expectedID := null.Int{}
switch {
case i <= 2:
expectedID = null.IntFrom(int64(team1.ID))
case i <= 5:
expectedID = null.IntFrom(int64(team2.ID))
}
assert.Equal(t, expectedID, host.TeamID)
}
require.NoError(t, ds.AddHostsToTeam(nil, []uint{1, 2, 3, 4}))
require.NoError(t, ds.AddHostsToTeam(&team1.ID, []uint{5, 6, 7, 8, 9, 10}))
for i := 1; i <= 10; i++ {
host, err := ds.Host(uint(i))
require.NoError(t, err)
expectedID := null.Int{}
switch {
case i >= 5:
expectedID = null.IntFrom(int64(team1.ID))
}
assert.Equal(t, expectedID, host.TeamID)
}
}

View file

@ -165,3 +165,24 @@ func (d *Datastore) ListTeams(opt kolide.ListOptions) ([]*kolide.Team, error) {
return teams, nil
}
func (d *Datastore) AddHostsToTeam(teamID *uint, hostIDs []uint) error {
if len(hostIDs) == 0 {
return nil
}
sql := `
UPDATE hosts SET team_id = ?
WHERE id IN (?)
`
sql, args, err := sqlx.In(sql, teamID, hostIDs)
if err != nil {
return errors.Wrap(err, "sqlx.In AddHostsToTeam")
}
if _, err := d.db.Exec(sql, args...); err != nil {
return errors.Wrap(err, "exec AddHostsToTeam")
}
return nil
}

View file

@ -20,6 +20,9 @@ type TeamStore interface {
// ListTeams lists teams with the ordering and filters in the provided
// options.
ListTeams(opt ListOptions) ([]*Team, error)
// AddHostsToTeam adds hosts to an existing team, clearing their team
// settings if teamID is nil.
AddHostsToTeam(teamID *uint, hostIDs []uint) error
}
type TeamService interface {
@ -40,6 +43,9 @@ type TeamService interface {
ListTeams(ctx context.Context, opt ListOptions) ([]*Team, error)
// ListTeams lists users on the team with the provided list options.
ListTeamUsers(ctx context.Context, teamID uint, opt ListOptions) ([]*User, error)
// AddHostsToTeam adds hosts to an existing team, clearing their team
// settings if teamID is nil.
AddHostsToTeam(ctx context.Context, teamID *uint, hostIDs []uint) error
}
type TeamPayload struct {

View file

@ -20,6 +20,8 @@ type TeamByNameFunc func(name string) (*kolide.Team, error)
type ListTeamsFunc func(opt kolide.ListOptions) ([]*kolide.Team, error)
type AddHostsToTeamFunc func(teamID *uint, hostIDs []uint) error
type TeamStore struct {
NewTeamFunc NewTeamFunc
NewTeamFuncInvoked bool
@ -38,6 +40,9 @@ type TeamStore struct {
ListTeamsFunc ListTeamsFunc
ListTeamsFuncInvoked bool
AddHostsToTeamFunc AddHostsToTeamFunc
AddHostsToTeamFuncInvoked bool
}
func (s *TeamStore) NewTeam(team *kolide.Team) (*kolide.Team, error) {
@ -69,3 +74,8 @@ func (s *TeamStore) ListTeams(opt kolide.ListOptions) ([]*kolide.Team, error) {
s.ListTeamsFuncInvoked = true
return s.ListTeamsFunc(opt)
}
func (s *TeamStore) AddHostsToTeam(teamID *uint, hostIDs []uint) error {
s.AddHostsToTeamFuncInvoked = true
return s.AddHostsToTeamFunc(teamID, hostIDs)
}

View file

@ -193,3 +193,28 @@ func makeDeleteTeamUsersEndpoint(svc kolide.Service) endpoint.Endpoint {
return teamResponse{Team: team}, err
}
}
////////////////////////////////////////////////////////////////////////////////
// Add Hosts to Team
////////////////////////////////////////////////////////////////////////////////
type addHostsToTeamRequest struct {
TeamID uint // From request path
HostIDs []uint `json:"hosts"`
}
type addHostsToTeamResponse struct {
Err error `json:"error,omitempty"`
}
func makeAddHostsToTeamEndpoint(svc kolide.Service) endpoint.Endpoint {
return func(ctx context.Context, request interface{}) (interface{}, error) {
req := request.(addHostsToTeamRequest)
err := svc.AddHostsToTeam(ctx, &req.TeamID, req.HostIDs)
if err != nil {
return addHostsToTeamResponse{Err: err}, nil
}
return addHostsToTeamResponse{}, err
}
}

View file

@ -116,6 +116,7 @@ type KolideEndpoints struct {
ListTeamUsers endpoint.Endpoint
AddTeamUsers endpoint.Endpoint
DeleteTeamUsers endpoint.Endpoint
AddHostsToTeam endpoint.Endpoint
}
// MakeKolideServerEndpoints creates the Kolide API endpoints.
@ -225,6 +226,7 @@ func MakeKolideServerEndpoints(svc kolide.Service, jwtKey, urlPrefix string, lim
ListTeamUsers: authenticatedUser(jwtKey, svc, makeListTeamUsersEndpoint(svc)),
AddTeamUsers: authenticatedUser(jwtKey, svc, makeAddTeamUsersEndpoint(svc)),
DeleteTeamUsers: authenticatedUser(jwtKey, svc, makeDeleteTeamUsersEndpoint(svc)),
AddHostsToTeam: authenticatedUser(jwtKey, svc, makeAddHostsToTeamEndpoint(svc)),
// Authenticated status endpoints
StatusResultStore: authenticatedUser(jwtKey, svc, makeStatusResultStoreEndpoint(svc)),
@ -340,6 +342,7 @@ type kolideHandlers struct {
ListTeamUsers http.Handler
AddTeamUsers http.Handler
DeleteTeamUsers http.Handler
AddHostsToTeam http.Handler
}
func makeKolideKitHandlers(e KolideEndpoints, opts []kithttp.ServerOption) *kolideHandlers {
@ -442,6 +445,7 @@ func makeKolideKitHandlers(e KolideEndpoints, opts []kithttp.ServerOption) *koli
ListTeamUsers: newServer(e.ListTeamUsers, decodeListTeamUsersRequest),
AddTeamUsers: newServer(e.AddTeamUsers, decodeModifyTeamUsersRequest),
DeleteTeamUsers: newServer(e.DeleteTeamUsers, decodeModifyTeamUsersRequest),
AddHostsToTeam: newServer(e.AddHostsToTeam, decodeAddHostsToTeamRequest),
}
}
@ -659,6 +663,7 @@ func attachKolideAPIRoutes(r *mux.Router, h *kolideHandlers) {
r.Handle("/api/v1/fleet/teams/{id}/users", h.ListTeamUsers).Methods("GET").Name("team_users")
r.Handle("/api/v1/fleet/teams/{id}/users", h.AddTeamUsers).Methods("PATCH").Name("add_team_users")
r.Handle("/api/v1/fleet/teams/{id}/users", h.DeleteTeamUsers).Methods("DELETE").Name("delete_team_users")
r.Handle("/api/v1/fleet/teams/{id}/hosts", h.AddHostsToTeam).Methods("POST").Name("add_hosts_to_team")
r.Handle("/api/v1/osquery/enroll", h.EnrollAgent).Methods("POST").Name("enroll_agent")
r.Handle("/api/v1/osquery/config", h.GetClientConfig).Methods("POST").Name("get_client_config")

View file

@ -133,3 +133,7 @@ func (svc service) ListTeams(ctx context.Context, opt kolide.ListOptions) ([]*ko
func (svc service) DeleteTeam(ctx context.Context, tid uint) error {
return svc.ds.DeleteTeam(tid)
}
func (svc service) AddHostsToTeam(ctx context.Context, teamID *uint, hostIDs []uint) error {
return svc.ds.AddHostsToTeam(teamID, hostIDs)
}

View file

@ -81,3 +81,16 @@ func decodeModifyTeamUsersRequest(ctx context.Context, r *http.Request) (interfa
}
return req, nil
}
func decodeAddHostsToTeamRequest(ctx context.Context, r *http.Request) (interface{}, error) {
id, err := idFromRequest(r, "id")
if err != nil {
return nil, err
}
req := addHostsToTeamRequest{TeamID: id}
err = json.NewDecoder(r.Body).Decode(&req)
if err != nil {
return nil, err
}
return req, nil
}