✨ Add scripts command

cmd/proxy.go

@@ -11,7 +11,7 @@ var proxyCmd = &cobra.Command{
 	Use:   "proxy",
 	Short: "Open the web UI (front-end) in the browser via proxy/port-forward",
 	RunE: func(cmd *cobra.Command, args []string) error {
-		runProxy()
+		runProxy(true)
 		return nil
 	},
 }

cmd/proxyRunner.go

@@ -14,7 +14,7 @@ import (
 	"github.com/rs/zerolog/log"
 )
 
-func runProxy() {
+func runProxy(block bool) {
 	kubernetesProvider, err := getKubernetesProviderForCli()
 	if err != nil {
 		return
@@ -64,7 +64,7 @@ func runProxy() {
 	var establishedProxy bool
 
 	hubUrl := kubernetes.GetLocalhostOnPort(config.Config.Tap.Proxy.Hub.SrcPort)
-	response, err := http.Get(fmt.Sprintf("%s/", hubUrl))
+	response, err := http.Get(fmt.Sprintf("%s/echo", hubUrl))
 	if err == nil && response.StatusCode == 200 {
 		log.Info().
 			Str("service", kubernetes.HubServiceName).
@@ -123,7 +123,7 @@ func runProxy() {
 		okToOpen("Kubeshark", frontUrl, false)
 	}
 
-	if establishedProxy {
+	if establishedProxy && block {
 		utils.WaitForTermination(ctx, cancel)
 	}
 }

cmd/scripts.go

@@ -0,0 +1,144 @@
+package cmd
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+
+	"github.com/creasty/defaults"
+	"github.com/fsnotify/fsnotify"
+	"github.com/kubeshark/kubeshark/config"
+	"github.com/kubeshark/kubeshark/config/configStructs"
+	"github.com/kubeshark/kubeshark/internal/connect"
+	"github.com/kubeshark/kubeshark/kubernetes"
+	"github.com/kubeshark/kubeshark/misc"
+	"github.com/kubeshark/kubeshark/utils"
+	"github.com/rs/zerolog/log"
+	"github.com/spf13/cobra"
+)
+
+var scriptsCmd = &cobra.Command{
+	Use:   "scripts",
+	Short: "Watch the `scripting.source` directory for changes and update the scripts.",
+	RunE: func(cmd *cobra.Command, args []string) error {
+		runScripts()
+		return nil
+	},
+}
+
+func init() {
+	rootCmd.AddCommand(scriptsCmd)
+
+	defaultTapConfig := configStructs.TapConfig{}
+	if err := defaults.Set(&defaultTapConfig); err != nil {
+		log.Debug().Err(err).Send()
+	}
+
+	scriptsCmd.Flags().Uint16(configStructs.ProxyHubPortLabel, defaultTapConfig.Proxy.Hub.SrcPort, "Provide a custom port for the Hub.")
+	scriptsCmd.Flags().String(configStructs.ProxyHostLabel, defaultTapConfig.Proxy.Host, "Provide a custom host for the Hub.")
+}
+
+func runScripts() {
+	if config.Config.Scripting.Source == "" {
+		log.Error().Msg("`scripting.source` field is empty.")
+		return
+	}
+
+	hubUrl := kubernetes.GetLocalhostOnPort(config.Config.Tap.Proxy.Hub.SrcPort)
+	response, err := http.Get(fmt.Sprintf("%s/echo", hubUrl))
+	if err != nil || response.StatusCode != 200 {
+		log.Info().Msg(fmt.Sprintf(utils.Yellow, "Couldn't connect to Hub. Establishing proxy..."))
+		runProxy(false)
+	}
+
+	files := make(map[string]int64)
+
+	connector = connect.NewConnector(kubernetes.GetLocalhostOnPort(config.Config.Tap.Proxy.Hub.SrcPort), connect.DefaultRetries, connect.DefaultTimeout)
+
+	scripts, err := config.Config.Scripting.GetScripts()
+	if err != nil {
+		log.Error().Err(err).Send()
+		return
+	}
+
+	for _, script := range scripts {
+		index, err := connector.PostScript(script)
+		if err != nil {
+			log.Error().Err(err).Send()
+			return
+		}
+
+		files[script.Path] = index
+	}
+
+	watcher, err := fsnotify.NewWatcher()
+	if err != nil {
+		log.Error().Err(err).Send()
+		return
+	}
+	defer watcher.Close()
+
+	go func() {
+		for {
+			select {
+			// watch for events
+			case event := <-watcher.Events:
+				switch event.Op {
+				case fsnotify.Create:
+					script, err := misc.ReadScriptFile(event.Name)
+					if err != nil {
+						log.Error().Err(err).Send()
+						return
+					}
+
+					index, err := connector.PostScript(script)
+					if err != nil {
+						log.Error().Err(err).Send()
+						return
+					}
+
+					files[script.Path] = index
+
+				case fsnotify.Write:
+					index := files[event.Name]
+					script, err := misc.ReadScriptFile(event.Name)
+					if err != nil {
+						log.Error().Err(err).Send()
+						return
+					}
+
+					err = connector.PutScript(script, index)
+					if err != nil {
+						log.Error().Err(err).Send()
+						return
+					}
+
+				case fsnotify.Rename:
+					index := files[event.Name]
+					err := connector.DeleteScript(index)
+					if err != nil {
+						log.Error().Err(err).Send()
+						return
+					}
+
+				default:
+					// pass
+				}
+
+			// watch for errors
+			case err := <-watcher.Errors:
+				log.Error().Err(err).Send()
+			}
+		}
+	}()
+
+	if err := watcher.Add(config.Config.Scripting.Source); err != nil {
+		log.Error().Err(err).Send()
+	}
+
+	log.Info().Str("directory", config.Config.Scripting.Source).Msg("Watching files against changes:")
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	utils.WaitForTermination(ctx, cancel)
+}

cmd/tapRunner.go

@@ -443,14 +443,17 @@ func postHubStarted(ctx context.Context, kubernetesProvider *kubernetes.Provider
 	// Scripting
 	connector.PostConsts(config.Config.Scripting.Consts)
 
-	var scripts []*configStructs.Script
+	var scripts []*misc.Script
 	scripts, err = config.Config.Scripting.GetScripts()
 	if err != nil {
 		log.Error().Err(err).Send()
 	}
 
 	for _, script := range scripts {
-		connector.PostScript(script)
+		_, err = connector.PostScript(script)
+		if err != nil {
+			log.Error().Err(err).Send()
+		}
 	}
 
 	connector.PostScriptDone()

config/configStructs/scriptingConfig.go

@@ -3,12 +3,9 @@ package configStructs
 import (
 	"io/fs"
 	"io/ioutil"
-	"os"
 	"path/filepath"
 
-	"github.com/robertkrimen/otto/ast"
+	"github.com/kubeshark/kubeshark/misc"
-	"github.com/robertkrimen/otto/file"
-	"github.com/robertkrimen/otto/parser"
 	"github.com/rs/zerolog/log"
 )
 
@@ -17,12 +14,7 @@ type ScriptingConfig struct {
 	Source string                 `yaml:"source" default:""`
 }
 
-type Script struct {
+func (config *ScriptingConfig) GetScripts() (scripts []*misc.Script, err error) {
-	Title string `json:"title"`
-	Code  string `json:"code"`
-}
-
-func (config *ScriptingConfig) GetScripts() (scripts []*Script, err error) {
 	if config.Source == "" {
 		return
 	}
@@ -38,38 +30,13 @@ func (config *ScriptingConfig) GetScripts() (scripts []*Script, err error) {
 			continue
 		}
 
-		filename := f.Name()
+		var script *misc.Script
-		var body []byte
+		path := filepath.Join(config.Source, f.Name())
-		path := filepath.Join(config.Source, filename)
+		script, err = misc.ReadScriptFile(path)
-		body, err = os.ReadFile(path)
 		if err != nil {
 			return
 		}
-		content := string(body)
+		scripts = append(scripts, script)
-
-		var program *ast.Program
-		program, err = parser.ParseFile(nil, filename, content, parser.StoreComments)
-		if err != nil {
-			return
-		}
-
-		var title string
-		code := content
-
-		var idx0 file.Idx
-		for node, comments := range program.Comments {
-			if (idx0 > 0 && node.Idx0() > idx0) || len(comments) == 0 {
-				continue
-			}
-			idx0 = node.Idx0()
-
-			title = comments[0].Text
-		}
-
-		scripts = append(scripts, &Script{
-			Title: title,
-			Code:  code,
-		})
 
 		log.Info().Str("path", path).Msg("Found script:")
 	}

go.mod

@@ -7,6 +7,7 @@ require (
 	github.com/docker/docker v20.10.22+incompatible
 	github.com/docker/go-connections v0.4.0
 	github.com/docker/go-units v0.4.0
+	github.com/fsnotify/fsnotify v1.5.1
 	github.com/google/go-github/v37 v37.0.0
 	github.com/gorilla/websocket v1.4.2
 	github.com/kubeshark/base v0.6.3

internal/connect/hub.go

@@ -3,12 +3,13 @@ package connect
 import (
 	"bytes"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"net/http"
 	"net/url"
 	"time"
 
-	"github.com/kubeshark/kubeshark/config/configStructs"
+	"github.com/kubeshark/kubeshark/misc"
 	"github.com/kubeshark/kubeshark/utils"
 
 	"github.com/rs/zerolog/log"
@@ -186,13 +187,13 @@ func (connector *Connector) PostConsts(consts map[string]interface{}) {
 
 	postConstsUrl := fmt.Sprintf("%s/scripts/consts", connector.url)
 
-	if payloadMarshalled, err := json.Marshal(consts); err != nil {
+	if constsMarshalled, err := json.Marshal(consts); err != nil {
-		log.Error().Err(err).Msg("Failed to marshal the payload:")
+		log.Error().Err(err).Msg("Failed to marshal the consts:")
 	} else {
 		ok := false
 		for !ok {
 			var resp *http.Response
-			if resp, err = utils.Post(postConstsUrl, "application/json", bytes.NewBuffer(payloadMarshalled), connector.client); err != nil || resp.StatusCode != http.StatusOK {
+			if resp, err = utils.Post(postConstsUrl, "application/json", bytes.NewBuffer(constsMarshalled), connector.client); err != nil || resp.StatusCode != http.StatusOK {
 				if _, ok := err.(*url.Error); ok {
 					break
 				}
@@ -206,27 +207,124 @@ func (connector *Connector) PostConsts(consts map[string]interface{}) {
 	}
 }
 
-func (connector *Connector) PostScript(script *configStructs.Script) {
+func (connector *Connector) PostScript(script *misc.Script) (index int64, err error) {
 	postScriptUrl := fmt.Sprintf("%s/scripts", connector.url)
 
-	if payloadMarshalled, err := json.Marshal(script); err != nil {
+	var scriptMarshalled []byte
-		log.Error().Err(err).Msg("Failed to marshal the payload:")
+	if scriptMarshalled, err = json.Marshal(script); err != nil {
+		log.Error().Err(err).Msg("Failed to marshal the script:")
 	} else {
 		ok := false
 		for !ok {
 			var resp *http.Response
-			if resp, err = utils.Post(postScriptUrl, "application/json", bytes.NewBuffer(payloadMarshalled), connector.client); err != nil || resp.StatusCode != http.StatusOK {
+			if resp, err = utils.Post(postScriptUrl, "application/json", bytes.NewBuffer(scriptMarshalled), connector.client); err != nil || resp.StatusCode != http.StatusOK {
 				if _, ok := err.(*url.Error); ok {
 					break
 				}
-				log.Warn().Err(err).Msg("Failed sending the script to Hub:")
+				log.Warn().Err(err).Msg("Failed creating script Hub:")
 			} else {
 				ok = true
-				log.Info().Interface("script", script).Msg("Reported script to Hub:")
+
+				var j map[string]interface{}
+				err = json.NewDecoder(resp.Body).Decode(&j)
+				if err != nil {
+					return
+				}
+
+				val, ok := j["key"]
+				if !ok {
+					err = errors.New("Response does not contain `key` field!")
+					return
+				}
+
+				index = int64(val.(float64))
+
+				log.Info().Int("index", int(index)).Interface("script", script).Msg("Created script on Hub:")
 			}
 			time.Sleep(time.Second)
 		}
 	}
+
+	return
+}
+
+func (connector *Connector) PutScript(script *misc.Script, index int64) (err error) {
+	putScriptUrl := fmt.Sprintf("%s/scripts/%d", connector.url, index)
+
+	var scriptMarshalled []byte
+	if scriptMarshalled, err = json.Marshal(script); err != nil {
+		log.Error().Err(err).Msg("Failed to marshal the script:")
+	} else {
+		ok := false
+		for !ok {
+			client := &http.Client{}
+
+			var req *http.Request
+			req, err = http.NewRequest(http.MethodPut, putScriptUrl, bytes.NewBuffer(scriptMarshalled))
+			if err != nil {
+				log.Error().Err(err).Send()
+				return
+			}
+			req.Header.Set("Content-Type", "application/json")
+
+			var resp *http.Response
+			resp, err = client.Do(req)
+			if err != nil {
+				log.Error().Err(err).Send()
+				return
+			}
+
+			if resp.StatusCode != http.StatusOK {
+				if _, ok := err.(*url.Error); ok {
+					break
+				}
+				log.Warn().Err(err).Msg("Failed updating script on Hub:")
+			} else {
+				ok = true
+				log.Info().Int("index", int(index)).Interface("script", script).Msg("Updated script on Hub:")
+			}
+			time.Sleep(time.Second)
+		}
+	}
+
+	return
+}
+
+func (connector *Connector) DeleteScript(index int64) (err error) {
+	deleteScriptUrl := fmt.Sprintf("%s/scripts/%d", connector.url, index)
+
+	ok := false
+	for !ok {
+		client := &http.Client{}
+
+		var req *http.Request
+		req, err = http.NewRequest(http.MethodDelete, deleteScriptUrl, nil)
+		if err != nil {
+			log.Error().Err(err).Send()
+			return
+		}
+		req.Header.Set("Content-Type", "application/json")
+
+		var resp *http.Response
+		resp, err = client.Do(req)
+		if err != nil {
+			log.Error().Err(err).Send()
+			return
+		}
+
+		if resp.StatusCode != http.StatusOK {
+			if _, ok := err.(*url.Error); ok {
+				break
+			}
+			log.Warn().Err(err).Msg("Failed deleting script on Hub:")
+		} else {
+			ok = true
+			log.Info().Int("index", int(index)).Msg("Deleted script on Hub:")
+		}
+		time.Sleep(time.Second)
+	}
+
+	return
 }
 
 func (connector *Connector) PostScriptDone() {

misc/scripting.go

@@ -0,0 +1,53 @@
+package misc
+
+import (
+	"os"
+	"path/filepath"
+
+	"github.com/robertkrimen/otto/ast"
+	"github.com/robertkrimen/otto/file"
+	"github.com/robertkrimen/otto/parser"
+)
+
+type Script struct {
+	Path  string `json:"path"`
+	Title string `json:"title"`
+	Code  string `json:"code"`
+}
+
+func ReadScriptFile(path string) (script *Script, err error) {
+	filename := filepath.Base(path)
+	var body []byte
+	body, err = os.ReadFile(path)
+	if err != nil {
+		return
+	}
+	content := string(body)
+
+	var program *ast.Program
+	program, err = parser.ParseFile(nil, filename, content, parser.StoreComments)
+	if err != nil {
+		return
+	}
+
+	var title string
+	code := content
+
+	var idx0 file.Idx
+	for node, comments := range program.Comments {
+		if (idx0 > 0 && node.Idx0() > idx0) || len(comments) == 0 {
+			continue
+		}
+		idx0 = node.Idx0()
+
+		title = comments[0].Text
+	}
+
+	script = &Script{
+		Path:  path,
+		Title: title,
+		Code:  code,
+	}
+
+	return
+}