Skip to content

Instantly share code, notes, and snippets.

@ivanvc
Last active June 24, 2023 06:11
Show Gist options
  • Save ivanvc/cb4580c02e8bb04d17ef286c9183098b to your computer and use it in GitHub Desktop.
Save ivanvc/cb4580c02e8bb04d17ef286c9183098b to your computer and use it in GitHub Desktop.
Wish program with Bubbletea middleware executing a remote command
package main
import (
"context"
"errors"
"os"
"os/exec"
"os/signal"
"syscall"
"time"
tea "github.com/charmbracelet/bubbletea"
"github.com/charmbracelet/log"
"github.com/charmbracelet/ssh"
"github.com/charmbracelet/wish"
bm "github.com/charmbracelet/wish/bubbletea"
lm "github.com/charmbracelet/wish/logging"
)
func main() {
s, err := wish.NewServer(
wish.WithAddress(":23234"),
wish.WithMiddleware(
bm.Middleware(teaHandler),
lm.Middleware(),
),
)
if err != nil {
log.Error("could not start server", "error", err)
}
done := make(chan os.Signal, 1)
signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
log.Info("Starting SSH server")
go func() {
if err = s.ListenAndServe(); err != nil && !errors.Is(err, ssh.ErrServerClosed) {
log.Error("could not start server", "error", err)
done <- nil
}
}()
<-done
log.Info("Stopping SSH server")
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer func() { cancel() }()
if err := s.Shutdown(ctx); err != nil && !errors.Is(err, ssh.ErrServerClosed) {
log.Error("could not stop server", "error", err)
}
}
func teaHandler(s ssh.Session) (tea.Model, []tea.ProgramOption) {
_, _, active := s.Pty()
if !active {
wish.Fatalln(s, "no active terminal, skipping")
return nil, nil
}
m := model{s}
return m, []tea.ProgramOption{
tea.WithAltScreen(),
}
}
type model struct {
session ssh.Session
}
func (m model) Init() tea.Cmd {
return runCommand(m.session)
}
func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
return m, nil
}
func (m model) View() string {
return ""
}
func runCommand(s ssh.Session) tea.Cmd {
c := exec.Command("journalctl", "-f")
return tea.ExecProcess(c, func(err error) tea.Msg {
if err != nil {
log.Error(err)
}
return tea.Quit()
})
}
package main
import (
"context"
"errors"
"fmt"
"io"
"os"
"os/exec"
"os/signal"
"syscall"
"time"
"unsafe"
tea "github.com/charmbracelet/bubbletea"
"github.com/charmbracelet/log"
"github.com/charmbracelet/ssh"
"github.com/charmbracelet/wish"
bm "github.com/charmbracelet/wish/bubbletea"
lm "github.com/charmbracelet/wish/logging"
"github.com/creack/pty"
)
func main() {
s, err := wish.NewServer(
wish.WithAddress(":23234"),
wish.WithMiddleware(
bm.Middleware(teaHandler),
lm.Middleware(),
),
)
if err != nil {
log.Error("could not start server", "error", err)
}
done := make(chan os.Signal, 1)
signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
log.Info("Starting SSH server")
go func() {
if err = s.ListenAndServe(); err != nil && !errors.Is(err, ssh.ErrServerClosed) {
log.Error("could not start server", "error", err)
done <- nil
}
}()
<-done
log.Info("Stopping SSH server")
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer func() { cancel() }()
if err := s.Shutdown(ctx); err != nil && !errors.Is(err, ssh.ErrServerClosed) {
log.Error("could not stop server", "error", err)
}
}
func teaHandler(s ssh.Session) (tea.Model, []tea.ProgramOption) {
_, _, active := s.Pty()
if !active {
wish.Fatalln(s, "no active terminal, skipping")
return nil, nil
}
m := model{s}
return m, []tea.ProgramOption{tea.WithAltScreen()}
}
type model struct {
session ssh.Session
}
func (m model) Init() tea.Cmd {
return runCommand(m.session)
}
func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
return m, nil
}
func (m model) View() string {
return ""
}
func runCommand(s ssh.Session) tea.Cmd {
c := exec.Command("/bin/bash", "-i")
pty, winChan, _ := s.Pty()
c.Env = append(c.Env, fmt.Sprintf("TERM=%s", pty.Term))
return tea.Exec(&execProc{cmd: c, winChan: winChan, pty: pty}, func(err error) tea.Msg {
if err != nil {
log.Error(err)
}
return tea.Quit()
})
}
type execProc struct {
cmd *exec.Cmd
winChan <-chan ssh.Window
stdin io.Reader
stdout io.Writer
pty ssh.Pty
}
func (e *execProc) SetStdin(r io.Reader) {
e.stdin = r
}
func (e *execProc) SetStdout(w io.Writer) {
e.stdout = w
}
func (e *execProc) SetStderr(w io.Writer) {}
func (e *execProc) Run() error {
f, err := pty.Start(e.cmd)
if err != nil {
return err
}
setWinsize(f, e.pty.Window.Width, e.pty.Window.Height)
go func() {
for win := range e.winChan {
setWinsize(f, win.Width, win.Height)
}
}()
go func() {
io.Copy(f, e.stdin)
}()
io.Copy(e.stdout, f)
return e.cmd.Wait()
}
func setWinsize(f *os.File, w, h int) {
syscall.Syscall(syscall.SYS_IOCTL, f.Fd(), uintptr(syscall.TIOCSWINSZ),
uintptr(unsafe.Pointer(&struct{ h, w, x, y uint16 }{uint16(h), uint16(w), 0, 0})))
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment