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