session.go

  1package ssh
  2
  3import (
  4	"context"
  5	"strings"
  6
  7	tea "github.com/charmbracelet/bubbletea"
  8	"github.com/charmbracelet/lipgloss"
  9	"github.com/charmbracelet/log"
 10	"github.com/charmbracelet/soft-serve/server/access"
 11	"github.com/charmbracelet/soft-serve/server/auth"
 12	"github.com/charmbracelet/soft-serve/server/backend"
 13	"github.com/charmbracelet/soft-serve/server/config"
 14	"github.com/charmbracelet/soft-serve/server/errors"
 15	"github.com/charmbracelet/soft-serve/server/sshutils"
 16	"github.com/charmbracelet/soft-serve/server/ui"
 17	"github.com/charmbracelet/soft-serve/server/ui/common"
 18	"github.com/charmbracelet/ssh"
 19	"github.com/charmbracelet/wish"
 20	bm "github.com/charmbracelet/wish/bubbletea"
 21	"github.com/muesli/termenv"
 22	"github.com/prometheus/client_golang/prometheus"
 23	"github.com/prometheus/client_golang/prometheus/promauto"
 24)
 25
 26var (
 27	tuiSessionCounter = promauto.NewCounterVec(prometheus.CounterOpts{
 28		Namespace: "soft_serve",
 29		Subsystem: "ssh",
 30		Name:      "tui_session_total",
 31		Help:      "The total number of TUI sessions",
 32	}, []string{"key", "user", "repo", "term"})
 33)
 34
 35// SessionHandler is the soft-serve bubbletea ssh session handler.
 36func SessionHandler(ctx context.Context) bm.ProgramHandler {
 37	be := backend.FromContext(ctx)
 38	cfg := config.FromContext(ctx)
 39	logger := log.FromContext(ctx).WithPrefix("ssh-pty")
 40	return func(s ssh.Session) *tea.Program {
 41		ctx := backend.WithContext(s.Context(), be)
 42		ctx = config.WithContext(ctx, cfg)
 43		ctx = log.WithContext(ctx, logger)
 44
 45		ak := sshutils.MarshalAuthorizedKey(s.PublicKey())
 46		pty, _, active := s.Pty()
 47		if !active {
 48			return nil
 49		}
 50
 51		cmd := s.Command()
 52		var initialRepo string
 53		if len(cmd) == 1 {
 54			user, _ := be.Authenticate(ctx, auth.NewPublicKey(s.PublicKey()))
 55			initialRepo = cmd[0]
 56			auth, _ := be.AccessLevel(ctx, initialRepo, user)
 57			if auth < access.ReadOnlyAccess {
 58				wish.Fatalln(s, errors.ErrUnauthorized)
 59				return nil
 60			}
 61		}
 62
 63		envs := &sessionEnv{s}
 64		output := lipgloss.NewRenderer(s, termenv.WithColorCache(true), termenv.WithEnvironment(envs))
 65		// FIXME: detect color profile and dark background from ssh.Session
 66		output.SetColorProfile(termenv.ANSI256)
 67		output.SetHasDarkBackground(true)
 68		c := common.NewCommon(ctx, output, pty.Window.Width, pty.Window.Height)
 69		m := ui.New(c, initialRepo)
 70		p := tea.NewProgram(m,
 71			tea.WithInput(s),
 72			tea.WithOutput(s),
 73			tea.WithAltScreen(),
 74			tea.WithoutCatchPanics(),
 75			tea.WithMouseCellMotion(),
 76		)
 77
 78		tuiSessionCounter.WithLabelValues(ak, s.User(), initialRepo, pty.Term).Inc()
 79
 80		return p
 81	}
 82}
 83
 84var _ termenv.Environ = &sessionEnv{}
 85
 86type sessionEnv struct {
 87	ssh.Session
 88}
 89
 90func (s *sessionEnv) Environ() []string {
 91	pty, _, _ := s.Pty()
 92	return append(s.Session.Environ(), "TERM="+pty.Term)
 93}
 94
 95func (s *sessionEnv) Getenv(key string) string {
 96	for _, env := range s.Environ() {
 97		if strings.HasPrefix(env, key+"=") {
 98			return strings.TrimPrefix(env, key+"=")
 99		}
100	}
101	return ""
102}