MuyueWorkspace/internal/api/terminal.go

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 {
		origin := r.Header.Get("Origin")
		if origin == "" {
			return true
		}
		switch {
		case strings.HasPrefix(origin, "http://127.0.0.1"),
			strings.HasPrefix(origin, "http://localhost"),
			strings.HasPrefix(origin, "http://[::1]"),
			strings.HasPrefix(origin, "https://127.0.0.1"),
			strings.HasPrefix(origin, "https://localhost"),
			strings.HasPrefix(origin, "https://[::1]"):
			return true
		default:
			return false
		}
	},
}

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()
		} else {
			if path, err := exec.LookPath(shell); err == nil {
				shell = path
			}
		}

		if _, err := os.Stat(shell); err != nil {
			conn.WriteJSON(wsMessage{Type: "error", Data: fmt.Sprintf("shell not found: %s", shell)})
			return
		}

		if strings.Contains(shell, "wsl") {
			cmd = exec.Command(shell, "--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) {
	if r.Method != "DELETE" {
		writeError(w, "DELETE only", http.StatusMethodNotAllowed)
		return
	}
	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 path, err := exec.LookPath(s); err == nil {
			return path
		}
	}
	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
}