Merge pull request kosmos-io#660 from village-way/release-0.4.0

Cherry-Pick: node-agent update

cmd/kubenest/node-agent/app_test.go → cmd/kubenest/node-agent/app/client/client_test.go

@@ -1,9 +1,7 @@
-package main
+package client
 
 import (
-	"bufio"
 	"crypto/tls"
-	"encoding/base64"
 	"fmt"
 	"net/http"
 	"net/url"
@@ -13,12 +11,9 @@ import (
 	"testing"
 	"time"
 
-	"github.com/gorilla/websocket"
+	"github.com/kosmos.io/kosmos/cmd/kubenest/node-agent/app/serve"
 )
 
-// create dialer
-var dialer = *websocket.DefaultDialer
-
 // test addr user pass
 var testAddr, username, pass string
 var headers http.Header
@@ -37,17 +32,19 @@ func init() {
 	username = os.Getenv("WEB_USER")
 	pass = os.Getenv("WEB_PASS")
 	testAddr = "127.0.0.1:5678"
+
 	headers = http.Header{
-		"Authorization": {"Basic " + basicAuth(username, pass)},
+		"Authorization": {"Basic " + BasicAuth(username, pass)},
 	}
-	go start(":5678", "cert.pem", "key.pem", username, pass)
+	go func() {
+		err := serve.Start(":5678", "cert.pem", "key.pem", username, pass)
+		if err != nil {
+			log.Fatal(err)
+		}
+	}()
 	time.Sleep(10 * time.Second)
 }
-func wsRespClose(resp *http.Response) {
-	if resp != nil && resp.Body != nil {
-		_ = resp.Body.Close()
-	}
-}
+
 func TestCmd(t *testing.T) {
 	fmt.Println("Command test")
 	command := url.QueryEscape("ls -l")
@@ -106,63 +103,3 @@ func TestPyScript(t *testing.T) {
 	sendFile(ws, filepath.Join(parentDir, "count.py"))
 	handleMessages(ws)
 }
-
-func basicAuth(username, password string) string {
-	auth := username + ":" + password
-	return base64.StdEncoding.EncodeToString([]byte(auth))
-}
-
-func handleMessages(ws *websocket.Conn) {
-	defer ws.Close()
-	for {
-		_, message, err := ws.ReadMessage()
-		if err != nil {
-			log.Println("Read message end :", err)
-			return
-		}
-		fmt.Printf("Received message: %s\n", message)
-	}
-}
-
-func sendFile(ws *websocket.Conn, filePath string) {
-	//if file not exists, close connection
-	if _, err := os.Stat(filePath); os.IsNotExist(err) {
-		log.Printf("File not exists: %v", err)
-		err := ws.WriteMessage(websocket.BinaryMessage, []byte("EOF"))
-		if err != nil {
-			log.Printf("Write message error: %v", err)
-		}
-		return
-	}
-
-	file, err := os.Open(filePath)
-	if err != nil {
-		log.Printf("File open error: %v", err)
-	}
-	defer file.Close()
-	// 指定每次读取的数据块大小
-	bufferSize := 1024 // 例如每次读取 1024 字节
-	buffer := make([]byte, bufferSize)
-
-	reader := bufio.NewReader(file)
-	for {
-		n, err := reader.Read(buffer)
-		if err != nil {
-			// check if EOF
-			if err.Error() == "EOF" {
-				break
-			}
-			log.Printf("failed to read file %v:", err)
-			return
-		}
-		dataToSend := buffer[:n]
-
-		_ = ws.WriteMessage(websocket.BinaryMessage, dataToSend)
-	}
-
-	err = ws.WriteMessage(websocket.BinaryMessage, []byte("EOF"))
-	log.Printf("send EOF ----")
-	if err != nil {
-		log.Printf("Write message error: %v", err)
-	}
-}

cmd/kubenest/node-agent/app/logger/logger.go

@@ -0,0 +1,27 @@
+package logger
+
+import (
+	"io"
+	"os"
+
+	"github.com/sirupsen/logrus"
+)
+
+var log *logrus.Logger
+
+func init() {
+	log = logrus.New()
+	// setup log
+	log.Out = os.Stdout
+	logFile, err := os.OpenFile("app.log", os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
+	if err == nil {
+		log.SetOutput(io.MultiWriter(os.Stdout, logFile))
+	} else {
+		log.Info("Failed to log to file, using default stderr")
+	}
+	log.SetLevel(logrus.InfoLevel)
+}
+
+func GetLogger() *logrus.Logger {
+	return log
+}

cmd/kubenest/node-agent/app/root.go

@@ -0,0 +1,74 @@
+package app
+
+import (
+	"os"
+
+	"github.com/spf13/cobra"
+	"github.com/spf13/viper"
+
+	"github.com/kosmos.io/kosmos/cmd/kubenest/node-agent/app/client"
+	"github.com/kosmos.io/kosmos/cmd/kubenest/node-agent/app/logger"
+	"github.com/kosmos.io/kosmos/cmd/kubenest/node-agent/app/serve"
+)
+
+var (
+	user     string // username for authentication
+	password string // password for authentication
+	log      = logger.GetLogger()
+)
+
+var RootCmd = &cobra.Command{
+	Use:   "node-agent",
+	Short: "node-agent is a tool for node to start websocket server and client",
+	Long: `node-agent client for connect server to execute command and upload file to node
+            node-agent serve for start websocket server to receive message from client and download file from client`,
+	Run: func(cmd *cobra.Command, args []string) {
+		_ = cmd.Help()
+	},
+}
+
+func initConfig() {
+	// Tell Viper to automatically look for a .env file
+	viper.SetConfigFile("agent.env")
+	currentDir, _ := os.Getwd()
+	viper.AddConfigPath(currentDir)
+	viper.AddConfigPath("/srv/node-agent/agent.env")
+	// If a agent.env file is found, read it in.
+	if err := viper.ReadInConfig(); err != nil {
+		log.Warnf("Load config file error, %s", err)
+	}
+	// set default value from agent.env
+	if len(user) == 0 {
+		user = viper.GetString("WEB_USER")
+	}
+	if len(password) == 0 {
+		password = viper.GetString("WEB_PASS")
+	}
+}
+
+func init() {
+	cobra.OnInitialize(initConfig)
+
+	RootCmd.PersistentFlags().StringVarP(&user, "user", "u", "", "Username for authentication")
+	RootCmd.PersistentFlags().StringVarP(&password, "password", "p", "", "Password for authentication")
+	// bind flags to viper
+	err := viper.BindPFlag("WEB_USER", RootCmd.PersistentFlags().Lookup("user"))
+	if err != nil {
+		log.Fatal(err)
+	}
+	err = viper.BindPFlag("WEB_PASS", RootCmd.PersistentFlags().Lookup("password"))
+	if err != nil {
+		log.Fatal(err)
+	}
+	// bind environment variables
+	err = viper.BindEnv("WEB_USER", "WEB_USER")
+	if err != nil {
+		log.Fatal(err)
+	}
+	err = viper.BindEnv("WEB_PASS", "WEB_PASS")
+	if err != nil {
+		log.Fatal(err)
+	}
+	RootCmd.AddCommand(client.ClientCmd)
+	RootCmd.AddCommand(serve.ServeCmd)
+}