hs9001/main.go

package main

import (
	"bufio"
	"database/sql"
	"flag"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	_ "modernc.org/sqlite"
)

func databaseLocation() string {
	envOverride := os.Getenv("HS9001_DB_PATH")
	if envOverride != "" {
		return envOverride
	}
	return filepath.Join(xdgOrFallback("XDG_DATA_HOME", filepath.Join(os.Getenv("HOME"), ".local/share")), "hs9001/db.sqlite")
}

func createConnection() *sql.DB {

	db, err := sql.Open("sqlite", databaseLocation())
	if err != nil {
		log.Panic(err)
	}
	return db
}

func initDatabase(conn *sql.DB) {
	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)
	}
}

func importFromStdin(conn *sql.DB) {
	scanner := bufio.NewScanner(os.Stdin)

	_, err := conn.Exec("BEGIN;")
	if err != nil {
		log.Panic(err)
	}

	for scanner.Scan() {
		add(conn, scanner.Text())
	}

	_, err = conn.Exec("END;")
	if err != nil {
		log.Panic(err)
	}
}

func search(conn *sql.DB, q string) {
	queryStmt := "SELECT command FROM history WHERE command LIKE ? ORDER BY timestamp ASC"

	rows, err := conn.Query(queryStmt, "%"+q+"%")
	if err != nil {
		log.Panic(err)
	}

	for rows.Next() {
		var command string
		err = rows.Scan(&command)
		if err != nil {
			log.Panic(err)
		}
		fmt.Printf("%s\n", command)
	}
}

func add(conn *sql.DB, cmd string) {
	user := os.Getenv("USER")
	hostname, err := os.Hostname()
	if err != nil {
		log.Panic(err)
	}

	stmt, err := conn.Prepare("INSERT INTO history (user, command, hostname) VALUES (?, ?, ?)")
	if err != nil {
		log.Panic(err)
	}

	_, err = stmt.Exec(user, cmd, hostname)
	if err != nil {
		log.Panic(err)
	}

}

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
}

func printUsage() {
	fmt.Fprintf(os.Stderr, "Usage:   ./hs9001 <add/search/init/import>\n")
}

func main() {
	var ret int
	flag.IntVar(&ret, "ret", 0, "Return value of the command to add")
	flag.Parse()
	args := flag.Args()
	argslen := len(args)

	if argslen < 1 {
		printUsage()
		return
	}

	cmd := args[0]

	conn := createConnection()

	if cmd == "add" {
		if ret == 23 { // 23 is our secret do not log status code
			return
		}
		if argslen < 2 {
			fmt.Fprint(os.Stderr, "Error: You need to provide the command to be added")

		}
		historycmd := args[1]
		var rgx = regexp.MustCompile("\\s+\\d+\\s+(.*)")
		rs := rgx.FindStringSubmatch(historycmd)
		if len(rs) == 2 {
			add(conn, rs[1])
		}
	} else if cmd == "search" {
		if argslen < 2 {
			fmt.Fprint(os.Stderr, "Please provide the search query\n")
		}
		q := strings.Join(args[1:], " ")
		search(conn, q)
		os.Exit(23)
	} else if cmd == "init" {
		err := os.MkdirAll(filepath.Dir(databaseLocation()), 0755)
		if err != nil {
			log.Panic(err)
		}
		initDatabase(conn)
	} else if cmd == "import" {
		importFromStdin(conn)
	} else {
		fmt.Fprint(os.Stderr, "Error: Unknown command supplied\n\n")
		printUsage()
		return
	}
}