Weetile/sshidentity
Weetile/sshidentity
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

SQLite support for user info

Browse Source
This commit is contained in:
Weetile 2025-04-06 22:18:09 +01:00
parent ae3081dee0
commit dc727699eb
4 changed files with 81 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
.gitignore vendored
View File

@ -23,3 +23,6 @@ go.work.sum
# SSH
.ssh/
# SQLite
*.db

1
go.mod
View File

@ -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 (

2
go.sum
View File

@ -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=

85
main.go
View File

@ -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"
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: