main pico / pkg / apps / pico / ssh.go
Antonio Mika  ·  2026-10-11
  1package pico
  2
  3import (
  4	"context"
  5	"os"
  6	"os/signal"
  7	"syscall"
  8	"time"
  9
 10	pgsdb "github.com/picosh/pico/pkg/apps/pgs/db"
 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	"go.rockorager.dev/vaxis"
 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 := shared.GetEnv("PICO_HOST", "0.0.0.0")
 44	port := shared.GetEnv("PICO_SSH_PORT", "2222")
 45	promPort := shared.GetEnv("PICO_PROM_PORT", "9222")
 46	cfg := NewConfig(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 func() {
 54		_ = dbpool.Close()
 55	}()
 56
 57	pgsDB := pgsdb.NewDBWithConn(dbpool.Db, cfg.Logger)
 58
 59	handler := NewUploadHandler(
 60		dbpool,
 61		cfg,
 62	)
 63
 64	cliHandler := &CliHandler{
 65		Logger:        logger,
 66		DBPool:        dbpool,
 67		DomainTargets: cfg.CustomDomainTargets,
 68	}
 69
 70	sshAuth := shared.NewSshAuthHandler(dbpool, logger, "pico")
 71
 72	// Create a new SSH server
 73	server, err := pssh.NewSSHServerWithConfig(
 74		ctx,
 75		logger,
 76		appName,
 77		host,
 78		port,
 79		promPort,
 80		"ssh_data/term_info_ed25519",
 81		func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
 82			perms, err := sshAuth.PubkeyAuthHandler(conn, key)
 83			logger.Warn("pubkey auth handler", "err", err)
 84			if perms == nil {
 85				perms = &ssh.Permissions{
 86					Extensions: map[string]string{
 87						"pubkey": shared.KeyForKeyText(key),
 88					},
 89				}
 90			}
 91
 92			return perms, nil
 93		},
 94		[]pssh.SSHServerMiddleware{
 95			pipe.Middleware(handler, ""),
 96			list.Middleware(handler),
 97			scp.Middleware(handler),
 98			rsync.Middleware(handler),
 99			auth.Middleware(handler),
100			func(next pssh.SSHServerHandler) pssh.SSHServerHandler {
101				return func(sesh *pssh.SSHServerConnSession) error {
102					shrd := &tui.SharedModel{
103						Session:    sesh,
104						TunsSecret: cfg.TunsSecret,
105						Dbpool:     handler.DBPool,
106						PgsDB:      pgsDB,
107						Logger:     pssh.GetLogger(sesh),
108					}
109					return pssh.PtyMdw(createTui(shrd), 200*time.Millisecond)(next)(sesh)
110				}
111			},
112			Middleware(cliHandler),
113			pssh.LogMiddleware(handler, dbpool),
114		},
115		[]pssh.SSHServerMiddleware{
116			sftp.Middleware(handler),
117			pssh.LogMiddleware(handler, dbpool),
118		},
119		nil,
120	)
121
122	if err != nil {
123		logger.Error("failed to create ssh server", "err", err.Error())
124		os.Exit(1)
125	}
126
127	done := make(chan os.Signal, 1)
128	signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
129	logger.Info("Starting SSH server", "addr", server.Config.ListenAddr)
130	go func() {
131		if err = server.ListenAndServe(); err != nil {
132			logger.Error("serve", "err", err.Error())
133			os.Exit(1)
134		}
135	}()
136
137	exit := func() {
138		logger.Info("stopping ssh server")
139		cancel()
140	}
141
142	<-done
143	exit()
144}