SQLite support for user info
This commit is contained in:
parent
ae3081dee0
commit
dc727699eb
3
.gitignore
vendored
3
.gitignore
vendored
@ -23,3 +23,6 @@ go.work.sum
|
|||||||
|
|
||||||
# SSH
|
# SSH
|
||||||
.ssh/
|
.ssh/
|
||||||
|
|
||||||
|
# SQLite
|
||||||
|
*.db
|
||||||
|
1
go.mod
1
go.mod
@ -9,6 +9,7 @@ require (
|
|||||||
github.com/charmbracelet/log v0.4.1
|
github.com/charmbracelet/log v0.4.1
|
||||||
github.com/charmbracelet/ssh v0.0.0-20250213143314-8712ec3ff3ef
|
github.com/charmbracelet/ssh v0.0.0-20250213143314-8712ec3ff3ef
|
||||||
github.com/charmbracelet/wish v1.4.7
|
github.com/charmbracelet/wish v1.4.7
|
||||||
|
github.com/mattn/go-sqlite3 v1.14.27
|
||||||
)
|
)
|
||||||
|
|
||||||
require (
|
require (
|
||||||
|
2
go.sum
2
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-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 h1:E5ScNMtiwvlvB5paMFdw9p4kSQzbXFikJ5SQO6TULQc=
|
||||||
github.com/mattn/go-runewidth v0.0.16/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
|
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 h1:ZK8zHtRHOkbHy6Mmr5D264iyp3TiX5OmNcI5cIARiQI=
|
||||||
github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6/go.mod h1:CJlz5H+gyd6CUWT45Oy4q24RdLyn7Md9Vj2/ldJBSIo=
|
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=
|
github.com/muesli/cancelreader v0.2.2 h1:3I4Kt4BQjOR54NavqnDogx/MIoWBFa0StPA8ELUXHmA=
|
||||||
|
89
main.go
89
main.go
@ -3,6 +3,7 @@ package main
|
|||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"crypto/md5"
|
"crypto/md5"
|
||||||
|
"database/sql"
|
||||||
"encoding/hex"
|
"encoding/hex"
|
||||||
"errors"
|
"errors"
|
||||||
"fmt"
|
"fmt"
|
||||||
@ -21,22 +22,45 @@ import (
|
|||||||
"github.com/charmbracelet/wish"
|
"github.com/charmbracelet/wish"
|
||||||
"github.com/charmbracelet/wish/bubbletea"
|
"github.com/charmbracelet/wish/bubbletea"
|
||||||
"github.com/charmbracelet/wish/logging"
|
"github.com/charmbracelet/wish/logging"
|
||||||
|
_ "github.com/mattn/go-sqlite3"
|
||||||
)
|
)
|
||||||
|
|
||||||
const (
|
const (
|
||||||
host = "localhost"
|
host = "localhost"
|
||||||
port = "23234"
|
port = "23234"
|
||||||
|
dbFile = "users.db"
|
||||||
)
|
)
|
||||||
|
|
||||||
var (
|
var (
|
||||||
// knownUsers maps public key fingerprints to nicknames
|
|
||||||
knownUsers = make(map[string]string)
|
|
||||||
usersMutex sync.Mutex
|
|
||||||
errorStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("9"))
|
errorStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("9"))
|
||||||
successStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("10"))
|
successStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("10"))
|
||||||
|
db *sql.DB
|
||||||
|
dbMutex sync.Mutex
|
||||||
)
|
)
|
||||||
|
|
||||||
func main() {
|
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(
|
s, err := wish.NewServer(
|
||||||
wish.WithAddress(net.JoinHostPort(host, port)),
|
wish.WithAddress(net.JoinHostPort(host, port)),
|
||||||
wish.WithHostKeyPath(".ssh/id_ed25519"),
|
wish.WithHostKeyPath(".ssh/id_ed25519"),
|
||||||
@ -78,6 +102,37 @@ func getFingerprint(pubKey ssh.PublicKey) string {
|
|||||||
return hex.EncodeToString(hash[:])
|
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
|
// teaHandler creates a Bubble Tea program for each SSH session
|
||||||
func teaHandler(s ssh.Session) (tea.Model, []tea.ProgramOption) {
|
func teaHandler(s ssh.Session) (tea.Model, []tea.ProgramOption) {
|
||||||
_, _, active := s.Pty()
|
_, _, active := s.Pty()
|
||||||
@ -95,11 +150,13 @@ func teaHandler(s ssh.Session) (tea.Model, []tea.ProgramOption) {
|
|||||||
fingerprint := getFingerprint(pubKey)
|
fingerprint := getFingerprint(pubKey)
|
||||||
|
|
||||||
// Check if we know this user
|
// Check if we know this user
|
||||||
usersMutex.Lock()
|
nickname, err := getUserNickname(fingerprint)
|
||||||
nickname, known := knownUsers[fingerprint]
|
if err != nil {
|
||||||
usersMutex.Unlock()
|
wish.Fatalln(s, fmt.Sprintf("database error: %v", err))
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
|
||||||
if known {
|
if nickname != "" {
|
||||||
// Known user - send greeting and close connection
|
// Known user - send greeting and close connection
|
||||||
wish.Println(s, fmt.Sprintf("Hello, %s!", nickname))
|
wish.Println(s, fmt.Sprintf("Hello, %s!", nickname))
|
||||||
return nil, nil
|
return nil, nil
|
||||||
@ -115,6 +172,7 @@ func teaHandler(s ssh.Session) (tea.Model, []tea.ProgramOption) {
|
|||||||
m := model{
|
m := model{
|
||||||
textInput: ti,
|
textInput: ti,
|
||||||
fingerprint: fingerprint,
|
fingerprint: fingerprint,
|
||||||
|
pubKey: pubKey,
|
||||||
session: s,
|
session: s,
|
||||||
}
|
}
|
||||||
return m, []tea.ProgramOption{tea.WithAltScreen()}
|
return m, []tea.ProgramOption{tea.WithAltScreen()}
|
||||||
@ -124,6 +182,7 @@ func teaHandler(s ssh.Session) (tea.Model, []tea.ProgramOption) {
|
|||||||
type model struct {
|
type model struct {
|
||||||
textInput textinput.Model
|
textInput textinput.Model
|
||||||
fingerprint string
|
fingerprint string
|
||||||
|
pubKey ssh.PublicKey
|
||||||
nickname string
|
nickname string
|
||||||
known bool
|
known bool
|
||||||
err string
|
err string
|
||||||
@ -159,13 +218,15 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
|
|||||||
return m, nil
|
return m, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// Store the nickname
|
// Store the user in the database
|
||||||
usersMutex.Lock()
|
err := addUser(m.fingerprint, m.pubKey, nick)
|
||||||
knownUsers[m.fingerprint] = nick
|
if err != nil {
|
||||||
usersMutex.Unlock()
|
m.err = fmt.Sprintf("Could not save nickname: %v", err)
|
||||||
|
return m, nil
|
||||||
|
}
|
||||||
|
|
||||||
// Send greeting and close connection
|
// 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
|
return m, tea.Quit
|
||||||
|
|
||||||
case tea.KeyCtrlC, tea.KeyEsc:
|
case tea.KeyCtrlC, tea.KeyEsc:
|
||||||
|
Loading…
x
Reference in New Issue
Block a user