remote-core/shell/shell_http.go

package shell

// 移植自 remote-backend-server internal/handlers/shell(去 DB 依赖, ConPTY 终端 + asciinema 录制)。
// 标准库 http.HandlerFunc 形态, 路由注册经 webhttp.Router 适配(去 gin 化)。

import (
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"github.com/gorilla/websocket"
	"github.com/runletapp/go-console"
	"github.com/sirupsen/logrus"

	"git.zeroonesoft.cn/golib/remote-core/common"
	"git.zeroonesoft.cn/golib/remote-core/webhttp"
)

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"`
}

// recStoreDir 会话录制目录, Register 时注入; 录制写入与记录列表两处共用。
// 原实现硬编码 ./rec, recDir 形参是死参数(remote-service 传 RunDir/rec 不生效)。
var recStoreDir = "./rec"

// RegisterRoutersWebShell 注册路由。recDir 为会话录制目录, 传空退回 ./rec。
func RegisterRoutersWebShell(r webhttp.Router, recDir string) {
	if recDir != "" {
		recStoreDir = recDir
	}
	r.GET("/web/shell", getWebShell)
	r.GET("/web/shell/record", RecoderList)
}

func getWebShell(w http.ResponseWriter, r *http.Request) {

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

	proc, err := console.New(180, 30)
	if err != nil {
		logrus.Errorf("创建终端失败:%s", err.Error())
		conn.Close()
		return
	}

	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())
		conn.WriteMessage(websocket.TextMessage, []byte(err.Error()))
		conn.Close()
		return
	}
	defer func() {
		proc.Close()
		conn.Close()
	}()

	var recorder *Recorder

	os.MkdirAll(recStoreDir, 0766)
	fileName := filepath.Join(recStoreDir, 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]
			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 {

				}

			} 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(w http.ResponseWriter, _ *http.Request) {
	files, err := os.ReadDir(recStoreDir)
	if err != nil {
		webhttp.WriteJSON(w, http.StatusOK, map[string]any{"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())
		}
	}
	webhttp.WriteJSON(w, http.StatusOK, filesName)
}