add WriteFile

Author: sp-yduck

proxmox/websocket.go

@@ -2,10 +2,12 @@ package proxmox
 
 import (
 	"context"
+	"encoding/base64"
 	"fmt"
 	"regexp"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/pkg/errors"
 
@@ -19,7 +21,8 @@ const (
 )
 
 type VNCWebSocketClient struct {
-	conn *websocket.Conn
+	conn   *websocket.Conn
+	ticker *time.Ticker
 }
 
 func (s *Service) NewNodeVNCWebSocketConnection(ctx context.Context, nodeName string) (*VNCWebSocketClient, error) {
@@ -32,11 +35,22 @@ func (s *Service) NewNodeVNCWebSocketConnection(ctx context.Context, nodeName st
 		return nil, err
 	}
 
-	return &VNCWebSocketClient{conn: conn}, nil
+	ticker := time.NewTicker(30 * time.Second)
+	go func() {
+		for {
+			select {
+			case <-ticker.C:
+				conn.WriteMessage(websocket.BinaryMessage, []byte("2"))
+			}
+		}
+	}()
+
+	return &VNCWebSocketClient{conn: conn, ticker: ticker}, nil
 }
 
 func (c *VNCWebSocketClient) Close() {
 	c.conn.Close()
+	c.ticker.Stop()
 }
 
 func (c *VNCWebSocketClient) Write(cmd string) error {
@@ -49,6 +63,38 @@ func (c *VNCWebSocketClient) Write(cmd string) error {
 	return c.sendFinMessage()
 }
 
+func (c *VNCWebSocketClient) WriteFile(ctx context.Context, content, path string) error {
+	c.Exec(ctx, fmt.Sprintf("rm %s", path))
+	if _, _, err := c.Exec(ctx, fmt.Sprintf("touch %s", path)); err != nil {
+		return err
+	}
+	chunks := chunkString(content, 3000)
+	for _, chunk := range chunks {
+		b64chunk := base64.StdEncoding.EncodeToString([]byte(chunk))
+		_, _, err := c.Exec(ctx, fmt.Sprintf("echo %s | base64 -d >> %s", b64chunk, path))
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func chunkString(s string, chunkSize int) []string {
+	var chunks []string
+	runes := []rune(s)
+	if len(runes) == 0 {
+		return []string{s}
+	}
+	for i := 0; i < len(runes); i += chunkSize {
+		nn := i + chunkSize
+		if nn > len(runes) {
+			nn = len(runes)
+		}
+		chunks = append(chunks, string(runes[i:nn]))
+	}
+	return chunks
+}
+
 func (c *VNCWebSocketClient) sendFinMessage() error {
 	b := []byte(fmt.Sprintf(`echo "%s$?"%s`, finMessage, "\n"))
 	bheader := []byte(fmt.Sprintf("0:%d:", len(b)))

proxmox/websocket_test.go

@@ -43,6 +43,21 @@ func (s *TestSuite) TestExec() {
 	s.T().Logf("exec command : %s : %d", out, code)
 }
 
+func (s *TestSuite) TestWriteFile() {
+	testNode := s.getTestNode()
+	client, err := s.service.NewNodeVNCWebSocketConnection(context.TODO(), testNode.Node)
+	if err != nil {
+		s.T().Fatalf("failed to create new vnc client: %v", err)
+	}
+	defer client.Close()
+
+	ctx, _ := context.WithTimeout(context.TODO(), 15*time.Second)
+	err = client.WriteFile(ctx, "this is a file content", "~/test-write-file.txt")
+	if err != nil {
+		s.T().Fatalf("failed to exec command: %v", err)
+	}
+}
+
 func TestParseFinMessage(t *testing.T) {
 	testMsg := " daf" + finMessage + "123\n"
 	if parseFinMessage(testMsg) == "" {