Skip to content

Instantly share code, notes, and snippets.

@tluyben

tluyben/catchout.go

Created September 12, 2024 13:38
Show Gist options
  • Save tluyben/95c300739f4a5aa12b97bbb83ca2b514 to your computer and use it in GitHub Desktop.
Save tluyben/95c300739f4a5aa12b97bbb83ca2b514 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
catchout.go
package main
import (
"fmt"
"io"
"os"
"os/exec"
"os/signal"
"strconv"
"strings"
"sync"
"syscall"
"github.com/creack/pty"
"golang.org/x/term"
)
type Screen struct {
content [][]rune
width int
height int
cursorX int
cursorY int
}
func NewScreen(width, height int) *Screen {
content := make([][]rune, height)
for i := range content {
content[i] = make([]rune, width)
for j := range content[i] {
content[i][j] = ' '
}
}
return &Screen{content: content, width: width, height: height}
}
func (s *Screen) Write(p []byte) (n int, err error) {
for i := 0; i < len(p); i++ {
if p[i] == 0x1b && i+1 < len(p) && p[i+1] == '[' {
i += 2
cmd := ""
for i < len(p) && ((p[i] >= 0x30 && p[i] <= 0x3f) || (p[i] >= 0x20 && p[i] <= 0x2f)) {
cmd += string(p[i])
i++
}
if i < len(p) {
cmd += string(p[i])
}
s.handleEscapeSequence(cmd)
} else if p[i] == '\n' {
s.cursorY++
s.cursorX = 0
} else if p[i] == '\r' {
s.cursorX = 0
} else {
if s.cursorX < s.width && s.cursorY < s.height {
s.content[s.cursorY][s.cursorX] = rune(p[i])
s.cursorX++
}
}
if s.cursorX >= s.width {
s.cursorY++
s.cursorX = 0
}
if s.cursorY >= s.height {
s.scrollUp()
}
}
return len(p), nil
}
func (s *Screen) handleEscapeSequence(cmd string) {
switch {
case cmd == "H":
s.cursorX, s.cursorY = 0, 0
case cmd == "2J":
for i := range s.content {
for j := range s.content[i] {
s.content[i][j] = ' '
}
}
case strings.HasSuffix(cmd, "H"):
parts := strings.Split(strings.TrimSuffix(cmd, "H"), ";")
if len(parts) == 2 {
y, _ := strconv.Atoi(parts[0])
x, _ := strconv.Atoi(parts[1])
s.cursorY = y - 1
s.cursorX = x - 1
}
}
}
func (s *Screen) scrollUp() {
copy(s.content, s.content[1:])
s.content[s.height-1] = make([]rune, s.width)
for i := range s.content[s.height-1] {
s.content[s.height-1][i] = ' '
}
s.cursorY = s.height - 1
}
func (s *Screen) String() string {
var result strings.Builder
for _, row := range s.content {
result.WriteString(string(row))
result.WriteString("\n")
}
return result.String()
}
func main() {
if len(os.Args) < 2 {
fmt.Println("Usage: go run main.go <command>")
os.Exit(1)
}
cmd := exec.Command(os.Args[1], os.Args[2:]...)
ptmx, err := pty.Start(cmd)
if err != nil {
fmt.Println("Error creating pseudo-terminal:", err)
os.Exit(1)
}
defer ptmx.Close()
width, height, err := term.GetSize(int(os.Stdout.Fd()))
if err != nil {
fmt.Println("Error getting terminal size:", err)
os.Exit(1)
}
ch := make(chan os.Signal, 1)
signal.Notify(ch, syscall.SIGWINCH)
go func() {
for range ch {
if w, h, err := term.GetSize(int(os.Stdout.Fd())); err == nil {
pty.Setsize(ptmx, &pty.Winsize{Rows: uint16(h), Cols: uint16(w)})
}
}
}()
ch <- syscall.SIGWINCH
oldState, err := term.MakeRaw(int(os.Stdin.Fd()))
if err != nil {
fmt.Println("Error setting raw mode:", err)
os.Exit(1)
}
defer term.Restore(int(os.Stdin.Fd()), oldState)
screen := NewScreen(width, height)
var wg sync.WaitGroup
// Handle input
wg.Add(1)
go func() {
defer wg.Done()
io.Copy(ptmx, os.Stdin)
}()
// Handle output
wg.Add(1)
go func() {
defer wg.Done()
io.Copy(io.MultiWriter(os.Stdout, screen), ptmx)
}()
// Wait for the command to finish
cmd.Wait()
// Close the pty to signal the goroutines to stop
ptmx.Close()
// Wait for goroutines to finish
wg.Wait()
// Restore terminal state
term.Restore(int(os.Stdin.Fd()), oldState)
// Print captured output
fmt.Println("\nCaptured output:")
fmt.Print(screen.String())
fmt.Printf("\nTerminal size: %dx%d\n", width, height)
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment