sshidentity/main.go

package main

import (
	"context"
	"crypto/md5"
	"encoding/hex"
	"errors"
	"fmt"
	"net"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"github.com/charmbracelet/bubbles/textinput"
	tea "github.com/charmbracelet/bubbletea"
	"github.com/charmbracelet/lipgloss"
	"github.com/charmbracelet/log"
	"github.com/charmbracelet/ssh"
	"github.com/charmbracelet/wish"
	"github.com/charmbracelet/wish/bubbletea"
	"github.com/charmbracelet/wish/logging"
)

const (
	host = "localhost"
	port = "23234"
)

var (
	// knownUsers maps public key fingerprints to nicknames
	knownUsers   = make(map[string]string)
	usersMutex   sync.Mutex
	errorStyle   = lipgloss.NewStyle().Foreground(lipgloss.Color("9"))
	successStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("10"))
)

func main() {
	s, err := wish.NewServer(
		wish.WithAddress(net.JoinHostPort(host, port)),
		wish.WithHostKeyPath(".ssh/id_ed25519"),
		wish.WithPublicKeyAuth(func(ctx ssh.Context, key ssh.PublicKey) bool {
			// Accept any public key authentication
			return true
		}),
		wish.WithMiddleware(
			bubbletea.Middleware(teaHandler),
			logging.Middleware(),
		),
	)
	if err != nil {
		log.Error("Could not start server", "error", err)
	}

	done := make(chan os.Signal, 1)
	signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
	log.Info("Starting SSH server", "host", host, "port", port)
	go func() {
		if err = s.ListenAndServe(); err != nil && !errors.Is(err, ssh.ErrServerClosed) {
			log.Error("Could not start server", "error", err)
			done <- nil
		}
	}()

	<-done
	log.Info("Stopping SSH server")
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer func() { cancel() }()
	if err := s.Shutdown(ctx); err != nil && !errors.Is(err, ssh.ErrServerClosed) {
		log.Error("Could not stop server", "error", err)
	}
}

// getFingerprint generates a fingerprint from the public key
func getFingerprint(pubKey ssh.PublicKey) string {
	hash := md5.Sum(pubKey.Marshal())
	return hex.EncodeToString(hash[:])
}

// teaHandler creates a Bubble Tea program for each SSH session
func teaHandler(s ssh.Session) (tea.Model, []tea.ProgramOption) {
	_, _, active := s.Pty()
	if !active {
		wish.Fatalln(s, "no active terminal, skipping")
		return nil, nil
	}

	// Get the public key
	pubKey := s.PublicKey()
	if pubKey == nil {
		wish.Fatalln(s, "no public key found")
		return nil, nil
	}
	fingerprint := getFingerprint(pubKey)

	// Check if we know this user
	usersMutex.Lock()
	nickname, known := knownUsers[fingerprint]
	usersMutex.Unlock()

	if known {
		// Known user - send greeting and close connection
		wish.Println(s, fmt.Sprintf("Hello, %s!", nickname))
		return nil, nil
	}

	// New user - prompt for nickname
	ti := textinput.New()
	ti.Placeholder = "3-12 characters"
	ti.Focus()
	ti.CharLimit = 12
	ti.Width = 20

	m := model{
		textInput:   ti,
		fingerprint: fingerprint,
		session:     s,
	}
	return m, []tea.ProgramOption{tea.WithAltScreen()}
}

// model manages the UI state
type model struct {
	textInput   textinput.Model
	fingerprint string
	nickname    string
	known       bool
	err         string
	session     ssh.Session
}

// Init initializes the model
func (m model) Init() tea.Cmd {
	if m.known {
		return nil
	}
	return textinput.Blink
}

// Update handles input and updates the model
func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
	if m.known {
		return m, nil
	}

	switch msg := msg.(type) {
	case tea.KeyMsg:
		switch msg.Type {
		case tea.KeyEnter:
			// Validate nickname
			nick := m.textInput.Value()
			if len(nick) < 3 {
				m.err = "Nickname must be at least 3 characters"
				return m, nil
			}
			if len(nick) > 12 {
				m.err = "Nickname must be no more than 12 characters"
				return m, nil
			}

			// Store the nickname
			usersMutex.Lock()
			knownUsers[m.fingerprint] = nick
			usersMutex.Unlock()

			// Send greeting and close connection
			wish.Println(m.session, fmt.Sprintf("Hello, %s!", nick))
			return m, tea.Quit

		case tea.KeyCtrlC, tea.KeyEsc:
			return m, tea.Quit
		}
	}

	var cmd tea.Cmd
	m.textInput, cmd = m.textInput.Update(msg)
	return m, cmd
}

// View renders the UI
func (m model) View() string {
	if m.known {
		return ""
	}

	// New user - show input prompt
	var errMsg string
	if m.err != "" {
		errMsg = errorStyle.Render("Error: " + m.err + "\n\n")
	}

	return fmt.Sprintf(
		"Welcome! Please choose a nickname (3-12 characters)\n\n%s%s\n\n%s",
		errMsg,
		m.textInput.View(),
		"(esc to quit)",
	) + "\n"
}