feat: add wish.Command and wish.Cmd

cmd.go

@@ -0,0 +1,71 @@
+package wish
+
+import (
+	"context"
+	"io"
+	"os/exec"
+
+	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 {
+	cmd := exec.CommandContext(ctx, name, args...)
+	return &Cmd{s, cmd}
+}
+
+// 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 {
+	sess ssh.Session
+	cmd  *exec.Cmd
+}
+
+// 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 {
+	ppty, winCh, ok := c.sess.Pty()
+	if !ok {
+		c.cmd.Stdin, c.cmd.Stdout, c.cmd.Stderr = c.sess, c.sess, c.sess
+		return c.cmd.Run()
+	}
+	return c.doRun(ppty, winCh)
+}
+
+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,147 @@
+package wish
+
+import (
+	"bytes"
+	"runtime"
+	"strings"
+	"testing"
+	"time"
+
+	"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) {
+	if runtime.GOOS == "windows" {
+		t.Skip()
+	}
+	tmp := t.TempDir()
+	srv := &ssh.Server{
+		Handler: func(s ssh.Session) {
+			runEcho(s, "hello")
+			runEnv(s, []string{"HELLO=world"})
+			runPwd(s, tmp)
+			// for some reason sometimes on macos github action runners,
+			// it cuts parts of the output.
+			time.Sleep(100 * time.Millisecond)
+		},
+	}
+	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) {
+	if runtime.GOOS == "windows" {
+		t.Skip()
+	}
+	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)
+	}
+}

cmd_unix.go

@@ -0,0 +1,13 @@
+//go:build darwin || dragonfly || freebsd || linux || netbsd || openbsd || solaris
+// +build darwin dragonfly freebsd linux netbsd openbsd solaris
+
+package wish
+
+import "github.com/charmbracelet/ssh"
+
+func (c *Cmd) doRun(ppty ssh.Pty, _ <-chan ssh.Window) error {
+	if err := ppty.Start(c.cmd); err != nil {
+		return err
+	}
+	return c.cmd.Wait()
+}

cmd_windows.go

@@ -0,0 +1,29 @@
+//go:build windows
+// +build windows
+
+package wish
+
+import (
+	"fmt"
+	"time"
+
+	"github.com/charmbracelet/ssh"
+)
+
+func (c *Cmd) doRun(ppty ssh.Pty, _ <-chan ssh.Window) error {
+	if err := ppty.Start(c.cmd); err != nil {
+		return err
+	}
+
+	start := time.Now()
+	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())
+	}
+	return nil
+}

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

go.mod

@@ -7,7 +7,7 @@ require (
 	github.com/charmbracelet/keygen v0.5.0
 	github.com/charmbracelet/lipgloss v0.9.1
 	github.com/charmbracelet/log v0.3.1
-	github.com/charmbracelet/ssh v0.0.0-20240118173142-6d7cf11c8371
+	github.com/charmbracelet/ssh v0.0.0-20240129182809-006ab784ccc7
 	github.com/go-git/go-git/v5 v5.11.0
 	github.com/google/go-cmp v0.6.0
 	github.com/hashicorp/golang-lru/v2 v2.0.7

go.sum

@@ -19,8 +19,8 @@ github.com/charmbracelet/lipgloss v0.9.1 h1:PNyd3jvaJbg4jRHKWXnCj1akQm4rh8dbEzN1
 github.com/charmbracelet/lipgloss v0.9.1/go.mod h1:1mPmG4cxScwUQALAAnacHaigiiHB9Pmr+v1VEawJl6I=
 github.com/charmbracelet/log v0.3.1 h1:TjuY4OBNbxmHWSwO3tosgqs5I3biyY8sQPny/eCMTYw=
 github.com/charmbracelet/log v0.3.1/go.mod h1:OR4E1hutLsax3ZKpXbgUqPtTjQfrh1pG3zwHGWuuq8g=
-github.com/charmbracelet/ssh v0.0.0-20240118173142-6d7cf11c8371 h1:lgr2JbKDeq13Ar9gwmwELlJAF6R13pP3Kp37C9KLVM8=
-github.com/charmbracelet/ssh v0.0.0-20240118173142-6d7cf11c8371/go.mod h1:l/6/exIt0QMHEW5IiqoY9HyPlSuXAm6xkngFIJgiEYI=
+github.com/charmbracelet/ssh v0.0.0-20240129182809-006ab784ccc7 h1:DL71svOeh8OSFdPc8GOAX9RMKb9+Ss0ldENOOkrL8/s=
+github.com/charmbracelet/ssh v0.0.0-20240129182809-006ab784ccc7/go.mod h1:l/6/exIt0QMHEW5IiqoY9HyPlSuXAm6xkngFIJgiEYI=
 github.com/charmbracelet/x/errors v0.0.0-20240117030013-d31dba354651 h1:3RXpZWGWTOeVXCTv0Dnzxdv/MhNUkBfEcbaTY0zrTQI=
 github.com/charmbracelet/x/errors v0.0.0-20240117030013-d31dba354651/go.mod h1:2P0UgXMEa6TsToMSuFqKFQR+fZTO9CNGUNokkPatT/0=
 github.com/charmbracelet/x/exp/term v0.0.0-20240117031359-6e25c76a1efe h1:HeRgHWxOTu7l73rKsa5BRAeaUenmNyomiPCUHXv/y14=