remote-core/shell/shell_http.go

package shell

// 移植自 remote-backend-server internal/handlers/shell(去 DB 依赖, ConPTY 终端 + asciinema 录制)。

import (
	"encoding/json"
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"github.com/runletapp/go-console"
	
	"github.com/sirupsen/logrus"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"git.zeroonesoft.cn/golib/remote-core/common"
	"runtime"
	"strings"
	"time"
)

type windowSize struct {
	Rows uint16 `json:"rows"`
	Cols uint16 `json:"cols"`
	X    uint16
	Y    uint16
}

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	// 解决跨域问题
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

type OpDataType struct {
	Op   string      `json:"op" form:"op"`
	Data interface{} `json:"data" form:"data" `
}

type ResizeDataType struct {
	Cols int `json:"cols"`
	Rows int `json:"rows"`
}

// RegisterRoutersWebShell 注册路由
func RegisterRoutersWebShell(r *gin.RouterGroup, recDir string) {
	r.GET("/web/shell", getWebShell(recDir))
	r.GET("/web/shell/record", RecoderList)
}

func getWebShell(recDir string) gin.HandlerFunc {
	return func(c *gin.Context) {

		conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
		if err != nil {
			logrus.Errorf("Unable to upgrade connection")
			return
		}

		proc, err := console.New(180, 30)
		if err != nil {
			panic(err)
		}

		arr := make([]string, 0)
		if runtime.GOOS == "windows" {
			if ok, _ := common.PathExists("C:\\Windows\\System32\\WindowsPowerShell\\v1.0\\powershell.exe"); ok {
				arr = append(arr, "powershell.exe") // powershell.exe cmd.exe
			} else {
				arr = append(arr, "cmd.exe")
			}

		} else {
			arr = []string{"bash"} // powershell.exe cmd.exe
		}

		if err := proc.Start(arr); err != nil {
			logrus.Errorf("启动出错:%s", err.Error())
			panic(err)
		}

		if err != nil {
			logrus.Errorf("Unable to start pty/cmd")
			conn.WriteMessage(websocket.TextMessage, []byte(err.Error()))
			return
		}
		defer func() {
			proc.Close()
			conn.Close()
		}()

		var recorder *Recorder

		os.MkdirAll("./rec", 0766)
		fileName := path.Join("./rec/", fmt.Sprintf("%s.cast", time.Now().Format("20060102_150405")))
		logrus.Debugf("记录文件名:%s", fileName)
		f, err := os.OpenFile(fileName, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0766)
		if err != nil {
			conn.WriteMessage(websocket.TextMessage, []byte(err.Error()))
			logrus.Errorf("创建记录文件失败")
			return
		}

		defer f.Close()
		recorder = NewRecorder(f)

		if recorder != nil {
			recorder.Lock()
			recorder.WriteHeader(30, 150)
			recorder.Unlock()
		}

		go func() {
			for {
				buf := make([]byte, 0x10000)
				read, err := proc.Read(buf)
				if err != nil {
					conn.WriteMessage(websocket.TextMessage, []byte(err.Error()))
					logrus.Errorf("读取错误:%s", err.Error())
					return
				}
				conn.WriteMessage(websocket.BinaryMessage, buf[:read])
				if recorder != nil {
					recorder.Lock()
					recorder.WriteData(InputType, string(buf[:read]))
					recorder.Unlock()
				}
			}
		}()

		for {
			messageType, reader, err := conn.NextReader()
			if err != nil {
				logrus.Errorf("Unable to grab next reader")
				return
			}

			if messageType != websocket.TextMessage {
				logrus.Errorf("无效的消息")
				//conn.WriteMessage(websocket.TextMessage, []byte("Unexpected text message"))
				//continue
			}

			buf := make([]byte, 0x10000)
			n, err := reader.Read(buf)

			if err != nil {
				logrus.Errorf("Error after copying %d bytes", n)
			}

			if n > 0 {

				opData := OpDataType{}
				newBuf := buf[:n]
				fmt.Println("收到数据", string(newBuf))
				err = json.Unmarshal(newBuf, &opData)
				if err != nil {
					logrus.Errorf("解析错误:%s", err.Error())
				}
				if opData.Op == "stdin" {
					value, ok := opData.Data.(string)
					if ok {

						n, err = proc.Write([]byte(value))
						if err != nil {
							logrus.Errorf("写放数据错误g %d bytes", n)
						} else {
							fmt.Println("写入成功")
						}
					} else {

					}

				} else if opData.Op == "resize" {
					// Data 反序列化后是 map[string]interface{}, 不能直接断言 ResizeDataType
					// (此前断言恒失败, 前端 resize 被静默吞掉, PTY 固定 180x30)
					raw, merr := json.Marshal(opData.Data)
					var size ResizeDataType
					if merr == nil && json.Unmarshal(raw, &size) == nil && size.Cols > 0 && size.Rows > 0 {
						if serr := proc.SetSize(size.Cols, size.Rows); serr != nil {
							logrus.Errorf("重置大小错误: %s", serr.Error())
						}
					}

				}

			}

		}
	}

}

func RecoderList(c *gin.Context) {
	files, err := ioutil.ReadDir("./rec")
	if err != nil {
		c.AbortWithStatusJSON(200, gin.H{"ok": false, "msg": err.Error()})
		return
	}
	var filesName []string
	for _, f := range files {
		if f.IsDir() {
			continue
		}

		if strings.HasSuffix(f.Name(), ".cast") {
			filesName = append(filesName, f.Name())
		}
	}
	c.JSON(200, filesName)
}