610 lines
13 KiB
Go
610 lines
13 KiB
Go
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()
|
|
}
|