server.go

 1package main
 2
 3import (
 4	"fmt"
 5	"log"
 6	"path/filepath"
 7	"strings"
 8
 9	tea "github.com/charmbracelet/bubbletea"
10	"github.com/charmbracelet/charm/keygen"
11	"github.com/gliderlabs/ssh"
12	gossh "golang.org/x/crypto/ssh"
13)
14
15type Middleware func(ssh.Handler) ssh.Handler
16
17func LoggingMiddleware() Middleware {
18	return func(sh ssh.Handler) ssh.Handler {
19		return func(s ssh.Session) {
20			hpk := s.PublicKey() != nil
21			log.Printf("%s connect %v %v\n", s.RemoteAddr().String(), hpk, s.Command())
22			sh(s)
23			log.Printf("%s disconnect %v %v\n", s.RemoteAddr().String(), hpk, s.Command())
24		}
25	}
26}
27
28func logError(s ssh.Session, err error) {
29	log.Printf("%s error %v: %s\n", s.RemoteAddr().String(), s.Command(), err)
30}
31
32func BubbleTeaMiddleware(bth func(ssh.Session) tea.Model, opts ...tea.ProgramOption) Middleware {
33	return func(sh ssh.Handler) ssh.Handler {
34		return func(s ssh.Session) {
35			m := bth(s)
36			if m != nil {
37				opts = append(opts, tea.WithInput(s), tea.WithOutput(s))
38				p := tea.NewProgram(m, opts...)
39				err := p.Start()
40				if err != nil {
41					logError(s, err)
42				}
43			}
44			sh(s)
45		}
46	}
47}
48
49type Server struct {
50	server *ssh.Server
51	key    gossh.PublicKey
52}
53
54func NewServer(port int, keyPath string, mw ...Middleware) (*Server, error) {
55	s := &Server{server: &ssh.Server{}}
56	s.server.Version = "OpenSSH_7.6p1"
57	s.server.Addr = fmt.Sprintf(":%d", port)
58	s.server.PasswordHandler = s.passHandler
59	s.server.PublicKeyHandler = s.authHandler
60	kps := strings.Split(keyPath, string(filepath.Separator))
61	kp := strings.Join(kps[:len(kps)-1], string(filepath.Separator))
62	n := strings.TrimRight(kps[len(kps)-1], "_ed25519")
63	_, err := keygen.NewSSHKeyPair(kp, n, nil, "ed25519")
64	if err != nil {
65		return nil, err
66	}
67	k := ssh.HostKeyFile(keyPath)
68	err = s.server.SetOption(k)
69	if err != nil {
70		return nil, err
71	}
72	h := func(s ssh.Session) {}
73	for _, m := range mw {
74		h = m(h)
75	}
76	s.server.Handler = h
77	return s, nil
78}
79
80func (srv *Server) sessionHandler(s ssh.Session) {
81}
82
83func (srv *Server) authHandler(ctx ssh.Context, key ssh.PublicKey) bool {
84	return true
85}
86
87func (srv *Server) passHandler(ctx ssh.Context, pass string) bool {
88	return true
89}
90
91func (srv *Server) Start() error {
92	log.Printf("Starting SSH server on %s\n", srv.server.Addr)
93	return srv.server.ListenAndServe()
94}