From f4c6613e6e42d7b66c1eb5e830fbc30078814d75 Mon Sep 17 00:00:00 2001
From: "M. Mert Yildiran" <me@mertyildiran.com>
Date: Thu, 9 Feb 2023 15:09:52 +0300
Subject: [PATCH] :sparkles: Add `console` command

---
 cmd/console.go   | 95 ++++++++++++++++++++++++++++++++++++++++++++++++
 config/config.go |  4 +-
 go.mod           |  1 +
 go.sum           |  1 +
 4 files changed, 100 insertions(+), 1 deletion(-)
 create mode 100644 cmd/console.go

diff --git a/cmd/console.go b/cmd/console.go
new file mode 100644
index 000000000..a4afc544b
--- /dev/null
+++ b/cmd/console.go
@@ -0,0 +1,95 @@
+package cmd
+
+import (
+	"fmt"
+	"net/url"
+	"os"
+	"os/signal"
+	"time"
+
+	"github.com/creasty/defaults"
+	"github.com/gorilla/websocket"
+	"github.com/kubeshark/kubeshark/config"
+	"github.com/kubeshark/kubeshark/config/configStructs"
+	"github.com/kubeshark/kubeshark/utils"
+	"github.com/rs/zerolog/log"
+	"github.com/spf13/cobra"
+)
+
+var consoleCmd = &cobra.Command{
+	Use:   "console",
+	Short: "Stream the scripting console logs into shell",
+	RunE: func(cmd *cobra.Command, args []string) error {
+		runConsole()
+		return nil
+	},
+}
+
+func init() {
+	rootCmd.AddCommand(consoleCmd)
+
+	defaultTapConfig := configStructs.TapConfig{}
+	if err := defaults.Set(&defaultTapConfig); err != nil {
+		log.Debug().Err(err).Send()
+	}
+
+	consoleCmd.Flags().Uint16(configStructs.ProxyHubPortLabel, defaultTapConfig.Proxy.Hub.SrcPort, "Provide a custom port for the Hub.")
+	consoleCmd.Flags().String(configStructs.ProxyHostLabel, defaultTapConfig.Proxy.Host, "Provide a custom host for the Hub.")
+}
+
+func runConsole() {
+	interrupt := make(chan os.Signal, 1)
+	signal.Notify(interrupt, os.Interrupt)
+
+	log.Info().Str("host", config.Config.Tap.Proxy.Host).Uint16("port", config.Config.Tap.Proxy.Hub.SrcPort).Msg("Connecting to:")
+	u := url.URL{
+		Scheme: "ws",
+		Host:   fmt.Sprintf("%s:%d", config.Config.Tap.Proxy.Host, config.Config.Tap.Proxy.Hub.SrcPort),
+		Path:   "/scripts/logs",
+	}
+
+	c, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
+	if err != nil {
+		log.Error().Err(err).Send()
+	}
+	defer c.Close()
+
+	done := make(chan struct{})
+
+	go func() {
+		defer close(done)
+		for {
+			_, message, err := c.ReadMessage()
+			if err != nil {
+				log.Error().Err(err).Send()
+				return
+			}
+
+			log.Info().Msg(string(message))
+		}
+	}()
+
+	ticker := time.NewTicker(time.Second)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-done:
+			return
+		case <-interrupt:
+			log.Warn().Msg(fmt.Sprintf(utils.Yellow, "Received interrupt, exiting..."))
+
+			err := c.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
+			if err != nil {
+				log.Error().Err(err).Send()
+				return
+			}
+
+			select {
+			case <-done:
+			case <-time.After(time.Second):
+			}
+			return
+		}
+	}
+}
diff --git a/config/config.go b/config/config.go
index 10ab35742..e377180fe 100644
--- a/config/config.go
+++ b/config/config.go
@@ -50,7 +50,9 @@ func InitConfig(cmd *cobra.Command) error {
 		return nil
 	}
 
-	go version.CheckNewerVersion()
+	if cmd.Use != "console" {
+		go version.CheckNewerVersion()
+	}
 
 	Config = CreateDefaultConfig()
 	cmdName = cmd.Name()
diff --git a/go.mod b/go.mod
index 37852c74e..51b17c415 100644
--- a/go.mod
+++ b/go.mod
@@ -8,6 +8,7 @@ require (
 	github.com/docker/go-connections v0.4.0
 	github.com/docker/go-units v0.4.0
 	github.com/google/go-github/v37 v37.0.0
+	github.com/gorilla/websocket v1.4.2
 	github.com/kubeshark/base v0.6.3
 	github.com/robertkrimen/otto v0.2.1
 	github.com/rs/zerolog v1.28.0
diff --git a/go.sum b/go.sum
index ab2587c71..5434bbeba 100644
--- a/go.sum
+++ b/go.sum
@@ -331,6 +331,7 @@ github.com/googleapis/gnostic v0.5.5/go.mod h1:7+EbHbldMins07ALC74bsA81Ovc97Dwqy
 github.com/gopherjs/gopherjs v0.0.0-20181017120253-0766667cb4d1/go.mod h1:wJfORRmW1u3UXTncJ5qlYoELFm8eSnnEO6hX4iZ3EWY=
 github.com/gorilla/mux v1.8.0/go.mod h1:DVbg23sWSpFRCP0SfiEN6jmj59UnW/n46BH5rLB71So=
 github.com/gorilla/websocket v1.4.0/go.mod h1:E7qHFY5m1UJ88s3WnNqhKjPHQ0heANvMoAMk2YaljkQ=
+github.com/gorilla/websocket v1.4.2 h1:+/TMaTYc4QFitKJxsQ7Yye35DkWvkdLcvGKqM+x0Ufc=
 github.com/gorilla/websocket v1.4.2/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/gregjones/httpcache v0.0.0-20180305231024-9cad4c3443a7/go.mod h1:FecbI9+v66THATjSRHfNgh1IVFe/9kFxbXtjV0ctIMA=
 github.com/gregjones/httpcache v0.0.0-20190611155906-901d90724c79 h1:+ngKgrYPPJrOjhax5N+uePQ0Fh1Z7PheYoUI/0nzkPA=