spacewalk/main.go

package main

import (
	"bufio"
	"bytes"
	"crypto/md5"
	"crypto/tls"
	"encoding/csv"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/url"
	"os"
	"os/user"
	"path/filepath"
	"sort"
	"strings"
	"time"
)

var swPath, fPath string
var timeout time.Duration = time.Duration(5) * time.Second

///////////////////////////////
// Startup & helpers
//////////////////////////////

func ifErrExit(e error, msg string) {
	if e != nil {
		fmt.Fprintf(os.Stderr, "%s: %s\n", msg, e.Error())
		os.Exit(1)
	}
}

func earlyExit(msg string) {
	fmt.Fprintf(os.Stderr, "Error: %s\n", msg)
	os.Exit(1)
}

// validateDataPaths makes sure that the necessary data folders exist
// for the user running spacewalk. It builds them if they are not
// present and does an error exit if there is an error.
func validateDataPaths() {
	xdgDataHome := os.Getenv("XDG_DATA_HOME")
	if xdgDataHome == "" {
		usr, _ := user.Current()
		xdgDataHome = filepath.Join(usr.HomeDir, "/.local/share")
	}
	neededFolders := filepath.Join(xdgDataHome, "spacewalk")
	neededFolders, _ = filepath.Abs(neededFolders)
	swPath = neededFolders
	neededFolders = filepath.Join(neededFolders, "flights")
	fPath = neededFolders
	err := os.MkdirAll(neededFolders, 0755)
	ifErrExit(err, "Error validation data paths")
	p := filepath.Join(swPath, "flight-manifest")
	f, err := os.OpenFile(p, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	ifErrExit(err, "Could not create flight manifest")
	f.Close()
}

func expandTilde(path string) string {
	usr, _ := user.Current()
	newpath := strings.Replace(path, "~", usr.HomeDir, -1)
	return strings.TrimSpace(newpath)
}

func displayUsage() {
	fmt.Println("spacewalk \033[3mcommand\033[23m [\033[3moptions...\033[23m]\n\nAvailable commands:\n\tcreate\n\tupdate\n\tremove\n\tdelete\n\tadd\n\tshow\n\tlaunch\n\thelp")
}

// parseArgs is the main entry point into scbm and will run the appropriate functions
// based on the arguments passed in at run time.
func parseArgs() {
	a := os.Args
	if len(a) == 1 {
		displayUsage()
		os.Exit(0)
	}

	switch a[1] {
	case "create":
		if len(a) == 3 {
			create(a[2])
		} else {
			earlyExit("Missing flight name.\n-   spacewalk create \033[3mflight\033[0m")
		}
	case "update":
		if len(a) == 5 {
			update(a[2], a[3], a[4])
		} else {
			earlyExit("Incorrect arguments.\n-   spacewalk update \033[3mflight item value\033[0m")
		}
	case "remove":
		if len(a) == 3 {
			remove(a[2])
		} else {
			earlyExit("Missing flight name.\n-   spacewalk remove \033[3mflight\033[0m")
		}
	case "add":
		if len(a) == 5 {
			add(a[2], a[3], a[4])
		} else {
			earlyExit("Incorrect syntax.\n-   spacewalk add \033[3mflight url title\033[0m")
		}
	case "delete":
		if len(a) == 4 {
			del(a[2], a[3])
		} else {
			earlyExit("Incorrect syntax.\n-   spacewalk delete \033[3mflight url|title\033[0m")
		}
	case "launch":
		if len(a) == 2 {
			launchFlights()
		} else if len(a) == 3 {
			launchFlight(a[2])
		} else {
			earlyExit("Incorrect syntax.\n-   spacewalk launch [\033[3mflight\033[0m]")
		}
	case "show":
		if len(a) == 2 {
			showFlights("")
		} else if len(a) == 3 {
			showFlights(a[2])
		} else {
			earlyExit("Incorrect syntax.\n-   spacewalk show [\033[3mflight\033[0m]")
		}
	case "help":
		if len(a) == 2 {
			displayHelp("")
		} else if len(a) == 3 {
			displayHelp(a[2])
		} else {
			earlyExit("Incorrect syntax.\n-   spacewalk help [\033[3mcommand\033[0m]")
		}
	default:
		displayUsage()
		earlyExit(fmt.Sprintf("Unknown command %q.\nTry: `spacewalk help`", a[1]))
	}
}

func getLine(prefix string) (string, error) {
	reader := bufio.NewReader(os.Stdin)
	fmt.Print(prefix)
	text, err := reader.ReadString('\n')
	if err != nil {
		return "", err
	}

	return text[:len(text)-1], nil
}

func readManifest() [][]string {
	p := filepath.Join(swPath, "flight-manifest")
	f, err := os.Open(p)
	ifErrExit(err, "Could not open flight manifest")
	defer f.Close()
	r := csv.NewReader(f)
	records, err := r.ReadAll()
	ifErrExit(err, "Could not read from flight manifest")
	return records
}

func writeFlight(flight string, records [][]string) error {
	p := filepath.Join(fPath, flight)
	f, err := os.OpenFile(p, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
	if err != nil {
		return err
	}
	defer f.Close()
	w := csv.NewWriter(f)
	w.WriteAll(records)
	if err := w.Error(); err != nil {
		return err
	}
	return nil
}

func writeManifest(data [][]string) error {
	p := filepath.Join(swPath, "flight-manifest")
	f, err := os.OpenFile(p, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
	if err != nil {
		return err
	}
	defer f.Close()
	w := csv.NewWriter(f)
	w.WriteAll(data)
	if err := w.Error(); err != nil {
		return err
	}
	return nil
}

func readFlightLog(flight string) [][]string {
	p := filepath.Join(fPath, flight)
	f, err := os.Open(p)
	ifErrExit(err, "Could not open flight log")
	defer f.Close()
	r := csv.NewReader(f)
	records, err := r.ReadAll()
	ifErrExit(err, "Could not read from flight log")
	return records
}

func findItemRow(name string, position int, csv [][]string) []string {
	for _, row := range csv {
		if len(row)-1 < position {
			earlyExit("Invalid array position passed to findItemRow")
		}
		if row[position] == name {
			return row
		}
	}
	earlyExit(fmt.Sprintf("Item %q was not found", name))
	return []string{}
}

func retrieve(addr string) (string, error) {
	addr = strings.TrimSpace(addr)
	u, err := url.Parse(addr)
	if err != nil {
		return "", fmt.Errorf("Retrieve error, parse url: %s", err.Error())
	}
	if u.Port() == "" {
		u.Host = u.Host + ":1965"
	}
	conf := &tls.Config{
		InsecureSkipVerify: true,
		MinVersion:         tls.VersionTLS12,
	}

	conn, err := tls.DialWithDialer(&net.Dialer{Timeout: timeout}, "tcp", u.Host, conf)
	if err != nil {
		return "", fmt.Errorf("TLS Dial Error: %s", err.Error())
	}

	defer conn.Close()

	send := u.String() + "\r\n"

	_, err = conn.Write([]byte(send))
	if err != nil {
		fmt.Println("Write error:")
		return "", err
	}

	result, err := ioutil.ReadAll(conn)
	if err != nil {
		fmt.Println("Read error:")
		return "", err
	}

	resp := strings.SplitN(string(result), "\r\n", 2)
	if len(resp[0]) <= 0 || resp[0][0] != '2' || len(resp) < 2 {
		fmt.Println(string(result))
		return "", fmt.Errorf("Invalid server response")
	}
	h := md5.New()
	io.WriteString(h, resp[1])

	return string(h.Sum(nil)), nil
}

func checkForUpdate(c []string, r chan []string) {
	chksum, err := retrieve(c[1])
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error updating capsule %q: %s\n", c[0], err.Error())
		r <- c
		return
	}

	if chksum == c[2] {
		fmt.Printf("%s has no changes\n", c[0])
		r <- c
		return
	}

	fmt.Printf("\033[1mUPDATED: %s\033[0m\n", c[0])
	c[2] = chksum
	currentTime := time.Now()
	c[3] = currentTime.Format("2006-01-02")
	r <- c
}

func getHeaderFooter(addr string) []byte {
	addr = strings.TrimSpace(addr)
	if addr == "none" {
		return make([]byte, 0)
	}

	content, err := ioutil.ReadFile(addr)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error reading %q\n", addr)
		return make([]byte, 0)
	}

	return content
}

///////////////////////////////
// Command Execution
//////////////////////////////

func create(flight string) {
	var lp, hp, fp string
	var err error
	records := readManifest()
	for _, r := range records {
		if r[0] == flight {
			earlyExit(fmt.Sprintf("There is already a flight with the name %q", flight))
		}
	}

	fmt.Printf("Creating flight: %q\n", flight)

	happy := false
	for !happy {
		lp, err = getLine("Enter the launch path (output path including file name): ")
		ifErrExit(err, "Error reading from stdin")
		if lp == "" {
			fmt.Println("Launch path cannot be empty")
			continue
		} else {
			lp = expandTilde(lp)
		}

		hp, err = getLine("Enter the header path, or leave blank for none: ")
		ifErrExit(err, "Error reading from stdin")
		if hp == "" {
			hp = "none"
		} else {
			hp = expandTilde(hp)
		}
		fp, err = getLine("Enter the footer path, or leave blank for none: ")
		ifErrExit(err, "Error reading from stdin")
		if fp == "" {
			fp = "none"
		} else {
			fp = expandTilde(fp)
		}
		fmt.Println("Are you happy with the following:")
		fmt.Printf("Launch path: %s\nHeader path: %s\nFooter path: %s\n", lp, hp, fp)
		yesNo, err := getLine("Type 'yes' to accept, anything else to redo: ")
		ifErrExit(err, "Error reading from stdin")

		if strings.ToLower(yesNo) == "yes" {
			happy = true
		}
	}

	p := filepath.Join(swPath, "flight-manifest")
	f, err := os.OpenFile(p, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	ifErrExit(err, "Could not open flight manifest")
	defer f.Close()
	ln := fmt.Sprintf("%s, %s, %s, %s\n", flight, lp, hp, fp)
	_, err = f.WriteString(ln)
	ifErrExit(err, "Unable to save new flight to data file")
	fpath := filepath.Join(fPath, flight)
	ff, err := os.OpenFile(fpath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	ifErrExit(err, fmt.Sprintf("Could not create flight log for %s", flight))
	ff.Close()
}

func displayHelp(command string) {
	switch command {
	case "launch":
		fmt.Println("Format: spacewalk launch [\033[3mflight\033[23m\n\nWhen run without a flight, spacewalk will launch all flights. When passed a flight it will just launch the given flight. A launch is defined as checking all of the remote pages for updates and generating a new version of the output file.")
	case "create":
		fmt.Println("Format: spacewalk create \033[3mflight\033[23m\n\nCreates a new flight with the given name. A flight can have capsules added to it, a header and footer assigned, and have it all launched to an output file.")
	case "add":
		fmt.Println("Format: spacewalk add \033[3mflight url title\033[23m\n\nAdds the given capsule url to the given flight. When the flight is launched, the given capsule will be named with the given title.")
	case "remove":
		fmt.Println("Format: spacewalk remove \033[3mflight\033[23m\n\nRemoves the given flight from the system. This is permanent and cannot be undone. Use with caution.")
	case "delete":
		fmt.Println("Format: spacewalk delete \033[3mflight title\033[23m\n\nRemoves the capsule with the given title from the given flight (remove just one capsule from a flight).")
	case "show":
		fmt.Println("Format: spacewalk show [\033[3mflight\033[23m]\n\nWhen run without a flight, will show basic information for all flights. If a flight is passed then detailed information, including a capsule list, will be shown.")
	case "update":
		fmt.Println("Format: spacewalk update \033[3mflight item value\033[23m\n\nWill update a flights manifest. Valid items:\n\tlaunch\n\theader\n\tfooter\n\nAll items take a filepath as their value.")
	case "help":
		fmt.Println("Format: spacewalk help \033[3mcommand\033[0m\n\nWill provide a brief detailed message about the given command. For more detailed information see the man page.")
	default:
		fmt.Println("Unknown command. Please run: spacewalk help \033[3mcommand\033[23m\nWhere \033[3mcommand\033[23m is any of the following:\n\tcreate\n\tlaunch\n\tadd\n\tremove\n\tdelete\n\tshow\n\tupdate\n\thelp")
	}

}

func update(flight, item, value string) {
	manifest := readManifest()
	row := -1
	for i, r := range manifest {
		if r[0] == flight {
			row = i
			break
		}
	}

	if row < 0 {
		earlyExit(fmt.Sprintf("Could not find flight \033[1m%s\033[0m", flight))
	}

	value = expandTilde(value)

	switch item {
	case "launch":
		manifest[row][1] = value
	case "header":
		manifest[row][2] = value
	case "footer":
		manifest[row][3] = value
	case "name":
		earlyExit("Changing the name of a flight is not allowed, you can remove the flight and make a new one?")
	default:
		earlyExit(fmt.Sprintf("No flight manifest item %q. Available items: launch, header, footer", item))
	}
	ifErrExit(writeManifest(manifest), "Update error")
	fmt.Printf("Successful update of %s's %s to %s\n", flight, item, value)
}

func remove(flight string) {
	manifest := readManifest()
	rowFound := false
	updated := make([][]string, 0, len(manifest)-1)
	for _, r := range manifest {
		if r[0] == flight {
			rowFound = true
		} else {
			updated = append(updated, r)
		}
	}

	if !rowFound {
		earlyExit(fmt.Sprintf("Could not find flight \033[1m%s\033[0m", flight))
	}
	ifErrExit(writeManifest(manifest), "Removal error")
	fmt.Printf("Successful removal of %s\n", flight)
}

func add(flight, url, title string) {
	p := filepath.Join(fPath, flight)
	f, err := os.OpenFile(p, os.O_APPEND|os.O_WRONLY, 0644)
	ifErrExit(err, fmt.Sprintf("Could not open flight log for %s", flight))
	defer f.Close()
	ln := fmt.Sprintf("%s, %s, %s, %s\n", title, url, "nil", "0")
	_, err = f.WriteString(ln)
	ifErrExit(err, "Unable to save new capsule to data file")
	fmt.Printf("Capsule %q added to %s's flight log\n", title, flight)
}

func del(flight, item string) {
	data := readFlightLog(flight)
	updatedData := make([][]string, 0, len(data))
	removed := false
	for _, row := range data {
		if row[0] == item {
			removed = true
		} else {
			updatedData = append(updatedData, row)
		}
	}

	if !removed {
		fmt.Printf("%q was not found in the flight log for \033[1m%s\033[0m\n", item, flight)
		return
	}

	ifErrExit(writeFlight(flight, updatedData), fmt.Sprintf("ERROR: Unable to write %s to file. Continuing with launch, but data will be out of date at next launch.", flight))
	fmt.Printf("\033[1mSuccess\033[0m, %q was removed from \033[1m%s\033[0m\n", item, flight)
}

func launchFlights() {
	fmt.Println("Pre-flight check")
	records := readManifest()
	fmt.Printf("Launching all flights (%d)\n", len(records))
	for _, fl := range records {
		launchFlight(fl[0])
	}
	fmt.Println("All flight launch procedures have been completed")
}

func launchFlight(flight string) {
	fmt.Printf("Launching %s --->\n", flight)
	manifest := findItemRow(flight, 0, readManifest())
	data := readFlightLog(flight)
	ch := make(chan []string)
	count := 0

	for _, capsule := range data {
		go checkForUpdate(capsule, ch)
		count++
	}

	updatedData := make([][]string, 0, count)

	for count > 0 {
		row := <-ch
		updatedData = append(updatedData, row)
		count--
	}

	sort.SliceStable(updatedData, func(i, j int) bool { return updatedData[i][3] > updatedData[j][3] })

	err := writeFlight(flight, updatedData)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: Unable to write %s to file. Continuing with launch, but data will be out of date at next launch. %s\n", flight, err.Error())
	}

	header := getHeaderFooter(manifest[2])
	footer := getHeaderFooter(manifest[3])

	var out bytes.Buffer
	out.Write(header)

	for _, capsule := range updatedData {
		ln := fmt.Sprintf("=> %s %s - %s\n", capsule[1], capsule[3], capsule[0])
		out.WriteString(ln)
	}

	out.Write(footer)

	f, err := os.OpenFile(strings.TrimSpace(manifest[1]), os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		fmt.Fprintf(os.Stderr, "\033[1mLAUNCH ABORTED\033[0m for %q: %s\n", flight, err.Error())
		return
	}
	defer f.Close()
	f.Write(out.Bytes())
	fmt.Printf("---> %s has been launched.\n", flight)
}

func showFlights(flight string) {
	records := readManifest()
	found := false
	for _, row := range records {
		if flight != "" && row[0] != flight {
			continue
		}
		found = true
		fmt.Printf("\033[1mFlight Name:\033[0m %s\n", row[0])
		fmt.Printf("\033[1mLaunch Path:\033[0m %s\n", row[1])
		fmt.Printf("\033[1mHeader Path:\033[0m %s\n", row[2])
		fmt.Printf("\033[1mFooter Path:\033[0m %s\n--\n", row[3])
	}

	if flight != "" {
		if !found {
			earlyExit(fmt.Sprintf("Could not find flight \033[1m%s\033[0m", flight))
		}
		fmt.Println()
		r := readFlightLog(flight)
		for _, fl := range r {
			fmt.Printf("\033[1m%s\033[0m: %s\n", fl[0], fl[1])
		}

	}
}

func main() {
	validateDataPaths()
	parseArgs()
}