Antonio Mika
·
08 Oct 24
ssh.go
1package pastes
2
3import (
4 "context"
5 "fmt"
6 "os"
7 "os/signal"
8 "syscall"
9 "time"
10
11 "github.com/charmbracelet/promwish"
12 "github.com/charmbracelet/ssh"
13 "github.com/charmbracelet/wish"
14 "github.com/picosh/pico/db/postgres"
15 "github.com/picosh/pico/filehandlers"
16 "github.com/picosh/pico/shared"
17 "github.com/picosh/pico/shared/storage"
18 wsh "github.com/picosh/pico/wish"
19 "github.com/picosh/send/auth"
20 "github.com/picosh/send/list"
21 "github.com/picosh/send/pipe"
22 wishrsync "github.com/picosh/send/protocols/rsync"
23 "github.com/picosh/send/protocols/scp"
24 "github.com/picosh/send/protocols/sftp"
25 "github.com/picosh/send/proxy"
26 "github.com/picosh/utils"
27)
28
29func createRouter(handler *filehandlers.FileHandlerRouter) proxy.Router {
30 return func(sh ssh.Handler, s ssh.Session) []wish.Middleware {
31 return []wish.Middleware{
32 pipe.Middleware(handler, ""),
33 list.Middleware(handler),
34 scp.Middleware(handler),
35 wishrsync.Middleware(handler),
36 auth.Middleware(handler),
37 wsh.PtyMdw(wsh.DeprecatedNotice()),
38 wsh.LogMiddleware(handler.GetLogger()),
39 }
40 }
41}
42
43func withProxy(handler *filehandlers.FileHandlerRouter, otherMiddleware ...wish.Middleware) ssh.Option {
44 return func(server *ssh.Server) error {
45 err := sftp.SSHOption(handler)(server)
46 if err != nil {
47 return err
48 }
49
50 return proxy.WithProxy(createRouter(handler), otherMiddleware...)(server)
51 }
52}
53
54func StartSshServer() {
55 host := utils.GetEnv("PASTES_HOST", "0.0.0.0")
56 port := utils.GetEnv("PASTES_SSH_PORT", "2222")
57 promPort := utils.GetEnv("PASTES_PROM_PORT", "9222")
58 cfg := NewConfigSite()
59 logger := cfg.Logger
60 dbh := postgres.NewDB(cfg.DbURL, cfg.Logger)
61 defer dbh.Close()
62 hooks := &FileHooks{
63 Cfg: cfg,
64 Db: dbh,
65 }
66
67 var st storage.StorageServe
68 var err error
69 if cfg.MinioURL == "" {
70 st, err = storage.NewStorageFS(cfg.StorageDir)
71 } else {
72 st, err = storage.NewStorageMinio(cfg.MinioURL, cfg.MinioUser, cfg.MinioPass)
73 }
74
75 if err != nil {
76 logger.Error(err.Error())
77 return
78 }
79
80 fileMap := map[string]filehandlers.ReadWriteHandler{
81 "fallback": filehandlers.NewScpPostHandler(dbh, cfg, hooks, st),
82 }
83 handler := filehandlers.NewFileHandlerRouter(cfg, dbh, fileMap)
84 sshAuth := shared.NewSshAuthHandler(dbh, logger, cfg)
85 s, err := wish.NewServer(
86 wish.WithAddress(fmt.Sprintf("%s:%s", host, port)),
87 wish.WithHostKeyPath("ssh_data/term_info_ed25519"),
88 wish.WithPublicKeyAuth(sshAuth.PubkeyAuthHandler),
89 withProxy(
90 handler,
91 promwish.Middleware(fmt.Sprintf("%s:%s", host, promPort), "pastes-ssh"),
92 ),
93 )
94 if err != nil {
95 logger.Error(err.Error())
96 return
97 }
98
99 done := make(chan os.Signal, 1)
100 signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
101 logger.Info("Starting SSH server", "host", host, "port", port)
102 go func() {
103 if err = s.ListenAndServe(); err != nil {
104 logger.Error(err.Error())
105 }
106 }()
107
108 <-done
109 logger.Info("Stopping SSH server")
110 ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
111 defer func() { cancel() }()
112 if err := s.Shutdown(ctx); err != nil {
113 logger.Error(err.Error())
114 }
115}