feat: Initial commit
This commit is contained in:
@@ -0,0 +1,249 @@
|
||||
package controllers
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"sync"
|
||||
"unicode/utf8"
|
||||
|
||||
"baihu/internal/constant"
|
||||
"baihu/internal/utils"
|
||||
|
||||
"github.com/creack/pty"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/gorilla/websocket"
|
||||
"golang.org/x/text/encoding/simplifiedchinese"
|
||||
"golang.org/x/text/transform"
|
||||
)
|
||||
|
||||
type TerminalController struct{}
|
||||
|
||||
func NewTerminalController() *TerminalController {
|
||||
return &TerminalController{}
|
||||
}
|
||||
|
||||
var upgrader = websocket.Upgrader{
|
||||
CheckOrigin: func(r *http.Request) bool {
|
||||
return true
|
||||
},
|
||||
}
|
||||
|
||||
// toUTF8 将可能是 GBK 编码的字节转换为 UTF-8
|
||||
func toUTF8(data []byte) string {
|
||||
if utf8.Valid(data) {
|
||||
return string(data)
|
||||
}
|
||||
// 尝试从 GBK 转换
|
||||
reader := transform.NewReader(
|
||||
bufio.NewReader(
|
||||
&byteReader{data: data},
|
||||
),
|
||||
simplifiedchinese.GBK.NewDecoder(),
|
||||
)
|
||||
result, err := io.ReadAll(reader)
|
||||
if err != nil {
|
||||
return string(data)
|
||||
}
|
||||
return string(result)
|
||||
}
|
||||
|
||||
type byteReader struct {
|
||||
data []byte
|
||||
pos int
|
||||
}
|
||||
|
||||
func (r *byteReader) Read(p []byte) (n int, err error) {
|
||||
if r.pos >= len(r.data) {
|
||||
return 0, io.EOF
|
||||
}
|
||||
n = copy(p, r.data[r.pos:])
|
||||
r.pos += n
|
||||
return n, nil
|
||||
}
|
||||
|
||||
func (tc *TerminalController) HandleWebSocket(c *gin.Context) {
|
||||
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer conn.Close()
|
||||
|
||||
// Windows 使用 pipe 模式,Unix 使用 PTY 模式
|
||||
if runtime.GOOS == "windows" {
|
||||
tc.handlePipeMode(conn)
|
||||
} else {
|
||||
tc.handlePtyMode(conn)
|
||||
}
|
||||
}
|
||||
|
||||
// handlePtyMode 使用 PTY 处理终端(Unix/macOS)
|
||||
func (tc *TerminalController) handlePtyMode(conn *websocket.Conn) {
|
||||
// 发送 PTY 模式标识
|
||||
conn.WriteMessage(websocket.TextMessage, []byte("__PTY_MODE__"))
|
||||
|
||||
cmd := utils.NewShellCmd()
|
||||
|
||||
if absDir, err := filepath.Abs(constant.ScriptsWorkDir); err == nil {
|
||||
cmd.Dir = absDir
|
||||
}
|
||||
|
||||
cmd.Env = append(os.Environ(), "TERM=xterm-256color")
|
||||
|
||||
ptmx, err := pty.Start(cmd)
|
||||
if err != nil {
|
||||
conn.WriteMessage(websocket.TextMessage, []byte("Error starting shell: "+err.Error()))
|
||||
return
|
||||
}
|
||||
defer ptmx.Close()
|
||||
|
||||
pty.Setsize(ptmx, &pty.Winsize{Rows: 24, Cols: 80})
|
||||
|
||||
var wg sync.WaitGroup
|
||||
var connMu sync.Mutex
|
||||
|
||||
writeMessage := func(data []byte) {
|
||||
connMu.Lock()
|
||||
defer connMu.Unlock()
|
||||
conn.WriteMessage(websocket.TextMessage, data)
|
||||
}
|
||||
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
buf := make([]byte, 4096)
|
||||
for {
|
||||
n, err := ptmx.Read(buf)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if n > 0 {
|
||||
text := toUTF8(buf[:n])
|
||||
writeMessage([]byte(text))
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
for {
|
||||
_, message, err := conn.ReadMessage()
|
||||
if err != nil {
|
||||
break
|
||||
}
|
||||
if _, err := ptmx.Write(message); err != nil {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
cmd.Process.Kill()
|
||||
cmd.Wait()
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
// handlePipeMode 使用 pipe 处理终端(Windows)
|
||||
func (tc *TerminalController) handlePipeMode(conn *websocket.Conn) {
|
||||
// 发送 pipe 模式标识
|
||||
conn.WriteMessage(websocket.TextMessage, []byte("__PIPE_MODE__"))
|
||||
|
||||
cmd := utils.NewShellCmd()
|
||||
|
||||
if absDir, err := filepath.Abs(constant.ScriptsWorkDir); err == nil {
|
||||
cmd.Dir = absDir
|
||||
}
|
||||
|
||||
stdin, err := cmd.StdinPipe()
|
||||
if err != nil {
|
||||
conn.WriteMessage(websocket.TextMessage, []byte("Error: "+err.Error()))
|
||||
return
|
||||
}
|
||||
|
||||
stdout, err := cmd.StdoutPipe()
|
||||
if err != nil {
|
||||
conn.WriteMessage(websocket.TextMessage, []byte("Error: "+err.Error()))
|
||||
return
|
||||
}
|
||||
|
||||
stderr, err := cmd.StderrPipe()
|
||||
if err != nil {
|
||||
conn.WriteMessage(websocket.TextMessage, []byte("Error: "+err.Error()))
|
||||
return
|
||||
}
|
||||
|
||||
if err := cmd.Start(); err != nil {
|
||||
conn.WriteMessage(websocket.TextMessage, []byte("Error: "+err.Error()))
|
||||
return
|
||||
}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
var connMu sync.Mutex
|
||||
|
||||
writeMessage := func(data []byte) {
|
||||
connMu.Lock()
|
||||
defer connMu.Unlock()
|
||||
conn.WriteMessage(websocket.TextMessage, data)
|
||||
}
|
||||
|
||||
readOutput := func(reader io.Reader) {
|
||||
defer wg.Done()
|
||||
defer func() { recover() }()
|
||||
buf := make([]byte, 4096)
|
||||
for {
|
||||
n, err := reader.Read(buf)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if n > 0 {
|
||||
text := toUTF8(buf[:n])
|
||||
writeMessage([]byte(text))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
wg.Add(2)
|
||||
go readOutput(stdout)
|
||||
go readOutput(stderr)
|
||||
|
||||
for {
|
||||
_, message, err := conn.ReadMessage()
|
||||
if err != nil {
|
||||
break
|
||||
}
|
||||
if _, err := stdin.Write(message); err != nil {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
stdin.Close()
|
||||
cmd.Process.Kill()
|
||||
cmd.Wait()
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
// ExecuteShellCommand 执行单个命令并返回结果
|
||||
func (tc *TerminalController) ExecuteShellCommand(c *gin.Context) {
|
||||
var req struct {
|
||||
Command string `json:"command" binding:"required"`
|
||||
}
|
||||
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
utils.BadRequest(c, err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
cmd := utils.NewShellCommandCmd(req.Command)
|
||||
output, err := cmd.CombinedOutput()
|
||||
|
||||
if err != nil {
|
||||
utils.Success(c, gin.H{
|
||||
"output": string(output),
|
||||
"error": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
utils.Success(c, gin.H{
|
||||
"output": string(output),
|
||||
})
|
||||
}
|
||||
Reference in New Issue
Block a user