Propagate user's prefered timezone for export tables (Velocidex#2232)

When the user sets their preferred timezone in the user preferences the UI adjusts the output. This PR also ensures this preference is propagated to the download table API so the encoded timestamps are also presented in the correct timezone. The user may always adjust the timezone back to UTC to receive their times in UTC again.
mirwitch · Nov 10, 2022 · 01979cd · 01979cd
1 parent 99c67c0
commit 01979cd
Show file tree

Hide file tree

Showing 19 changed files with 145 additions and 69 deletions.
diff --git a/api/csv.go b/api/csv.go
@@ -93,6 +93,8 @@ func getTable(
 		return nil, err
 	}
 
+	opts := getJsonOptsForTimezone(in.Timezone)
+
 	// Unpack the rows into the output protobuf
 	for row := range rs_reader.Rows(ctx) {
 		if result.Columns == nil {
@@ -102,7 +104,7 @@ func getTable(
 		row_data := make([]string, 0, len(result.Columns))
 		for _, key := range result.Columns {
 			value, _ := row.Get(key)
-			row_data = append(row_data, csv.AnyToString(value))
+			row_data = append(row_data, csv.AnyToString(value, opts))
 		}
 		result.Rows = append(result.Rows, &api_proto.Row{
 			Cell: row_data,
@@ -262,6 +264,8 @@ func getEventTableWithPathManager(
 		rs_reader.SetMaxTime(time.Unix(int64(in.EndTime), 0))
 	}
 
+	opts := getJsonOptsForTimezone(in.Timezone)
+
 	// Unpack the rows into the output protobuf
 	for row := range rs_reader.Rows(ctx) {
 		if result.Columns == nil {
@@ -271,7 +275,7 @@ func getEventTableWithPathManager(
 		row_data := make([]string, 0, len(result.Columns))
 		for _, key := range result.Columns {
 			value, _ := row.Get(key)
-			row_data = append(row_data, csv.AnyToString(value))
+			row_data = append(row_data, csv.AnyToString(value, opts))
 		}
 		result.Rows = append(result.Rows, &api_proto.Row{
 			Cell: row_data,
@@ -315,6 +319,7 @@ func getTimeline(
 	}
 
 	rows := uint64(0)
+	opts := getJsonOptsForTimezone(in.Timezone)
 	for item := range reader.Read(ctx) {
 		if result.StartTime == 0 {
 			result.StartTime = item.Time.UnixNano()
@@ -323,8 +328,8 @@ func getTimeline(
 		result.Rows = append(result.Rows, &api_proto.Row{
 			Cell: []string{
 				item.Source,
-				csv.AnyToString(item.Time),
-				csv.AnyToString(item.Row)},
+				csv.AnyToString(item.Time, opts),
+				csv.AnyToString(item.Row, opts)},
 		})
 
 		rows += 1

diff --git a/api/download.go b/api/download.go
@@ -52,6 +52,7 @@ import (
 	"www.velocidex.com/golang/velociraptor/file_store/path_specs"
 	"www.velocidex.com/golang/velociraptor/flows"
 	"www.velocidex.com/golang/velociraptor/json"
+	vjson "www.velocidex.com/golang/velociraptor/json"
 	"www.velocidex.com/golang/velociraptor/logging"
 	"www.velocidex.com/golang/velociraptor/paths"
 	"www.velocidex.com/golang/velociraptor/paths/artifacts"
@@ -400,6 +401,7 @@ func downloadTable() http.Handler {
 			return
 		}
 
+		opts := getJsonOptsForTimezone(request.Timezone)
 		switch request.DownloadFormat {
 		case "csv":
 			download_name = strings.TrimSuffix(download_name, ".json")
@@ -421,7 +423,8 @@ func downloadTable() http.Handler {
 
 			scope := vql_subsystem.MakeScope()
 			csv_writer := csv.GetCSVAppender(
-				org_config_obj, scope, w, true /* write_headers */)
+				org_config_obj, scope, w,
+				csv.WriteHeaders, opts)
 			for row := range row_chan {
 				csv_writer.Write(
 					filterColumns(request.Columns, transform(row)))
@@ -449,8 +452,9 @@ func downloadTable() http.Handler {
 			}).Info("DownloadTable")
 
 			for row := range row_chan {
-				serialized, err := json.Marshal(
-					filterColumns(request.Columns, transform(row)))
+				serialized, err := json.MarshalWithOptions(
+					filterColumns(request.Columns, transform(row)),
+					getJsonOptsForTimezone(request.Timezone))
 				if err != nil {
 					return
 				}
@@ -540,3 +544,25 @@ func filterColumns(columns []string, row *ordereddict.Dict) *ordereddict.Dict {
 	}
 	return new_row
 }
+
+func getJsonOptsForTimezone(timezone string) *json.EncOpts {
+	if timezone == "" {
+		return vjson.NoEncOpts
+	}
+
+	loc := time.UTC
+	if timezone != "" {
+		loc, _ = time.LoadLocation(timezone)
+	}
+
+	return vjson.NewEncOpts().
+		WithCallback(time.Time{},
+			func(v interface{}, opts *json.EncOpts) ([]byte, error) {
+				switch t := v.(type) {
+				case time.Time:
+					return t.In(loc).MarshalJSON()
+				}
+				return nil, json.EncoderCallbackSkip
+			})
+
+}
diff --git a/api/hunts.go b/api/hunts.go
@@ -14,6 +14,7 @@ import (
 	api_proto "www.velocidex.com/golang/velociraptor/api/proto"
 	"www.velocidex.com/golang/velociraptor/file_store/csv"
 	"www.velocidex.com/golang/velociraptor/json"
+	vjson "www.velocidex.com/golang/velociraptor/json"
 	"www.velocidex.com/golang/velociraptor/logging"
 	"www.velocidex.com/golang/velociraptor/services"
 	"www.velocidex.com/golang/velociraptor/services/hunt_dispatcher"
@@ -72,11 +73,12 @@ func (self *ApiServer) GetHuntFlows(
 			flow.Context.ClientId,
 			services.GetHostname(ctx, org_config_obj, flow.Context.ClientId),
 			flow.Context.SessionId,
-			csv.AnyToString(flow.Context.StartTime / 1000),
+			csv.AnyToString(flow.Context.StartTime/1000, vjson.NoEncOpts),
 			flow.Context.State.String(),
-			csv.AnyToString(flow.Context.ExecutionDuration / 1000000000),
-			csv.AnyToString(flow.Context.TotalUploadedBytes),
-			csv.AnyToString(flow.Context.TotalCollectedRows)}
+			csv.AnyToString(flow.Context.ExecutionDuration/1000000000,
+				vjson.NoEncOpts),
+			csv.AnyToString(flow.Context.TotalUploadedBytes, vjson.NoEncOpts),
+			csv.AnyToString(flow.Context.TotalCollectedRows, vjson.NoEncOpts)}
 
 		result.Rows = append(result.Rows, &api_proto.Row{Cell: row_data})
 

diff --git a/api/proto/csv.pb.go b/api/proto/csv.pb.go
diff --git a/api/proto/csv.proto b/api/proto/csv.proto
@@ -62,6 +62,9 @@ message GetTableRequest {
     // The org id may be specified in the query string - The protobuf
     // is normally parsed from the query string directly.
     string org_id = 23;
+
+    // The required timezone to export in.
+    string timezone = 24;
 }
 
 message Row {

diff --git a/api/vql.go b/api/vql.go
@@ -23,6 +23,7 @@ import (
 	api_proto "www.velocidex.com/golang/velociraptor/api/proto"
 	config_proto "www.velocidex.com/golang/velociraptor/config/proto"
 	"www.velocidex.com/golang/velociraptor/file_store/csv"
+	"www.velocidex.com/golang/velociraptor/json"
 	"www.velocidex.com/golang/velociraptor/logging"
 	"www.velocidex.com/golang/velociraptor/services"
 	"www.velocidex.com/golang/velociraptor/vql/acl_managers"
@@ -69,7 +70,8 @@ func RunVQL(
 			if !pres {
 				value = ""
 			}
-			new_row.Cell = append(new_row.Cell, csv.AnyToString(value))
+			new_row.Cell = append(new_row.Cell,
+				csv.AnyToString(value, json.NoEncOpts))
 		}
 
 		result.Rows = append(result.Rows, new_row)

diff --git a/bin/query.go b/bin/query.go
@@ -127,7 +127,7 @@ func outputCSV(ctx context.Context,
 		10, *max_wait)
 
 	csv_writer := csv.GetCSVAppender(config_obj,
-		scope, &StdoutWrapper{out}, true /* write_headers */)
+		scope, &StdoutWrapper{out}, csv.WriteHeaders, json.NoEncOpts)
 	defer csv_writer.Close()
 
 	for result := range result_chan {
@@ -241,7 +241,8 @@ func doRemoteQuery(
 			scope := vql_subsystem.MakeScope()
 
 			csv_writer := csv.GetCSVAppender(config_obj,
-				scope, &StdoutWrapper{os.Stdout}, true /* write_headers */)
+				scope, &StdoutWrapper{os.Stdout},
+				csv.WriteHeaders, json.NoEncOpts)
 			defer csv_writer.Close()
 
 			for _, row := range rows {

diff --git a/config/proto/config.pb.go b/config/proto/config.pb.go
diff --git a/file_store/csv/utils.go b/file_store/csv/utils.go
@@ -29,9 +29,15 @@ import (
 	"github.com/Velocidex/ordereddict"
 	config_proto "www.velocidex.com/golang/velociraptor/config/proto"
 	"www.velocidex.com/golang/velociraptor/file_store/api"
+	"www.velocidex.com/golang/velociraptor/json"
 	"www.velocidex.com/golang/vfilter"
 )
 
+const (
+	WriteHeaders     = true
+	DoNotWriteHeader = false
+)
+
 type CSVWriter struct {
 	row_chan chan vfilter.Row
 	wg       sync.WaitGroup
@@ -89,7 +95,8 @@ func GetCSVReader(ctx context.Context, fd api.FileReader) CSVReader {
 
 func GetCSVAppender(
 	config_obj *config_proto.Config,
-	scope vfilter.Scope, fd io.Writer, write_headers bool) *CSVWriter {
+	scope vfilter.Scope, fd io.Writer,
+	write_headers bool, opts *json.EncOpts) *CSVWriter {
 	result := &CSVWriter{
 		row_chan: make(chan vfilter.Row),
 		wg:       sync.WaitGroup{},
@@ -142,7 +149,7 @@ func GetCSVAppender(
 					item, _ := scope.Associative(row, column)
 					csv_row = append(csv_row, item)
 				}
-				err := w.WriteAny(csv_row)
+				err := w.WriteAny(csv_row, opts)
 				if err != nil {
 					return
 				}
@@ -160,25 +167,27 @@ func GetCSVAppender(
 
 func GetCSVWriter(
 	config_obj *config_proto.Config,
-	scope vfilter.Scope, fd api.FileWriter) (*CSVWriter, error) {
+	scope vfilter.Scope, fd api.FileWriter,
+	opts *json.EncOpts) (*CSVWriter, error) {
 	// Seek to the end of the file.
 	length, err := fd.Size()
 	if err != nil {
 		return nil, err
 	}
-	return GetCSVAppender(config_obj, scope, fd, length == 0), nil
+	return GetCSVAppender(config_obj, scope, fd, length == 0, opts), nil
 }
 
 func EncodeToCSV(
 	config_obj *config_proto.Config,
-	scope vfilter.Scope, v interface{}) (string, error) {
+	scope vfilter.Scope, v interface{},
+	opts *json.EncOpts) (string, error) {
 	slice := reflect.ValueOf(v)
 	if slice.Type().Kind() != reflect.Slice {
 		return "", errors.New("EncodeToCSV - should be a list of rows")
 	}
 
 	buffer := &bytes.Buffer{}
-	writer := GetCSVAppender(config_obj, scope, buffer, true)
+	writer := GetCSVAppender(config_obj, scope, buffer, true, opts)
 
 	for i := 0; i < slice.Len(); i++ {
 		value := slice.Index(i).Interface()