Added prompt and admin check to collector. (Velocidex#634)

actions/vql.go

@@ -95,7 +95,13 @@ func (self VQLClientAction) StartQuery(
 
 	// Clients do not have a copy of artifacts so they need to be
 	// sent all artifacts from the server.
-	repository := services.GetRepositoryManager().NewRepository()
+	manager, err := services.GetRepositoryManager()
+	if err != nil {
+		responder.RaiseError(fmt.Sprintf("%v", err))
+		return
+	}
+
+	repository := manager.NewRepository()
 	for _, artifact := range arg.Artifacts {
 		_, err := repository.LoadProto(artifact, false /* validate */)
 		if err != nil {
@@ -123,7 +129,7 @@ func (self VQLClientAction) StartQuery(
 		builder.Env.Set(env_spec.Key, env_spec.Value)
 	}
 
-	scope := services.GetRepositoryManager().BuildScope(builder)
+	scope := manager.BuildScope(builder)
 	defer scope.Close()
 
 	scope.Log("Starting query execution.")

api/api.go

@@ -149,7 +149,12 @@ func (self *ApiServer) GetReport(
 
 	acl_manager := vql_subsystem.NewServerACLManager(self.config, user_name)
 
-	global_repo, err := services.GetRepositoryManager().GetGlobalRepository(self.config)
+	manager, err := services.GetRepositoryManager()
+	if err != nil {
+		return nil, err
+	}
+
+	global_repo, err := manager.GetGlobalRepository(self.config)
 	if err != nil {
 		return nil, err
 	}
@@ -185,7 +190,12 @@ func (self *ApiServer) CollectArtifact(
 		}
 	}
 
-	repository, err := services.GetRepositoryManager().GetGlobalRepository(self.config)
+	manager, err := services.GetRepositoryManager()
+	if err != nil {
+		return nil, err
+	}
+
+	repository, err := manager.GetGlobalRepository(self.config)
 	if err != nil {
 		return nil, err
 	}
@@ -667,7 +677,12 @@ func (self *ApiServer) GetArtifacts(
 
 	if len(in.Names) > 0 {
 		result := &artifacts_proto.ArtifactDescriptors{}
-		repository, err := services.GetRepositoryManager().GetGlobalRepository(self.config)
+		manager, err := services.GetRepositoryManager()
+		if err != nil {
+			return nil, err
+		}
+
+		repository, err := manager.GetGlobalRepository(self.config)
 		if err != nil {
 			return nil, err
 		}
@@ -725,7 +740,12 @@ func (self *ApiServer) SetArtifactFile(
 	permissions := acls.ARTIFACT_WRITER
 
 	// First ensure that the artifact is correct.
-	tmp_repository := services.GetRepositoryManager().NewRepository()
+	manager, err := services.GetRepositoryManager()
+	if err != nil {
+		return nil, err
+	}
+
+	tmp_repository := manager.NewRepository()
 	artifact_definition, err := tmp_repository.LoadYaml(
 		in.Artifact, true /* validate */)
 	if err != nil {
@@ -1008,7 +1028,12 @@ func (self *ApiServer) CreateDownloadFile(ctx context.Context,
 
 	}
 
-	scope := services.GetRepositoryManager().BuildScope(
+	manager, err := services.GetRepositoryManager()
+	if err != nil {
+		return nil, err
+	}
+
+	scope := manager.BuildScope(
 		services.ScopeBuilder{
 			Config:     self.config,
 			Env:        env,

api/artifacts.go

@@ -79,7 +79,12 @@ func getArtifactFile(
 	config_obj *config_proto.Config,
 	name string) (string, error) {
 
-	repository, err := services.GetRepositoryManager().GetGlobalRepository(config_obj)
+	manager, err := services.GetRepositoryManager()
+	if err != nil {
+		return "", err
+	}
+
+	repository, err := manager.GetGlobalRepository(config_obj)
 	if err != nil {
 		return "", err
 	}
@@ -116,13 +121,16 @@ func setArtifactFile(config_obj *config_proto.Config,
 	required_prefix string) (
 	*artifacts_proto.Artifact, error) {
 
-	repository_manager := services.GetRepositoryManager()
+	manager, err := services.GetRepositoryManager()
+	if err != nil {
+		return nil, err
+	}
 
 	switch in.Op {
 	case api_proto.SetArtifactRequest_DELETE:
 
 		// First ensure that the artifact is correct.
-		tmp_repository := services.GetRepositoryManager().NewRepository()
+		tmp_repository := manager.NewRepository()
 		artifact_definition, err := tmp_repository.LoadYaml(
 			in.Artifact, true /* validate */)
 		if err != nil {
@@ -135,11 +143,11 @@ func setArtifactFile(config_obj *config_proto.Config,
 					required_prefix + "'")
 		}
 
-		return artifact_definition, repository_manager.DeleteArtifactFile(config_obj,
+		return artifact_definition, manager.DeleteArtifactFile(config_obj,
 			artifact_definition.Name)
 
 	case api_proto.SetArtifactRequest_SET:
-		return repository_manager.SetArtifactFile(
+		return manager.SetArtifactFile(
 			config_obj, in.Artifact, required_prefix)
 	}
 
@@ -156,7 +164,11 @@ func getReportArtifacts(
 		number_of_results = 100
 	}
 
-	repository, err := services.GetRepositoryManager().GetGlobalRepository(config_obj)
+	manager, err := services.GetRepositoryManager()
+	if err != nil {
+		return nil, err
+	}
+	repository, err := manager.GetGlobalRepository(config_obj)
 	if err != nil {
 		return nil, err
 	}
@@ -230,7 +242,11 @@ func searchArtifact(
 		return true
 	}
 
-	repository, err := services.GetRepositoryManager().GetGlobalRepository(config_obj)
+	manager, err := services.GetRepositoryManager()
+	if err != nil {
+		return nil, err
+	}
+	repository, err := manager.GetGlobalRepository(config_obj)
 	if err != nil {
 		return nil, err
 	}

api/notebooks.go

@@ -486,7 +486,11 @@ func (self *ApiServer) UpdateNotebookCell(
 
 	acl_manager := vql_subsystem.NewServerACLManager(self.config, user_name)
 
-	global_repo, err := services.GetRepositoryManager().GetGlobalRepository(self.config)
+	manager, err := services.GetRepositoryManager()
+	if err != nil {
+		return nil, err
+	}
+	global_repo, err := manager.GetGlobalRepository(self.config)
 	if err != nil {
 		return nil, err
 	}

api/query.go

@@ -73,7 +73,11 @@ func streamQuery(
 	scope_logger := MakeLogger(response_channel)
 
 	// Add extra artifacts to the query from the global repository.
-	repository, err := services.GetRepositoryManager().GetGlobalRepository(config_obj)
+	manager, err := services.GetRepositoryManager()
+	if err != nil {
+		return err
+	}
+	repository, err := manager.GetGlobalRepository(config_obj)
 	if err != nil {
 		return err
 	}
@@ -91,7 +95,7 @@ func streamQuery(
 	}
 
 	// Now execute the query.
-	scope := services.GetRepositoryManager().BuildScope(builder)
+	scope := manager.BuildScope(builder)
 	defer scope.Close()
 
 	// Throttle the query if required.

api/reflect.go

@@ -74,7 +74,11 @@ func (self *ApiServer) GetKeywordCompletions(
 		})
 	}
 
-	repository, err := services.GetRepositoryManager().GetGlobalRepository(self.config)
+	manager, err := services.GetRepositoryManager()
+	if err != nil {
+		return nil, err
+	}
+	repository, err := manager.GetGlobalRepository(self.config)
 	if err != nil {
 		return nil, err
 	}

api/vql.go

@@ -38,7 +38,11 @@ func RunVQL(
 
 	result := &api_proto.GetTableResponse{}
 
-	scope := services.GetRepositoryManager().BuildScope(services.ScopeBuilder{
+	manager, err := services.GetRepositoryManager()
+	if err != nil {
+		return nil, err
+	}
+	scope := manager.BuildScope(services.ScopeBuilder{
 		Config:     config_obj,
 		Env:        env,
 		ACLManager: vql_subsystem.NewServerACLManager(config_obj, principal),

artifacts/definitions/Generic/System/Pstree.yaml

@@ -12,6 +12,10 @@ parameters:
   - name: ProcessNameRegex
     default: .
 
+  - name: PidFilter
+    description: Filter pids by this regex
+    default: .
+
   - name: CallChainFilter
     default: .
 
@@ -64,4 +68,6 @@ sources:
       SELECT Name, Pid, Ppid,
              join(array=pstree(LookupPid=Pid).Name, sep=CallChainSep) AS CallChain
       FROM processes
-      WHERE Name =~ ProcessNameRegex AND CallChain =~ CallChainFilter
+      WHERE Name =~ ProcessNameRegex
+        AND CallChain =~ CallChainFilter
+        AND str(str=Pid) =~ str(str=PidFilter)

artifacts/definitions/Server/Utils/CreateCollector.yaml

@@ -71,6 +71,26 @@ parameters:
     type: json
     default: "{}"
 
+  - name: opt_verbose
+    default: Y
+    type: bool
+    description: Show verbose progress.
+
+  - name: opt_banner
+    default: Y
+    type: bool
+    description: Show Velociraptor banner.
+
+  - name: opt_prompt
+    default: Y
+    type: bool
+    description: Wait for a prompt before closing.
+
+  - name: opt_admin
+    default: Y
+    type: bool
+    description: Require administrator privilege when running.
+
   - name: StandardCollection
     type: hidden
     default: |
@@ -300,11 +320,18 @@ sources:
             ) AS sources FROM scope()  }
       )
 
+      LET optional_cmdline = SELECT * FROM chain(
+        a={ SELECT "-v" AS Opt FROM scope() WHERE opt_verbose = "Y"},
+        b={ SELECT "--nobanner" AS Opt FROM scope() WHERE opt_banner != "Y"},
+        c={ SELECT "--require_admin" AS Opt FROM scope() WHERE opt_admin = "Y"},
+        d={ SELECT "--prompt" AS Opt FROM scope() WHERE opt_prompt = "Y"}
+      )
+
       // Build the autoexec config file depending on the user's
       // collection type choices.
       LET autoexec <= dict(autoexec=dict(
-          argv=["artifacts", "collect", "-v", "Collector",
-                "--logfile", CollectorName + ".log"],
+          argv=("artifacts", "collect", "Collector",
+                "--logfile", CollectorName + ".log") + optional_cmdline.Opt,
           artifact_definitions=definitions)
       )
 

bin/admin_linux.go

@@ -0,0 +1,17 @@
+//+ build: !windows
+
+package main
+
+import (
+	"fmt"
+	"os"
+	"syscall"
+)
+
+func checkAdmin() {
+	if *artificat_command_collect_admin_flag && syscall.Geteuid() != 0 {
+		fmt.Println("Velociraptor requires administrator level access. Use a 'Run as administrator' command shell to launch the binary.")
+		os.Exit(-1)
+	}
+
+}

bin/admin_windows.go

@@ -0,0 +1,57 @@
+//+ build: windows
+
+package main
+
+import (
+	"fmt"
+	"os"
+
+	"golang.org/x/sys/windows"
+)
+
+func checkAdmin() {
+	if *artificat_command_collect_admin_flag && !IsAdmin() {
+		fmt.Println("Velociraptor requires administrator level access. Use a 'Run as administrator' command shell to launch the binary.")
+		os.Exit(-1)
+	}
+}
+
+// https://github.com/golang/go/issues/28804
+func IsAdmin() bool {
+	var sid *windows.SID
+
+	// Although this looks scary, it is directly copied from the
+	// official windows documentation. The Go API for this is a
+	// direct wrap around the official C++ API.
+	// See https://docs.microsoft.com/en-us/windows/desktop/api/securitybaseapi/nf-securitybaseapi-checktokenmembership
+	err := windows.AllocateAndInitializeSid(
+		&windows.SECURITY_NT_AUTHORITY,
+		2,
+		windows.SECURITY_BUILTIN_DOMAIN_RID,
+		windows.DOMAIN_ALIAS_RID_ADMINS,
+		0, 0, 0, 0, 0, 0,
+		&sid)
+	if err != nil {
+		fmt.Printf("SID Error: %s\n", err)
+		return false
+	}
+	defer windows.FreeSid(sid)
+
+	// This appears to cast a null pointer so I'm not sure why this
+	// works, but this guy says it does and it Works for Me™:
+	// https://github.com/golang/go/issues/28804#issuecomment-438838144
+	token := windows.Token(0)
+
+	member, err := token.IsMember(sid)
+	if err != nil {
+		fmt.Printf("Token Membership Error: %s", err)
+		return false
+	}
+
+	// Also note that an admin is _not_ necessarily considered
+	// elevated.
+	// For elevation see https://github.com/mozey/run-as-admin
+	fmt.Printf("IsElevated %v\n", token.IsElevated())
+	fmt.Printf("Member %v\n", member)
+	return token.IsElevated() || member
+}