diff --git a/cmd/proxy.go b/cmd/proxy.go index 030c85205..544ec0326 100644 --- a/cmd/proxy.go +++ b/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 }, } diff --git a/cmd/proxyRunner.go b/cmd/proxyRunner.go index 7ea95a2ab..809dfb2ff 100644 --- a/cmd/proxyRunner.go +++ b/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) } } diff --git a/cmd/scripts.go b/cmd/scripts.go new file mode 100644 index 000000000..edf046dc6 --- /dev/null +++ b/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) +} diff --git a/cmd/tapRunner.go b/cmd/tapRunner.go index 403f67dd1..1314fe480 100644 --- a/cmd/tapRunner.go +++ b/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() diff --git a/config/configStructs/scriptingConfig.go b/config/configStructs/scriptingConfig.go index 31245ba44..34b0021b4 100644 --- a/config/configStructs/scriptingConfig.go +++ b/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/robertkrimen/otto/file" - "github.com/robertkrimen/otto/parser" + "github.com/kubeshark/kubeshark/misc" "github.com/rs/zerolog/log" ) @@ -17,12 +14,7 @@ type ScriptingConfig struct { Source string `yaml:"source" default:""` } -type Script struct { - Title string `json:"title"` - Code string `json:"code"` -} - -func (config *ScriptingConfig) GetScripts() (scripts []*Script, err error) { +func (config *ScriptingConfig) GetScripts() (scripts []*misc.Script, err error) { if config.Source == "" { return } @@ -38,38 +30,13 @@ func (config *ScriptingConfig) GetScripts() (scripts []*Script, err error) { continue } - filename := f.Name() - var body []byte - path := filepath.Join(config.Source, filename) - body, err = os.ReadFile(path) + var script *misc.Script + path := filepath.Join(config.Source, f.Name()) + script, err = misc.ReadScriptFile(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 - } - - scripts = append(scripts, &Script{ - Title: title, - Code: code, - }) + scripts = append(scripts, script) log.Info().Str("path", path).Msg("Found script:") } diff --git a/go.mod b/go.mod index 51b17c415..2b2f27145 100644 --- a/go.mod +++ b/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 diff --git a/internal/connect/hub.go b/internal/connect/hub.go index 9b42e0ff6..1fa1040e0 100644 --- a/internal/connect/hub.go +++ b/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 { - log.Error().Err(err).Msg("Failed to marshal the payload:") + if constsMarshalled, err := json.Marshal(consts); err != nil { + 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 { - log.Error().Err(err).Msg("Failed to marshal the payload:") + 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 { 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() { diff --git a/misc/scripting.go b/misc/scripting.go new file mode 100644 index 000000000..03b22b222 --- /dev/null +++ b/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 +}