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 SessionHandler func(ssh.Session) (tea.Model, error)
16
17type Server struct {
18 server *ssh.Server
19 key gossh.PublicKey
20 handler SessionHandler
21}
22
23func NewServer(port int, keyPath string, handler SessionHandler) (*Server, error) {
24 s := &Server{
25 server: &ssh.Server{},
26 handler: handler,
27 }
28 s.server.Version = "OpenSSH_7.6p1"
29 s.server.Addr = fmt.Sprintf(":%d", port)
30 s.server.Handler = s.sessionHandler
31 s.server.PasswordHandler = s.passHandler
32 s.server.PublicKeyHandler = s.authHandler
33 kps := strings.Split(keyPath, string(filepath.Separator))
34 kp := strings.Join(kps[:len(kps)-1], string(filepath.Separator))
35 n := strings.TrimRight(kps[len(kps)-1], "_ed25519")
36 _, err := keygen.NewSSHKeyPair(kp, n, nil, "ed25519")
37 if err != nil {
38 return nil, err
39 }
40 k := ssh.HostKeyFile(keyPath)
41 err = s.server.SetOption(k)
42 if err != nil {
43 return nil, err
44 }
45 return s, nil
46}
47
48func (srv *Server) sessionHandler(s ssh.Session) {
49 hpk := s.PublicKey() != nil
50 log.Printf("%s connect %v %v\n", s.RemoteAddr().String(), hpk, s.Command())
51 m, err := srv.handler(s)
52 if err != nil {
53 log.Printf("%s error %v %s\n", s.RemoteAddr().String(), hpk, err)
54 s.Exit(1)
55 return
56 }
57 if m != nil {
58 p := tea.NewProgram(m, tea.WithInput(s), tea.WithOutput(s))
59 err = p.Start()
60 if err != nil {
61 log.Printf("%s error %v %s\n", s.RemoteAddr().String(), hpk, err)
62 s.Exit(1)
63 return
64 }
65 }
66 log.Printf("%s disconnect %v %v\n", s.RemoteAddr().String(), hpk, s.Command())
67}
68
69func (srv *Server) authHandler(ctx ssh.Context, key ssh.PublicKey) bool {
70 return true
71}
72
73func (srv *Server) passHandler(ctx ssh.Context, pass string) bool {
74 return true
75}
76
77func (srv *Server) Start() error {
78 return srv.server.ListenAndServe()
79}