datagraph/script.go

package main

import (
	// External
	werr "git.gibonuddevalla.se/go/wrappederror"
	"github.com/jmoiron/sqlx"

	// Standard
	"crypto/md5"
	"database/sql"
	"encoding/hex"
	"encoding/json"
	"strings"
	"time"
)

type Script struct {
	ID      int
	Group   string
	Name    string
	Source  string
	Updated time.Time
}

type Hook struct {
	ID     int
	Node   Node
	Script Script
	SSH    string
	Env    map[string]string
}

func GetScript(scriptID int) (script Script, err error) { // {{{
	row := db.QueryRowx(`SELECT * FROM script WHERE id=$1`, scriptID)

	err = row.StructScan(&script)
	if err != nil {
		err = werr.Wrap(err)
		return
	}

	return
} // }}}
func GetScripts() (scripts []Script, err error) { // {{{
	scripts = []Script{}

	var rows *sqlx.Rows
	rows, err = db.Queryx(`
		SELECT *
		FROM script
		ORDER BY
			"group" ASC,
		name ASC
	`)
	if err != nil {
		err = werr.Wrap(err)
		return
	}
	defer rows.Close()

	for rows.Next() {
		var script Script
		err = rows.StructScan(&script)
		if err != nil {
			err = werr.Wrap(err)
			return
		}
		scripts = append(scripts, script)
	}

	return
} // }}}
func UpdateScript(scriptID int, data []byte) (script Script, err error) { // {{{
	err = json.Unmarshal(data, &script)
	if err != nil {
		err = werr.Wrap(err)
		return
	}

	script.ID = scriptID
	script.Group = strings.TrimSpace(script.Group)
	script.Name = strings.TrimSpace(script.Name)

	if script.Group == "" || script.Name == "" {
		err = werr.New("Group and name must be provided.")
		return
	}

	if script.ID < 1 {
		row := db.QueryRowx(`
			INSERT INTO script("group", "name", "source")
			VALUES($1, $2, $3)
			RETURNING
				id
			`,
			strings.TrimSpace(script.Group),
			strings.TrimSpace(script.Name),
			script.Source,
		)
		err = row.Scan(&script.ID)
	} else {
		_, err = db.Exec(`
			UPDATE script
			SET
				"group" = $2,
				"name" = $3,
				"source" = $4
			WHERE
				id = $1
			`,
			scriptID,
			strings.TrimSpace(script.Group),
			strings.TrimSpace(script.Name),
			script.Source,
		)
	}

	if err != nil {
		err = werr.Wrap(err)
		return
	}

	return
} // }}}
func DeleteScript(scriptID int) (err error) { // {{{
	_, err = db.Exec(`DELETE FROM script WHERE id = $1`, scriptID)
	if err != nil {
		err = werr.Wrap(err)
		return
	}
	return
} // }}}
func SearchScripts(search string) (scripts []Script, err error) { // {{{
	scripts = []Script{}

	row := db.QueryRow(`
		SELECT
			json_agg(script) AS scripts
		FROM public.script
		WHERE
			name ILIKE $1
		ORDER BY
			"group" ASC,
			name ASC
		`,
		search,
	)

	var jsonBody []byte
	err = row.Scan(&jsonBody)
	if err != nil {
		err = werr.Wrap(err)
		return
	}

	err = json.Unmarshal(jsonBody, &scripts)
	if err != nil {
		err = werr.Wrap(err)
		return
	}

	return
} // }}}
func HookScript(nodeID, scriptID int) (err error) { // {{{
	_, err = db.Exec(`INSERT INTO hook(node_id, script_id, ssh) VALUES($1, $2, '<host>')`, nodeID, scriptID)
	return
} // }}}

func GetHook(hookID int) (hook Hook, err error) { // {{{
	row := db.QueryRow(`
		SELECT
			to_json(res)
		FROM (
			SELECT
			   h.id,
			   h.ssh,
			   h.env,
			   (SELECT to_json(node) FROM node WHERE id = h.node_id) AS node,
			   (SELECT to_json(script) FROM script WHERE id = h.script_id) AS script
			FROM hook h
			WHERE
				h.id = $1
		) res
		`,
		hookID,
	)
	var data []byte
	if err = row.Scan(&data); err != nil {
		err = werr.Wrap(err)
	}

	err = json.Unmarshal(data, &hook)
	if err != nil {
		err = werr.Wrap(err)
	}

	return
} // }}}
func UpdateHook(hook Hook) (err error) { // {{{
	j, _ := json.Marshal(hook.Env)
	_, err = db.Exec(`UPDATE hook SET ssh=$2, env=$3 WHERE id=$1`, hook.ID, strings.TrimSpace(hook.SSH), j)
	if err != nil {
		err = werr.Wrap(err)
		return
	}
	return
} // }}}
func DeleteHook(hookID int) (err error) { // {{{
	_, err = db.Exec(`DELETE FROM hook WHERE id=$1`, hookID)
	if err != nil {
		err = werr.Wrap(err)
		return
	}
	return
} // }}}
func ScheduleHook(hookID int) (err error) { // {{{
	/* Script source is needed to preserve the full execution
	   against changing of the script at a later date.
	   To not waste db disk, the md5sum of the script is
	   calculated and the changed version of the script is just
	   stored once.
	*/
	hook, err := GetHook(hookID)
	if err != nil {
		err = werr.Wrap(err)
		return
	}

	scriptLogID, err := ScriptPreservedID(hook.Script.Name, hook.Script.Source)
	if err != nil {
		err = werr.Wrap(err)
		return
	}

	// The node tree data is retrieved and sent as input to the script.
	var topNode *Node
	topNode, err = GetNodeTree(hook.Node.ID, 8, true)
	if err != nil {
		err = werr.Wrap(err)
		return
	}

	nodeData, err := json.Marshal(topNode)
	if err != nil {
		err = werr.Wrap(err)
		return
	}

	j, _ := json.Marshal(hook.Env)
	_, err = db.Exec(`INSERT INTO execution(script_log_id, data, ssh, env) VALUES($1, $2, $3, $4)`, scriptLogID, nodeData, hook.SSH, j)
	if err != nil {
		err = werr.Wrap(err)
		return
	}

	go func() {
		scriptScheduler.EventQueue <- "SCRIPT_SCHEDULED"
	}()

	return
} // }}}

func ScriptPreservedID(name, source string) (id int, err error) { // {{{
	sum := md5.Sum([]byte(source))
	md5sum := hex.EncodeToString(sum[:])

	row := db.QueryRow(`SELECT id FROM script_log WHERE md5sum=$1`, md5sum)
	if err = row.Scan(&id); err == nil {
		return
	}

	if err != sql.ErrNoRows {
		err = werr.Wrap(err)
		return
	}

	row = db.QueryRow(`INSERT INTO script_log(md5sum, name, source) VALUES($1, $2, $3) RETURNING id`, md5sum, name, source)
	err = row.Scan(&id)
	if err != nil {
		err = werr.Wrap(err)
	}
	return
} // }}}