ScuttleBot

scuttlebot / internal / store / store.go
Source Blame History 249 lines
0e78954… lmata 1 // Package store provides a thin database/sql wrapper for scuttlebot's
0e78954… lmata 2 // persistent state: agent registry, admin accounts, and policies.
0e78954… lmata 3 // Both SQLite and PostgreSQL are supported.
0e78954… lmata 4 package store
0e78954… lmata 5
0e78954… lmata 6 import (
0e78954… lmata 7 "database/sql"
0e78954… lmata 8 "encoding/base64"
0e78954… lmata 9 "fmt"
0e78954… lmata 10 "strconv"
0e78954… lmata 11 "time"
0e78954… lmata 12
0e78954… lmata 13 _ "github.com/lib/pq"
0e78954… lmata 14 _ "modernc.org/sqlite"
0e78954… lmata 15 )
0e78954… lmata 16
0e78954… lmata 17 // AgentRow is the flat database representation of a registered agent.
0e78954… lmata 18 type AgentRow struct {
0e78954… lmata 19 Nick string
0e78954… lmata 20 Type string
0e78954… lmata 21 Config []byte // JSON-encoded EngagementConfig
0e78954… lmata 22 CreatedAt time.Time
0e78954… lmata 23 Revoked bool
66d18d7… lmata 24 LastSeen *time.Time
0e78954… lmata 25 }
0e78954… lmata 26
0e78954… lmata 27 // AdminRow is the flat database representation of an admin account.
0e78954… lmata 28 type AdminRow struct {
0e78954… lmata 29 Username string
0e78954… lmata 30 Hash []byte // bcrypt hash
0e78954… lmata 31 CreatedAt time.Time
0e78954… lmata 32 }
0e78954… lmata 33
0e78954… lmata 34 // Store wraps a sql.DB with scuttlebot-specific CRUD operations.
0e78954… lmata 35 type Store struct {
0e78954… lmata 36 db *sql.DB
0e78954… lmata 37 driver string
0e78954… lmata 38 }
0e78954… lmata 39
0e78954… lmata 40 // Open opens a database connection, runs schema migrations, and returns a Store.
0e78954… lmata 41 // driver must be "sqlite" or "postgres". dsn is the connection string.
0e78954… lmata 42 func Open(driver, dsn string) (*Store, error) {
0e78954… lmata 43 db, err := sql.Open(driver, dsn)
0e78954… lmata 44 if err != nil {
0e78954… lmata 45 return nil, fmt.Errorf("store: open %s: %w", driver, err)
0e78954… lmata 46 }
0e78954… lmata 47 if err := db.Ping(); err != nil {
0e78954… lmata 48 db.Close()
0e78954… lmata 49 return nil, fmt.Errorf("store: ping %s: %w", driver, err)
0e78954… lmata 50 }
0e78954… lmata 51 s := &Store{db: db, driver: driver}
0e78954… lmata 52 if err := s.migrate(); err != nil {
0e78954… lmata 53 db.Close()
0e78954… lmata 54 return nil, fmt.Errorf("store: migrate: %w", err)
0e78954… lmata 55 }
0e78954… lmata 56 return s, nil
0e78954… lmata 57 }
0e78954… lmata 58
0e78954… lmata 59 // Close closes the underlying database connection.
0e78954… lmata 60 func (s *Store) Close() error { return s.db.Close() }
0e78954… lmata 61
0e78954… lmata 62 // ph returns the query placeholder for argument n (1-indexed).
0e78954… lmata 63 // SQLite uses "?"; PostgreSQL uses "$1", "$2", …
0e78954… lmata 64 func (s *Store) ph(n int) string {
0e78954… lmata 65 if s.driver == "postgres" {
0e78954… lmata 66 return "$" + strconv.Itoa(n)
0e78954… lmata 67 }
0e78954… lmata 68 return "?"
0e78954… lmata 69 }
0e78954… lmata 70
0e78954… lmata 71 func (s *Store) migrate() error {
0e78954… lmata 72 stmts := []string{
0e78954… lmata 73 `CREATE TABLE IF NOT EXISTS agents (
0e78954… lmata 74 nick TEXT PRIMARY KEY,
0e78954… lmata 75 type TEXT NOT NULL,
0e78954… lmata 76 config TEXT NOT NULL,
0e78954… lmata 77 created_at TEXT NOT NULL,
0e78954… lmata 78 revoked INTEGER NOT NULL DEFAULT 0
0e78954… lmata 79 )`,
0e78954… lmata 80 `CREATE TABLE IF NOT EXISTS admins (
0e78954… lmata 81 username TEXT PRIMARY KEY,
0e78954… lmata 82 hash TEXT NOT NULL,
0e78954… lmata 83 created_at TEXT NOT NULL
0e78954… lmata 84 )`,
0e78954… lmata 85 `CREATE TABLE IF NOT EXISTS policies (
0e78954… lmata 86 id INTEGER PRIMARY KEY,
0e78954… lmata 87 data TEXT NOT NULL
0e78954… lmata 88 )`,
0e78954… lmata 89 }
66d18d7… lmata 90 // Run base schema.
0e78954… lmata 91 for _, stmt := range stmts {
0e78954… lmata 92 if _, err := s.db.Exec(stmt); err != nil {
0e78954… lmata 93 return fmt.Errorf("migrate: %w", err)
0e78954… lmata 94 }
66d18d7… lmata 95 }
66d18d7… lmata 96 // Additive migrations — safe to re-run.
66d18d7… lmata 97 addColumns := []string{
66d18d7… lmata 98 `ALTER TABLE agents ADD COLUMN last_seen TEXT`,
66d18d7… lmata 99 }
66d18d7… lmata 100 for _, stmt := range addColumns {
66d18d7… lmata 101 _, _ = s.db.Exec(stmt) // ignore "column already exists"
0e78954… lmata 102 }
0e78954… lmata 103 return nil
0e78954… lmata 104 }
d924aea… lmata 105
0e78954… lmata 106 // AgentUpsert inserts or updates an agent row by nick.
0e78954… lmata 107 func (s *Store) AgentUpsert(r *AgentRow) error {
0e78954… lmata 108 revoked := 0
0e78954… lmata 109 if r.Revoked {
0e78954… lmata 110 revoked = 1
0e78954… lmata 111 }
66d18d7… lmata 112 var lastSeen string
66d18d7… lmata 113 if r.LastSeen != nil {
66d18d7… lmata 114 lastSeen = r.LastSeen.UTC().Format(time.RFC3339Nano)
66d18d7… lmata 115 }
0e78954… lmata 116 q := fmt.Sprintf(
66d18d7… lmata 117 `INSERT INTO agents (nick, type, config, created_at, revoked, last_seen)
66d18d7… lmata 118 VALUES (%s, %s, %s, %s, %s, %s)
0e78954… lmata 119 ON CONFLICT(nick) DO UPDATE SET
0e78954… lmata 120 type=EXCLUDED.type, config=EXCLUDED.config,
66d18d7… lmata 121 created_at=EXCLUDED.created_at, revoked=EXCLUDED.revoked,
66d18d7… lmata 122 last_seen=EXCLUDED.last_seen`,
66d18d7… lmata 123 s.ph(1), s.ph(2), s.ph(3), s.ph(4), s.ph(5), s.ph(6),
0e78954… lmata 124 )
0e78954… lmata 125 _, err := s.db.Exec(q,
0e78954… lmata 126 r.Nick, r.Type, string(r.Config),
66d18d7… lmata 127 r.CreatedAt.UTC().Format(time.RFC3339), revoked, lastSeen,
0e78954… lmata 128 )
0e78954… lmata 129 return err
0e78954… lmata 130 }
0e78954… lmata 131
0e78954… lmata 132 // AgentDelete removes an agent row entirely.
0e78954… lmata 133 func (s *Store) AgentDelete(nick string) error {
0e78954… lmata 134 _, err := s.db.Exec(
0e78954… lmata 135 fmt.Sprintf(`DELETE FROM agents WHERE nick=%s`, s.ph(1)),
0e78954… lmata 136 nick,
0e78954… lmata 137 )
0e78954… lmata 138 return err
0e78954… lmata 139 }
0e78954… lmata 140
0e78954… lmata 141 // AgentList returns all agent rows, including revoked ones.
0e78954… lmata 142 func (s *Store) AgentList() ([]*AgentRow, error) {
66d18d7… lmata 143 rows, err := s.db.Query(`SELECT nick, type, config, created_at, revoked, COALESCE(last_seen,'') FROM agents`)
0e78954… lmata 144 if err != nil {
0e78954… lmata 145 return nil, err
0e78954… lmata 146 }
0e78954… lmata 147 defer rows.Close()
0e78954… lmata 148
0e78954… lmata 149 var out []*AgentRow
0e78954… lmata 150 for rows.Next() {
0e78954… lmata 151 var r AgentRow
66d18d7… lmata 152 var cfg, ts, lastSeenStr string
0e78954… lmata 153 var revokedInt int
66d18d7… lmata 154 if err := rows.Scan(&r.Nick, &r.Type, &cfg, &ts, &revokedInt, &lastSeenStr); err != nil {
0e78954… lmata 155 return nil, err
0e78954… lmata 156 }
0e78954… lmata 157 r.Config = []byte(cfg)
0e78954… lmata 158 r.Revoked = revokedInt != 0
0e78954… lmata 159 r.CreatedAt, err = time.Parse(time.RFC3339, ts)
0e78954… lmata 160 if err != nil {
0e78954… lmata 161 return nil, fmt.Errorf("store: agent %s timestamp: %w", r.Nick, err)
66d18d7… lmata 162 }
66d18d7… lmata 163 if lastSeenStr != "" {
66d18d7… lmata 164 if t, err := time.Parse(time.RFC3339Nano, lastSeenStr); err == nil {
66d18d7… lmata 165 r.LastSeen = &t
66d18d7… lmata 166 }
0e78954… lmata 167 }
0e78954… lmata 168 out = append(out, &r)
0e78954… lmata 169 }
0e78954… lmata 170 return out, rows.Err()
0e78954… lmata 171 }
0e78954… lmata 172
0e78954… lmata 173 // AdminUpsert inserts or updates an admin row. The bcrypt hash is stored as base64.
0e78954… lmata 174 func (s *Store) AdminUpsert(r *AdminRow) error {
0e78954… lmata 175 q := fmt.Sprintf(
0e78954… lmata 176 `INSERT INTO admins (username, hash, created_at)
0e78954… lmata 177 VALUES (%s, %s, %s)
0e78954… lmata 178 ON CONFLICT(username) DO UPDATE SET hash=EXCLUDED.hash, created_at=EXCLUDED.created_at`,
0e78954… lmata 179 s.ph(1), s.ph(2), s.ph(3),
0e78954… lmata 180 )
0e78954… lmata 181 _, err := s.db.Exec(q,
0e78954… lmata 182 r.Username,
0e78954… lmata 183 base64.StdEncoding.EncodeToString(r.Hash),
0e78954… lmata 184 r.CreatedAt.UTC().Format(time.RFC3339),
0e78954… lmata 185 )
0e78954… lmata 186 return err
0e78954… lmata 187 }
0e78954… lmata 188
0e78954… lmata 189 // AdminDelete removes an admin row.
0e78954… lmata 190 func (s *Store) AdminDelete(username string) error {
0e78954… lmata 191 _, err := s.db.Exec(
0e78954… lmata 192 fmt.Sprintf(`DELETE FROM admins WHERE username=%s`, s.ph(1)),
0e78954… lmata 193 username,
0e78954… lmata 194 )
0e78954… lmata 195 return err
0e78954… lmata 196 }
0e78954… lmata 197
0e78954… lmata 198 // AdminList returns all admin rows.
0e78954… lmata 199 func (s *Store) AdminList() ([]*AdminRow, error) {
0e78954… lmata 200 rows, err := s.db.Query(`SELECT username, hash, created_at FROM admins`)
0e78954… lmata 201 if err != nil {
0e78954… lmata 202 return nil, err
0e78954… lmata 203 }
0e78954… lmata 204 defer rows.Close()
0e78954… lmata 205
0e78954… lmata 206 var out []*AdminRow
0e78954… lmata 207 for rows.Next() {
0e78954… lmata 208 var r AdminRow
0e78954… lmata 209 var hashB64, ts string
0e78954… lmata 210 if err := rows.Scan(&r.Username, &hashB64, &ts); err != nil {
0e78954… lmata 211 return nil, err
0e78954… lmata 212 }
0e78954… lmata 213 r.Hash, err = base64.StdEncoding.DecodeString(hashB64)
0e78954… lmata 214 if err != nil {
0e78954… lmata 215 return nil, fmt.Errorf("store: admin %s hash decode: %w", r.Username, err)
0e78954… lmata 216 }
0e78954… lmata 217 r.CreatedAt, err = time.Parse(time.RFC3339, ts)
0e78954… lmata 218 if err != nil {
0e78954… lmata 219 return nil, fmt.Errorf("store: admin %s timestamp: %w", r.Username, err)
0e78954… lmata 220 }
0e78954… lmata 221 out = append(out, &r)
0e78954… lmata 222 }
0e78954… lmata 223 return out, rows.Err()
0e78954… lmata 224 }
0e78954… lmata 225
0e78954… lmata 226 // PolicyGet returns the raw JSON blob for the singleton policy record.
0e78954… lmata 227 // Returns nil, nil if no policies have been saved yet.
0e78954… lmata 228 func (s *Store) PolicyGet() ([]byte, error) {
0e78954… lmata 229 var data string
0e78954… lmata 230 err := s.db.QueryRow(`SELECT data FROM policies WHERE id=1`).Scan(&data)
0e78954… lmata 231 if err == sql.ErrNoRows {
0e78954… lmata 232 return nil, nil
0e78954… lmata 233 }
0e78954… lmata 234 if err != nil {
0e78954… lmata 235 return nil, err
0e78954… lmata 236 }
0e78954… lmata 237 return []byte(data), nil
0e78954… lmata 238 }
0e78954… lmata 239
0e78954… lmata 240 // PolicySet writes the raw JSON blob for the singleton policy record.
0e78954… lmata 241 func (s *Store) PolicySet(data []byte) error {
0e78954… lmata 242 q := fmt.Sprintf(
0e78954… lmata 243 `INSERT INTO policies (id, data) VALUES (1, %s)
0e78954… lmata 244 ON CONFLICT(id) DO UPDATE SET data=EXCLUDED.data`,
0e78954… lmata 245 s.ph(1),
0e78954… lmata 246 )
0e78954… lmata 247 _, err := s.db.Exec(q, string(data))
0e78954… lmata 248 return err
0e78954… lmata 249 }

Keyboard Shortcuts

Open search /
Next entry (timeline) j
Previous entry (timeline) k
Open focused entry Enter
Show this help ?
Toggle theme Top nav button