porterm/main.go at main · ryu-ryuk/porterm · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
// package main

// import (
// 	"log"

// 	tea "github.com/charmbracelet/bubbletea"
// 	"github.com/charmbracelet/ssh"
// 	"github.com/charmbracelet/wish"
// 	bubbletea "github.com/charmbracelet/wish/bubbletea"
// 	"github.com/charmbracelet/wish/logging"
// 	"github.com/charmbracelet/wish/recover"

// 	"porterm/model"
// )

// func main() {
// 	s, err := wish.NewServer(
// 		wish.WithAddress(":22"),
// 		wish.WithMiddleware(
// 			recover.Middleware(), // catch panics
// 			logging.Middleware(), // log connections
// 			bubbletea.Middleware(
// 				func(s ssh.Session) (tea.Model, []tea.ProgramOption) {
// 					return model.New(), nil
// 				},
// 			),
// 		),
// 	)
// 	if err != nil {
// 		log.Fatalf("failed to create server: %s", err)
// 	}

// 	log.Println("serving at :22")
// 	if err := s.ListenAndServe(); err != nil {
// 		log.Fatalf("failed to start server: %s", err)
// 	}
// }

// // -------------------------------------------
package main

import (
	"context"
	"fmt"
	"log"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"porterm/model"

	tea "github.com/charmbracelet/bubbletea"
	"github.com/charmbracelet/ssh"
	"github.com/charmbracelet/wish"
)

const (
	host        = "0.0.0.0"
	port        = 22
	hostKeyPath = ".ssh/term_host_rsa"
)

func main() {
	log.SetOutput(os.Stderr)

	s, err := wish.NewServer(
		wish.WithAddress(fmt.Sprintf("%s:%d", host, port)),
		wish.WithHostKeyPath(hostKeyPath),
		wish.WithMiddleware(handleSSHConnection()),
	)
	if err != nil {
		log.Fatalf("could not start ssh server: %v", err)
	}

	done := make(chan os.Signal, 1)
	signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)

	log.Printf("starting porterm ssh server on %s:%d", host, port)
	go func() {
		if err = s.ListenAndServe(); err != nil && err != ssh.ErrServerClosed {
			log.Fatalf("ssh server failed: %v", err)
		}
	}()

	<-done
	log.Println("stopping porterm ssh server gracefully...")
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	if err := s.Shutdown(ctx); err != nil {
		log.Fatalf("ssh server shutdown failed: %v", err)
	}
	log.Println("porterm ssh server stopped.")
}

func handleSSHConnection() wish.Middleware {
	return func(next ssh.Handler) ssh.Handler {
		return func(s ssh.Session) {
			// manually set TERM
			_ = os.Setenv("TERM", "xterm-256color")

			// apply all session env vars (important for TERM, LANG, etc)
			for _, env := range s.Environ() {
				parts := strings.SplitN(env, "=", 2)
				if len(parts) == 2 {
					_ = os.Setenv(parts[0], parts[1])
				}
			}

			p := tea.NewProgram(
				model.New(),
				tea.WithInput(s),
				tea.WithOutput(s),
				tea.WithAltScreen(),
				tea.WithMouseCellMotion(),
			)

			if _, err := p.Run(); err != nil {
				log.Printf("bubbletea error in session %s: %v", s.RemoteAddr(), err)
			}

			next(s)
		}
	}
}