Store output of notebook in filestore. (Velocidex#354)

Previously output was stored in the notebook datastore object but this object is limited to 64kb in the mysql data store implementation. This change makes the output of VQL queries write to the filestore the output, keeping the notebook cell small.
tlinman67 · May 10, 2020 · d0e06ae · d0e06ae
1 parent cee9c79
commit d0e06ae
Show file tree

Hide file tree

Showing 15 changed files with 487 additions and 298 deletions.
diff --git a/api/csv.go b/api/csv.go
@@ -23,6 +23,7 @@ import (
 	"www.velocidex.com/golang/velociraptor/file_store/api"
 	"www.velocidex.com/golang/velociraptor/file_store/csv"
 	"www.velocidex.com/golang/velociraptor/paths"
+	"www.velocidex.com/golang/velociraptor/reporting"
 	"www.velocidex.com/golang/velociraptor/result_sets"
 
 	api_proto "www.velocidex.com/golang/velociraptor/api/proto"
@@ -57,6 +58,9 @@ func getTable(
 		}
 	} else if in.HuntId != "" && in.Type == "clients" {
 		path_manager = paths.NewHuntPathManager(in.HuntId).Clients()
+	} else if in.NotebookId != "" && in.CellId != "" {
+		path_manager = reporting.NewNotebookPathManager(in.NotebookId).Cell(
+			in.CellId).QueryStorage(in.TableId)
 	}
 
 	result := &api_proto.GetTableResponse{}

diff --git a/api/notebooks.go b/api/notebooks.go
@@ -54,9 +54,10 @@ func (self *ApiServer) GetNotebooks(
 
 	// We want a single notebook metadata.
 	if in.NotebookId != "" {
+		notebook_path_manager := reporting.NewNotebookPathManager(
+			in.NotebookId)
 		notebook := &api_proto.NotebookMetadata{}
-		err := db.GetSubject(self.config,
-			reporting.GetNotebookPath(in.NotebookId),
+		err := db.GetSubject(self.config, notebook_path_manager.Path(),
 			notebook)
 		if err != nil {
 			logging.GetLogger(
@@ -81,10 +82,9 @@ func (self *ApiServer) GetNotebooks(
 	}
 
 	// List all available notebooks
+
 	notebook_urns, err := db.ListChildren(
-		self.config,
-		path.Dir(reporting.GetNotebookPath("X")),
-		in.Offset, in.Count)
+		self.config, reporting.NotebookDir(), in.Offset, in.Count)
 	if err != nil {
 		return nil, err
 	}
@@ -179,8 +179,8 @@ func (self *ApiServer) NewNotebook(
 		return nil, err
 	}
 
-	err = db.SetSubject(self.config, reporting.GetNotebookPath(
-		in.NotebookId), in)
+	notebook_path_manager := reporting.NewNotebookPathManager(in.NotebookId)
+	err = db.SetSubject(self.config, notebook_path_manager.Path(), in)
 
 	// Add a new single cell to the notebook.
 	new_cell_request := &api_proto.NotebookCellRequest{
@@ -222,8 +222,8 @@ func (self *ApiServer) NewNotebookCell(
 	}
 
 	notebook := &api_proto.NotebookMetadata{}
-	err = db.GetSubject(self.config, reporting.GetNotebookPath(
-		in.NotebookId), notebook)
+	notebook_path_manager := reporting.NewNotebookPathManager(in.NotebookId)
+	err = db.GetSubject(self.config, notebook_path_manager.Path(), notebook)
 	if err != nil {
 		return nil, err
 	}
@@ -259,8 +259,7 @@ func (self *ApiServer) NewNotebookCell(
 
 	notebook.CellMetadata = new_cell_md
 
-	err = db.SetSubject(self.config, reporting.GetNotebookPath(
-		in.NotebookId), notebook)
+	err = db.SetSubject(self.config, notebook_path_manager.Path(), notebook)
 	if err != nil {
 		return nil, err
 	}
@@ -305,8 +304,8 @@ func (self *ApiServer) UpdateNotebook(
 	}
 
 	old_notebook := &api_proto.NotebookMetadata{}
-	err = db.GetSubject(self.config, reporting.GetNotebookPath(
-		in.NotebookId), old_notebook)
+	notebook_path_manager := reporting.NewNotebookPathManager(in.NotebookId)
+	err = db.GetSubject(self.config, notebook_path_manager.Path(), old_notebook)
 	if err != nil {
 		return nil, err
 	}
@@ -317,8 +316,7 @@ func (self *ApiServer) UpdateNotebook(
 
 	in.ModifiedTime = time.Now().Unix()
 
-	err = db.SetSubject(self.config, reporting.GetNotebookPath(
-		in.NotebookId), in)
+	err = db.SetSubject(self.config, notebook_path_manager.Path(), in)
 
 	return in, err
 }
@@ -354,8 +352,9 @@ func (self *ApiServer) GetNotebookCell(
 	}
 
 	notebook := &api_proto.NotebookCell{}
+	notebook_path_manager := reporting.NewNotebookPathManager(in.NotebookId)
 	err = db.GetSubject(self.config,
-		reporting.GetNotebookCellPath(in.NotebookId, in.CellId),
+		notebook_path_manager.Cell(in.CellId).Path(),
 		notebook)
 
 	// Cell does not exist, make it a default cell.
@@ -370,7 +369,7 @@ func (self *ApiServer) GetNotebookCell(
 
 		// And store it for next time.
 		err = db.SetSubject(self.config,
-			reporting.GetNotebookCellPath(in.NotebookId, in.CellId),
+			notebook_path_manager.Cell(in.CellId).Path(),
 			notebook)
 		if err != nil {
 			return nil, err
@@ -407,9 +406,10 @@ func (self *ApiServer) UpdateNotebookCell(
 		return nil, status.Error(codes.PermissionDenied,
 			"User is not allowed to edit notebooks.")
 	}
-
+	notebook_path_manager := reporting.NewNotebookPathManager(in.NotebookId)
 	tmpl, err := reporting.NewGuiTemplateEngine(
 		self.config, ctx, user_name, /* principal */
+		notebook_path_manager.Cell(in.CellId),
 		"Server.Internal.ArtifactDescription")
 	if err != nil {
 		return nil, err
@@ -460,16 +460,14 @@ func (self *ApiServer) UpdateNotebookCell(
 
 	// And store it for next time.
 	err = db.SetSubject(self.config,
-		reporting.GetNotebookCellPath(in.NotebookId, in.CellId),
+		notebook_path_manager.Cell(in.CellId).Path(),
 		notebook_cell)
 	if err != nil {
 		return nil, err
 	}
 
 	notebook := &api_proto.NotebookMetadata{}
-	err = db.GetSubject(self.config,
-		reporting.GetNotebookPath(in.NotebookId),
-		notebook)
+	err = db.GetSubject(self.config, notebook_path_manager.Path(), notebook)
 	if err != nil {
 		return nil, err
 	}
@@ -487,9 +485,7 @@ func (self *ApiServer) UpdateNotebookCell(
 	}
 	notebook.CellMetadata = new_cell_md
 
-	err = db.SetSubject(self.config, reporting.GetNotebookPath(
-		in.NotebookId), notebook)
-
+	err = db.SetSubject(self.config, notebook_path_manager.Path(), notebook)
 	return notebook_cell, err
 }
 
@@ -558,17 +554,14 @@ func (self *ApiServer) CreateNotebookDownloadFile(
 	}
 
 	notebook := &api_proto.NotebookMetadata{}
-	err = db.GetSubject(self.config,
-		reporting.GetNotebookPath(in.NotebookId),
-		notebook)
+	notebook_path_manager := reporting.NewNotebookPathManager(in.NotebookId)
+	err = db.GetSubject(self.config, notebook_path_manager.Path(), notebook)
 	if err != nil {
 		return nil, err
 	}
 
 	file_store_factory := file_store.GetFileStore(self.config)
-	filename := path.Join("/downloads/notebooks", notebook.NotebookId,
-		fmt.Sprintf("%s-%s.html", notebook.NotebookId,
-			time.Now().Format("20060102150405Z")))
+	filename := notebook_path_manager.HtmlExport()
 
 	lock_file, err := file_store_factory.WriteFile(filename + ".lock")
 	if err != nil {
@@ -581,12 +574,16 @@ func (self *ApiServer) CreateNotebookDownloadFile(
 		return nil, err
 	}
 
+	// Allow 1 hour to export the notebook.
+	sub_ctx, cancel := context.WithTimeout(context.Background(), time.Hour)
+
 	go func() {
 		defer file_store_factory.Delete(filename + ".lock")
 		defer writer.Close()
+		defer cancel()
 
 		err := reporting.ExportNotebookToHTML(
-			self.config, notebook.NotebookId, writer)
+			sub_ctx, self.config, notebook.NotebookId, writer)
 		if err != nil {
 			logger := logging.GetLogger(self.config, &logging.GUIComponent)
 			logger.WithFields(logrus.Fields{

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
@@ -28,6 +28,10 @@ message GetTableRequest {
     string type = 7;
 
     string hunt_id = 8;
+
+    string notebook_id = 9;
+    string cell_id = 10;
+    int64 table_id = 11;
 }
 
 message Row {

diff --git a/api/reports.go b/api/reports.go
@@ -18,11 +18,11 @@ func getReport(ctx context.Context,
 	*api_proto.GetReportResponse, error) {
 
 	template_engine, err := reporting.NewGuiTemplateEngine(
-		config_obj, ctx, principal, in.Artifact)
+		config_obj, ctx, principal, nil, in.Artifact)
 	if err != nil {
 		if strings.HasPrefix(in.Artifact, "Custom.") {
 			template_engine, err = reporting.NewGuiTemplateEngine(
-				config_obj, ctx, principal,
+				config_obj, ctx, principal, nil,
 				strings.TrimPrefix(in.Artifact, "Custom."))
 		}
 		if err != nil {

diff --git a/bin/notebook.go b/bin/notebook.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 
 	kingpin "gopkg.in/alecthomas/kingpin.v2"
@@ -20,9 +21,10 @@ func doExportNotebook() {
 	config_obj, err := get_server_config(*config_path)
 	kingpin.FatalIfError(err, "Unable to load config file")
 
+	ctx := context.Background()
 	writer := &bytes.Buffer{}
 	err = reporting.ExportNotebookToHTML(
-		config_obj, *notebook_command_notebook_id, writer)
+		ctx, config_obj, *notebook_command_notebook_id, writer)
 	kingpin.FatalIfError(err, "Generating report")
 
 	fmt.Println(writer.String())