chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
+383
@@ -0,0 +1,383 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"os"
|
||||
"os/signal"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/creasty/defaults"
|
||||
"github.com/fsnotify/fsnotify"
|
||||
"github.com/kubeshark/kubeshark/config"
|
||||
"github.com/kubeshark/kubeshark/config/configStructs"
|
||||
"github.com/kubeshark/kubeshark/kubernetes"
|
||||
"github.com/kubeshark/kubeshark/misc"
|
||||
"github.com/rs/zerolog/log"
|
||||
"github.com/spf13/cobra"
|
||||
k8serrors "k8s.io/apimachinery/pkg/api/errors"
|
||||
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
|
||||
"k8s.io/apimachinery/pkg/watch"
|
||||
)
|
||||
|
||||
var scriptsCmd = &cobra.Command{
|
||||
Use: "scripts",
|
||||
Short: "Watch the `scripting.source` and/or `scripting.sources` folders 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.ProxyFrontPortLabel, defaultTapConfig.Proxy.Front.Port, "Provide a custom port for the Kubeshark")
|
||||
scriptsCmd.Flags().String(configStructs.ProxyHostLabel, defaultTapConfig.Proxy.Host, "Provide a custom host for the Kubeshark")
|
||||
scriptsCmd.Flags().StringP(configStructs.ReleaseNamespaceLabel, "s", defaultTapConfig.Release.Namespace, "Release namespace of Kubeshark")
|
||||
}
|
||||
|
||||
func runScripts() {
|
||||
if config.Config.Scripting.Source == "" && len(config.Config.Scripting.Sources) == 0 {
|
||||
log.Error().Msg("Both `scripting.source` and `scripting.sources` fields are empty.")
|
||||
return
|
||||
}
|
||||
|
||||
kubernetesProvider, err := getKubernetesProviderForCli(false, false)
|
||||
if err != nil {
|
||||
log.Error().Err(err).Send()
|
||||
return
|
||||
}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
signalChan := make(chan os.Signal, 1)
|
||||
signal.Notify(signalChan, os.Interrupt)
|
||||
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
watchConfigMap(ctx, kubernetesProvider)
|
||||
}()
|
||||
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
watchScripts(ctx, kubernetesProvider, true)
|
||||
}()
|
||||
|
||||
go func() {
|
||||
<-signalChan
|
||||
log.Debug().Msg("Received interrupt, stopping watchers.")
|
||||
cancel()
|
||||
}()
|
||||
|
||||
wg.Wait()
|
||||
|
||||
}
|
||||
|
||||
func createScript(provider *kubernetes.Provider, script misc.ConfigMapScript) (index int64, err error) {
|
||||
const maxRetries = 5
|
||||
var scripts map[int64]misc.ConfigMapScript
|
||||
|
||||
for i := 0; i < maxRetries; i++ {
|
||||
scripts, err = kubernetes.ConfigGetScripts(provider)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
script.Active = kubernetes.IsActiveScript(provider, script.Title)
|
||||
index = 0
|
||||
if script.Title != "New Script" {
|
||||
for i, v := range scripts {
|
||||
if index <= i {
|
||||
index = i + 1
|
||||
}
|
||||
if v.Title == script.Title {
|
||||
index = int64(i)
|
||||
}
|
||||
}
|
||||
}
|
||||
scripts[index] = script
|
||||
|
||||
log.Info().Str("title", script.Title).Bool("Active", script.Active).Int64("Index", index).Msg("Creating script")
|
||||
var data []byte
|
||||
data, err = json.Marshal(scripts)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
_, err = kubernetes.SetConfig(provider, kubernetes.CONFIG_SCRIPTING_SCRIPTS, string(data))
|
||||
if err == nil {
|
||||
return index, nil
|
||||
}
|
||||
|
||||
if k8serrors.IsConflict(err) {
|
||||
log.Debug().Err(err).Msg("Conflict detected, retrying update...")
|
||||
time.Sleep(500 * time.Millisecond)
|
||||
continue
|
||||
}
|
||||
|
||||
return 0, err
|
||||
}
|
||||
|
||||
log.Error().Msg("Max retries reached for creating script due to conflicts.")
|
||||
return 0, errors.New("max retries reached due to conflicts while creating script")
|
||||
}
|
||||
|
||||
func updateScript(provider *kubernetes.Provider, index int64, script misc.ConfigMapScript) (err error) {
|
||||
var scripts map[int64]misc.ConfigMapScript
|
||||
scripts, err = kubernetes.ConfigGetScripts(provider)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
script.Active = kubernetes.IsActiveScript(provider, script.Title)
|
||||
scripts[index] = script
|
||||
|
||||
var data []byte
|
||||
data, err = json.Marshal(scripts)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
_, err = kubernetes.SetConfig(provider, kubernetes.CONFIG_SCRIPTING_SCRIPTS, string(data))
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
func deleteScript(provider *kubernetes.Provider, index int64) (err error) {
|
||||
var scripts map[int64]misc.ConfigMapScript
|
||||
scripts, err = kubernetes.ConfigGetScripts(provider)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
err = kubernetes.DeleteActiveScriptByTitle(provider, scripts[index].Title)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
delete(scripts, index)
|
||||
|
||||
var data []byte
|
||||
data, err = json.Marshal(scripts)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
_, err = kubernetes.SetConfig(provider, kubernetes.CONFIG_SCRIPTING_SCRIPTS, string(data))
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
func watchScripts(ctx context.Context, provider *kubernetes.Provider, block bool) {
|
||||
files := make(map[string]int64)
|
||||
|
||||
scripts, err := config.Config.Scripting.GetScripts()
|
||||
if err != nil {
|
||||
log.Error().Err(err).Send()
|
||||
return
|
||||
}
|
||||
|
||||
for _, script := range scripts {
|
||||
index, err := createScript(provider, script.ConfigMap())
|
||||
if err != nil {
|
||||
log.Error().Err(err).Send()
|
||||
continue
|
||||
}
|
||||
|
||||
files[script.Path] = index
|
||||
}
|
||||
|
||||
watcher, err := fsnotify.NewWatcher()
|
||||
if err != nil {
|
||||
log.Error().Err(err).Send()
|
||||
return
|
||||
}
|
||||
if block {
|
||||
defer watcher.Close()
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
defer cancel()
|
||||
|
||||
signalChan := make(chan os.Signal, 1)
|
||||
signal.Notify(signalChan, os.Interrupt)
|
||||
|
||||
go func() {
|
||||
<-signalChan
|
||||
log.Debug().Msg("Received interrupt, stopping script watch.")
|
||||
cancel()
|
||||
watcher.Close()
|
||||
}()
|
||||
|
||||
if err := watcher.Add(config.Config.Scripting.Source); err != nil {
|
||||
log.Error().Err(err).Msg("Failed to add scripting source to watcher")
|
||||
return
|
||||
}
|
||||
|
||||
go func() {
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
log.Debug().Msg("Script watcher exiting gracefully.")
|
||||
return
|
||||
|
||||
// watch for events
|
||||
case event := <-watcher.Events:
|
||||
if !strings.HasSuffix(event.Name, "js") {
|
||||
log.Info().Str("file", event.Name).Msg("Ignoring file")
|
||||
continue
|
||||
}
|
||||
switch event.Op {
|
||||
case fsnotify.Create:
|
||||
script, err := misc.ReadScriptFile(event.Name)
|
||||
if err != nil {
|
||||
log.Error().Err(err).Send()
|
||||
continue
|
||||
}
|
||||
|
||||
index, err := createScript(provider, script.ConfigMap())
|
||||
if err != nil {
|
||||
log.Error().Err(err).Send()
|
||||
continue
|
||||
}
|
||||
|
||||
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()
|
||||
continue
|
||||
}
|
||||
|
||||
err = updateScript(provider, index, script.ConfigMap())
|
||||
if err != nil {
|
||||
log.Error().Err(err).Send()
|
||||
continue
|
||||
}
|
||||
|
||||
case fsnotify.Rename:
|
||||
index := files[event.Name]
|
||||
err := deleteScript(provider, index)
|
||||
if err != nil {
|
||||
log.Error().Err(err).Send()
|
||||
continue
|
||||
}
|
||||
|
||||
default:
|
||||
// pass
|
||||
}
|
||||
|
||||
case err, ok := <-watcher.Errors:
|
||||
if !ok {
|
||||
log.Info().Msg("Watcher errors channel closed.")
|
||||
return
|
||||
}
|
||||
log.Error().Err(err).Msg("Watcher error encountered")
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
if err := watcher.Add(config.Config.Scripting.Source); err != nil {
|
||||
log.Error().Err(err).Send()
|
||||
}
|
||||
|
||||
for _, source := range config.Config.Scripting.Sources {
|
||||
if err := watcher.Add(source); err != nil {
|
||||
log.Error().Err(err).Send()
|
||||
}
|
||||
}
|
||||
|
||||
log.Info().Str("folder", config.Config.Scripting.Source).Interface("folders", config.Config.Scripting.Sources).Msg("Watching scripts against changes:")
|
||||
|
||||
if block {
|
||||
<-ctx.Done()
|
||||
}
|
||||
}
|
||||
|
||||
func watchConfigMap(ctx context.Context, provider *kubernetes.Provider) {
|
||||
clientset := provider.GetClientSet()
|
||||
configMapName := kubernetes.SELF_RESOURCES_PREFIX + kubernetes.SUFFIX_CONFIG_MAP
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
log.Info().Msg("ConfigMap watcher exiting gracefully.")
|
||||
return
|
||||
|
||||
default:
|
||||
watcher, err := clientset.CoreV1().ConfigMaps(config.Config.Tap.Release.Namespace).Watch(context.TODO(), metav1.ListOptions{
|
||||
FieldSelector: "metadata.name=" + configMapName,
|
||||
})
|
||||
if err != nil {
|
||||
log.Warn().Err(err).Msg("ConfigMap not found, retrying in 5 seconds...")
|
||||
time.Sleep(5 * time.Second)
|
||||
continue
|
||||
}
|
||||
|
||||
// Create a goroutine to process events
|
||||
watcherClosed := make(chan struct{})
|
||||
go func() {
|
||||
defer close(watcherClosed)
|
||||
for event := range watcher.ResultChan() {
|
||||
if event.Type == watch.Added {
|
||||
log.Info().Msg("ConfigMap created or modified")
|
||||
runScriptsSync(provider)
|
||||
} else if event.Type == watch.Deleted {
|
||||
log.Warn().Msg("ConfigMap deleted, waiting for recreation...")
|
||||
break
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
// Wait for either context cancellation or watcher completion
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
watcher.Stop()
|
||||
log.Info().Msg("ConfigMap watcher stopping due to context cancellation")
|
||||
return
|
||||
case <-watcherClosed:
|
||||
log.Info().Msg("Watcher closed, restarting...")
|
||||
}
|
||||
|
||||
time.Sleep(5 * time.Second)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func runScriptsSync(provider *kubernetes.Provider) {
|
||||
files := make(map[string]int64)
|
||||
|
||||
scripts, err := config.Config.Scripting.GetScripts()
|
||||
if err != nil {
|
||||
log.Error().Err(err).Send()
|
||||
return
|
||||
}
|
||||
|
||||
for _, script := range scripts {
|
||||
index, err := createScript(provider, script.ConfigMap())
|
||||
if err != nil {
|
||||
log.Error().Err(err).Send()
|
||||
continue
|
||||
}
|
||||
files[script.Path] = index
|
||||
}
|
||||
log.Info().Msg("Synchronized scripts with ConfigMap.")
|
||||
}
|
||||
Reference in New Issue
Block a user