repos / pico

pico services - prose.sh, pastes.sh, imgs.sh, feeds.sh, pgs.sh
git clone https://github.com/picosh/pico.git

pico / pastes
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}