Added a delete on exit option for the offline collector. (#3974)

If this option is set, the offline collector will delete the collection
zip when exiting. This is only useful if the zip will be uploaded over
the network.

api/query.go

@@ -213,6 +213,12 @@ type logWriter struct {
 }
 
 func (self *logWriter) Write(b []byte) (int, error) {
+	select {
+	case <-self.ctx.Done():
+		return 0, io.EOF
+	default:
+	}
+
 	select {
 	case <-self.ctx.Done():
 		return 0, io.EOF

artifacts/definitions/Server/Internal/ArtifactDescription.yaml

@@ -43,7 +43,7 @@ reports:
       ### Tools
 
       {{ range $artifact.Tools -}}
-      * <grr-tool-viewer name="{{.Name}}" version="{{.Version}}"></grr-tool-viewer>
+      * <velo-tool-viewer name="{{.Name}}" version="{{.Version}}"></velo-tool-viewer>
       {{ end }}
 
       {{ end }}

artifacts/definitions/Server/Utils/CreateCollector.yaml

@@ -144,6 +144,16 @@ parameters:
       because the packed builtin artifacts are only compatible with
       the current release version.
 
+  - name: opt_delete_at_exit
+    type: bool
+    default: N
+    description: |
+      If specified the collection will be deleted at exit. This only
+      makes sense when uploading to the cloud or a remote
+      location. NOTE: There is no way to check that the upload
+      actually worked so this flag deletes the collection regardless
+      of upload success.
+
   - name: StandardCollection
     type: hidden
     default: |
@@ -267,6 +277,12 @@ parameters:
       // The log is always written to the executable path
       LET log_filename <= pathspec(parse= baseline[0].Exe + ".log")
 
+      -- Remove the zip file and log file when done if the user asked for it.
+      LET _ <= if(condition=DeleteOnExit, then=atexit(query={
+         SELECT rm(filename=zip_filename), rm(filename=log_filename) FROM scope()
+         WHERE log(message="Removed Zip file %v", args=zip_filename)
+      }, env=dict(zip_filename=zip_filename, log_filename=log_filename)))
+
       -- Make a random hex string as a random password
       LET RandomPassword <= SELECT format(format="%02x",
             args=rand(range=255)) AS A
@@ -470,6 +486,7 @@ sources:
                     dict(name="ProgressTimeout", type="int",
                          default=opt_progress_timeout),
                     dict(name="Timeout", default=opt_timeout, type="int"),
+                    dict(name="DeleteOnExit", default=opt_delete_at_exit, type="bool"),
                     dict(name="target_args",
                          default=serialize(format='json', item=target_args),
                          type="json"),

artifacts/testdata/server/testcases/atexit.out.yaml

@@ -12,7 +12,7 @@ SELECT * FROM query(query=''' LET Foo = SELECT log(message="I Ran at the end") F
   "Log": "Velociraptor: DEFAULT:Second!\n"
  },
  {
-  "Log": "Velociraptor: Running AtExit query\n"
+  "Log": "Velociraptor: Running AtExit query SELECT log(message=\"I Ran at the end\") FROM scope()\n"
  },
  {
   "Log": "Velociraptor: DEFAULT:I Ran at the end\n"

bin/offline.go

@@ -90,6 +90,13 @@ OptTimeout: 0
 # the packed builtin artifacts are only compatible with the current
 # release version.
 OptVersion: ""
+
+# If specified the collection will be deleted at exit. This only
+# makes sense when uploading to the cloud or a remote
+# location. NOTE: There is no way to check that the upload
+# actually worked so this flag deletes the collection regardless
+# of upload success.
+OptDeleteAtExit: N
 `
 
 func doCollector() error {
@@ -219,7 +226,8 @@ SELECT * FROM Artifact.Server.Utils.CreateCollector(
    opt_cpu_limit=Spec.OptCpuLimit,
    opt_progress_timeout=Spec.OptProgressTimeout,
    opt_timeout=Spec.OptTimeout,
-   opt_version=Spec.OptVersion
+   opt_version=Spec.OptVersion,
+   opt_delete_at_exit=Spec.OptDeleteAtExit
    )
 `
 	return runQueryWithEnv(query, builder, "json")

docs/offline_collector/offline_collector.sh

@@ -68,7 +68,8 @@ SELECT * FROM Artifact.Server.Utils.CreateCollector(
    opt_cpu_limit=Spec.OptCpuLimit,
    opt_progress_timeout=Spec.OptProgressTimeout,
    opt_timeout=Spec.OptTimeout,
-   opt_version=Spec.OptVersion
+   opt_version=Spec.OptVersion,
+   opt_delete_at_exit=Spec.OptDeleteAtExit
    )
 EOF
 )

docs/offline_collector/sample.spec.yaml

@@ -60,3 +60,10 @@ OptTimeout: 0
 # the packed builtin artifacts are only compatible with the current
 # release version.
 OptVersion: ""
+
+# If specified the collection will be deleted at exit. This only
+# makes sense when uploading to the cloud or a remote
+# location. NOTE: There is no way to check that the upload
+# actually worked so this flag deletes the collection regardless
+# of upload success.
+OptDeleteAtExit: N

gui/velociraptor/package.json

@@ -6,7 +6,7 @@
     "dependencies": {
         "@babel/runtime": "^7.26.0",
         "@fortawesome/fontawesome-svg-core": "6.7.1",
-        "@fortawesome/free-regular-svg-icons": "6.6.0",
+        "@fortawesome/free-regular-svg-icons": "6.7.1",
         "@fortawesome/free-solid-svg-icons": "^6.7.1",
         "@fortawesome/react-fontawesome": "0.2.2",
         "@popperjs/core": "^2.11.8",

gui/velociraptor/src/components/flows/offline-collector.jsx

@@ -563,13 +563,14 @@ class OfflineCollectorParameters  extends React.Component {
                     <Form.Label column sm="3">{T("Pause For Prompt")}</Form.Label>
                     <Col sm="8">
                       <Form.Check
-                        type="checkbox"
+                        type="switch"
+                        label={T("Pause For Prompt")}
                         onChange={(e) => {
+                            let value = "N";
                             if (e.currentTarget.checked) {
-                                this.props.parameters.opt_prompt = "Y";
-                            } else {
-                                this.props.parameters.opt_prompt = "N";
-                            }
+                                value = "Y";
+                            };
+                            this.props.parameters.opt_prompt = value;
                             this.props.setParameters(this.props.parameters);
                         }}
                         checked={this.props.parameters.opt_prompt === "Y"}
@@ -622,6 +623,24 @@ class OfflineCollectorParameters  extends React.Component {
                       />
                     </Col>
                   </Form.Group>
+                  <Form.Group as={Row}>
+                    <Form.Label column sm="3">{T("Delete Collection at Exit")}</Form.Label>
+                    <Col sm="8">
+                      <Form.Check
+                        type="switch"
+                        label={T("Delete Collection at Exit")}
+                        onChange={(e) => {
+                            let value = "N";
+                            if (e.currentTarget.checked) {
+                                value = "Y";
+                            }
+                            this.props.parameters.opt_delete_at_exit = value;
+                            this.props.setParameters(this.props.parameters);
+                        }}
+                        checked={this.props.parameters.opt_delete_at_exit === "Y"}
+                      />
+                    </Col>
+                  </Form.Group>
 
                 </Form>
               </Modal.Body>
@@ -788,6 +807,7 @@ function getDefaultCollectionParameters() {
         opt_tempdir: undefined,
         opt_filename_template: "Collection-%FQDN%-%TIMESTAMP%",
         opt_collector_filename: undefined,
+        opt_delete_at_exit: "N",
         opt_format: "jsonl",
         opt_prompt: "N",
     };
@@ -872,6 +892,9 @@ export default class OfflineCollectorWizard extends React.Component {
                 case "opt_prompt":
                     collector_parameters.opt_prompt = value;
                     break;
+                case "opt_delete_at_exit":
+                    collector_parameters.opt_delete_at_exit = value;
+                    break;
                 case "opt_tempdir":
                     collector_parameters.opt_tempdir = value;
                     break;
@@ -890,7 +913,6 @@ export default class OfflineCollectorWizard extends React.Component {
                 case "opt_collector_filename":
                     collector_parameters.opt_collector_filename = value;
                     break;
-
                 case "opt_progress_timeout":
                     resources.progress_timeout =  JSONparse(value);
                     break;
@@ -958,6 +980,7 @@ export default class OfflineCollectorWizard extends React.Component {
         env.push({key: "opt_output_directory", value: str(params.opt_output_directory)});
         env.push({key: "opt_filename_template", value: str(params.opt_filename_template)});
         env.push({key: "opt_collector_filename", value: str(params.opt_collector_filename)});
+        env.push({key: "opt_delete_at_exit", value: str(params.opt_delete_at_exit)});
         env.push({key: "opt_progress_timeout", value: str(this.state.resources.progress_timeout)});
         env.push({key: "opt_timeout", value: str(this.state.resources.timeout)});
         env.push({key: "opt_cpu_limit", value: str( this.state.resources.cpu_limit)});

vql/tools/atexit.go

@@ -40,28 +40,28 @@ func (self AtExitFunction) Call(
 
 	switch t := arg.Query.(type) {
 	case types.StoredQuery:
+		subscope := scope.Copy()
+		subscope.AppendVars(arg.Env)
+
 		vql_subsystem.GetRootScope(scope).AddDestructor(func() {
-			scope.Log("Running AtExit query")
+			scope.Log("Running AtExit query %v", vfilter.FormatToString(scope, t))
 
 			// We need to create a new context to run the
-			// desctructors in because the main context
+			// destructors in because the main context
 			// may already be cancelled.
 			ctx, cancel := context.WithTimeout(
 				context.Background(),
 				time.Duration(timeout)*time.Second)
 			defer cancel()
 
-			subscope := scope.Copy()
-			subscope.AppendVars(arg.Env)
-
 			for _ = range t.Eval(ctx, subscope) {
 			}
 		})
 	default:
 		scope.Log("atexit: Query type %T not supported.", arg.Query)
 	}
 
-	return vfilter.Null{}
+	return true
 }
 
 func (self AtExitFunction) Info(scope vfilter.Scope, type_map *vfilter.TypeMap) *vfilter.FunctionInfo {