Initial commit

main.go

@@ -0,0 +1,199 @@
+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"
+}