httpprocprobed/main.go

/*
	httpprocwatchd provides HTTP interface to a list of process' statuses
	formatted as JSON.
	Copyright (c) 2021 Alexander "Arav" Andreev <me@arav.top>

	This program 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.

    This program 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 this program.  If not, see <https://www.gnu.org/licenses/>.
*/

package main

import (
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"syscall"
	"time"
)

// ErrPgrepNotFound occurs when pgrep program is not found.
var ErrPgrepNotFound = errors.New("pgrep not found")
var ErrNotFound = errors.New("process is not on list")
var ErrIsOnList = errors.New("process is already on list")

// Configuration holds a list of process names to be tracked and a listen address.
type Configuration struct {
	ListenAddress string   `json:"listen_address"`
	Processes     []string `json:"processes"`
}

// LoadConfiguration loads configuration from a JSON file.
func LoadConfiguration(path string) (conf *Configuration, err error) {
	conf = &Configuration{}

	file, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, err
	}

	if err := json.Unmarshal(file, conf); err != nil {
		return nil, err
	}

	return conf, nil
}

func StoreConfiguration(path string, conf *Configuration) (err error) {
	config, err := json.Marshal(*conf)
	if err != nil {
		return err
	}

	err = ioutil.WriteFile(path, config, 0644)
	if err != nil {
		return err
	}

	return nil
}

// ProcessList is a map of processes' statuses.
type ProcessList map[string]bool

// NewProcessList returns a ProcessList with initialised map.
func NewProcessList() *ProcessList {
	pl := make(ProcessList)
	return &pl
}

// AddProcess appends a process to a ProcessList
func (pl *ProcessList) AddProcess(name string, isup bool) {
	(*pl)[name] = isup
}

// IsProcessUp returns true if process is up. Uses pgrep to get PID of a process
// and if a process is not working, then pgrep returns nothing.
func IsProcessUp(name string) (bool, error) {
	pgrep := exec.Command("pgrep", "-nf", name)

	out, err := pgrep.CombinedOutput()
	if err != nil {
		return false, err
	}

	if len(out) == 0 {
		return false, nil
	}

	return true, nil
}

// AreProcessesUp handles a GET /processes request and sends back status of given
// processes.
func AreProcessesUp(processes *[]string) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method == http.MethodGet {
			proclist := NewProcessList()

			for _, proc := range *processes {
				isup, _ := IsProcessUp(proc)

				proclist.AddProcess(proc, isup)
			}

			w.Header().Add("Content-Type", "application/json")
			json.NewEncoder(w).Encode(proclist)
		} else {
			w.WriteHeader(http.StatusMethodNotAllowed)
			w.Header().Add("Allow", "GET")
		}
	}
}

func AddProcessToList(process string, conf *Configuration, configPath string) error {
	for _, v := range conf.Processes {
		if v == process {
			return ErrIsOnList
		}
	}

	conf.Processes = append(conf.Processes, process)
	if err := StoreConfiguration(configPath, conf); err != nil {
		return err
	}

	return nil
}

func RemoveProcessFromList(process string, conf *Configuration, configPath string) error {
	for k, v := range conf.Processes {
		if v == process {
			copy(conf.Processes[:k], conf.Processes[k+1:])

			if err := StoreConfiguration(configPath, conf); err != nil {
				return err
			}

			return nil
		}
	}

	return ErrNotFound
}

func ListWatchedProcesses(processes *[]string) {
	lastidx := len(*processes) - 1
	for k, v := range *processes {
		if k == lastidx {
			fmt.Printf("%s\n", v)
		} else {
			fmt.Printf("%s, ", v)
		}
	}
}

func version() {
	fmt.Println("httpprocwatchd ver. 1.1")
	fmt.Println("Copyright (c) 2021 Alexander \"Arav\" Andreev <me@arav.top>")
	fmt.Println("License GPLv3+: GNU GPL version 3 or later <https://gnu.org/licenses/gpl.html>.")
	fmt.Println("This is free software, and you are welcome to change and redistribute it.")
	fmt.Println("There is NO WARRANTY, to the extent permitted by law.")
}

var oConfigPath string
var oShowVersion bool
var oListProcesses bool

var oAddProcess string
var oRemoveProcess string

func main() {
	log.SetFlags(0)

	flag.StringVar(&oConfigPath, "config", "processes.json", "path to configuration file")
	flag.StringVar(&oConfigPath, "c", "processes.json", "path to configuration file (shorthand)")
	flag.BoolVar(&oShowVersion, "version", false, "show version")
	flag.BoolVar(&oShowVersion, "v", false, "show version (shorthand)")

	flag.BoolVar(&oListProcesses, "list", false, "list watched processes")
	flag.BoolVar(&oListProcesses, "l", false, "list watched processes (shorthand)")

	flag.StringVar(&oAddProcess, "add", "", "add process to list")
	flag.StringVar(&oAddProcess, "a", "", "add process to list (shorthand)")

	flag.StringVar(&oRemoveProcess, "remove", "", "remove process from list")
	flag.StringVar(&oRemoveProcess, "r", "", "remove process from list (shorthand)")

	flag.Parse()

	if oShowVersion {
		version()
		return
	}

	if _, err := exec.LookPath("pgrep"); err != nil {
		log.Fatalln(ErrPgrepNotFound)
	}

	conf, err := LoadConfiguration(oConfigPath)
	if err != nil {
		log.Fatalf("Cannot load configuration file: %s\n", err)
	}

	if oListProcesses {
		ListWatchedProcesses(&conf.Processes)
		return
	}

	if oAddProcess != "" {
		err := AddProcessToList(oAddProcess, conf, oConfigPath)
		if err != nil {
			log.Fatalf("Cannot add process: %s\n", err)
		}
		return
	}

	if oRemoveProcess != "" {
		err := RemoveProcessFromList(oRemoveProcess, conf, oConfigPath)
		if err != nil {
			log.Fatalf("Cannot remove process: %s\n", err)
		}
		return
	}
	router := http.NewServeMux()
	router.HandleFunc("/processes", AreProcessesUp(&conf.Processes))

	srv := &http.Server{
		Addr:    conf.ListenAddress,
		Handler: router,
	}

	done := make(chan os.Signal, 1)
	signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)

	go func() {
		if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			log.Fatalf("ListenAndServe: %s\n", err)
		}
	}()

	log.Printf("httpprocwatchd is running on \"%s\".", conf.ListenAddress)

	<-done

	log.Printf("Shutting down... ")

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	if err := srv.Shutdown(ctx); err != nil {
		log.Fatalf("%s\n", err)
	}

	log.Println("Done.")
}