diff --git a/.gitignore b/.gitignore index 9182466..d710ae1 100644 --- a/.gitignore +++ b/.gitignore @@ -23,3 +23,6 @@ go.work.sum # SSH .ssh/ + +# SQLite +*.db diff --git a/go.mod b/go.mod index 15841fc..6552776 100644 --- a/go.mod +++ b/go.mod @@ -9,6 +9,7 @@ require ( github.com/charmbracelet/log v0.4.1 github.com/charmbracelet/ssh v0.0.0-20250213143314-8712ec3ff3ef github.com/charmbracelet/wish v1.4.7 + github.com/mattn/go-sqlite3 v1.14.27 ) require ( diff --git a/go.sum b/go.sum index f722e5e..8a91f12 100644 --- a/go.sum +++ b/go.sum @@ -52,6 +52,8 @@ github.com/mattn/go-localereader v0.0.1 h1:ygSAOl7ZXTx4RdPYinUpg6W99U8jWvWi9Ye2J github.com/mattn/go-localereader v0.0.1/go.mod h1:8fBrzywKY7BI3czFoHkuzRoWE9C+EiG4R1k4Cjx5p88= github.com/mattn/go-runewidth v0.0.16 h1:E5ScNMtiwvlvB5paMFdw9p4kSQzbXFikJ5SQO6TULQc= github.com/mattn/go-runewidth v0.0.16/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w= +github.com/mattn/go-sqlite3 v1.14.27 h1:drZCnuvf37yPfs95E5jd9s3XhdVWLal+6BOK6qrv6IU= +github.com/mattn/go-sqlite3 v1.14.27/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y= github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 h1:ZK8zHtRHOkbHy6Mmr5D264iyp3TiX5OmNcI5cIARiQI= github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6/go.mod h1:CJlz5H+gyd6CUWT45Oy4q24RdLyn7Md9Vj2/ldJBSIo= github.com/muesli/cancelreader v0.2.2 h1:3I4Kt4BQjOR54NavqnDogx/MIoWBFa0StPA8ELUXHmA= diff --git a/main.go b/main.go index ef641d7..b9b7db6 100644 --- a/main.go +++ b/main.go @@ -3,6 +3,7 @@ package main import ( "context" "crypto/md5" + "database/sql" "encoding/hex" "errors" "fmt" @@ -21,22 +22,45 @@ import ( "github.com/charmbracelet/wish" "github.com/charmbracelet/wish/bubbletea" "github.com/charmbracelet/wish/logging" + _ "github.com/mattn/go-sqlite3" ) const ( - host = "localhost" - port = "23234" + host = "localhost" + port = "23234" + dbFile = "users.db" ) 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")) + db *sql.DB + dbMutex sync.Mutex ) func main() { + // Initialize database + var err error + db, err = sql.Open("sqlite3", dbFile) + if err != nil { + log.Fatal("Could not open database", "error", err) + } + defer db.Close() + + // Create users table if it doesn't exist + _, err = db.Exec(` + CREATE TABLE IF NOT EXISTS users ( + fingerprint TEXT PRIMARY KEY, + public_key BLOB NOT NULL, + nickname TEXT NOT NULL, + created_at DATETIME NOT NULL, + CHECK(length(nickname) >= 3 AND length(nickname) <= 12) + ) + `) + if err != nil { + log.Fatal("Could not create users table", "error", err) + } + s, err := wish.NewServer( wish.WithAddress(net.JoinHostPort(host, port)), wish.WithHostKeyPath(".ssh/id_ed25519"), @@ -78,6 +102,37 @@ func getFingerprint(pubKey ssh.PublicKey) string { return hex.EncodeToString(hash[:]) } +// getUserNickname retrieves a user's nickname from the database +func getUserNickname(fingerprint string) (string, error) { + dbMutex.Lock() + defer dbMutex.Unlock() + + var nickname string + err := db.QueryRow("SELECT nickname FROM users WHERE fingerprint = ?", fingerprint).Scan(&nickname) + if err != nil { + if err == sql.ErrNoRows { + return "", nil + } + return "", err + } + return nickname, nil +} + +// addUser adds a new user to the database +func addUser(fingerprint string, pubKey ssh.PublicKey, nickname string) error { + dbMutex.Lock() + defer dbMutex.Unlock() + + _, err := db.Exec( + "INSERT INTO users (fingerprint, public_key, nickname, created_at) VALUES (?, ?, ?, ?)", + fingerprint, + pubKey.Marshal(), + nickname, + time.Now().UTC(), + ) + return err +} + // teaHandler creates a Bubble Tea program for each SSH session func teaHandler(s ssh.Session) (tea.Model, []tea.ProgramOption) { _, _, active := s.Pty() @@ -95,11 +150,13 @@ func teaHandler(s ssh.Session) (tea.Model, []tea.ProgramOption) { fingerprint := getFingerprint(pubKey) // Check if we know this user - usersMutex.Lock() - nickname, known := knownUsers[fingerprint] - usersMutex.Unlock() + nickname, err := getUserNickname(fingerprint) + if err != nil { + wish.Fatalln(s, fmt.Sprintf("database error: %v", err)) + return nil, nil + } - if known { + if nickname != "" { // Known user - send greeting and close connection wish.Println(s, fmt.Sprintf("Hello, %s!", nickname)) return nil, nil @@ -115,6 +172,7 @@ func teaHandler(s ssh.Session) (tea.Model, []tea.ProgramOption) { m := model{ textInput: ti, fingerprint: fingerprint, + pubKey: pubKey, session: s, } return m, []tea.ProgramOption{tea.WithAltScreen()} @@ -124,6 +182,7 @@ func teaHandler(s ssh.Session) (tea.Model, []tea.ProgramOption) { type model struct { textInput textinput.Model fingerprint string + pubKey ssh.PublicKey nickname string known bool err string @@ -159,13 +218,15 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { return m, nil } - // Store the nickname - usersMutex.Lock() - knownUsers[m.fingerprint] = nick - usersMutex.Unlock() + // Store the user in the database + err := addUser(m.fingerprint, m.pubKey, nick) + if err != nil { + m.err = fmt.Sprintf("Could not save nickname: %v", err) + return m, nil + } // Send greeting and close connection - wish.Println(m.session, fmt.Sprintf("Hello, %s!", nick)) + wish.Println(m.session, fmt.Sprintf("Hello, %s! Your account has been created.", nick)) return m, tea.Quit case tea.KeyCtrlC, tea.KeyEsc: