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