repos / pico

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

pico / feeds
Antonio Mika · 08 Oct 24

ssh.go

  1package feeds
  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, ".txt"),
 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("LISTS_HOST", "0.0.0.0")
 56	port := utils.GetEnv("LISTS_SSH_PORT", "2222")
 57	promPort := utils.GetEnv("LISTS_PROM_PORT", "9222")
 58	cfg := NewConfigSite()
 59	logger := cfg.Logger
 60	dbh := postgres.NewDB(cfg.DbURL, cfg.Logger)
 61	defer dbh.Close()
 62
 63	hooks := &FeedHooks{
 64		Cfg: cfg,
 65		Db:  dbh,
 66	}
 67
 68	var st storage.StorageServe
 69	var err error
 70	if cfg.MinioURL == "" {
 71		st, err = storage.NewStorageFS(cfg.StorageDir)
 72	} else {
 73		st, err = storage.NewStorageMinio(cfg.MinioURL, cfg.MinioUser, cfg.MinioPass)
 74	}
 75
 76	if err != nil {
 77		logger.Error(err.Error())
 78		return
 79	}
 80
 81	fileMap := map[string]filehandlers.ReadWriteHandler{
 82		"fallback": filehandlers.NewScpPostHandler(dbh, cfg, hooks, st),
 83	}
 84	handler := filehandlers.NewFileHandlerRouter(cfg, dbh, fileMap)
 85
 86	sshAuth := shared.NewSshAuthHandler(dbh, logger, cfg)
 87	s, err := wish.NewServer(
 88		wish.WithAddress(fmt.Sprintf("%s:%s", host, port)),
 89		wish.WithHostKeyPath("ssh_data/term_info_ed25519"),
 90		wish.WithPublicKeyAuth(sshAuth.PubkeyAuthHandler),
 91		withProxy(
 92			handler,
 93			promwish.Middleware(fmt.Sprintf("%s:%s", host, promPort), "feeds-ssh"),
 94		),
 95	)
 96	if err != nil {
 97		logger.Error(err.Error())
 98		return
 99	}
100
101	done := make(chan os.Signal, 1)
102	signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
103	logger.Info("Starting SSH server", "host", host, "port", port)
104	go func() {
105		if err = s.ListenAndServe(); err != nil {
106			logger.Error(err.Error())
107		}
108	}()
109
110	<-done
111	logger.Info("Stopping SSH server")
112	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
113	defer func() { cancel() }()
114	if err := s.Shutdown(ctx); err != nil {
115		logger.Error(err.Error())
116	}
117}