Added vql functions to manipulate notebooks (Velocidex#3325)

This allows for automations around creating custom notebooks. Also added tests and refactored notebook code: - Notebooks uploads have been moved to the cells so they can be versioned. - Notebook attachments are not versions and exist in the notebook top level.
Genio40921 · Mar 6, 2024 · 913367d · 913367d
1 parent 23fb6fb
commit 913367d
Show file tree

Hide file tree

Showing 49 changed files with 1,959 additions and 919 deletions.
diff --git a/api/api.go b/api/api.go
@@ -880,8 +880,12 @@ func (self *ApiServer) GetServerMonitoringState(
 			fmt.Sprintf("User is not allowed to read results (%v).", permissions))
 	}
 
-	result, err := getServerMonitoringState(org_config_obj)
-	return result, Status(self.verbose, err)
+	server_event_manager, err := services.GetServerEventManager(org_config_obj)
+	if err != nil {
+		return nil, Status(self.verbose, err)
+	}
+
+	return server_event_manager.Get(), nil
 }
 
 func (self *ApiServer) SetServerMonitoringState(
@@ -907,7 +911,12 @@ func (self *ApiServer) SetServerMonitoringState(
 			fmt.Sprintf("User is not allowed to modify artifacts (%v).", permissions))
 	}
 
-	err = setServerMonitoringState(ctx, org_config_obj, principal, in)
+	server_event_manager, err := services.GetServerEventManager(org_config_obj)
+	if err != nil {
+		return nil, Status(self.verbose, err)
+	}
+
+	err = server_event_manager.Update(ctx, org_config_obj, principal, in)
 	return in, Status(self.verbose, err)
 }
 

diff --git a/api/clients.go b/api/clients.go
@@ -1,19 +1,19 @@
 /*
-   Velociraptor - Dig Deeper
-   Copyright (C) 2019-2024 Rapid7 Inc.
+Velociraptor - Dig Deeper
+Copyright (C) 2019-2024 Rapid7 Inc.
 
-   This program is free software: you can redistribute it and/or modify
-   it under the terms of the GNU Affero General Public License as published
-   by the Free Software Foundation, either version 3 of the License, or
-   (at your option) any later version.
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU Affero General Public License as published
+by the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
 
-   This program is distributed in the hope that it will be useful,
-   but WITHOUT ANY WARRANTY; without even the implied warranty of
-   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-   GNU Affero General Public License for more details.
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU Affero General Public License for more details.
 
-   You should have received a copy of the GNU Affero General Public License
-   along with this program.  If not, see <https://www.gnu.org/licenses/>.
+You should have received a copy of the GNU Affero General Public License
+along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */
 package api
 
@@ -27,8 +27,6 @@ import (
 	"google.golang.org/protobuf/types/known/emptypb"
 	"www.velocidex.com/golang/velociraptor/acls"
 	api_proto "www.velocidex.com/golang/velociraptor/api/proto"
-	"www.velocidex.com/golang/velociraptor/datastore"
-	"www.velocidex.com/golang/velociraptor/paths"
 	"www.velocidex.com/golang/velociraptor/services"
 )
 
@@ -54,19 +52,36 @@ func (self *ApiServer) GetClientMetadata(
 			"User is not allowed to view clients.")
 	}
 
-	client_path_manager := paths.NewClientPathManager(in.ClientId)
-	db, err := datastore.GetDB(org_config_obj)
+	client_info_manager, err := services.GetClientInfoManager(org_config_obj)
 	if err != nil {
 		return nil, Status(self.verbose, err)
 	}
 
-	result := &api_proto.ClientMetadata{}
-	err = db.GetSubject(org_config_obj, client_path_manager.Metadata(), result)
+	result := &api_proto.ClientMetadata{
+		ClientId: in.ClientId,
+	}
+
+	client_metadata, err := client_info_manager.GetMetadata(ctx, in.ClientId)
 	if errors.Is(err, os.ErrNotExist) {
 		// Metadata not set, start with empty set.
 		err = nil
 	}
-	return result, err
+	if err != nil {
+		return nil, Status(self.verbose, err)
+	}
+
+	for _, k := range client_metadata.Keys() {
+		v, _ := client_metadata.GetString(k)
+		if v != "" {
+			result.Items = append(result.Items,
+				&api_proto.ClientMetadataItem{
+					Key:   k,
+					Value: v,
+				})
+		}
+	}
+
+	return result, nil
 }
 
 func (self *ApiServer) SetClientMetadata(

diff --git a/api/notebooks.go b/api/notebooks.go
@@ -1,30 +1,20 @@
 package api
 
 import (
-	"crypto/md5"
-	"crypto/sha256"
-	"encoding/hex"
 	"os"
 	"strings"
 	"sync"
 	"time"
 
 	"github.com/Velocidex/ordereddict"
 	errors "github.com/go-errors/errors"
-	"github.com/sirupsen/logrus"
 	context "golang.org/x/net/context"
 	"google.golang.org/protobuf/types/known/emptypb"
 	"www.velocidex.com/golang/velociraptor/acls"
 	api_proto "www.velocidex.com/golang/velociraptor/api/proto"
-	config_proto "www.velocidex.com/golang/velociraptor/config/proto"
-	"www.velocidex.com/golang/velociraptor/datastore"
-	file_store "www.velocidex.com/golang/velociraptor/file_store"
-	"www.velocidex.com/golang/velociraptor/file_store/path_specs"
 	"www.velocidex.com/golang/velociraptor/logging"
-	"www.velocidex.com/golang/velociraptor/paths"
-	"www.velocidex.com/golang/velociraptor/reporting"
 	"www.velocidex.com/golang/velociraptor/services"
-	"www.velocidex.com/golang/velociraptor/utils"
+	"www.velocidex.com/golang/velociraptor/vql/server/notebooks"
 )
 
 const (
@@ -62,7 +52,8 @@ func (self *ApiServer) GetNotebooks(
 
 	// We want a single notebook metadata.
 	if in.NotebookId != "" {
-		notebook_metadata, err := notebook_manager.GetNotebook(ctx, in.NotebookId, in.IncludeUploads)
+		notebook_metadata, err := notebook_manager.GetNotebook(
+			ctx, in.NotebookId, in.IncludeUploads)
 		// Handle the EOF especially: it means there is no such
 		// notebook and return an empty result set.
 		if errors.Is(err, os.ErrNotExist) ||
@@ -430,7 +421,13 @@ func (self *ApiServer) UploadNotebookAttachment(
 	if err != nil {
 		return nil, Status(self.verbose, err)
 	}
-	return notebook_manager.UploadNotebookAttachment(ctx, in)
+	res, err := notebook_manager.UploadNotebookAttachment(ctx, in)
+	if err != nil {
+		return nil, Status(self.verbose, err)
+	}
+
+	res.MimeType = detectMime([]byte(in.Data), true)
+	return res, nil
 }
 
 func (self *ApiServer) CreateNotebookDownloadFile(
@@ -453,144 +450,22 @@ func (self *ApiServer) CreateNotebookDownloadFile(
 			"User is not allowed to export notebooks.")
 	}
 
+	wg := &sync.WaitGroup{}
+
 	switch in.Type {
 	case "zip":
-		return &emptypb.Empty{}, exportZipNotebook(
-			org_config_obj, in.NotebookId, principal)
-	default:
-		return &emptypb.Empty{}, exportHTMLNotebook(
-			org_config_obj, in.NotebookId, principal)
-	}
-}
+		_, err := notebooks.ExportNotebookToZip(ctx,
+			org_config_obj, wg, in.NotebookId,
+			principal, in.PreferredName)
 
-// Create a portable notebook into a zip file.
-func exportZipNotebook(
-	config_obj *config_proto.Config,
-	notebook_id, principal string) error {
-	db, err := datastore.GetDB(config_obj)
-	if err != nil {
-		return err
-	}
+		return &emptypb.Empty{}, Status(self.verbose, err)
 
-	notebook := &api_proto.NotebookMetadata{}
-	notebook_path_manager := paths.NewNotebookPathManager(notebook_id)
-	err = db.GetSubject(config_obj, notebook_path_manager.Path(), notebook)
-	if err != nil {
-		return err
-	}
-
-	notebook_manager, err := services.GetNotebookManager(config_obj)
-	if err != nil {
-		return err
-	}
-	if !notebook_manager.CheckNotebookAccess(notebook, principal) {
-		return InvalidStatus("Notebook is not shared with user.")
-	}
-
-	filename := notebook_path_manager.ZipExport()
-
-	// Allow 1 hour to export the notebook.
-	sub_ctx, cancel := context.WithTimeout(context.Background(), time.Hour)
-
-	go func() {
-		defer cancel()
-
-		wg := &sync.WaitGroup{}
-
-		err := reporting.ExportNotebookToZip(
-			sub_ctx, config_obj, wg, notebook_path_manager)
-		if err != nil {
-			logger := logging.GetLogger(config_obj, &logging.GUIComponent)
-			logger.WithFields(logrus.Fields{
-				"notebook_id": notebook.NotebookId,
-				"export_file": filename,
-				"error":       err.Error(),
-			}).Error("CreateNotebookDownloadFile")
-			return
-		}
-
-		// Wait for the export to finish before we return.
-		wg.Wait()
-	}()
-
-	return nil
-}
-
-func exportHTMLNotebook(config_obj *config_proto.Config,
-	notebook_id, principal string) error {
-	db, err := datastore.GetDB(config_obj)
-	if err != nil {
-		return err
-	}
-
-	notebook := &api_proto.NotebookMetadata{}
-	notebook_path_manager := paths.NewNotebookPathManager(notebook_id)
-	err = db.GetSubject(config_obj, notebook_path_manager.Path(), notebook)
-	if err != nil {
-		return err
-	}
-
-	notebook_manager, err := services.GetNotebookManager(config_obj)
-	if err != nil {
-		return err
-	}
-
-	if !notebook_manager.CheckNotebookAccess(notebook, principal) {
-		return InvalidStatus("Notebook is not shared with user.")
-	}
-
-	file_store_factory := file_store.GetFileStore(config_obj)
-	filename := notebook_path_manager.HtmlExport()
-
-	writer, err := file_store_factory.WriteFile(filename)
-	if err != nil {
-		return err
-	}
-
-	sha_sum := sha256.New()
-	md5_sum := md5.New()
-	tee_writer := utils.NewTee(writer, sha_sum, md5_sum)
-
-	stats := &api_proto.ContainerStats{
-		Timestamp:  uint64(time.Now().Unix()),
-		Type:       "html",
-		Components: path_specs.AsGenericComponentList(filename),
-	}
-	stats_path := notebook_path_manager.PathStats(filename)
-
-	err = db.SetSubject(config_obj, stats_path, stats)
-	if err != nil {
-		return err
+	default:
+		_, err := notebooks.ExportNotebookToHTML(
+			org_config_obj, wg, in.NotebookId,
+			principal, in.PreferredName)
+		return &emptypb.Empty{}, Status(self.verbose, err)
 	}
-
-	// Allow 1 hour to export the notebook.
-	sub_ctx, cancel := context.WithTimeout(context.Background(), time.Hour)
-
-	go func() {
-		defer writer.Close()
-		defer cancel()
-
-		defer func() {
-			stats.Hash = hex.EncodeToString(sha_sum.Sum(nil))
-			stats.TotalDuration = uint64(time.Now().Unix()) - stats.Timestamp
-
-			db.SetSubject(config_obj, stats_path, stats)
-		}()
-
-		err := reporting.ExportNotebookToHTML(
-			sub_ctx, config_obj, notebook.NotebookId, tee_writer)
-		if err != nil {
-			logger := logging.GetLogger(config_obj, &logging.GUIComponent)
-			logger.WithFields(logrus.Fields{
-				"notebook_id": notebook.NotebookId,
-				"export_file": filename,
-				"error":       err.Error(),
-			}).Error("CreateNotebookDownloadFile")
-			return
-		}
-	}()
-
-	return nil
 }
 
 func (self *ApiServer) RemoveNotebookAttachment(