github/kubeshark
1
0
Fork 0
mirror of https://github.com/kubeshark/kubeshark.git synced 2025-07-05 04:18:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

✨ Add scripts command

Browse Source
This commit is contained in:
M. Mert Yildiran 2023-02-14 20:23:25 +03:00
parent 85da7f71ac
commit 41ba509428
No known key found for this signature in database
GPG Key ID: DA5D6DCBB758A461
8 changed files with 321 additions and 55 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
cmd/proxy.go
View File

@ -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
},
}

6
cmd/proxyRunner.go
View File

@ -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)
}
}

144
cmd/scripts.go Normal file
View File

@ -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)
}

7
cmd/tapRunner.go
View File

@ -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()

45
config/configStructs/scriptingConfig.go
View File

@ -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:")
}

1
go.mod
View File

@ -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

118
internal/connect/hub.go
View File

@ -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() {

53
misc/scripting.go Normal file
View File

@ -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
}