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