tunnel-please/session/interaction/interaction.go

package interaction

import (
	"context"
	"log"
	"sync"
	"tunnel_pls/internal/config"
	"tunnel_pls/internal/random"
	"tunnel_pls/session/slug"
	"tunnel_pls/types"

	"charm.land/bubbles/v2/help"
	"charm.land/bubbles/v2/key"
	"charm.land/bubbles/v2/list"
	"charm.land/bubbles/v2/textinput"
	tea "charm.land/bubbletea/v2"
	"github.com/charmbracelet/colorprofile"
	"golang.org/x/crypto/ssh"
)

type Interaction interface {
	Mode() types.InteractiveMode
	SetChannel(channel ssh.Channel)
	SetMode(m types.InteractiveMode)
	SetWH(w, h int)
	Start()
	Redraw()
	Send(message string) error
}

type SessionRegistry interface {
	Update(user string, oldKey, newKey types.SessionKey) error
}

type Forwarder interface {
	Close() error
	TunnelType() types.TunnelType
	ForwardedPort() uint16
}

type CloseFunc func() error
type interaction struct {
	randomizer      random.Random
	config          config.Config
	channel         ssh.Channel
	slug            slug.Slug
	forwarder       Forwarder
	closeFunc       CloseFunc
	user            string
	sessionRegistry SessionRegistry
	program         *tea.Program
	ctx             context.Context
	cancel          context.CancelFunc
	mode            types.InteractiveMode
	programMu       sync.Mutex
	width           int
	height          int
}

func (i *interaction) SetMode(m types.InteractiveMode) {
	i.mode = m
}

func (i *interaction) Mode() types.InteractiveMode {
	return i.mode
}

func (i *interaction) Send(message string) error {
	if i.channel != nil {
		_, err := i.channel.Write([]byte(message))
		return err
	}
	return nil
}
func (i *interaction) SetWH(w, h int) {
	i.programMu.Lock()
	i.width = w
	i.height = h
	prog := i.program
	i.programMu.Unlock()

	if prog != nil {
		prog.Send(tea.WindowSizeMsg{
			Width:  w,
			Height: h,
		})
	}
}

func New(randomizer random.Random, config config.Config, slug slug.Slug, forwarder Forwarder, sessionRegistry SessionRegistry, user string, closeFunc CloseFunc) Interaction {
	ctx, cancel := context.WithCancel(context.Background())
	return &interaction{
		randomizer:      randomizer,
		config:          config,
		channel:         nil,
		slug:            slug,
		forwarder:       forwarder,
		closeFunc:       closeFunc,
		user:            user,
		sessionRegistry: sessionRegistry,
		program:         nil,
		ctx:             ctx,
		cancel:          cancel,
	}
}

func (i *interaction) SetChannel(channel ssh.Channel) {
	i.channel = channel
}

func (i *interaction) Stop() {
	if i.cancel != nil {
		i.cancel()
	}

	i.programMu.Lock()
	defer i.programMu.Unlock()

	if i.program != nil {
		i.program.Kill()
		i.program = nil
	}
}

func (m *model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {

	switch msg := msg.(type) {
	case tickMsg:
		m.showingComingSoon = false
		return m, textinput.Blink

	case tea.WindowSizeMsg:
		m.width = msg.Width
		m.height = msg.Height
		m.commandList.SetWidth(msg.Width)
		m.commandList.SetHeight(msg.Height - 4)

		if msg.Width < 80 {
			m.slugInput.SetWidth(msg.Width - 10)
		} else {
			m.slugInput.SetWidth(50)
		}
		return m, nil

	case tea.QuitMsg:
		m.quitting = true
		return m, tea.Batch(textinput.Blink, tea.Quit)

	case tea.KeyPressMsg:
		if m.showingComingSoon {
			return m.comingSoonUpdate(msg)
		}

		if m.editingSlug {
			return m.slugUpdate(msg)
		}

		if m.showingCommands {
			return m.commandsUpdate(msg)
		}

		return m.dashboardUpdate(msg)
	}

	return m, nil
}

func (i *interaction) Redraw() {
	i.programMu.Lock()
	prog := i.program
	w, h := i.width, i.height
	i.programMu.Unlock()

	if prog != nil {
		prog.Send(tea.WindowSizeMsg{Width: w, Height: h})
	}
}

func (m *model) View() tea.View {
	if m.quitting {
		return tea.NewView("")
	}

	if m.showingComingSoon {
		return tea.NewView(m.comingSoonView())
	}

	if m.editingSlug {
		return tea.NewView(m.slugView())
	}

	if m.showingCommands {
		return tea.NewView(m.commandsView())
	}

	v := tea.NewView(m.dashboardView())
	v.AltScreen = true
	v.MouseMode = tea.MouseModeCellMotion
	return v
}

func (i *interaction) Start() {
	if i.mode == types.InteractiveModeHEADLESS {
		return
	}

	protocol := "http"
	if i.config.TLSEnabled() {
		protocol = "https"
	}

	tunnelType := i.forwarder.TunnelType()
	port := i.forwarder.ForwardedPort()

	items := []list.Item{
		commandItem{name: "slug", desc: "Set custom subdomain"},
		commandItem{name: "tunnel-type", desc: "Change tunnel type (Coming Soon)"},
	}

	delegate := list.NewDefaultDelegate()
	delegate.ShowDescription = true
	delegate.SetHeight(2)

	commandList := list.New(items, delegate, 80, 20)
	commandList.Title = "Select a command"
	commandList.SetShowStatusBar(false)
	commandList.SetFilteringEnabled(false)
	commandList.SetShowHelp(false)

	ti := textinput.New()
	ti.Placeholder = "my-custom-slug"
	ti.CharLimit = 20
	ti.SetWidth(50)

	m := &model{
		randomizer:  i.randomizer,
		domain:      i.config.Domain(),
		protocol:    protocol,
		tunnelType:  tunnelType,
		port:        port,
		commandList: commandList,
		slugInput:   ti,
		interaction: i,
		keymap: keymap{
			quit: key.NewBinding(
				key.WithKeys("q", "ctrl+c"),
				key.WithHelp("q", "quit"),
			),
			command: key.NewBinding(
				key.WithKeys("c"),
				key.WithHelp("c", "commands"),
			),
			random: key.NewBinding(
				key.WithKeys("ctrl+r"),
				key.WithHelp("ctrl+r", "random"),
			),
		},
		help: help.New(),
	}

	i.programMu.Lock()
	w, h := i.width, i.height
	i.programMu.Unlock()

	opts := []tea.ProgramOption{
		tea.WithInput(i.channel),
		tea.WithOutput(i.channel),
		tea.WithColorProfile(colorprofile.TrueColor),
		tea.WithoutSignals(),
		tea.WithoutSignalHandler(),
		tea.WithFPS(30),
		tea.WithEnvironment([]string{"TERM=xterm-256color", "COLORTERM=truecolor"}),
	}
	if w > 0 && h > 0 {
		opts = append(opts, tea.WithWindowSize(w, h))
	}

	prog := tea.NewProgram(m, opts...)

	i.programMu.Lock()
	i.program = prog
	i.programMu.Unlock()

	_, err := prog.Run()
	if err != nil {
		log.Printf("Cannot close tea: %s \n", err)
	}

	i.programMu.Lock()
	if i.program != nil {
		i.program.Kill()
		i.program = nil
	}
	i.programMu.Unlock()

	if i.closeFunc != nil {
		_ = i.closeFunc()
	}
}