-
Notifications
You must be signed in to change notification settings - Fork 2
/
queries.go
412 lines (375 loc) · 12.2 KB
/
queries.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
// Copyright (c) 2015, Marios Andreopoulos.
//
// This file is part of bashistdb.
//
// Bashistdb is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// Bashistdb is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with bashistdb. If not, see <http://www.gnu.org/licenses/>.
/*
Package database handles a SQLite3 database and access methods for the
specific needs of bashistdb.
*/
package database
import (
"bytes"
"database/sql"
"errors"
"fmt"
"regexp"
"strconv"
"strings"
"time"
conf "github.com/andmarios/bashistdb/configuration"
"github.com/andmarios/bashistdb/result"
)
// TopK returns the k most frequent command lines in history
func (d Database) TopK(qp conf.QueryParams) ([]byte, error) {
rows, err := d.Query(`SELECT command, count(*) as count FROM history
WHERE user LIKE ? AND host LIKE ? AND command LIKE ? ESCAPE '\'
GROUP BY command ORDER BY count DESC LIMIT ?`,
qp.User, qp.Host, qp.Command, qp.Kappa)
if err != nil {
return []byte{}, err
}
defer rows.Close()
res := result.New("")
for rows.Next() {
var command string
var count int
rows.Scan(&command, &count)
res.AddCountRow(count, command)
}
return res.Formatted(), err
}
// LastK returns the k most recent command lines in history
func (d Database) LastK(qp conf.QueryParams) ([]byte, error) {
var rows *sql.Rows
var err error
switch qp.Unique {
case true:
rows, err = d.Query(`SELECT * FROM
(SELECT rowid, * FROM history
WHERE user LIKE ? AND host LIKE ? AND command LIKE ? ESCAPE '\'
GROUP BY command
ORDER BY datetime DESC LIMIT ?)
ORDER BY datetime ASC`,
qp.User, qp.Host, qp.Command, qp.Kappa)
default:
rows, err = d.Query(`SELECT * FROM
(SELECT rowid, * FROM history
WHERE user LIKE ? AND host LIKE ? AND command LIKE ? ESCAPE '\'
ORDER BY datetime DESC LIMIT ?)
ORDER BY datetime ASC`,
qp.User, qp.Host, qp.Command, qp.Kappa)
}
if err != nil {
return []byte{}, err
}
defer rows.Close()
res := result.New(qp.Format)
for rows.Next() {
var user, host, command string
var t time.Time
var row int
rows.Scan(&row, &user, &host, &command, &t)
res.AddRow(row, user, host, command, t)
}
return res.Formatted(), nil
}
// DefaultQuery returns history within the search criteria in the format requested
func (d Database) DefaultQuery(qp conf.QueryParams) ([]byte, error) {
// SQLite's regexp extension is problematic; most systems don't have it, loading
// it is extremely error prone (almost impossible to get right), even we manage
// to load it, it doesn't seem to work with our queries. Thus I use go's regexp
// library. This makes bashistdb slower when in regexp mode since it has to process
// all the history entries itself. Also it makes difficult to add pcre support
// to anything but the DefaultQuery
var regex *regexp.Regexp
var err error
commandQuery := "AND command LIKE ?" // This is used for normal searches. Fast.
if qp.Regex {
regex, err = regexp.Compile(qp.Command)
if err != nil {
return []byte{}, err
}
commandQuery = "" // For PCRE we do the search, so we want everything. Slow.
}
var rows *sql.Rows
switch qp.Unique {
case true:
rows, err = d.Query(`SELECT rowid, * FROM history
WHERE user LIKE ? AND host LIKE ? `+commandQuery+` ESCAPE '\'
GROUP BY command ORDER BY DATETIME ASC`,
qp.User, qp.Host, qp.Command)
default:
rows, err = d.Query(`SELECT rowid, * FROM history
WHERE user LIKE ? AND host LIKE ? `+commandQuery+` ESCAPE '\'`,
qp.User, qp.Host, qp.Command)
}
if err != nil {
return nil, err
}
defer rows.Close()
res := result.New(qp.Format)
for rows.Next() {
var user, host, command string
var t time.Time
var row int
rows.Scan(&row, &user, &host, &command, &t)
switch qp.Regex {
case true:
if regex.MatchString(command) {
res.AddRow(row, user, host, command, t)
}
default:
res.AddRow(row, user, host, command, t)
}
}
// Return the result without the newline at the end.
return res.Formatted(), nil
}
// RunQuery is a wrapper around various queries.
func (d Database) RunQuery(p conf.QueryParams) ([]byte, error) {
switch p.Type {
case conf.QUERY:
return d.DefaultQuery(p)
case conf.QUERY_LASTK:
return d.LastK(p)
case conf.QUERY_TOPK:
return d.TopK(p)
case conf.QUERY_USERS:
return d.Users(p)
case conf.QUERY_DEMO:
return d.Demo(p)
case conf.QUERY_ROW:
return d.ReturnRow(p)
case conf.DELETE:
return d.DeleteRows(p)
case conf.QUERY_CONTENT:
return d.ContentQuery(p)
}
return []byte{}, errors.New("Unknown query type.")
}
// Users returns unique user@host pairs from the database.
func (d Database) Users(qp conf.QueryParams) (res []byte, e error) {
var result bytes.Buffer
result.WriteString(fmt.Sprintf("Unique user-hosts pairs:"))
rows, e := d.Query(`SELECT distinct(user), host FROM history
WHERE user LIKE ? AND host LIKE ? AND command LIKE ? ESCAPE '\'`,
qp.User, qp.Host, qp.Command)
if e != nil {
return result.Bytes(), e
}
defer rows.Close()
for rows.Next() {
var user string
var host string
rows.Scan(&user, &host)
result.WriteString(fmt.Sprintf("\n%s@%s", user, host))
}
return result.Bytes(), e
}
// Demo returns some stats from the database to showcase bashistdb.
func (d Database) Demo(qp conf.QueryParams) (res []byte, e error) {
var result bytes.Buffer
var numUsers int
err := d.QueryRow("SELECT count(*) FROM (SELECT distinct(user), host FROM history)").Scan(&numUsers)
if err != nil {
return result.Bytes(), err
}
var numHosts int
err = d.QueryRow("SELECT count(distinct(host)) FROM history").Scan(&numHosts)
if err != nil {
return result.Bytes(), err
}
var numLines int
err = d.QueryRow("SELECT count(command) FROM history").Scan(&numLines)
if err != nil {
return result.Bytes(), err
}
var numUniqueLines int
err = d.QueryRow("SELECT count(distinct(command)) FROM history").Scan(&numUniqueLines)
if err != nil {
return result.Bytes(), err
}
qp.Kappa = 15
restop, err := d.TopK(qp)
if err != nil {
return result.Bytes(), err
}
qp.Kappa = 10
reslast, err := d.LastK(qp)
if err != nil {
return result.Bytes(), err
}
result.WriteString(fmt.Sprintf("There are %d command lines (%d unique) in your database from %d users across %d hosts.\n\n", numLines, numUniqueLines, numUsers, numHosts))
result.WriteString(fmt.Sprintf("Top-15 commands for user %s@%s:\n", qp.User, qp.Host))
result.Write(restop)
result.WriteString(fmt.Sprintf("\n\nLast 10 commands user %s@%s ran:\n", qp.User, qp.Host))
result.Write(reslast)
return result.Bytes(), nil
}
// ReturnRow returns a single row with no other data.
// It is useful to pipe to bash.
func (d Database) ReturnRow(qp conf.QueryParams) ([]byte, error) {
var command string
err := d.QueryRow("SELECT command FROM history WHERE rowid = ?", qp.Kappa).Scan(&command)
if err != nil {
return []byte{}, err
}
return []byte(command), nil
}
// DeleteRows deletes a range of rows.
func (d Database) DeleteRows(qp conf.QueryParams) ([]byte, error) {
tx, err := d.Begin()
defer tx.Rollback()
if err != nil {
return []byte{}, err
}
stmt, err := tx.Prepare(`DELETE FROM history WHERE rowid=?`)
if err != nil {
return []byte{}, err
}
for i := len(qp.Rows) - 1; i >= 0; i-- {
_, err = stmt.Exec(qp.Rows[i])
if err != nil {
return []byte{}, err
}
}
tx.Commit()
return []byte("No errors during deletion."), nil
}
// ContentQuery returns matches of a row plus rows before or after the match.
// Think of it as grep -A(fter) / -B(efore) / -C(ontent)
// It works on 4 stages:
// 1. find matches and get their datetime
// 2. for each match, get the content asked by rowid
// 3. if the content of two matches overlap, join them
// 4. given the sets of rowids, get them from the database
func (d Database) ContentQuery(qp conf.QueryParams) ([]byte, error) {
// Using Goland regexp library. Check DefaultQuery() for more info.
var regex *regexp.Regexp
var err error
commandQuery := "AND command LIKE ?" // This is used for normal searches. Fast.
if qp.Regex {
regex, err = regexp.Compile(qp.Command)
if err != nil {
return []byte{}, err
}
commandQuery = "" // For PCRE we do the search, so we want everything. Slow.
}
// Stage 1: find matches and get an array with their datetime
var rows *sql.Rows
rows, err = d.Query(`SELECT datetime, command FROM history
WHERE user LIKE ? AND host LIKE ? `+commandQuery+` ESCAPE '\'`,
qp.User, qp.Host, qp.Command)
if err != nil {
return nil, err
}
defer rows.Close()
var hits []time.Time
for rows.Next() {
var t time.Time
var command string
rows.Scan(&t, &command)
switch qp.Regex {
case true:
if regex.MatchString(command) {
hits = append(hits, t)
}
default:
hits = append(hits, t)
}
}
// Stage 2: for each match create a slice with its content by rowid
var hitsContent [][]int
for _, v := range hits {
var content []int
// Before query also includes the current command, thus is always run.
rows, err = d.Query(`SELECT rowid, datetime FROM
(SELECT rowid, datetime FROM history
WHERE datetime <= ? AND user LIKE ? AND host LIKE ? ESCAPE '\'
ORDER BY datetime DESC LIMIT ?)
ORDER BY datetime ASC`,
v, qp.User, qp.Host, qp.BeforeContent+1) // Here we include current query to before
if err != nil {
return nil, err
}
defer rows.Close()
for rows.Next() {
var row int
var datetime time.Time
rows.Scan(&row, &datetime)
content = append(content, row)
}
// After runs only if needed.
if qp.AfterContent > 0 {
rows, err = d.Query(`SELECT rowid, datetime FROM history
WHERE datetime > ? AND user LIKE ? AND host LIKE ? ESCAPE '\'
ORDER BY datetime ASC LIMIT ?`,
v, qp.User, qp.Host, qp.AfterContent)
if err != nil {
return nil, err
}
defer rows.Close()
for rows.Next() {
var row int
var datetime time.Time
rows.Scan(&row, &datetime)
content = append(content, row)
}
}
hitsContent = append(hitsContent, content)
}
// Stage 3: let's merge the sets that overlap
for i := len(hitsContent) - 2; i >= 0; i-- {
for k, v := range hitsContent[i] {
if v == hitsContent[i+1][0] {
hitsContent[i] = append(hitsContent[i][:k], hitsContent[i+1]...)
if len(hitsContent) > i+2 {
hitsContent = append(hitsContent[:i+1], hitsContent[i+2:]...)
} else {
hitsContent = hitsContent[:i+1]
}
break
}
}
}
var out bytes.Buffer
// Stage 4: get the tuples for each set's rowids and add them formatted to the result
for i := 0; i < len(hitsContent); i++ {
var rowids []string
for _, v := range hitsContent[i] {
rowids = append(rowids, strconv.Itoa(v))
}
rows, err = d.Query(`SELECT rowid, * FROM history
WHERE rowid IN (` + strings.Join(rowids, ",") + `)
ORDER BY datetime ASC`)
if err != nil {
return nil, err
}
defer rows.Close()
res := result.New(qp.Format)
for rows.Next() {
var user, host, command string
var t time.Time
var row int
rows.Scan(&row, &user, &host, &command, &t)
res.AddRow(row, user, host, command, t)
}
out.Write(res.Formatted())
if i < len(hitsContent)-1 {
out.WriteString("\n------------------\n")
}
}
return out.Bytes(), nil
}