repos / pico

pico services mono repo
git clone https://github.com/picosh/pico.git

pico / pkg / apps / pico
Antonio Mika  ·  2025-04-10

ssh.go

  1package pico
  2
  3import (
  4	"context"
  5	"os"
  6	"os/signal"
  7	"syscall"
  8	"time"
  9
 10	"git.sr.ht/~rockorager/vaxis"
 11	"github.com/picosh/pico/pkg/db/postgres"
 12	"github.com/picosh/pico/pkg/pssh"
 13	"github.com/picosh/pico/pkg/send/auth"
 14	"github.com/picosh/pico/pkg/send/list"
 15	"github.com/picosh/pico/pkg/send/pipe"
 16	"github.com/picosh/pico/pkg/send/protocols/rsync"
 17	"github.com/picosh/pico/pkg/send/protocols/scp"
 18	"github.com/picosh/pico/pkg/send/protocols/sftp"
 19	"github.com/picosh/pico/pkg/shared"
 20	"github.com/picosh/pico/pkg/tui"
 21	"github.com/picosh/utils"
 22	"golang.org/x/crypto/ssh"
 23)
 24
 25func createTui(shrd *tui.SharedModel) pssh.SSHServerMiddleware {
 26	return func(next pssh.SSHServerHandler) pssh.SSHServerHandler {
 27		return func(sesh *pssh.SSHServerConnSession) error {
 28			vty, err := shared.NewVConsole(sesh)
 29			if err != nil {
 30				return err
 31			}
 32			opts := vaxis.Options{
 33				WithConsole: vty,
 34			}
 35			return tui.NewTui(opts, shrd)
 36		}
 37	}
 38}
 39
 40func StartSshServer() {
 41	appName := "pico-ssh"
 42
 43	host := utils.GetEnv("PICO_HOST", "0.0.0.0")
 44	port := utils.GetEnv("PICO_SSH_PORT", "2222")
 45	promPort := utils.GetEnv("PICO_PROM_PORT", "9222")
 46	cfg := NewConfigSite(appName)
 47	logger := cfg.Logger
 48
 49	ctx, cancel := context.WithCancel(context.Background())
 50	defer cancel()
 51
 52	dbpool := postgres.NewDB(cfg.DbURL, cfg.Logger)
 53	defer dbpool.Close()
 54
 55	handler := NewUploadHandler(
 56		dbpool,
 57		cfg,
 58	)
 59
 60	cliHandler := &CliHandler{
 61		Logger: logger,
 62		DBPool: dbpool,
 63	}
 64
 65	sshAuth := shared.NewSshAuthHandler(dbpool, logger)
 66
 67	// Create a new SSH server
 68	server, err := pssh.NewSSHServerWithConfig(
 69		ctx,
 70		logger,
 71		appName,
 72		host,
 73		port,
 74		promPort,
 75		"ssh_data/term_info_ed25519",
 76		func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
 77			perms, _ := sshAuth.PubkeyAuthHandler(conn, key)
 78			if perms == nil {
 79				perms = &ssh.Permissions{
 80					Extensions: map[string]string{
 81						"pubkey": utils.KeyForKeyText(key),
 82					},
 83				}
 84			}
 85
 86			return perms, nil
 87		},
 88		[]pssh.SSHServerMiddleware{
 89			pipe.Middleware(handler, ""),
 90			list.Middleware(handler),
 91			scp.Middleware(handler),
 92			rsync.Middleware(handler),
 93			auth.Middleware(handler),
 94			func(next pssh.SSHServerHandler) pssh.SSHServerHandler {
 95				return func(sesh *pssh.SSHServerConnSession) error {
 96					shrd := &tui.SharedModel{
 97						Session: sesh,
 98						Cfg:     cfg,
 99						Dbpool:  handler.DBPool,
100						Logger:  cfg.Logger,
101					}
102					return pssh.PtyMdw(createTui(shrd), 200*time.Millisecond)(next)(sesh)
103				}
104			},
105			Middleware(cliHandler),
106			pssh.LogMiddleware(handler, dbpool),
107		},
108		[]pssh.SSHServerMiddleware{
109			sftp.Middleware(handler),
110			pssh.LogMiddleware(handler, dbpool),
111		},
112		nil,
113	)
114
115	if err != nil {
116		logger.Error("failed to create ssh server", "err", err.Error())
117		os.Exit(1)
118	}
119
120	done := make(chan os.Signal, 1)
121	signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
122	logger.Info("Starting SSH server", "addr", server.Config.ListenAddr)
123	go func() {
124		if err = server.ListenAndServe(); err != nil {
125			logger.Error("serve", "err", err.Error())
126			os.Exit(1)
127		}
128	}()
129
130	exit := func() {
131		logger.Info("stopping ssh server")
132		cancel()
133	}
134
135	<-done
136	exit()
137}