package storage import ( "context" "database/sql" "fmt" "time" "github.com/google/uuid" _ "modernc.org/sqlite" ) // SQLiteStore implements Store using a SQLite database. type SQLiteStore struct { db *sql.DB } // NewSQLiteStore opens or creates a SQLite database at the given path, // runs pending migrations, and returns a ready-to-use store. func NewSQLiteStore(dbPath string) (*SQLiteStore, error) { dsn := dbPath + "?_pragma=journal_mode(wal)&_pragma=foreign_keys(on)&_pragma=busy_timeout(5000)" db, err := sql.Open("sqlite", dsn) if err != nil { return nil, fmt.Errorf("opening database: %w", err) } db.SetMaxOpenConns(1) if err := Migrate(db); err != nil { db.Close() return nil, fmt.Errorf("running migrations: %w", err) } return &SQLiteStore{db: db}, nil } func (s *SQLiteStore) RecordLoginAttempt(ctx context.Context, username, password, ip string) error { now := time.Now().UTC().Format(time.RFC3339) _, err := s.db.ExecContext(ctx, ` INSERT INTO login_attempts (username, password, ip, count, first_seen, last_seen) VALUES (?, ?, ?, 1, ?, ?) ON CONFLICT(username, password, ip) DO UPDATE SET count = count + 1, last_seen = ?`, username, password, ip, now, now, now) if err != nil { return fmt.Errorf("recording login attempt: %w", err) } return nil } func (s *SQLiteStore) CreateSession(ctx context.Context, ip, username, shellName string) (string, error) { id := uuid.New().String() now := time.Now().UTC().Format(time.RFC3339) _, err := s.db.ExecContext(ctx, ` INSERT INTO sessions (id, ip, username, shell_name, connected_at) VALUES (?, ?, ?, ?, ?)`, id, ip, username, shellName, now) if err != nil { return "", fmt.Errorf("creating session: %w", err) } return id, nil } func (s *SQLiteStore) EndSession(ctx context.Context, sessionID string, disconnectedAt time.Time) error { _, err := s.db.ExecContext(ctx, ` UPDATE sessions SET disconnected_at = ? WHERE id = ?`, disconnectedAt.UTC().Format(time.RFC3339), sessionID) if err != nil { return fmt.Errorf("ending session: %w", err) } return nil } func (s *SQLiteStore) UpdateHumanScore(ctx context.Context, sessionID string, score float64) error { _, err := s.db.ExecContext(ctx, ` UPDATE sessions SET human_score = ? WHERE id = ?`, score, sessionID) if err != nil { return fmt.Errorf("updating human score: %w", err) } return nil } func (s *SQLiteStore) AppendSessionLog(ctx context.Context, sessionID, input, output string) error { now := time.Now().UTC().Format(time.RFC3339) _, err := s.db.ExecContext(ctx, ` INSERT INTO session_logs (session_id, timestamp, input, output) VALUES (?, ?, ?, ?)`, sessionID, now, input, output) if err != nil { return fmt.Errorf("appending session log: %w", err) } return nil } func (s *SQLiteStore) DeleteRecordsBefore(ctx context.Context, cutoff time.Time) (int64, error) { cutoffStr := cutoff.UTC().Format(time.RFC3339) tx, err := s.db.BeginTx(ctx, nil) if err != nil { return 0, fmt.Errorf("begin transaction: %w", err) } defer tx.Rollback() var total int64 // Delete session logs for old sessions. res, err := tx.ExecContext(ctx, ` DELETE FROM session_logs WHERE session_id IN ( SELECT id FROM sessions WHERE connected_at < ? )`, cutoffStr) if err != nil { return 0, fmt.Errorf("deleting session logs: %w", err) } n, _ := res.RowsAffected() total += n // Delete old sessions. res, err = tx.ExecContext(ctx, `DELETE FROM sessions WHERE connected_at < ?`, cutoffStr) if err != nil { return 0, fmt.Errorf("deleting sessions: %w", err) } n, _ = res.RowsAffected() total += n // Delete old login attempts. res, err = tx.ExecContext(ctx, `DELETE FROM login_attempts WHERE last_seen < ?`, cutoffStr) if err != nil { return 0, fmt.Errorf("deleting login attempts: %w", err) } n, _ = res.RowsAffected() total += n if err := tx.Commit(); err != nil { return 0, fmt.Errorf("commit transaction: %w", err) } return total, nil } func (s *SQLiteStore) Close() error { return s.db.Close() }