diff --git a/server/datastore/datastore.go b/server/datastore/datastore.go index 171a02c85e..e91c828f5b 100644 --- a/server/datastore/datastore.go +++ b/server/datastore/datastore.go @@ -98,5 +98,6 @@ var TestFunctions = [...]func(*testing.T, kolide.Datastore){ testTeamUsers, testUserTeams, testUserCreateWithTeams, + testTeamAddHostsToTeam, testSaveHostSoftware, } diff --git a/server/datastore/datastore_teams.go b/server/datastore/datastore_teams.go index 91c2d6f7aa..3b211c2a58 100644 --- a/server/datastore/datastore_teams.go +++ b/server/datastore/datastore_teams.go @@ -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) + } +} diff --git a/server/datastore/mysql/teams.go b/server/datastore/mysql/teams.go index dd5134f8b7..88fd3a73e8 100644 --- a/server/datastore/mysql/teams.go +++ b/server/datastore/mysql/teams.go @@ -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 +} diff --git a/server/kolide/teams.go b/server/kolide/teams.go index c7ff55187c..3c6b0d4424 100644 --- a/server/kolide/teams.go +++ b/server/kolide/teams.go @@ -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 { diff --git a/server/mock/datastore_teams.go b/server/mock/datastore_teams.go index a6331069a9..690e0005b0 100644 --- a/server/mock/datastore_teams.go +++ b/server/mock/datastore_teams.go @@ -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) +} diff --git a/server/service/endpoint_teams.go b/server/service/endpoint_teams.go index ef1364ef7e..8c2a4b3164 100644 --- a/server/service/endpoint_teams.go +++ b/server/service/endpoint_teams.go @@ -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 + } +} diff --git a/server/service/handler.go b/server/service/handler.go index 0687e44692..087069e1e9 100644 --- a/server/service/handler.go +++ b/server/service/handler.go @@ -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") diff --git a/server/service/service_teams.go b/server/service/service_teams.go index 918443e6a8..06a52ff4ab 100644 --- a/server/service/service_teams.go +++ b/server/service/service_teams.go @@ -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) +} diff --git a/server/service/transport_teams.go b/server/service/transport_teams.go index ff4856027e..d73d97f101 100644 --- a/server/service/transport_teams.go +++ b/server/service/transport_teams.go @@ -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 +}