O_O
This commit is contained in:
609
internal/model.go
Normal file
609
internal/model.go
Normal file
@@ -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()
|
||||
}
|
||||
Reference in New Issue
Block a user