package api import ( "encoding/json" "fmt" "log" "net/http" "os" "os/exec" "runtime" "strings" "sync" "time" "github.com/creack/pty/v2" "github.com/gorilla/websocket" "github.com/muyue/muyue/internal/config" ) var upgrader = websocket.Upgrader{ CheckOrigin: func(r *http.Request) bool { return true }, } type wsMessage struct { Type string `json:"type"` Data string `json:"data"` Rows uint16 `json:"rows,omitempty"` Cols uint16 `json:"cols,omitempty"` } func (s *Server) handleTerminalWS(w http.ResponseWriter, r *http.Request) { conn, err := upgrader.Upgrade(w, r, nil) if err != nil { log.Printf("ws upgrade: %v", err) return } defer conn.Close() var initMsg wsMessage _, raw, err := conn.ReadMessage() if err != nil { conn.WriteJSON(wsMessage{Type: "error", Data: "failed to read init message"}) return } if err := json.Unmarshal(raw, &initMsg); err != nil { conn.WriteJSON(wsMessage{Type: "error", Data: "invalid init message"}) return } var cmd *exec.Cmd if initMsg.Type == "ssh" && initMsg.Data != "" { var sshConf struct { Host string `json:"host"` Port int `json:"port"` User string `json:"user"` KeyPath string `json:"key_path"` } if err := json.Unmarshal([]byte(initMsg.Data), &sshConf); err != nil { conn.WriteJSON(wsMessage{Type: "error", Data: "invalid ssh config"}) return } if sshConf.Port == 0 { sshConf.Port = 22 } sshArgs := []string{ "-o", "StrictHostKeyChecking=accept-new", "-o", "UserKnownHostsFile=/dev/null", "-o", "LogLevel=ERROR", } if sshConf.KeyPath != "" { sshArgs = append(sshArgs, "-i", sshConf.KeyPath) } if sshConf.Port != 22 { sshArgs = append(sshArgs, "-p", fmt.Sprintf("%d", sshConf.Port)) } sshArgs = append(sshArgs, fmt.Sprintf("%s@%s", sshConf.User, sshConf.Host)) cmd = exec.Command("ssh", sshArgs...) } else { shell := initMsg.Data if shell == "" { shell = detectShell() } if strings.Contains(shell, "wsl") { cmd = exec.Command("wsl", "--shell-type", "login") } else if strings.Contains(shell, "powershell") || strings.Contains(shell, "pwsh") { cmd = exec.Command(shell, "-NoLogo", "-NoProfile") } else { cmd = exec.Command(shell, "--login") } } cmd.Env = append(os.Environ(), "TERM=xterm-256color") ptmx, err := pty.Start(cmd) if err != nil { log.Printf("pty start: %v", err) conn.WriteJSON(wsMessage{Type: "error", Data: err.Error()}) return } defer func() { ptmx.Close() if cmd.Process != nil { cmd.Process.Kill() cmd.Wait() } }() var once sync.Once cleanup := func() { once.Do(func() { ptmx.Close() if cmd.Process != nil { cmd.Process.Kill() cmd.Wait() } }) } go func() { buf := make([]byte, 4096) for { n, err := ptmx.Read(buf) if err != nil { cleanup() conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")) return } if err := conn.WriteJSON(wsMessage{ Type: "output", Data: string(buf[:n]), }); err != nil { cleanup() return } } }() conn.SetReadLimit(1 << 20) conn.SetReadDeadline(time.Time{}) for { _, raw, err := conn.ReadMessage() if err != nil { cleanup() return } var msg wsMessage if err := json.Unmarshal(raw, &msg); err != nil { continue } switch msg.Type { case "input": if _, err := ptmx.Write([]byte(msg.Data)); err != nil { cleanup() return } case "resize": if msg.Rows > 0 && msg.Cols > 0 { pty.Setsize(ptmx, &pty.Winsize{ Rows: msg.Rows, Cols: msg.Cols, }) } } } } func (s *Server) handleTerminalSessions(w http.ResponseWriter, r *http.Request) { if r.Method == "GET" { writeJSON(w, map[string]interface{}{ "ssh": s.config.Terminal.SSH, "system": detectSystemTerminals(), }) return } if r.Method != "POST" { writeError(w, "POST only", http.StatusMethodNotAllowed) return } var body struct { Name string `json:"name"` Host string `json:"host"` Port int `json:"port"` User string `json:"user"` Password string `json:"password"` KeyPath string `json:"key_path"` } if err := json.NewDecoder(r.Body).Decode(&body); err != nil { writeError(w, err.Error(), http.StatusBadRequest) return } if body.Name == "" || body.Host == "" { writeError(w, "name and host required", http.StatusBadRequest) return } if body.Port == 0 { body.Port = 22 } conn := config.SSHConnection{ Name: body.Name, Host: body.Host, Port: body.Port, User: body.User, KeyPath: body.KeyPath, } if s.config.Terminal.SSH == nil { s.config.Terminal.SSH = []config.SSHConnection{} } s.config.Terminal.SSH = append(s.config.Terminal.SSH, conn) if err := config.Save(s.config); err != nil { writeError(w, err.Error(), http.StatusInternalServerError) return } writeJSON(w, map[string]string{"status": "ok"}) } func (s *Server) handleTerminalSessionsDelete(w http.ResponseWriter, r *http.Request) { name := strings.TrimPrefix(r.URL.Path, "/api/terminal/sessions/") if name == "" { writeError(w, "name required", http.StatusBadRequest) return } found := false for i, c := range s.config.Terminal.SSH { if c.Name == name { s.config.Terminal.SSH = append(s.config.Terminal.SSH[:i], s.config.Terminal.SSH[i+1:]...) found = true break } } if !found { writeError(w, "not found", http.StatusNotFound) return } if err := config.Save(s.config); err != nil { writeError(w, err.Error(), http.StatusInternalServerError) return } writeJSON(w, map[string]string{"status": "ok"}) } func detectShell() string { shells := []string{"zsh", "bash", "fish", "pwsh", "powershell"} for _, s := range shells { if _, err := exec.LookPath(s); err == nil { return s } } return "/bin/sh" } func detectSystemTerminals() []map[string]string { var terminals []map[string]string terminals = append(terminals, map[string]string{ "type": "local", "name": "Default Shell", "shell": detectShell(), }) if runtime.GOOS == "windows" { if _, err := exec.LookPath("wsl"); err == nil { terminals = append(terminals, map[string]string{ "type": "local", "name": "WSL", "shell": "wsl", }) } if _, err := exec.LookPath("powershell"); err == nil { terminals = append(terminals, map[string]string{ "type": "local", "name": "PowerShell", "shell": "powershell", }) } if _, err := exec.LookPath("pwsh"); err == nil { terminals = append(terminals, map[string]string{ "type": "local", "name": "PowerShell Core", "shell": "pwsh", }) } if _, err := exec.LookPath("cmd"); err == nil { terminals = append(terminals, map[string]string{ "type": "local", "name": "Command Prompt", "shell": "cmd", }) } } return terminals }