repos / pico

ssh.go

  1package pico
  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	bm "github.com/charmbracelet/wish/bubbletea"
 15	"github.com/muesli/termenv"
 16	"github.com/picosh/pico/db/postgres"
 17	"github.com/picosh/pico/shared"
 18	"github.com/picosh/pico/tui"
 19	wsh "github.com/picosh/pico/wish"
 20	"github.com/picosh/send/auth"
 21	"github.com/picosh/send/list"
 22	"github.com/picosh/send/pipe"
 23	wishrsync "github.com/picosh/send/protocols/rsync"
 24	"github.com/picosh/send/protocols/scp"
 25	"github.com/picosh/send/protocols/sftp"
 26	"github.com/picosh/send/proxy"
 27	"github.com/picosh/utils"
 28)
 29
 30func authHandler(ctx ssh.Context, key ssh.PublicKey) bool {
 31	shared.SetPublicKey(ctx, key)
 32	return true
 33}
 34
 35func createRouter(cfg *shared.ConfigSite, handler *UploadHandler, cliHandler *CliHandler) proxy.Router {
 36	return func(sh ssh.Handler, s ssh.Session) []wish.Middleware {
 37		return []wish.Middleware{
 38			pipe.Middleware(handler, ""),
 39			list.Middleware(handler),
 40			scp.Middleware(handler),
 41			wishrsync.Middleware(handler),
 42			auth.Middleware(handler),
 43			wsh.PtyMdw(bm.MiddlewareWithColorProfile(tui.CmsMiddleware(cfg), termenv.ANSI256)),
 44			WishMiddleware(cliHandler),
 45			wsh.LogMiddleware(handler.GetLogger()),
 46		}
 47	}
 48}
 49
 50func withProxy(cfg *shared.ConfigSite, handler *UploadHandler, cliHandler *CliHandler, otherMiddleware ...wish.Middleware) ssh.Option {
 51	return func(server *ssh.Server) error {
 52		err := sftp.SSHOption(handler)(server)
 53		if err != nil {
 54			return err
 55		}
 56
 57		return proxy.WithProxy(createRouter(cfg, handler, cliHandler), otherMiddleware...)(server)
 58	}
 59}
 60
 61func StartSshServer() {
 62	host := utils.GetEnv("PICO_HOST", "0.0.0.0")
 63	port := utils.GetEnv("PICO_SSH_PORT", "2222")
 64	promPort := utils.GetEnv("PICO_PROM_PORT", "9222")
 65	cfg := NewConfigSite()
 66	logger := cfg.Logger
 67	dbpool := postgres.NewDB(cfg.DbURL, cfg.Logger)
 68	defer dbpool.Close()
 69
 70	handler := NewUploadHandler(
 71		dbpool,
 72		cfg,
 73	)
 74	cliHandler := &CliHandler{
 75		Logger: logger,
 76		DBPool: dbpool,
 77	}
 78
 79	s, err := wish.NewServer(
 80		wish.WithAddress(fmt.Sprintf("%s:%s", host, port)),
 81		wish.WithHostKeyPath("ssh_data/term_info_ed25519"),
 82		wish.WithPublicKeyAuth(authHandler),
 83		withProxy(
 84			cfg,
 85			handler,
 86			cliHandler,
 87			promwish.Middleware(fmt.Sprintf("%s:%s", host, promPort), "pico-ssh"),
 88		),
 89	)
 90	if err != nil {
 91		logger.Error(err.Error())
 92		return
 93	}
 94
 95	done := make(chan os.Signal, 1)
 96	signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
 97	logger.Info("starting SSH server on", "host", host, "port", port)
 98	go func() {
 99		if err = s.ListenAndServe(); err != nil {
100			logger.Error("serve", "err", err.Error())
101			os.Exit(1)
102		}
103	}()
104
105	<-done
106	logger.Info("stopping SSH server")
107	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
108	defer func() { cancel() }()
109	if err := s.Shutdown(ctx); err != nil {
110		logger.Error("shutdown", "err", err.Error())
111		os.Exit(1)
112	}
113}