Torjus Håkestad
d655968216
Adds persistent storage using modernc.org/sqlite (pure Go). Login attempts are deduplicated by (username, password, ip) with counts. Sessions and session logs are tracked with UUID IDs. Includes embedded SQL migrations, configurable retention with background pruning, and an in-memory store for tests. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
144 lines
2.9 KiB
Go
144 lines
2.9 KiB
Go
package storage
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
)
|
|
|
|
// MemoryStore is an in-memory implementation of Store for use in tests.
|
|
type MemoryStore struct {
|
|
mu sync.Mutex
|
|
LoginAttempts []LoginAttempt
|
|
Sessions map[string]*Session
|
|
SessionLogs []SessionLog
|
|
}
|
|
|
|
// NewMemoryStore returns a new empty MemoryStore.
|
|
func NewMemoryStore() *MemoryStore {
|
|
return &MemoryStore{
|
|
Sessions: make(map[string]*Session),
|
|
}
|
|
}
|
|
|
|
func (m *MemoryStore) RecordLoginAttempt(_ context.Context, username, password, ip string) error {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
|
|
now := time.Now().UTC()
|
|
for i := range m.LoginAttempts {
|
|
a := &m.LoginAttempts[i]
|
|
if a.Username == username && a.Password == password && a.IP == ip {
|
|
a.Count++
|
|
a.LastSeen = now
|
|
return nil
|
|
}
|
|
}
|
|
|
|
m.LoginAttempts = append(m.LoginAttempts, LoginAttempt{
|
|
ID: int64(len(m.LoginAttempts) + 1),
|
|
Username: username,
|
|
Password: password,
|
|
IP: ip,
|
|
Count: 1,
|
|
FirstSeen: now,
|
|
LastSeen: now,
|
|
})
|
|
return nil
|
|
}
|
|
|
|
func (m *MemoryStore) CreateSession(_ context.Context, ip, username, shellName string) (string, error) {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
|
|
id := uuid.New().String()
|
|
now := time.Now().UTC()
|
|
m.Sessions[id] = &Session{
|
|
ID: id,
|
|
IP: ip,
|
|
Username: username,
|
|
ShellName: shellName,
|
|
ConnectedAt: now,
|
|
}
|
|
return id, nil
|
|
}
|
|
|
|
func (m *MemoryStore) EndSession(_ context.Context, sessionID string, disconnectedAt time.Time) error {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
|
|
if s, ok := m.Sessions[sessionID]; ok {
|
|
t := disconnectedAt.UTC()
|
|
s.DisconnectedAt = &t
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (m *MemoryStore) UpdateHumanScore(_ context.Context, sessionID string, score float64) error {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
|
|
if s, ok := m.Sessions[sessionID]; ok {
|
|
s.HumanScore = &score
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (m *MemoryStore) AppendSessionLog(_ context.Context, sessionID, input, output string) error {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
|
|
m.SessionLogs = append(m.SessionLogs, SessionLog{
|
|
ID: int64(len(m.SessionLogs) + 1),
|
|
SessionID: sessionID,
|
|
Timestamp: time.Now().UTC(),
|
|
Input: input,
|
|
Output: output,
|
|
})
|
|
return nil
|
|
}
|
|
|
|
func (m *MemoryStore) DeleteRecordsBefore(_ context.Context, cutoff time.Time) (int64, error) {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
|
|
var total int64
|
|
|
|
// Delete old login attempts.
|
|
kept := m.LoginAttempts[:0]
|
|
for _, a := range m.LoginAttempts {
|
|
if a.LastSeen.Before(cutoff) {
|
|
total++
|
|
} else {
|
|
kept = append(kept, a)
|
|
}
|
|
}
|
|
m.LoginAttempts = kept
|
|
|
|
// Delete old sessions and their logs.
|
|
for id, s := range m.Sessions {
|
|
if s.ConnectedAt.Before(cutoff) {
|
|
delete(m.Sessions, id)
|
|
total++
|
|
}
|
|
}
|
|
|
|
keptLogs := m.SessionLogs[:0]
|
|
for _, l := range m.SessionLogs {
|
|
if _, ok := m.Sessions[l.SessionID]; ok {
|
|
keptLogs = append(keptLogs, l)
|
|
} else {
|
|
total++
|
|
}
|
|
}
|
|
m.SessionLogs = keptLogs
|
|
|
|
return total, nil
|
|
}
|
|
|
|
func (m *MemoryStore) Close() error {
|
|
return nil
|
|
}
|