Added hunt_add() VQL function. (Velocidex#608)

This allows VQL to directly add a client to a hunt. The client will be
immediately scheduled regardless of label matches. This works even if
the hunt is paused or expired.

bin/binary_test.go

@@ -109,6 +109,52 @@ func (self *MainTestSuite) TestAutoexec() {
 	require.Contains(self.T(), string(out), "MySpecialInterface")
 }
 
+func (self *MainTestSuite) TestBuildDeb() {
+	// A temp file for the generated config.
+	config_file, err := ioutil.TempFile("", "config")
+	assert.NoError(self.T(), err)
+	defer os.Remove(config_file.Name())
+
+	cmd := exec.Command(
+		self.binary, "config", "generate", "--merge",
+		`{"Client": {"nonce": "Foo", "writeback_linux": "some_location"}}`)
+	out, err := cmd.Output()
+	require.NoError(self.T(), err)
+
+	// Write the config to the tmp file
+	config_file.Write(out)
+	config_file.Close()
+
+	// Create a tempfile for the binary executable.
+	binary_file, err := ioutil.TempFile("", "binary")
+	assert.NoError(self.T(), err)
+
+	defer os.Remove(binary_file.Name())
+	binary_file.Write([]byte("\x7f\x45\x4c\x46XXXXXXXXXX"))
+	binary_file.Close()
+
+	output_file, err := ioutil.TempFile("", "output*.deb")
+	assert.NoError(self.T(), err)
+	output_file.Close()
+	defer os.Remove(output_file.Name())
+
+	cmd = exec.Command(
+		self.binary, "--config", config_file.Name(),
+		"debian", "client", "--binary", binary_file.Name(),
+		"--output", output_file.Name())
+	_, err = cmd.Output()
+	require.NoError(self.T(), err)
+
+	// Make sure the file is written
+	fd, err := os.Open(output_file.Name())
+	assert.NoError(self.T(), err)
+
+	stat, err := fd.Stat()
+	assert.NoError(self.T(), err)
+
+	assert.Greater(self.T(), stat.Size(), int64(0))
+}
+
 func (self *MainTestSuite) TestGenerateConfigWithMerge() {
 	// A temp file for the generated config.
 	config_file, err := ioutil.TempFile("", "config")

bin/collector_test.go

@@ -69,6 +69,7 @@ func (self *CollectorTestSuite) SetupTest() {
 	self.config_obj.Datastore.Implementation = "FileBaseDataStore"
 	self.config_obj.Datastore.Location = self.tmpdir
 	self.config_obj.Datastore.FilestoreDirectory = self.tmpdir
+	self.config_obj.Frontend.DoNotCompressArtifacts = true
 
 	// Start a web server that serves the filesystem
 	self.test_server = httptest.NewServer(
@@ -166,16 +167,34 @@ reports:
 	fmt.Println(string(out))
 	require.NoError(self.T(), err)
 
-	for _, os_name := range []string{"Windows", "Windows_x86", "Linux", "Darwin"} {
+	var os_name string
+	for _, os_name = range []string{"Windows", "Windows_x86", "Linux", "Darwin"} {
 		cmd = exec.Command(self.binary, "--config", self.config_file,
 			"tools", "upload", "--name", "Velociraptor"+os_name,
-			self.test_server.URL+"/"+filepath.Base(self.binary),
+			self.config_file,
 			"--serve_remote")
 		out, err = cmd.CombinedOutput()
 		fmt.Println(string(out))
 		require.NoError(self.T(), err)
 	}
 
+	switch runtime.GOOS {
+	case "windows":
+		os_name = "Windows"
+	case "linux":
+		os_name = "Linux"
+	case "darwin":
+		os_name = "Darwin"
+	}
+
+	cmd = exec.Command(self.binary, "--config", self.config_file,
+		"tools", "upload", "--name", "Velociraptor"+os_name,
+		self.test_server.URL+"/"+filepath.Base(self.binary),
+		"--serve_remote")
+	out, err = cmd.CombinedOutput()
+	fmt.Println(string(out))
+	require.NoError(self.T(), err)
+
 	// Make sure the binary is proprly added.
 	assert.Regexp(self.T(), "name: Velociraptor", string(out))
 

bin/installer_windows.go

@@ -543,6 +543,7 @@ func NewVelociraptorService(name string) (*VelociraptorService, error) {
 	go func() {
 		for {
 			runOnce(result, elog)
+			time.Sleep(10 * time.Second)
 		}
 	}()
 

reporting/container.go

@@ -419,6 +419,9 @@ func (self *Container) maybeCollectSparseFile(
 }
 
 func (self *Container) Close() error {
+	self.Lock()
+	defer self.Unlock()
+
 	if self.current_writers != 0 {
 		for _, i := range self.backtraces {
 			fmt.Println(i)

services/hunt_manager/hunt_manager.go

@@ -71,6 +71,7 @@ type ParticipationRecord struct {
 	Fqdn        string `vfilter:"optional,field=Fqdn"`
 	FlowId      string `vfilter:"optional,field=FlowId"`
 	Participate bool   `vfilter:"required,field=Participate"`
+	Override    bool   `vfilter:"optional,field=Override"`
 	Timestamp   uint64 `vfilter:"optional,field=Timestamp"`
 	TS          uint64 `vfilter:"optional,field=_ts"`
 }
@@ -98,7 +99,11 @@ func (self *HuntManager) StartParticipation(
 	config_obj *config_proto.Config,
 	wg *sync.WaitGroup) error {
 
-	scope := vfilter.NewScope()
+	scope := services.GetRepositoryManager().BuildScope(
+		services.ScopeBuilder{
+			Config: config_obj,
+			Logger: logging.NewPlainLogger(config_obj, &logging.GenericComponent),
+		})
 	qm_chan, cancel := services.GetJournal().Watch("System.Hunt.Participation")
 
 	wg.Add(1)
@@ -200,7 +205,7 @@ func (self *HuntManager) ProcessRow(
 	participation_row := &ParticipationRecord{}
 	err := vfilter.ExtractArgs(scope, row, participation_row)
 	if err != nil {
-		scope.Log("ExtractArgs %v", err)
+		scope.Log("hunt_manager: %v", err)
 		return
 	}
 
@@ -248,6 +253,15 @@ func (self *HuntManager) ProcessRow(
 				hunt_obj.Stats = &api_proto.HuntStats{}
 			}
 
+			// The event may override the regular hunt
+			// logic.
+			if participation_row.Override {
+				proto.Merge(request, hunt_obj.StartRequest)
+				hunt_obj.Stats.TotalClientsScheduled += 1
+
+				return nil
+			}
+
 			// Ignore stopped hunts.
 			if hunt_obj.Stats.Stopped ||
 				hunt_obj.State != api_proto.Hunt_RUNNING {
@@ -280,13 +294,12 @@ func (self *HuntManager) ProcessRow(
 		})
 
 	if err != nil {
-		scope.Log("hunt manager: launching %v:  %v", participation_row, err)
 		return
 	}
 
 	repository, err := services.GetRepositoryManager().GetGlobalRepository(config_obj)
 	if err != nil {
-		scope.Log("hunt manager: launching %v:  %v", participation_row, err)
+		scope.Log("hunt manager: GetGlobalRepository: %v", err)
 		return
 	}
 

services/hunt_manager/hunt_manager_test.go

@@ -94,7 +94,6 @@ func (self *HuntTestSuite) TestHuntManager() {
 		[]*ordereddict.Dict{ordereddict.NewDict().
 			Set("HuntId", self.hunt_id).
 			Set("ClientId", self.client_id).
-			Set("Fqdn", "MyHost").
 			Set("Participate", true)},
 		"System.Hunt.Participation", self.client_id, "")
 
@@ -232,6 +231,58 @@ func (self *HuntTestSuite) TestHuntWithLabelClientHasLabelDifferentCase() {
 	assert.Equal(t, collection_context.Request.Artifacts, self.expected.Artifacts)
 }
 
+func (self *HuntTestSuite) TestHuntWithOverride() {
+	t := self.T()
+
+	services.GetLauncher().SetFlowIdForTests("F.1234")
+
+	// Hunt is paused so normally will not receive any clients.
+	hunt_obj := &api_proto.Hunt{
+		HuntId:       self.hunt_id,
+		StartRequest: self.expected,
+		State:        api_proto.Hunt_PAUSED,
+		Stats:        &api_proto.HuntStats{},
+		Expires:      uint64(time.Now().Add(7*24*time.Hour).UTC().UnixNano() / 1000),
+	}
+
+	db, err := datastore.GetDB(self.config_obj)
+	assert.NoError(t, err)
+
+	hunt_path_manager := paths.NewHuntPathManager(hunt_obj.HuntId)
+	err = db.SetSubject(self.config_obj, hunt_path_manager.Path(), hunt_obj)
+	assert.NoError(t, err)
+
+	services.GetHuntDispatcher().Refresh(self.config_obj)
+
+	// Simulate a System.Hunt.Participation event
+	path_manager := result_sets.NewArtifactPathManager(self.config_obj,
+		self.client_id, "", "System.Hunt.Participation")
+	services.GetJournal().PushRows(self.config_obj, path_manager,
+		[]*ordereddict.Dict{ordereddict.NewDict().
+			Set("HuntId", self.hunt_id).
+			Set("ClientId", self.client_id).
+			Set("Override", true).
+			Set("Participate", true)})
+
+	vtesting.WaitUntil(5*time.Second, self.T(), func() bool {
+		// The hunt index is updated since we have seen this client
+		// already (even if we decided not to launch on it).
+		err = db.CheckIndex(self.config_obj, constants.HUNT_INDEX,
+			self.client_id, []string{hunt_obj.HuntId})
+		if err != nil {
+			return false
+		}
+
+		_, err := LoadCollectionContext(self.config_obj, self.client_id, "F.1234")
+		return err == nil
+	})
+
+	collection_context, err := LoadCollectionContext(self.config_obj,
+		self.client_id, "F.1234")
+	assert.NoError(t, err)
+	assert.Equal(t, collection_context.Request.Artifacts, self.expected.Artifacts)
+}
+
 func (self *HuntTestSuite) TestHuntWithLabelClientHasLabel() {
 	t := self.T()
 

vql/acls.go

@@ -36,6 +36,7 @@ type ServerACLManager struct {
 	Token     *acl_proto.ApiClientACL
 }
 
+// Token must have *ALL* the specified permissions.
 func (self *ServerACLManager) CheckAccess(
 	permissions ...acls.ACL_PERMISSION) (bool, error) {
 	for _, permission := range permissions {

vql/server/hunt.go

@@ -29,6 +29,7 @@ import (
 	"www.velocidex.com/golang/velociraptor/artifacts"
 	flows_proto "www.velocidex.com/golang/velociraptor/flows/proto"
 	"www.velocidex.com/golang/velociraptor/json"
+	"www.velocidex.com/golang/velociraptor/services"
 	vql_subsystem "www.velocidex.com/golang/velociraptor/vql"
 
 	"www.velocidex.com/golang/velociraptor/grpc_client"
@@ -118,6 +119,66 @@ func (self ScheduleHuntFunction) Info(scope *vfilter.Scope, type_map *vfilter.Ty
 	}
 }
 
+type AddToHuntFunctionArg struct {
+	ClientId string `vfilter:"required,field=ClientId"`
+	HuntId   string `vfilter:"required,field=HuntId"`
+}
+
+type AddToHuntFunction struct{}
+
+func (self *AddToHuntFunction) Call(ctx context.Context,
+	scope *vfilter.Scope,
+	args *ordereddict.Dict) vfilter.Any {
+
+	err := vql_subsystem.CheckAccess(scope, acls.COLLECT_CLIENT)
+	if err != nil {
+		scope.Log("hunt_add: %s", err)
+		return vfilter.Null{}
+	}
+
+	arg := &AddToHuntFunctionArg{}
+	err = vfilter.ExtractArgs(scope, args, arg)
+	if err != nil {
+		scope.Log("hunt_add: %s", err.Error())
+		return vfilter.Null{}
+	}
+
+	config_obj, ok := artifacts.GetServerConfig(scope)
+	if !ok {
+		scope.Log("Command can only run on the server")
+		return vfilter.Null{}
+	}
+
+	journal := services.GetJournal()
+	if journal == nil {
+		return vfilter.Null{}
+	}
+
+	err = journal.PushRowsToArtifact(config_obj,
+		[]*ordereddict.Dict{ordereddict.NewDict().
+			Set("HuntId", arg.HuntId).
+			Set("ClientId", arg.ClientId).
+			Set("Override", true).
+			Set("Participate", true)},
+		"System.Hunt.Participation", arg.ClientId, "")
+	if err != nil {
+		scope.Log("hunt_add: %s", err.Error())
+		return vfilter.Null{}
+	}
+
+	return arg.ClientId
+}
+
+func (self AddToHuntFunction) Info(scope *vfilter.Scope,
+	type_map *vfilter.TypeMap) *vfilter.FunctionInfo {
+	return &vfilter.FunctionInfo{
+		Name:    "hunt_add",
+		Doc:     "Assign a client to a hunt.",
+		ArgType: type_map.AddType(scope, &AddToHuntFunctionArg{}),
+	}
+}
+
 func init() {
 	vql_subsystem.RegisterFunction(&ScheduleHuntFunction{})
+	vql_subsystem.RegisterFunction(&AddToHuntFunction{})
 }