Fixed bugs in VQL import_collection function (Velocidex#1244)

file_store/test_utils/testsuite.go

@@ -10,6 +10,7 @@ import (
 	"www.velocidex.com/golang/velociraptor/config"
 	config_proto "www.velocidex.com/golang/velociraptor/config/proto"
 	"www.velocidex.com/golang/velociraptor/services"
+	"www.velocidex.com/golang/velociraptor/services/client_info"
 	"www.velocidex.com/golang/velociraptor/services/inventory"
 	"www.velocidex.com/golang/velociraptor/services/journal"
 	"www.velocidex.com/golang/velociraptor/services/labels"
@@ -45,6 +46,7 @@ func (self *TestSuite) SetupTest() {
 	require.NoError(self.T(), self.Sm.Start(journal.StartJournalService))
 	require.NoError(self.T(), self.Sm.Start(notifications.StartNotificationService))
 	require.NoError(self.T(), self.Sm.Start(inventory.StartInventoryService))
+	require.NoError(self.T(), self.Sm.Start(client_info.StartClientInfoService))
 	require.NoError(self.T(), self.Sm.Start(launcher.StartLauncherService))
 	require.NoError(self.T(), self.Sm.Start(repository.StartRepositoryManager))
 	require.NoError(self.T(), self.Sm.Start(labels.StartLabelService))

vql/tools/collector_test.go

@@ -11,19 +11,11 @@ import (
 	"github.com/Velocidex/ordereddict"
 	"github.com/alecthomas/assert"
 	"github.com/sebdah/goldie"
-	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
-	"www.velocidex.com/golang/velociraptor/config"
-	config_proto "www.velocidex.com/golang/velociraptor/config/proto"
 	"www.velocidex.com/golang/velociraptor/file_store/test_utils"
 	"www.velocidex.com/golang/velociraptor/json"
 	"www.velocidex.com/golang/velociraptor/logging"
 	"www.velocidex.com/golang/velociraptor/services"
-	"www.velocidex.com/golang/velociraptor/services/inventory"
-	"www.velocidex.com/golang/velociraptor/services/journal"
-	"www.velocidex.com/golang/velociraptor/services/launcher"
-	"www.velocidex.com/golang/velociraptor/services/notifications"
-	"www.velocidex.com/golang/velociraptor/services/repository"
 	"www.velocidex.com/golang/velociraptor/utils"
 	vql_subsystem "www.velocidex.com/golang/velociraptor/vql"
 	"www.velocidex.com/golang/vfilter"
@@ -118,47 +110,18 @@ sources:
 )
 
 type TestSuite struct {
-	suite.Suite
-	config_obj *config_proto.Config
-	sm         *services.Service
-}
-
-func (self *TestSuite) SetupTest() {
-	var err error
-	self.config_obj, err = new(config.Loader).WithFileLoader(
-		"../../http_comms/test_data/server.config.yaml").
-		WithRequiredFrontend().WithWriteback().WithVerbose(true).
-		LoadAndValidate()
-	require.NoError(self.T(), err)
-
-	self.config_obj.Frontend.DoNotCompressArtifacts = true
-
-	// Start essential services.
-	ctx, _ := context.WithTimeout(context.Background(), time.Second*60)
-	self.sm = services.NewServiceManager(ctx, self.config_obj)
-
-	require.NoError(self.T(), self.sm.Start(journal.StartJournalService))
-	require.NoError(self.T(), self.sm.Start(notifications.StartNotificationService))
-	require.NoError(self.T(), self.sm.Start(inventory.StartInventoryService))
-	require.NoError(self.T(), self.sm.Start(repository.StartRepositoryManager))
-	require.NoError(self.T(), self.sm.Start(launcher.StartLauncherService))
-}
-
-func (self *TestSuite) TearDownTest() {
-	self.sm.Close()
-	test_utils.GetMemoryFileStore(self.T(), self.config_obj).Clear()
-	test_utils.GetMemoryDataStore(self.T(), self.config_obj).Clear()
+	test_utils.TestSuite
 }
 
 func (self *TestSuite) TestSimpleCollection() {
 	scope := vql_subsystem.MakeScope()
 
-	scope.SetLogger(logging.NewPlainLogger(self.config_obj, &logging.FrontendComponent))
+	scope.SetLogger(logging.NewPlainLogger(self.ConfigObj, &logging.FrontendComponent))
 
-	repository, err := getRepository(self.config_obj, nil)
+	repository, err := getRepository(self.ConfigObj, nil)
 	assert.NoError(self.T(), err)
 
-	request, err := getArtifactCollectorArgs(self.config_obj,
+	request, err := getArtifactCollectorArgs(self.ConfigObj,
 		repository, scope, simpleCollectorArgs)
 	assert.NoError(self.T(), err)
 
@@ -167,7 +130,7 @@ func (self *TestSuite) TestSimpleCollection() {
 
 	acl_manager := vql_subsystem.NullACLManager{}
 	vql_requests, err := launcher.CompileCollectorArgs(
-		context.Background(), self.config_obj, acl_manager, repository,
+		context.Background(), self.ConfigObj, acl_manager, repository,
 		services.CompilerOptions{}, request)
 
 	serialized, err := json.MarshalIndent(ordereddict.NewDict().
@@ -191,9 +154,9 @@ func (self *TestSuite) TestCollectionWithArtifacts() {
 	defer os.Remove(report_file.Name())
 
 	builder := services.ScopeBuilder{
-		Config:     self.config_obj,
+		Config:     self.ConfigObj,
 		ACLManager: vql_subsystem.NullACLManager{},
-		Logger:     logging.NewPlainLogger(self.config_obj, &logging.FrontendComponent),
+		Logger:     logging.NewPlainLogger(self.ConfigObj, &logging.FrontendComponent),
 		Env:        ordereddict.NewDict(),
 	}
 
@@ -237,9 +200,9 @@ func (self *TestSuite) TestCollectionWithTypes() {
 	defer os.Remove(output_file.Name())
 
 	builder := services.ScopeBuilder{
-		Config:     self.config_obj,
+		Config:     self.ConfigObj,
 		ACLManager: vql_subsystem.NullACLManager{},
-		Logger:     logging.NewPlainLogger(self.config_obj, &logging.FrontendComponent),
+		Logger:     logging.NewPlainLogger(self.ConfigObj, &logging.FrontendComponent),
 		Env:        ordereddict.NewDict(),
 	}
 
@@ -277,9 +240,9 @@ func (self *TestSuite) TestCollectionWithUpload() {
 	defer os.Remove(output_file.Name())
 
 	builder := services.ScopeBuilder{
-		Config:     self.config_obj,
+		Config:     self.ConfigObj,
 		ACLManager: vql_subsystem.NullACLManager{},
-		Logger:     logging.NewPlainLogger(self.config_obj, &logging.FrontendComponent),
+		Logger:     logging.NewPlainLogger(self.ConfigObj, &logging.FrontendComponent),
 		Env:        ordereddict.NewDict(),
 	}
 

vql/tools/import.go

@@ -13,6 +13,7 @@ import (
 	"github.com/Velocidex/ordereddict"
 	"www.velocidex.com/golang/velociraptor/acls"
 	actions_proto "www.velocidex.com/golang/velociraptor/actions/proto"
+	api_proto "www.velocidex.com/golang/velociraptor/api/proto"
 	config_proto "www.velocidex.com/golang/velociraptor/config/proto"
 	"www.velocidex.com/golang/velociraptor/datastore"
 	"www.velocidex.com/golang/velociraptor/file_store"
@@ -71,7 +72,8 @@ func (self ImportCollectionFunction) Call(ctx context.Context,
 	}
 
 	if arg.ClientId == "auto" {
-		arg.ClientId, err = makeNewClient(config_obj, arg.Hostname)
+		arg.ClientId, err = getExistingClientOrNewClient(
+			ctx, config_obj, arg.Hostname)
 		if err != nil {
 			scope.Log("import_collection: %v", err)
 			return vfilter.Null{}
@@ -254,7 +256,7 @@ func (self ImportCollectionFunction) Call(ctx context.Context,
 				}
 				defer out_fd.Close()
 
-				log("Copying file %v -> %v", file.Name, out_path)
+				log("Copying file %v -> %v", file.Name, out_path.AsClientPath())
 
 				_, err = utils.Copy(ctx, out_fd, fd)
 				if err != nil {
@@ -283,6 +285,14 @@ func (self ImportCollectionFunction) Call(ctx context.Context,
 		return vfilter.Null{}
 	}
 
+	// Write an empty request so we can show something in the GUI
+	err = db.SetSubject(config_obj, path_manager.Task(),
+		&api_proto.ApiFlowRequestDetails{})
+	if err != nil {
+		scope.Log("import_collection: %v", err)
+		return vfilter.Null{}
+	}
+
 	// Generate a fake System.Flow.Completion event for the
 	// uploaded flow in case there are any listeners who are
 	// interested.
@@ -324,6 +334,21 @@ func NewClientId() string {
 	return "C." + string(dst)
 }
 
+func getExistingClientOrNewClient(
+	ctx context.Context,
+	config_obj *config_proto.Config,
+	hostname string) (string, error) {
+
+	// Search for an existing client with the same hostname
+	search_resp, err := search.SearchClients(ctx, config_obj,
+		&api_proto.SearchClientsRequest{Query: hostname}, "")
+	if err == nil && len(search_resp.Items) > 0 {
+		return search_resp.Items[0].ClientId, nil
+	}
+
+	return makeNewClient(config_obj, hostname)
+}
+
 // Create a new client record
 func makeNewClient(
 	config_obj *config_proto.Config,
@@ -354,18 +379,20 @@ func makeNewClient(
 		return "", err
 	}
 	// Add the new client to the index.
-	keywords := []string{
+	for _, term := range []string{
 		"all", // This is used for "." search
 		client_id,
 		client_info.Hostname,
 		client_info.Fqdn,
 		"host:" + client_info.Hostname,
+	} {
+		err = search.SetIndex(config_obj, client_id, term)
+		if err != nil {
+			return client_id, err
+		}
 	}
 
-	return client_id, db.SetIndex(config_obj,
-		paths.CLIENT_INDEX_URN,
-		client_id, keywords,
-	)
+	return client_id, nil
 }
 
 func init() {

vql/tools/import_test.go

@@ -8,27 +8,25 @@ import (
 
 	"github.com/Velocidex/ordereddict"
 	"github.com/alecthomas/assert"
-	"github.com/stretchr/testify/require"
+	api_proto "www.velocidex.com/golang/velociraptor/api/proto"
 	"www.velocidex.com/golang/velociraptor/flows/proto"
 	flows_proto "www.velocidex.com/golang/velociraptor/flows/proto"
 	"www.velocidex.com/golang/velociraptor/logging"
+	"www.velocidex.com/golang/velociraptor/search"
 	"www.velocidex.com/golang/velociraptor/services"
-	"www.velocidex.com/golang/velociraptor/services/labels"
 	vql_subsystem "www.velocidex.com/golang/velociraptor/vql"
 )
 
 func (self *TestSuite) TestImportCollection() {
-	require.NoError(self.T(), self.sm.Start(labels.StartLabelService))
-
 	manager, _ := services.GetRepositoryManager()
-	repository, _ := manager.GetGlobalRepository(self.config_obj)
+	repository, _ := manager.GetGlobalRepository(self.ConfigObj)
 	_, err := repository.LoadYaml(CustomTestArtifactDependent, true)
 	assert.NoError(self.T(), err)
 
 	builder := services.ScopeBuilder{
-		Config:     self.config_obj,
+		Config:     self.ConfigObj,
 		ACLManager: vql_subsystem.NullACLManager{},
-		Logger:     logging.NewPlainLogger(self.config_obj, &logging.FrontendComponent),
+		Logger:     logging.NewPlainLogger(self.ConfigObj, &logging.FrontendComponent),
 		Env:        ordereddict.NewDict(),
 	}
 
@@ -59,4 +57,28 @@ func (self *TestSuite) TestImportCollection() {
 	assert.Equal(self.T(), uint64(1), context.TotalCollectedRows)
 	assert.Equal(self.T(), flows_proto.ArtifactCollectorContext_FINISHED,
 		context.State)
+
+	// Check the indexes are correct for the new client_id
+	search_resp, err := search.SearchClients(ctx, self.ConfigObj,
+		&api_proto.SearchClientsRequest{Query: "MyNewHost"}, "")
+	assert.NoError(self.T(), err)
+
+	// There is one hit - a new client is added to the index.
+	assert.Equal(self.T(), 1, len(search_resp.Items))
+	assert.Equal(self.T(), search_resp.Items[0].ClientId, context.ClientId)
+
+	// Importing the collection again and providing the same host name
+	// will reuse the client id
+
+	result2 := ImportCollectionFunction{}.Call(ctx, scope,
+		ordereddict.NewDict().
+			Set("client_id", "auto").
+			Set("hostname", "MyNewHost").
+			Set("accessor", "data").
+			Set("filename", data))
+	context2, ok := result2.(*proto.ArtifactCollectorContext)
+	assert.True(self.T(), ok)
+
+	// The new flow was created on the same client id as before.
+	assert.Equal(self.T(), context2.ClientId, context.ClientId)
 }