2021-03-20 19:39:18 +01:00
package main
import (
2021-03-20 20:47:05 +01:00
"bufio"
2021-04-05 12:49:13 +02:00
"container/list"
2021-03-20 19:39:18 +01:00
"database/sql"
"flag"
"fmt"
"log"
"os"
2021-03-20 20:18:27 +01:00
"path/filepath"
2021-03-20 19:39:18 +01:00
"regexp"
2021-03-22 19:43:53 +01:00
"strings"
2021-04-05 14:00:51 +02:00
"time"
2021-03-20 20:18:27 +01:00
2021-04-05 14:00:51 +02:00
"github.com/tj/go-naturaldate"
2021-03-20 20:18:27 +01:00
_ "modernc.org/sqlite"
2021-03-20 19:39:18 +01:00
)
2021-04-05 12:49:13 +02:00
type HistoryEntry struct {
2021-04-05 14:00:51 +02:00
id uint32
cmd string
cwd string
hostname string
user string
timestamp time . Time
2021-04-05 12:49:13 +02:00
}
2021-03-20 20:18:27 +01:00
func databaseLocation ( ) string {
2021-03-21 10:38:12 +01:00
envOverride := os . Getenv ( "HS9001_DB_PATH" )
if envOverride != "" {
return envOverride
}
2021-03-20 20:18:27 +01:00
return filepath . Join ( xdgOrFallback ( "XDG_DATA_HOME" , filepath . Join ( os . Getenv ( "HOME" ) , ".local/share" ) ) , "hs9001/db.sqlite" )
}
func createConnection ( ) * sql . DB {
2021-03-20 19:39:18 +01:00
2021-03-20 20:18:27 +01:00
db , err := sql . Open ( "sqlite" , databaseLocation ( ) )
2021-03-20 19:39:18 +01:00
if err != nil {
log . Panic ( err )
}
2021-04-05 11:44:07 +02:00
2021-03-20 19:39:18 +01:00
return db
}
2021-03-20 20:47:05 +01:00
func initDatabase ( conn * sql . DB ) {
2021-03-20 20:18:27 +01:00
queryStmt := "CREATE TABLE history(id INTEGER PRIMARY KEY, command varchar(512), timestamp datetime DEFAULT current_timestamp, user varchar(25), hostname varchar(32));\n" +
"CREATE VIEW count_by_date AS SELECT COUNT(id), STRFTIME('%Y-%m-%d', timestamp) FROM history GROUP BY strftime('%Y-%m-%d', timestamp)"
_ , err := conn . Exec ( queryStmt )
if err != nil {
log . Panic ( err )
}
}
2021-04-05 11:44:07 +02:00
func migrateDatabase ( conn * sql . DB , currentVersion int ) {
2021-04-05 11:54:09 +02:00
migrations := [ ] string {
2021-04-05 14:37:37 +02:00
"ALTER TABLE history add column workdir varchar(4096) DEFAULT ''" ,
2021-04-05 11:54:09 +02:00
}
2021-04-05 11:44:07 +02:00
if ! ( len ( migrations ) > currentVersion ) {
return
}
_ , err := conn . Exec ( "BEGIN;" )
if err != nil {
log . Panic ( err )
}
for _ , m := range migrations [ currentVersion : ] {
_ , err := conn . Exec ( m )
if err != nil {
log . Panic ( err )
}
}
setDBVersion ( conn , len ( migrations ) )
_ , err = conn . Exec ( "END;" )
if err != nil {
log . Panic ( err )
}
}
func fetchDBVersion ( conn * sql . DB ) int {
rows , err := conn . Query ( "PRAGMA user_version;" )
if err != nil {
log . Panic ( err )
}
defer rows . Close ( )
rows . Next ( )
var res int
rows . Scan ( & res )
return res
}
func setDBVersion ( conn * sql . DB , ver int ) {
_ , err := conn . Exec ( fmt . Sprintf ( "PRAGMA user_version=%d" , ver ) )
if err != nil {
log . Panic ( err )
}
}
2021-04-05 13:00:27 +02:00
func NewHistoryEntry ( cmd string ) HistoryEntry {
wd , err := os . Getwd ( )
if err != nil {
log . Panic ( err )
}
hostname , err := os . Hostname ( )
if err != nil {
log . Panic ( err )
}
return HistoryEntry {
2021-04-05 14:00:51 +02:00
user : os . Getenv ( "USER" ) ,
hostname : hostname ,
cmd : cmd ,
cwd : wd ,
timestamp : time . Now ( ) ,
2021-04-05 13:00:27 +02:00
}
}
2021-03-20 20:47:05 +01:00
func importFromStdin ( conn * sql . DB ) {
scanner := bufio . NewScanner ( os . Stdin )
_ , err := conn . Exec ( "BEGIN;" )
if err != nil {
log . Panic ( err )
}
for scanner . Scan ( ) {
2021-04-05 13:00:27 +02:00
entry := NewHistoryEntry ( scanner . Text ( ) )
entry . cwd = ""
2021-04-05 14:00:51 +02:00
entry . timestamp = time . Unix ( 0 , 0 )
2021-04-05 13:00:27 +02:00
add ( conn , entry )
2021-03-20 20:47:05 +01:00
}
2021-03-20 19:39:18 +01:00
2021-03-20 20:47:05 +01:00
_ , err = conn . Exec ( "END;" )
if err != nil {
log . Panic ( err )
}
}
2021-04-05 14:21:33 +02:00
func search ( conn * sql . DB , q string , workdir string , beginTime time . Time , endTime time . Time ) list . List {
queryStmt := "SELECT id, command, workdir, user, hostname FROM history WHERE timestamp BETWEEN datetime(?, 'unixepoch') AND datetime(?, 'unixepoch') AND command LIKE ? AND workdir LIKE ? ORDER BY timestamp ASC"
2021-03-20 19:39:18 +01:00
2021-04-05 14:21:33 +02:00
rows , err := conn . Query ( queryStmt , beginTime . Unix ( ) , endTime . Unix ( ) , q , workdir )
2021-03-20 19:39:18 +01:00
if err != nil {
log . Panic ( err )
}
2021-04-05 12:49:13 +02:00
var result list . List
2021-04-05 11:44:07 +02:00
defer rows . Close ( )
2021-03-20 19:39:18 +01:00
for rows . Next ( ) {
2021-04-05 12:49:13 +02:00
var entry HistoryEntry
err = rows . Scan ( & entry . id , & entry . cmd , & entry . cwd , & entry . user , & entry . hostname )
2021-03-20 19:39:18 +01:00
if err != nil {
log . Panic ( err )
}
2021-04-05 12:49:13 +02:00
result . PushBack ( & entry )
2021-03-20 19:39:18 +01:00
}
2021-04-05 12:49:13 +02:00
return result
2021-03-20 19:39:18 +01:00
}
2021-04-05 12:49:13 +02:00
func delete ( conn * sql . DB , entryId uint32 ) {
queryStmt := "DELETE FROM history WHERE id = ?"
2021-04-05 10:29:50 +02:00
2021-04-05 12:49:13 +02:00
_ , err := conn . Exec ( queryStmt , entryId )
2021-04-05 10:29:50 +02:00
if err != nil {
log . Panic ( err )
}
}
2021-04-05 13:15:20 +02:00
func add ( conn * sql . DB , entry HistoryEntry ) {
2021-04-05 14:00:51 +02:00
stmt , err := conn . Prepare ( "INSERT INTO history (user, command, hostname, workdir, timestamp) VALUES (?, ?, ?, ?, datetime(?, 'unixepoch'))" )
2021-03-20 19:39:18 +01:00
if err != nil {
log . Panic ( err )
}
2021-04-05 14:00:51 +02:00
_ , err = stmt . Exec ( entry . user , entry . cmd , entry . hostname , entry . cwd , entry . timestamp . Unix ( ) )
2021-03-20 19:39:18 +01:00
if err != nil {
log . Panic ( err )
}
}
2021-03-20 20:18:27 +01:00
func xdgOrFallback ( xdg string , fallback string ) string {
dir := os . Getenv ( xdg )
if dir != "" {
if ok , err := exists ( dir ) ; ok && err == nil {
return dir
}
}
return fallback
}
func exists ( path string ) ( bool , error ) {
_ , err := os . Stat ( path )
if err == nil {
return true , nil
}
if os . IsNotExist ( err ) {
return false , nil
}
return false , err
}
2021-03-20 19:39:18 +01:00
2021-03-22 19:50:27 +01:00
func printUsage ( ) {
2021-04-05 11:44:07 +02:00
fmt . Fprintf ( os . Stderr , "Usage: ./hs9001 <add/search/import>\n" )
2021-03-22 19:50:27 +01:00
}
2021-03-20 19:39:18 +01:00
func main ( ) {
2021-04-05 10:16:40 +02:00
addCmd := flag . NewFlagSet ( "add" , flag . ExitOnError )
searchCmd := flag . NewFlagSet ( "search" , flag . ExitOnError )
2021-03-20 20:22:38 +01:00
2021-04-05 10:43:54 +02:00
if len ( os . Args ) < 2 {
2021-03-22 19:50:27 +01:00
printUsage ( )
2021-03-20 20:22:38 +01:00
return
}
2021-04-05 10:16:40 +02:00
cmd := os . Args [ 1 ]
globalargs := os . Args [ 2 : ]
2021-03-20 19:39:18 +01:00
2021-04-05 11:44:07 +02:00
var conn * sql . DB
ok , _ := exists ( databaseLocation ( ) )
if ! ok {
err := os . MkdirAll ( filepath . Dir ( databaseLocation ( ) ) , 0755 )
if err != nil {
log . Panic ( err )
}
conn = createConnection ( )
initDatabase ( conn )
} else {
conn = createConnection ( )
}
migrateDatabase ( conn , fetchDBVersion ( conn ) )
2021-03-20 20:47:05 +01:00
2021-04-05 10:16:40 +02:00
switch cmd {
case "add" :
var ret int
addCmd . IntVar ( & ret , "ret" , 0 , "Return value of the command to add" )
addCmd . Parse ( globalargs )
args := addCmd . Args ( )
2021-03-21 12:45:26 +01:00
if ret == 23 { // 23 is our secret do not log status code
return
}
2021-04-05 10:16:40 +02:00
if len ( args ) < 1 {
2021-03-20 19:39:18 +01:00
fmt . Fprint ( os . Stderr , "Error: You need to provide the command to be added" )
}
2021-04-05 10:16:40 +02:00
historycmd := args [ 0 ]
2021-03-20 19:39:18 +01:00
var rgx = regexp . MustCompile ( "\\s+\\d+\\s+(.*)" )
rs := rgx . FindStringSubmatch ( historycmd )
2021-03-21 10:43:06 +01:00
if len ( rs ) == 2 {
2021-04-05 13:00:27 +02:00
add ( conn , NewHistoryEntry ( rs [ 1 ] ) )
2021-03-21 10:43:06 +01:00
}
2021-04-05 13:15:20 +02:00
case "search" :
fallthrough
2021-04-05 12:49:13 +02:00
case "delete" :
2021-04-05 13:15:20 +02:00
var workDir string
2021-04-05 14:21:33 +02:00
var beginTime string
var endTime string
2021-05-16 17:13:29 +02:00
var distinct bool = true
2021-04-05 14:21:33 +02:00
2021-04-05 13:15:20 +02:00
searchCmd . StringVar ( & workDir , "workdir" , "%" , "Search only within this workdir" )
2021-04-05 14:21:33 +02:00
searchCmd . StringVar ( & beginTime , "begin" , "50 years ago" , "Start searching from this timeframe" )
searchCmd . StringVar ( & endTime , "end" , "now" , "End searching from this timeframe" )
2021-05-16 17:13:29 +02:00
searchCmd . BoolVar ( & distinct , "distinct" , true , "Remove consecutive duplicate commands from output" )
2021-04-05 14:21:33 +02:00
2021-04-05 10:16:40 +02:00
searchCmd . Parse ( globalargs )
2021-04-05 13:15:20 +02:00
args := searchCmd . Args ( )
2021-04-05 10:16:40 +02:00
2021-04-05 14:21:33 +02:00
beginTimestamp , err := naturaldate . Parse ( beginTime , time . Now ( ) )
if err != nil {
fmt . Fprintf ( os . Stderr , "Failed to convert time string: %s\n" , err . Error ( ) )
}
endTimeStamp , err := naturaldate . Parse ( endTime , time . Now ( ) )
if err != nil {
fmt . Fprintf ( os . Stderr , "Failed to convert time string: %s\n" , err . Error ( ) )
}
2021-04-05 10:16:40 +02:00
q := strings . Join ( args , " " )
2021-04-05 14:21:33 +02:00
results := search ( conn , "%" + q + "%" , workDir , beginTimestamp , endTimeStamp )
2021-04-05 12:49:13 +02:00
2021-05-16 17:13:29 +02:00
previousCmd := ""
2021-04-05 12:49:13 +02:00
for e := results . Front ( ) ; e != nil ; e = e . Next ( ) {
entry , ok := e . Value . ( * HistoryEntry )
if ! ok {
log . Panic ( "Failed to retrieve entries" )
}
2021-05-16 17:13:29 +02:00
if ! distinct || previousCmd != entry . cmd {
fmt . Printf ( "%s\n" , entry . cmd )
}
previousCmd = entry . cmd
2021-04-05 10:29:50 +02:00
}
2021-04-05 12:49:13 +02:00
if cmd == "delete" {
2021-04-05 11:44:07 +02:00
2021-04-05 12:49:13 +02:00
_ , err := conn . Exec ( "BEGIN;" )
if err != nil {
log . Panic ( err )
}
for e := results . Front ( ) ; e != nil ; e = e . Next ( ) {
entry , ok := e . Value . ( * HistoryEntry )
if ! ok {
log . Panic ( "Failed to retrieve entries" )
}
delete ( conn , entry . id )
}
_ , err = conn . Exec ( "END;" )
if err != nil {
log . Panic ( err )
}
_ , err = conn . Exec ( "VACUUM" )
if err != nil {
log . Panic ( err )
}
}
os . Exit ( 23 )
2021-04-05 10:16:40 +02:00
case "import" :
2021-03-20 20:47:05 +01:00
importFromStdin ( conn )
2021-04-05 10:16:40 +02:00
default :
fmt . Fprintf ( os . Stderr , "Error: Unknown subcommand '%s' supplied\n\n" , cmd )
2021-03-22 19:50:27 +01:00
printUsage ( )
return
2021-03-20 19:39:18 +01:00
}
2021-04-05 10:16:40 +02:00
2021-03-20 19:39:18 +01:00
}