repos / pico

ssh.go

  1package prose
  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	uploadimgs "github.com/picosh/pico/filehandlers/imgs"
 17	"github.com/picosh/pico/filehandlers/util"
 18	"github.com/picosh/pico/shared"
 19	"github.com/picosh/pico/shared/storage"
 20	wsh "github.com/picosh/pico/wish"
 21	"github.com/picosh/send/list"
 22	"github.com/picosh/send/pipe"
 23	"github.com/picosh/send/proxy"
 24	"github.com/picosh/send/send/auth"
 25	wishrsync "github.com/picosh/send/send/rsync"
 26	"github.com/picosh/send/send/scp"
 27	"github.com/picosh/send/send/sftp"
 28)
 29
 30func createRouter(handler *filehandlers.FileHandlerRouter, cliHandler *CliHandler) proxy.Router {
 31	return func(sh ssh.Handler, s ssh.Session) []wish.Middleware {
 32		return []wish.Middleware{
 33			pipe.Middleware(handler, ".md"),
 34			list.Middleware(handler),
 35			scp.Middleware(handler),
 36			wishrsync.Middleware(handler),
 37			auth.Middleware(handler),
 38			wsh.PtyMdw(wsh.DeprecatedNotice()),
 39			WishMiddleware(cliHandler),
 40			wsh.LogMiddleware(handler.GetLogger()),
 41		}
 42	}
 43}
 44
 45func withProxy(handler *filehandlers.FileHandlerRouter, cliHandler *CliHandler, otherMiddleware ...wish.Middleware) ssh.Option {
 46	return func(server *ssh.Server) error {
 47		err := sftp.SSHOption(handler)(server)
 48		if err != nil {
 49			return err
 50		}
 51
 52		return proxy.WithProxy(createRouter(handler, cliHandler), otherMiddleware...)(server)
 53	}
 54}
 55
 56func StartSshServer() {
 57	host := shared.GetEnv("PROSE_HOST", "0.0.0.0")
 58	port := shared.GetEnv("PROSE_SSH_PORT", "2222")
 59	promPort := shared.GetEnv("PROSE_PROM_PORT", "9222")
 60	cfg := NewConfigSite()
 61	logger := cfg.Logger
 62	dbh := postgres.NewDB(cfg.DbURL, cfg.Logger)
 63	defer dbh.Close()
 64	hooks := &MarkdownHooks{
 65		Cfg: cfg,
 66		Db:  dbh,
 67	}
 68
 69	var st storage.StorageServe
 70	var err error
 71	if cfg.MinioURL == "" {
 72		st, err = storage.NewStorageFS(cfg.StorageDir)
 73	} else {
 74		st, err = storage.NewStorageMinio(cfg.MinioURL, cfg.MinioUser, cfg.MinioPass)
 75	}
 76
 77	if err != nil {
 78		logger.Error(err.Error())
 79		return
 80	}
 81
 82	fileMap := map[string]filehandlers.ReadWriteHandler{
 83		".md":      filehandlers.NewScpPostHandler(dbh, cfg, hooks, st),
 84		".css":     filehandlers.NewScpPostHandler(dbh, cfg, hooks, st),
 85		"fallback": uploadimgs.NewUploadImgHandler(dbh, cfg, st),
 86	}
 87	handler := filehandlers.NewFileHandlerRouter(cfg, dbh, fileMap)
 88	handler.Spaces = []string{cfg.Space, "imgs"}
 89
 90	cliHandler := &CliHandler{
 91		Logger: logger,
 92		DBPool: dbh,
 93	}
 94
 95	sshAuth := util.NewSshAuthHandler(dbh, logger, cfg)
 96	s, err := wish.NewServer(
 97		wish.WithAddress(fmt.Sprintf("%s:%s", host, port)),
 98		wish.WithHostKeyPath("ssh_data/term_info_ed25519"),
 99		wish.WithPublicKeyAuth(sshAuth.PubkeyAuthHandler),
100		withProxy(
101			handler,
102			cliHandler,
103			promwish.Middleware(fmt.Sprintf("%s:%s", host, promPort), "prose-ssh"),
104		),
105	)
106	if err != nil {
107		logger.Error(err.Error())
108		return
109	}
110
111	done := make(chan os.Signal, 1)
112	signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
113	logger.Info("Starting SSH server", "host", host, "port", port)
114	go func() {
115		if err = s.ListenAndServe(); err != nil {
116			logger.Error(err.Error())
117		}
118	}()
119
120	<-done
121	logger.Info("Stopping SSH server")
122	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
123	defer func() { cancel() }()
124	if err := s.Shutdown(ctx); err != nil {
125		logger.Error(err.Error())
126	}
127}