feat: add wish.Command and wish.Cmd

cmd.go

@@ -0,0 +1,99 @@
+package wish
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"os/exec"
+	"runtime"
+	"time"
+
+	tea "github.com/charmbracelet/bubbletea"
+	"github.com/charmbracelet/ssh"
+)
+
+// CommandContext is like Command but includes a context.
+//
+// If the current session does not have a PTY, it sets them to the session
+// itself.
+func CommandContext(ctx context.Context, s ssh.Session, name string, args ...string) *Cmd {
+	c := exec.CommandContext(ctx, name, args...)
+	pty, _, ok := s.Pty()
+	if !ok {
+		c.Stdin, c.Stdout, c.Stderr = s, s, s
+		return &Cmd{cmd: c}
+	}
+
+	return &Cmd{c, &pty}
+}
+
+// Command sets stdin, stdout, and stderr to the current session's PTY slave.
+//
+// If the current session does not have a PTY, it sets them to the session
+// itself.
+//
+// This will call CommandContext using the session's Context.
+func Command(s ssh.Session, name string, args ...string) *Cmd {
+	return CommandContext(s.Context(), s, name, args...)
+}
+
+// Cmd wraps a *exec.Cmd and a ssh.Pty so a command can be properly run.
+type Cmd struct {
+	cmd *exec.Cmd
+	pty *ssh.Pty
+}
+
+// SetDir set the underlying exec.Cmd env.
+func (c *Cmd) SetEnv(env []string) {
+	c.cmd.Env = env
+}
+
+// Environ returns the underlying exec.Cmd environment.
+func (c *Cmd) Environ() []string {
+	return c.cmd.Environ()
+}
+
+// SetDir set the underlying exec.Cmd dir.
+func (c *Cmd) SetDir(dir string) {
+	c.cmd.Dir = dir
+}
+
+// Run runs the program and waits for it to finish.
+func (c *Cmd) Run() error {
+	if c.pty == nil {
+		return c.cmd.Run()
+	}
+
+	if err := c.pty.Start(c.cmd); err != nil {
+		return err
+	}
+	start := time.Now()
+	if runtime.GOOS == "windows" {
+		for c.cmd.ProcessState == nil {
+			if time.Since(start) > time.Second*10 {
+				return fmt.Errorf("could not start process")
+			}
+			time.Sleep(100 * time.Millisecond)
+		}
+
+		if !c.cmd.ProcessState.Success() {
+			return fmt.Errorf("process failed: exit %d", c.cmd.ProcessState.ExitCode())
+		}
+	} else {
+		if err := c.cmd.Wait(); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+var _ tea.ExecCommand = &Cmd{}
+
+// SetStderr conforms with tea.ExecCommand.
+func (*Cmd) SetStderr(io.Writer) {}
+
+// SetStdin conforms with tea.ExecCommand.
+func (*Cmd) SetStdin(io.Reader) {}
+
+// SetStdout conforms with tea.ExecCommand.
+func (*Cmd) SetStdout(io.Writer) {}

cmd_test.go

@@ -0,0 +1,137 @@
+package wish
+
+import (
+	"bytes"
+	"runtime"
+	"strings"
+	"testing"
+
+	"github.com/charmbracelet/ssh"
+	"github.com/charmbracelet/wish/testsession"
+)
+
+func TestCommandNoPty(t *testing.T) {
+	tmp := t.TempDir()
+	sess := testsession.New(t, &ssh.Server{
+		Handler: func(s ssh.Session) {
+			runEcho(s, "hello")
+			runEnv(s, []string{"HELLO=world"})
+			runPwd(s, tmp)
+		},
+	}, nil)
+	var stdout bytes.Buffer
+	var stderr bytes.Buffer
+	sess.Stdout = &stdout
+	sess.Stderr = &stderr
+	if err := sess.Run(""); err != nil {
+		t.Errorf("expected no error, got %v: %s", err, stderr.String())
+	}
+	out := stdout.String()
+	expectContains(t, out, "hello")
+	expectContains(t, out, "HELLO=world")
+	expectContains(t, out, tmp)
+}
+
+func TestCommandPty(t *testing.T) {
+	tmp := t.TempDir()
+	srv := &ssh.Server{
+		Handler: func(s ssh.Session) {
+			runEcho(s, "hello")
+			runEnv(s, []string{"HELLO=world"})
+			runPwd(s, tmp)
+		},
+	}
+	if err := ssh.AllocatePty()(srv); err != nil {
+		t.Fatalf("expected no error, got %v", err)
+	}
+
+	sess := testsession.New(t, srv, nil)
+	if err := sess.RequestPty("xterm", 500, 200, nil); err != nil {
+		t.Fatalf("expected no error, got %v", err)
+	}
+
+	var stdout bytes.Buffer
+	var stderr bytes.Buffer
+	sess.Stdout = &stdout
+	sess.Stderr = &stderr
+	if err := sess.Run(""); err != nil {
+		t.Errorf("expected no error, got %v: %s", err, stderr.String())
+	}
+	out := stdout.String()
+	expectContains(t, out, "hello")
+	expectContains(t, out, "HELLO=world")
+	expectContains(t, out, tmp)
+}
+
+func TestCommandPtyError(t *testing.T) {
+	srv := &ssh.Server{
+		Handler: func(s ssh.Session) {
+			if err := Command(s, "nopenopenope").Run(); err != nil {
+				Fatal(s, err)
+			}
+		},
+	}
+	if err := ssh.AllocatePty()(srv); err != nil {
+		t.Fatalf("expected no error, got %v", err)
+	}
+
+	sess := testsession.New(t, srv, nil)
+	if err := sess.RequestPty("xterm", 500, 200, nil); err != nil {
+		t.Fatalf("expected no error, got %v", err)
+	}
+
+	var stderr bytes.Buffer
+	sess.Stderr = &stderr
+	if err := sess.Run(""); err == nil {
+		t.Errorf("expected an error, got nil")
+	}
+	expect := `exec: "nopenopenope"`
+	if s := stderr.String(); !strings.Contains(s, expect) {
+		t.Errorf("expected output to contain %q, got %q", expect, s)
+	}
+}
+
+func runEcho(s ssh.Session, str string) {
+	cmd := Command(s, "echo", str)
+	if runtime.GOOS == "windows" {
+		cmd = Command(s, "cmd", "/C", "echo", str)
+	}
+	// these should do nothing...
+	cmd.SetStderr(nil)
+	cmd.SetStdin(nil)
+	cmd.SetStdout(nil)
+	if err := cmd.Run(); err != nil {
+		Fatal(s, err)
+	}
+}
+
+func runEnv(s ssh.Session, env []string) {
+	cmd := Command(s, "env")
+	if runtime.GOOS == "windows" {
+		cmd = Command(s, "cmd", "/C", "set")
+	}
+	cmd.SetEnv(env)
+	if err := cmd.Run(); err != nil {
+		Fatal(s, err)
+	}
+	if len(cmd.Environ()) == 0 {
+		Fatal(s, "cmd.Environ() should not be empty")
+	}
+}
+
+func runPwd(s ssh.Session, dir string) {
+	cmd := Command(s, "pwd")
+	if runtime.GOOS == "windows" {
+		cmd = Command(s, "cmd", "/C", "cd")
+	}
+	cmd.SetDir(dir)
+	if err := cmd.Run(); err != nil {
+		Fatal(s, err)
+	}
+}
+
+func expectContains(tb testing.TB, s, substr string) {
+	if !strings.Contains(s, substr) {
+		tb.Errorf("expected output %q to contain %q", s, substr)
+	}
+}

examples/wish-exec/main.go

@@ -5,7 +5,6 @@ import (
 	"errors"
 	"fmt"
 	"os"
-	"os/exec"
 	"os/signal"
 	"syscall"
 	"time"
@@ -77,26 +76,35 @@ func (m model) Init() tea.Cmd {
 	return nil
 }
 
-type vimFinishedMsg struct{ err error }
+type cmdFinishedMsg struct{ err error }
 
 func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 	switch msg := msg.(type) {
 	case tea.KeyMsg:
+		// PS: the execs won't work on windows.
 		switch msg.String() {
 		case "e":
-			// PS: this does not work on Windows.
-			c := exec.Command("vim", "file.txt")
-			cmd := tea.ExecProcess(c, func(err error) tea.Msg {
+			c := wish.Command(m.sess, "vim", "file.txt")
+			cmd := tea.Exec(c, func(err error) tea.Msg {
 				if err != nil {
 					log.Error("vim finished", "error", err)
 				}
-				return vimFinishedMsg{err: err}
+				return cmdFinishedMsg{err: err}
+			})
+			return m, cmd
+		case "s":
+			c := wish.Command(m.sess, "bash", "-im")
+			cmd := tea.Exec(c, func(err error) tea.Msg {
+				if err != nil {
+					log.Error("shell finished", "error", err)
+				}
+				return cmdFinishedMsg{err: err}
 			})
 			return m, cmd
 		case "q", "ctrl+c":
 			return m, tea.Quit
 		}
-	case vimFinishedMsg:
+	case cmdFinishedMsg:
 		m.err = msg.err
 		return m, nil
 	}
@@ -108,5 +116,5 @@ func (m model) View() string {
 	if m.err != nil {
 		return m.errStyle.Render(m.err.Error() + "\n")
 	}
-	return m.style.Render("Press 'e' to edit or 'q' to quit...\n")
+	return m.style.Render("Press 'e' to edit, 's' to hop into a shell, or 'q' to quit...\n")
 }