O_O

internal/broadcaster.go

@@ -0,0 +1,42 @@
+package internal
+
+import "sync"
+
+type broadcaster struct {
+	clients map[chan bool]bool
+	mu      sync.RWMutex
+}
+
+func newBroadcaster() *broadcaster {
+	return &broadcaster{
+		clients: make(map[chan bool]bool),
+	}
+}
+
+func (b *broadcaster) subscribe() chan bool {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	ch := make(chan bool, 10)
+	b.clients[ch] = true
+	return ch
+}
+
+func (b *broadcaster) unsubscribe(ch chan bool) {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	delete(b.clients, ch)
+	close(ch)
+}
+
+func (b *broadcaster) broadcast() {
+	b.mu.RLock()
+	defer b.mu.RUnlock()
+	for ch := range b.clients {
+		select {
+		case ch <- true:
+		default:
+		}
+	}
+}
+
+var updates = newBroadcaster()

internal/crypto.go

@@ -0,0 +1,154 @@
+package internal
+
+import (
+	"crypto/aes"
+	"crypto/cipher"
+	"crypto/ed25519"
+	"crypto/rand"
+	"crypto/sha256"
+	"io"
+
+	"golang.org/x/crypto/argon2"
+	"golang.org/x/crypto/curve25519"
+)
+
+const (
+	saltSize  = 16
+	keySize   = 32
+	nonceSize = 12
+)
+
+// IdentityKeyPair represents a user's long-term identity key
+type IdentityKeyPair struct {
+	PublicKey  [32]byte
+	PrivateKey [32]byte
+}
+
+// PrekeyBundle contains public keys for establishing a session
+type PrekeyBundle struct {
+	IdentityKey     [32]byte
+	Prekey          [32]byte
+	PrekeySignature []byte
+}
+
+// generateIdentityKeyPair creates a new Curve25519 key pair for identity
+func generateIdentityKeyPair() (*IdentityKeyPair, error) {
+	var privateKey [32]byte
+	if _, err := io.ReadFull(rand.Reader, privateKey[:]); err != nil {
+		return nil, err
+	}
+
+	var publicKey [32]byte
+	curve25519.ScalarBaseMult(&publicKey, &privateKey)
+
+	return &IdentityKeyPair{
+		PublicKey:  publicKey,
+		PrivateKey: privateKey,
+	}, nil
+}
+
+// generateSignedPrekey creates a prekey signed with an Ed25519 signing key
+func generateSignedPrekey() ([32]byte, []byte, error) {
+	var prekey [32]byte
+	var prekeyPriv [32]byte
+	if _, err := io.ReadFull(rand.Reader, prekeyPriv[:]); err != nil {
+		return prekey, nil, err
+	}
+	curve25519.ScalarBaseMult(&prekey, &prekeyPriv)
+
+	// Generate signing key
+	signingPub, signingPriv, err := ed25519.GenerateKey(rand.Reader)
+	if err != nil {
+		return prekey, nil, err
+	}
+
+	// Sign the prekey
+	signature := ed25519.Sign(signingPriv, prekey[:])
+
+	// In practice, we'd store signingPub to verify later. For now, we include it in signature
+	combinedSig := append(signingPub, signature...)
+
+	return prekey, combinedSig, nil
+}
+
+// performDH does an ECDH key exchange
+func performDH(privateKey, publicKey [32]byte) ([32]byte, error) {
+	var sharedSecret [32]byte
+	curve25519.ScalarMult(&sharedSecret, &privateKey, &publicKey)
+	return sharedSecret, nil
+}
+
+// deriveSharedSecret combines multiple DH outputs to create a shared secret (simplified X3DH)
+func deriveSharedSecret(dh1, dh2, dh3 [32]byte) []byte {
+	// KDF: hash all DH outputs together
+	h := sha256.New()
+	h.Write(dh1[:])
+	h.Write(dh2[:])
+	h.Write(dh3[:])
+	return h.Sum(nil)
+}
+
+// encryptUserKey encrypts a user's private key with their passphrase
+func encryptUserKey(privateKey [32]byte, passphrase string) ([]byte, []byte, []byte, error) {
+	salt := make([]byte, saltSize)
+	if _, err := io.ReadFull(rand.Reader, salt); err != nil {
+		return nil, nil, nil, err
+	}
+
+	key := deriveKey(passphrase, salt)
+	ciphertext, nonce, err := encryptMsg(string(privateKey[:]), key)
+	if err != nil {
+		return nil, nil, nil, err
+	}
+
+	return ciphertext, nonce, salt, nil
+}
+
+// decryptUserKey decrypts a user's private key with their passphrase
+func decryptUserKey(ciphertext, nonce, salt []byte, passphrase string) ([32]byte, error) {
+	var privateKey [32]byte
+	key := deriveKey(passphrase, salt)
+	plain, err := decryptMsg(ciphertext, nonce, key)
+	if err != nil {
+		return privateKey, err
+	}
+	copy(privateKey[:], plain)
+	return privateKey, nil
+}
+
+func deriveKey(passphrase string, salt []byte) []byte {
+	return argon2.IDKey([]byte(passphrase), salt, 1, 64*1024, 4, keySize)
+}
+
+func encryptMsg(plain string, key []byte) ([]byte, []byte, error) {
+	block, err := aes.NewCipher(key)
+	if err != nil {
+		return nil, nil, err
+	}
+	gcm, err := cipher.NewGCM(block)
+	if err != nil {
+		return nil, nil, err
+	}
+	nonce := make([]byte, gcm.NonceSize())
+	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
+		return nil, nil, err
+	}
+	ciphertext := gcm.Seal(nil, nonce, []byte(plain), nil)
+	return ciphertext, nonce, nil
+}
+
+func decryptMsg(ciphertext, nonce, key []byte) (string, error) {
+	block, err := aes.NewCipher(key)
+	if err != nil {
+		return "", err
+	}
+	gcm, err := cipher.NewGCM(block)
+	if err != nil {
+		return "", err
+	}
+	plain, err := gcm.Open(nil, nonce, ciphertext, nil)
+	if err != nil {
+		return "", err
+	}
+	return string(plain), nil
+}

internal/database.go

@@ -0,0 +1,268 @@
+package internal
+
+import (
+	"crypto/rand"
+	"database/sql"
+	"encoding/base64"
+	"io"
+	"log"
+	"time"
+
+	_ "github.com/mattn/go-sqlite3"
+)
+
+const dbFile = "teachat.db"
+
+var db *sql.DB
+
+func InitDB() {
+	var err error
+	db, err = sql.Open("sqlite3", dbFile)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	query := `
+    CREATE TABLE IF NOT EXISTS config (key TEXT PRIMARY KEY, val TEXT);
+    CREATE TABLE IF NOT EXISTS users (
+        pubkey TEXT PRIMARY KEY,
+        username TEXT UNIQUE,
+        identity_key BLOB,
+        prekey BLOB,
+        prekey_signature BLOB,
+        created_at DATETIME
+    );
+    CREATE TABLE IF NOT EXISTS rooms (
+        id TEXT PRIMARY KEY,
+        name TEXT,
+        creator TEXT,
+        is_dm BOOLEAN,
+        created_at DATETIME
+    );
+    CREATE TABLE IF NOT EXISTS room_members (
+        room_id TEXT,
+        username TEXT,
+        shared_secret BLOB,
+        joined_at DATETIME,
+        PRIMARY KEY (room_id, username),
+        FOREIGN KEY (room_id) REFERENCES rooms(id),
+        FOREIGN KEY (username) REFERENCES users(username)
+    );
+    CREATE TABLE IF NOT EXISTS messages (
+        id INTEGER PRIMARY KEY,
+        room_id TEXT,
+        timestamp DATETIME,
+        sender TEXT,
+        ciphertext BLOB,
+        nonce BLOB,
+        FOREIGN KEY (room_id) REFERENCES rooms(id)
+    );
+    `
+	if _, err := db.Exec(query); err != nil {
+		log.Fatal(err)
+	}
+
+	// Migration for older DBs
+	db.Exec("ALTER TABLE rooms ADD COLUMN creator TEXT")
+}
+
+func ensureSalt() []byte {
+	row := db.QueryRow("SELECT val FROM config WHERE key='global_salt'")
+	var b64 string
+	if err := row.Scan(&b64); err == sql.ErrNoRows {
+		salt := make([]byte, saltSize)
+		if _, err := io.ReadFull(rand.Reader, salt); err != nil {
+			log.Fatal(err)
+		}
+		b64 = base64.StdEncoding.EncodeToString(salt)
+		db.Exec("INSERT INTO config (key, val) VALUES ('global_salt', ?)", b64)
+		return salt
+	} else if err != nil {
+		log.Fatal(err)
+	}
+	salt, _ := base64.StdEncoding.DecodeString(b64)
+	return salt
+}
+
+func authorizeUser(pubkey, username string, identityKey, prekey, prekeySignature []byte) error {
+	_, err := db.Exec("INSERT INTO users (pubkey, username, identity_key, prekey, prekey_signature, created_at) VALUES (?, ?, ?, ?, ?, ?)",
+		pubkey, username, identityKey, prekey, prekeySignature, time.Now())
+	return err
+}
+
+func getUsername(pubkey string) (string, error) {
+	var username string
+	err := db.QueryRow("SELECT username FROM users WHERE pubkey = ?", pubkey).Scan(&username)
+	return username, err
+}
+
+func getUserKeys(username string) (identityKey, prekey, prekeySignature []byte, err error) {
+	err = db.QueryRow("SELECT identity_key, prekey, prekey_signature FROM users WHERE username = ?", username).Scan(&identityKey, &prekey, &prekeySignature)
+	return
+}
+
+func listUsers() ([]string, error) {
+	rows, err := db.Query("SELECT username FROM users ORDER BY username ASC")
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+
+	var users []string
+	for rows.Next() {
+		var username string
+		if err := rows.Scan(&username); err != nil {
+			continue
+		}
+		users = append(users, username)
+	}
+	return users, nil
+}
+
+func createRoom(roomID, name, creator string, isDM bool) error {
+	_, err := db.Exec("INSERT INTO rooms (id, name, creator, is_dm, created_at) VALUES (?, ?, ?, ?, ?)",
+		roomID, name, creator, isDM, time.Now())
+	return err
+}
+
+func getRoomIDByName(name string) (string, error) {
+	var id string
+	err := db.QueryRow("SELECT id FROM rooms WHERE name = ? AND is_dm = 0 LIMIT 1", name).Scan(&id)
+	return id, err
+}
+
+// getAnyRoomSecret gets the shared secret from ANY member of the room.
+// In a real decentralized system this wouldn't be possible, but here the server acts as the key distributor.
+func getAnyRoomSecret(roomID string) ([]byte, error) {
+	var secret []byte
+	err := db.QueryRow("SELECT shared_secret FROM room_members WHERE room_id = ? LIMIT 1", roomID).Scan(&secret)
+	return secret, err
+}
+
+func deleteRoom(roomID string) error {
+	tx, err := db.Begin()
+	if err != nil {
+		return err
+	}
+	if _, err := tx.Exec("DELETE FROM messages WHERE room_id = ?", roomID); err != nil {
+		tx.Rollback()
+		return err
+	}
+	if _, err := tx.Exec("DELETE FROM room_members WHERE room_id = ?", roomID); err != nil {
+		tx.Rollback()
+		return err
+	}
+	if _, err := tx.Exec("DELETE FROM rooms WHERE id = ?", roomID); err != nil {
+		tx.Rollback()
+		return err
+	}
+	return tx.Commit()
+}
+
+func leaveRoom(roomID, username string) error {
+	_, err := db.Exec("DELETE FROM room_members WHERE room_id = ? AND username = ?", roomID, username)
+	return err
+}
+
+func joinRoom(roomID, username string, sharedSecret []byte) error {
+	// Check if already joined
+	var count int
+	db.QueryRow("SELECT COUNT(*) FROM room_members WHERE room_id = ? AND username = ?", roomID, username).Scan(&count)
+	if count > 0 {
+		return nil
+	}
+	_, err := db.Exec("INSERT INTO room_members (room_id, username, shared_secret, joined_at) VALUES (?, ?, ?, ?)",
+		roomID, username, sharedSecret, time.Now())
+	return err
+}
+
+func getMemberJoinedAt(roomID, username string) (time.Time, error) {
+	var joinedAt time.Time
+	err := db.QueryRow("SELECT joined_at FROM room_members WHERE room_id = ? AND username = ?", roomID, username).Scan(&joinedAt)
+	return joinedAt, err
+}
+
+func listUserRooms(username string) ([]struct {
+	ID      string
+	Name    string
+	Creator string
+	IsDM    bool
+}, error) {
+	rows, err := db.Query(`
+		SELECT r.id, r.name, COALESCE(r.creator, ''), r.is_dm
+		FROM rooms r
+		JOIN room_members rm ON r.id = rm.room_id
+		WHERE rm.username = ?
+		ORDER BY r.created_at DESC`, username)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+
+	var rooms []struct {
+		ID      string
+		Name    string
+		Creator string
+		IsDM    bool
+	}
+	for rows.Next() {
+		var room struct {
+			ID      string
+			Name    string
+			Creator string
+			IsDM    bool
+		}
+		if err := rows.Scan(&room.ID, &room.Name, &room.Creator, &room.IsDM); err != nil {
+			continue
+		}
+		rooms = append(rooms, room)
+	}
+	return rooms, nil
+}
+
+func getRoomSharedSecret(roomID, username string) ([]byte, error) {
+	var secret []byte
+	err := db.QueryRow("SELECT shared_secret FROM room_members WHERE room_id = ? AND username = ?", roomID, username).Scan(&secret)
+	return secret, err
+}
+
+func saveMessage(roomID, sender string, ciphertext, nonce []byte) error {
+	_, err := db.Exec("INSERT INTO messages (room_id, timestamp, sender, ciphertext, nonce) VALUES (?, ?, ?, ?, ?)",
+		roomID, time.Now(), sender, ciphertext, nonce)
+	return err
+}
+
+func loadMessages(roomID string) ([]struct {
+	Timestamp  time.Time
+	Sender     string
+	Ciphertext []byte
+	Nonce      []byte
+}, error) {
+	rows, err := db.Query("SELECT timestamp, sender, ciphertext, nonce FROM messages WHERE room_id = ? ORDER BY timestamp ASC", roomID)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+
+	var messages []struct {
+		Timestamp  time.Time
+		Sender     string
+		Ciphertext []byte
+		Nonce      []byte
+	}
+
+	for rows.Next() {
+		var msg struct {
+			Timestamp  time.Time
+			Sender     string
+			Ciphertext []byte
+			Nonce      []byte
+		}
+		if err := rows.Scan(&msg.Timestamp, &msg.Sender, &msg.Ciphertext, &msg.Nonce); err != nil {
+			continue
+		}
+		messages = append(messages, msg)
+	}
+
+	return messages, nil
+}

internal/model.go

@@ -0,0 +1,609 @@
+package internal
+
+import (
+	"crypto/rand"
+	"database/sql"
+	"encoding/base64"
+	"fmt"
+	"io"
+	"strings"
+	"time"
+
+	"github.com/charmbracelet/bubbles/textinput"
+	"github.com/charmbracelet/bubbles/viewport"
+	tea "github.com/charmbracelet/bubbletea"
+	"github.com/charmbracelet/lipgloss"
+	"github.com/charmbracelet/ssh"
+)
+
+var (
+	msgStyle     = lipgloss.NewStyle().Foreground(lipgloss.Color("252")).PaddingLeft(1)
+	senderStyle  = lipgloss.NewStyle().Foreground(lipgloss.Color("205")).Bold(true)
+	errStyle     = lipgloss.NewStyle().Foreground(lipgloss.Color("196")).Bold(true)
+	sysStyle     = lipgloss.NewStyle().Foreground(lipgloss.Color("117")).Italic(true)
+	roomStyle    = lipgloss.NewStyle().Foreground(lipgloss.Color("86")).Bold(true)
+	dmStyle      = lipgloss.NewStyle().Foreground(lipgloss.Color("141")).Bold(true)
+	noteStyle    = lipgloss.NewStyle().Foreground(lipgloss.Color("228")).Bold(true)
+	commandStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("39")).Italic(true)
+	timeStyle    = lipgloss.NewStyle().Foreground(lipgloss.Color("246"))
+	lockedStyle  = lipgloss.NewStyle().Foreground(lipgloss.Color("240")).Italic(true)
+)
+
+type chatMsg struct {
+	Timestamp time.Time
+	Sender    string
+	Content   string
+}
+
+type model struct {
+	sshSession        ssh.Session
+	pubkey            string
+	username          string
+	authorized        bool
+	needsRegistration bool
+
+	state            int
+	passphrase       string
+	identityKey      *IdentityKeyPair
+	encryptedPrivKey []byte
+	privKeySalt      []byte
+	privKeyNonce     []byte
+
+	currentRoomID       string
+	currentRoomKey      []byte
+	currentRoomJoinedAt time.Time
+	rooms               []struct {
+		ID      string
+		Name    string
+		Creator string
+		IsDM    bool
+	}
+	availableUsers []string
+
+	err error
+
+	input      textinput.Model
+	viewport   viewport.Model
+	messages   []chatMsg
+	updateChan chan bool
+}
+
+func initialModel(s ssh.Session) model {
+	pk := s.PublicKey()
+	var pubkeyStr string
+	var auth bool
+	var user string
+	var needsRegistration bool
+
+	if pk == nil {
+		auth = false
+		user = ""
+		pubkeyStr = ""
+	} else {
+		pubkeyStr = fmt.Sprintf("%s %s", pk.Type(), base64.StdEncoding.EncodeToString(pk.Marshal()))
+
+		username, err := getUsername(pubkeyStr)
+		if err == sql.ErrNoRows {
+			needsRegistration = true
+			auth = true
+			user = ""
+		} else {
+			auth = true
+			needsRegistration = false
+			user = username
+		}
+	}
+
+	ti := textinput.New()
+	if needsRegistration {
+		ti.Placeholder = "Choose a username..."
+	} else {
+		ti.Placeholder = "Enter Passphrase..."
+		ti.EchoMode = textinput.EchoPassword
+	}
+	ti.Focus()
+	ti.CharLimit = 156
+	ti.Width = 50
+
+	vp := viewport.New(80, 20)
+
+	return model{
+		sshSession:        s,
+		pubkey:            pubkeyStr,
+		username:          user,
+		authorized:        auth,
+		needsRegistration: needsRegistration,
+		state:             0,
+		input:             ti,
+		viewport:          vp,
+		updateChan:        updates.subscribe(),
+	}
+}
+
+func (m model) Init() tea.Cmd {
+	return tea.Batch(textinput.Blink, m.waitForUpdate())
+}
+
+func (m model) waitForUpdate() tea.Cmd {
+	return func() tea.Msg {
+		<-m.updateChan
+		return struct{}{}
+	}
+}
+
+func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
+	var cmd tea.Cmd
+
+	switch msg := msg.(type) {
+	case tea.WindowSizeMsg:
+		m.viewport.Width = msg.Width
+		m.viewport.Height = msg.Height - 4
+		return m, tea.ClearScreen
+
+	case tea.KeyMsg:
+		switch msg.Type {
+		case tea.KeyCtrlC, tea.KeyEsc:
+			return m, tea.Quit
+		case tea.KeyEnter:
+			if !m.authorized {
+				return m, tea.Quit
+			}
+
+			if m.needsRegistration {
+				username := m.input.Value()
+				if len(username) > 0 {
+					identityKey, err := generateIdentityKeyPair()
+					if err != nil {
+						m.err = err
+						return m, nil
+					}
+
+					prekey, prekeySignature, err := generateSignedPrekey()
+					if err != nil {
+						m.err = err
+						return m, nil
+					}
+
+					if err := authorizeUser(m.pubkey, username, identityKey.PublicKey[:], prekey[:], prekeySignature); err != nil {
+						m.err = err
+					} else {
+						m.username = username
+						m.identityKey = identityKey
+						m.needsRegistration = false
+
+						m.input.Reset()
+						m.input.Placeholder = "Enter Passphrase..."
+						m.input.EchoMode = textinput.EchoPassword
+						m.viewport.SetContent("")
+						return m, tea.ClearScreen
+					}
+				}
+				return m, nil
+			}
+
+			if m.state == 0 {
+				pass := m.input.Value()
+				if len(pass) > 0 {
+					if m.identityKey != nil {
+						ciphertext, nonce, salt, err := encryptUserKey(m.identityKey.PrivateKey, pass)
+						if err != nil {
+							m.err = err
+							return m, nil
+						}
+						m.encryptedPrivKey = ciphertext
+						m.privKeyNonce = nonce
+						m.privKeySalt = salt
+					}
+
+					m.passphrase = pass
+					m.state = 1
+
+					m.input.Reset()
+					m.input.EchoMode = textinput.EchoNormal
+					m.input.Placeholder = "Loading rooms..."
+					m.viewport.SetContent("")
+
+					m.loadRooms()
+					return m, tea.ClearScreen
+				}
+			} else if m.state == 1 {
+				val := m.input.Value()
+				if len(val) > 0 {
+					m.handleRoomListInput(val)
+					m.input.Reset()
+					if m.state == 2 {
+						return m, tea.ClearScreen
+					}
+				}
+			} else if m.state == 2 {
+				val := m.input.Value()
+				if len(val) > 0 {
+					if val == "/back" {
+						m.exitChat()
+						return m, tea.ClearScreen
+					}
+
+					if val == "/leave" {
+						leaveRoom(m.currentRoomID, m.username)
+						m.exitChat()
+						return m, tea.ClearScreen
+					}
+
+					if val == "/delete" {
+						isCreator := false
+						for _, r := range m.rooms {
+							if r.ID == m.currentRoomID {
+								if r.Creator == m.username {
+									isCreator = true
+								}
+								break
+							}
+						}
+						if isCreator {
+							deleteRoom(m.currentRoomID)
+							m.exitChat()
+							return m, tea.ClearScreen
+						} else {
+							m.saveMessage("System", "Error: Only the creator can delete this room.")
+							m.input.Reset()
+							return m, nil
+						}
+					}
+					m.handleChatInput(val)
+					m.input.Reset()
+				}
+			}
+		}
+
+	case struct{}:
+		if m.state == 2 {
+			m.loadMessages()
+			return m, m.waitForUpdate()
+		}
+	}
+
+	m.input, cmd = m.input.Update(msg)
+	return m, cmd
+}
+
+func (m *model) exitChat() {
+	m.state = 1
+	m.currentRoomID = ""
+	m.currentRoomKey = nil
+	m.messages = nil
+	m.viewport.SetContent("")
+	m.loadRooms()
+	m.input.Reset()
+}
+
+func (m *model) loadRooms() {
+	rooms, err := listUserRooms(m.username)
+	if err != nil {
+		m.err = err
+		return
+	}
+	m.rooms = rooms
+	m.viewport.SetContent("")
+	m.input.SetValue("")
+	m.input.Placeholder = "Enter room # or: /join <name>, /dm <username>, /list"
+}
+
+func (m *model) handleRoomListInput(text string) {
+	if text == "/list" {
+		m.loadRooms()
+		return
+	}
+
+	if strings.HasPrefix(text, "/join ") {
+		roomName := strings.TrimPrefix(text, "/join ")
+		m.handleJoinRoom(roomName)
+		return
+	}
+	// Keep /new as an alias for /join
+	if strings.HasPrefix(text, "/new ") {
+		roomName := strings.TrimPrefix(text, "/new ")
+		m.handleJoinRoom(roomName)
+		return
+	}
+
+	if strings.HasPrefix(text, "/dm ") {
+		targetUser := strings.TrimPrefix(text, "/dm ")
+		m.handleSelectUserForDM(targetUser)
+		return
+	}
+
+	if idx := parseInt(text); idx > 0 && idx <= len(m.rooms) {
+		room := m.rooms[idx-1]
+		m.enterRoom(room.ID, room.Name)
+	}
+}
+
+func (m *model) handleJoinRoom(roomName string) {
+	// Check if room exists
+	existingID, err := getRoomIDByName(roomName)
+
+	// Room Exists: Join it
+	if err == nil && existingID != "" {
+		// Need key
+		secret, err := getAnyRoomSecret(existingID)
+		if err != nil {
+			m.err = fmt.Errorf("could not join: room key unreachable")
+			return
+		}
+
+		if err := joinRoom(existingID, m.username, secret); err != nil {
+			m.err = err
+			return
+		}
+		m.enterRoom(existingID, roomName)
+		return
+	}
+
+	// Room Does Not Exist: Create it
+	roomID := generateRoomID()
+	if err := createRoom(roomID, roomName, m.username, false); err != nil {
+		m.err = err
+		return
+	}
+
+	sharedSecret := make([]byte, 32)
+	if _, err := io.ReadFull(rand.Reader, sharedSecret); err != nil {
+		m.err = err
+		return
+	}
+
+	if err := joinRoom(roomID, m.username, sharedSecret); err != nil {
+		m.err = err
+		return
+	}
+
+	m.enterRoom(roomID, roomName)
+}
+
+func (m *model) handleSelectUserForDM(username string) {
+	if username == "/cancel" {
+		m.state = 1
+		m.input.Placeholder = "Enter room # or: /join <name>, /dm <username>, /list"
+		return
+	}
+
+	if username == m.username {
+		m.createNoteToSelf()
+		return
+	}
+
+	identityKey, prekey, prekeySignature, err := getUserKeys(username)
+	if err != nil {
+		m.err = fmt.Errorf("user not found: %w", err)
+		return
+	}
+
+	var theirIdentityKey, theirPrekey [32]byte
+	copy(theirIdentityKey[:], identityKey)
+	copy(theirPrekey[:], prekey)
+
+	ephemeralKey, err := generateIdentityKeyPair()
+	if err != nil {
+		m.err = err
+		return
+	}
+
+	dh1, _ := performDH(m.identityKey.PrivateKey, theirPrekey)
+	dh2, _ := performDH(ephemeralKey.PrivateKey, theirIdentityKey)
+	dh3, _ := performDH(ephemeralKey.PrivateKey, theirPrekey)
+
+	sharedSecret := deriveSharedSecret(dh1, dh2, dh3)
+
+	roomID := generateRoomID()
+	roomName := fmt.Sprintf("DM: %s <-> %s", m.username, username)
+
+	if err := createRoom(roomID, roomName, m.username, true); err != nil {
+		m.err = err
+		return
+	}
+
+	if err := joinRoom(roomID, m.username, sharedSecret); err != nil {
+		m.err = err
+		return
+	}
+
+	m.enterRoom(roomID, roomName)
+	_ = prekeySignature
+}
+
+func (m *model) createNoteToSelf() {
+	roomID := generateRoomID()
+	roomName := "[Note to Self]"
+
+	sharedSecret := make([]byte, 32)
+	if _, err := io.ReadFull(rand.Reader, sharedSecret); err != nil {
+		m.err = err
+		return
+	}
+
+	if err := createRoom(roomID, roomName, m.username, true); err != nil {
+		m.err = err
+		return
+	}
+
+	if err := joinRoom(roomID, m.username, sharedSecret); err != nil {
+		m.err = err
+		return
+	}
+
+	m.enterRoom(roomID, roomName)
+}
+
+func (m *model) enterRoom(roomID, roomName string) {
+	secret, err := getRoomSharedSecret(roomID, m.username)
+	if err != nil {
+		m.err = err
+		return
+	}
+
+	joinedAt, err := getMemberJoinedAt(roomID, m.username)
+	if err != nil {
+		joinedAt = time.Now() // Fallback
+	}
+
+	m.currentRoomID = roomID
+	m.currentRoomKey = secret
+	m.currentRoomJoinedAt = joinedAt
+	m.state = 2
+	m.input.Placeholder = fmt.Sprintf("[%s] /back to menu, /leave to quit room", roomName)
+	m.loadMessages()
+}
+
+func (m *model) handleChatInput(text string) {
+	m.saveMessage(m.username, text)
+}
+
+func parseInt(s string) int {
+	var num int
+	fmt.Sscanf(s, "%d", &num)
+	return num
+}
+
+func generateRoomID() string {
+	b := make([]byte, 16)
+	io.ReadFull(rand.Reader, b)
+	return base64.URLEncoding.EncodeToString(b)
+}
+
+func (m *model) saveMessage(sender, text string) {
+	ct, nonce, err := encryptMsg(text, m.currentRoomKey)
+	if err != nil {
+		m.err = err
+		return
+	}
+
+	err = saveMessage(m.currentRoomID, sender, ct, nonce)
+	if err != nil {
+		m.err = err
+		return
+	}
+
+	go updates.broadcast()
+}
+
+func (m *model) loadMessages() {
+	if m.currentRoomID == "" {
+		return
+	}
+
+	dbMessages, err := loadMessages(m.currentRoomID)
+	if err != nil {
+		m.err = err
+		return
+	}
+
+	var msgs []chatMsg
+	for _, dbMsg := range dbMessages {
+		plain, err := decryptMsg(dbMsg.Ciphertext, dbMsg.Nonce, m.currentRoomKey)
+		if err != nil {
+			plain = "[Decryption Failed]"
+		}
+		msgs = append(msgs, chatMsg{
+			Timestamp: dbMsg.Timestamp,
+			Sender:    dbMsg.Sender,
+			Content:   plain,
+		})
+	}
+	m.messages = msgs
+	m.updateViewport()
+}
+
+func (m *model) updateViewport() {
+	var b strings.Builder
+	for _, msg := range m.messages {
+		dateStr := msg.Timestamp.Format("Jan 2 15:04")
+
+		// Mask history before join time
+		// Add 1 second buffer to avoid hiding immediate first messages
+		if msg.Timestamp.Before(m.currentRoomJoinedAt.Add(-1 * time.Second)) {
+			b.WriteString(fmt.Sprintf("%s %s\n",
+				timeStyle.Render(dateStr),
+				lockedStyle.Render("[Encrypted message - History hidden]"),
+			))
+		} else {
+			b.WriteString(fmt.Sprintf("%s %s: %s\n",
+				timeStyle.Render(dateStr),
+				senderStyle.Render(msg.Sender),
+				msgStyle.Render(msg.Content),
+			))
+		}
+	}
+	m.viewport.SetContent(b.String())
+	m.viewport.GotoBottom()
+}
+
+func (m model) View() string {
+	if !m.authorized {
+		return fmt.Sprintf(`
+%s
+ACCESS DENIED
+No public key authentication provided.
+
+(Press Ctrl+C to quit)
+`, errStyle.Render("SECURITY ALERT"))
+	}
+
+	if m.needsRegistration {
+		return fmt.Sprintf(`
+%s
+Welcome! You're a new user.
+Please choose a username to register.
+
+%s
+`, sysStyle.Render("SECURE TUI CHAT - REGISTRATION"), m.input.View())
+	}
+
+	if m.state == 0 {
+		return fmt.Sprintf(`
+%s
+Welcome back, %s.
+This environment is encrypted at rest.
+Please enter your passphrase to unlock your keys.
+
+%s
+`, sysStyle.Render("SECURE TUI CHAT"), senderStyle.Render(m.username), m.input.View())
+	}
+
+	if m.state == 1 {
+		var b strings.Builder
+		b.WriteString(sysStyle.Render("=== YOUR ROOMS ===") + "\n\n")
+		if len(m.rooms) == 0 {
+			b.WriteString(sysStyle.Render("No rooms yet.") + "\n")
+			b.WriteString(commandStyle.Render("  /join <name>") + " - Join/Create room\n")
+			b.WriteString(commandStyle.Render("  /dm <user>") + " - Start a DM\n")
+		} else {
+			for i, room := range m.rooms {
+				var style lipgloss.Style
+				var prefix string
+				if strings.Contains(room.Name, "Note to Self") {
+					style = noteStyle
+					prefix = "[*]"
+				} else if room.IsDM {
+					style = dmStyle
+					prefix = "[DM]"
+				} else {
+					style = roomStyle
+					prefix = "[#]"
+				}
+				b.WriteString(fmt.Sprintf("%s %s %s\n",
+					senderStyle.Render(fmt.Sprintf("%d.", i+1)),
+					prefix,
+					style.Render(room.Name)))
+			}
+			b.WriteString("\n" + sysStyle.Render("Commands: ") + commandStyle.Render("/join /dm /list") + "\n")
+		}
+		b.WriteString("\n" + m.input.View())
+		return b.String()
+	}
+
+	if m.state == 2 {
+		return fmt.Sprintf("%s\n%s", m.viewport.View(), m.input.View())
+	}
+
+	return m.input.View()
+}

internal/server.go

@@ -0,0 +1,45 @@
+package internal
+
+import (
+	"fmt"
+	"log"
+
+	tea "github.com/charmbracelet/bubbletea"
+	"github.com/charmbracelet/ssh"
+	"github.com/charmbracelet/wish"
+	lm "github.com/charmbracelet/wish/logging"
+)
+
+func NewServer(host string, port int) (*ssh.Server, error) {
+	srv, err := wish.NewServer(
+		wish.WithAddress(fmt.Sprintf("%s:%d", host, port)),
+		wish.WithHostKeyPath(".ssh/term_info_ed25519"),
+		wish.WithPublicKeyAuth(func(ctx ssh.Context, key ssh.PublicKey) bool {
+			return true
+		}),
+		wish.WithMiddleware(
+			appMiddleware,
+			lm.Middleware(),
+		),
+	)
+	return srv, err
+}
+
+func appMiddleware(sh ssh.Handler) ssh.Handler {
+	return func(s ssh.Session) {
+		_, _, active := s.Pty()
+		if !active {
+			wish.Fatalln(s, "no active terminal, skipping")
+			return
+		}
+
+		m := initialModel(s)
+		p := tea.NewProgram(m, tea.WithInput(s), tea.WithOutput(s), tea.WithAltScreen())
+
+		if _, err := p.Run(); err != nil {
+			log.Println("Error running program:", err)
+		}
+
+		updates.unsubscribe(m.updateChan)
+	}
+}