repos / pico

ssh.go

  1package pgs
  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"
 15	"github.com/picosh/pico/db/postgres"
 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/tunkit"
 27	"github.com/picosh/utils"
 28)
 29
 30func createRouter(handler *UploadAssetHandler) proxy.Router {
 31	return func(sh ssh.Handler, s ssh.Session) []wish.Middleware {
 32		return []wish.Middleware{
 33			pipe.Middleware(handler, ""),
 34			list.Middleware(handler),
 35			scp.Middleware(handler),
 36			wishrsync.Middleware(handler),
 37			auth.Middleware(handler),
 38			wsh.PtyMdw(wsh.DeprecatedNotice()),
 39			WishMiddleware(handler),
 40			wsh.LogMiddleware(handler.GetLogger()),
 41		}
 42	}
 43}
 44
 45func withProxy(handler *UploadAssetHandler, 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), otherMiddleware...)(server)
 53	}
 54}
 55
 56func StartSshServer() {
 57	host := utils.GetEnv("PGS_HOST", "0.0.0.0")
 58	port := utils.GetEnv("PGS_SSH_PORT", "2222")
 59	promPort := utils.GetEnv("PGS_PROM_PORT", "9222")
 60	cfg := NewConfigSite()
 61	logger := cfg.Logger
 62	dbpool := postgres.NewDB(cfg.DbURL, cfg.Logger)
 63	defer dbpool.Close()
 64
 65	var st storage.StorageServe
 66	var err error
 67	if cfg.MinioURL == "" {
 68		st, err = storage.NewStorageFS(cfg.StorageDir)
 69	} else {
 70		st, err = storage.NewStorageMinio(cfg.MinioURL, cfg.MinioUser, cfg.MinioPass)
 71	}
 72
 73	if err != nil {
 74		logger.Error(err.Error())
 75		return
 76	}
 77
 78	handler := NewUploadAssetHandler(
 79		dbpool,
 80		cfg,
 81		st,
 82	)
 83
 84	ch := make(chan *db.AnalyticsVisits)
 85	go shared.AnalyticsCollect(ch, dbpool, logger)
 86	apiConfig := &shared.ApiConfig{
 87		Cfg:            cfg,
 88		Dbpool:         dbpool,
 89		Storage:        st,
 90		AnalyticsQueue: ch,
 91	}
 92
 93	webTunnel := &tunkit.WebTunnelHandler{
 94		Logger:      logger,
 95		HttpHandler: createHttpHandler(apiConfig),
 96	}
 97
 98	sshAuth := shared.NewSshAuthHandler(dbpool, logger, cfg)
 99	s, err := wish.NewServer(
100		wish.WithAddress(fmt.Sprintf("%s:%s", host, port)),
101		wish.WithHostKeyPath("ssh_data/term_info_ed25519"),
102		wish.WithPublicKeyAuth(sshAuth.PubkeyAuthHandler),
103		tunkit.WithWebTunnel(webTunnel),
104		withProxy(
105			handler,
106			promwish.Middleware(fmt.Sprintf("%s:%s", host, promPort), "pgs-ssh"),
107		),
108	)
109	if err != nil {
110		logger.Error(err.Error())
111		return
112	}
113
114	done := make(chan os.Signal, 1)
115	signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
116	logger.Info("starting SSH server on", "host", host, "port", port)
117	go func() {
118		if err = s.ListenAndServe(); err != nil {
119			logger.Error("serve", "err", err.Error())
120			os.Exit(1)
121		}
122	}()
123
124	<-done
125	logger.Info("stopping SSH server")
126	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
127	defer func() { cancel() }()
128	if err := s.Shutdown(ctx); err != nil {
129		logger.Error("shutdown", "err", err.Error())
130		os.Exit(1)
131	}
132}