add.go

package goss

import (
	"fmt"
	"github.com/SimonBaeumer/goss/resource"
	"github.com/SimonBaeumer/goss/system"
	"github.com/SimonBaeumer/goss/util"
	"io"
	"os"
	"reflect"
	"strconv"
	"strings"
)

type Add struct {
	Writer            io.Writer
	ExcludeAttr       []string
	Timeout           int
	AllowInsecure     bool
	NoFollowRedirects bool
	Server            string
	Username          string
	Password          string
	Header            string
	Sys               *system.System
}

// AddResources is a sSimple wrapper to add multiple resources
func (a *Add) AddResources(fileName, resourceName string, keys []string) error {
	OutStoreFormat = getStoreFormatFromFileName(fileName)
	header := extractHeaderArgument(a.Header)

	config := util.Config{
		IgnoreList:        a.ExcludeAttr,
		Timeout:           a.Timeout,
		AllowInsecure:     a.AllowInsecure,
		NoFollowRedirects: a.NoFollowRedirects,
		Server:            a.Server,
		Username:          a.Username,
		Password:          a.Password,
		Header:            header,
	}

	var gossConfig GossConfig
	if _, err := os.Stat(fileName); err == nil {
		gossConfig = ReadJSON(fileName)
	} else {
		gossConfig = *NewGossConfig()
	}

	for _, key := range keys {
		if err := a.AddResource(fileName, gossConfig, resourceName, key, config); err != nil {
			return err
		}
	}
	WriteJSON(fileName, gossConfig)

	return nil
}

func extractHeaderArgument(headerArg string) map[string][]string {
	if headerArg == "" {
		return make(map[string][]string)
	}
	rawHeaders := strings.Split(headerArg, ":")
	headers := make(map[string][]string)
	headers[rawHeaders[0]] = []string{strings.TrimSpace(rawHeaders[1])}
	return headers
}

// AddResource adds a resource to the configuration file
func (a *Add) AddResource(fileName string, gossConfig GossConfig, resourceName, key string, config util.Config) error {
	// Need to figure out a good way to refactor this
	switch resourceName {
	case "addr":
		res, err := gossConfig.Addrs.AppendSysResource(key, a.Sys, config)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		a.resourcePrint(fileName, res)
	case "command":
		res, err := gossConfig.Commands.AppendSysResource(key, a.Sys, config)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		a.resourcePrint(fileName, res)
	case "dns":
		res, err := gossConfig.DNS.AppendSysResource(key, a.Sys, config)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		a.resourcePrint(fileName, res)
	case "file":
		res, err := gossConfig.Files.AppendSysResource(key, a.Sys, config)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		a.resourcePrint(fileName, res)
	case "Group":
		res, err := gossConfig.Groups.AppendSysResource(key, a.Sys, config)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		a.resourcePrint(fileName, res)
	case "package":
		res, err := gossConfig.Packages.AppendSysResource(key, a.Sys, config)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		a.resourcePrint(fileName, res)
	case "port":
		res, err := gossConfig.Ports.AppendSysResource(key, a.Sys, config)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		a.resourcePrint(fileName, res)
	case "process":
		res, err := gossConfig.Processes.AppendSysResource(key, a.Sys, config)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		a.resourcePrint(fileName, res)
	case "service":
		res, err := gossConfig.Services.AppendSysResource(key, a.Sys, config)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		a.resourcePrint(fileName, res)
	case "user":
		res, err := gossConfig.Users.AppendSysResource(key, a.Sys, config)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		a.resourcePrint(fileName, res)
	case "gossfile":
		res, err := gossConfig.Gossfiles.AppendSysResource(key, a.Sys, config)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		a.resourcePrint(fileName, res)
	case "kernel-param":
		res, err := gossConfig.KernelParams.AppendSysResource(key, a.Sys, config)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		a.resourcePrint(fileName, res)
	case "Mount":
		res, err := gossConfig.Mounts.AppendSysResource(key, a.Sys, config)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		a.resourcePrint(fileName, res)
	case "interface":
		res, err := gossConfig.Interfaces.AppendSysResource(key, a.Sys, config)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		a.resourcePrint(fileName, res)
	case "http":
		res, err := gossConfig.HTTPs.AppendSysResource(key, a.Sys, config)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		a.resourcePrint(fileName, res)
	default:
		panic("Undefined resource name: " + resourceName)
	}

	return nil
}

// Simple wrapper to add multiple resources
func (a *Add) AutoAddResources(fileName string, keys []string) error {
	OutStoreFormat = getStoreFormatFromFileName(fileName)

	var gossConfig GossConfig
	if _, err := os.Stat(fileName); err == nil {
		gossConfig = ReadJSON(fileName)
	} else {
		gossConfig = *NewGossConfig()
	}

	for _, key := range keys {
		if err := a.AutoAddResource(fileName, gossConfig, key); err != nil {
			return err
		}
	}
	if err := WriteJSON(fileName, gossConfig); err != nil {
		return err
	}

	return nil
}

// Autoadds all resources to the config file
func (a *Add) AutoAddResource(fileName string, gossConfig GossConfig, key string) error {
	// file
	if strings.Contains(key, "/") {
		if res, _, ok := gossConfig.Files.AppendSysResourceIfExists(key, a.Sys); ok == true {
			a.resourcePrint(fileName, res)
		}
	}

	// group
	if res, _, ok := gossConfig.Groups.AppendSysResourceIfExists(key, a.Sys); ok == true {
		a.resourcePrint(fileName, res)
	}

	// package
	if res, _, ok := gossConfig.Packages.AppendSysResourceIfExists(key, a.Sys); ok == true {
		a.resourcePrint(fileName, res)
	}

	// port
	if res, _, ok := gossConfig.Ports.AppendSysResourceIfExists(key, a.Sys); ok == true {
		a.resourcePrint(fileName, res)
	}

	// process
	if res, sysres, ok := gossConfig.Processes.AppendSysResourceIfExists(key, a.Sys); ok == true {
		a.resourcePrint(fileName, res)
		ports := system.GetPorts(true)
		pids, _ := sysres.Pids()
		for _, pid := range pids {
			pidS := strconv.Itoa(pid)
			for port, entries := range ports {
				for _, entry := range entries {
					if entry.Pid == pidS {
						// port
						if res, _, ok := gossConfig.Ports.AppendSysResourceIfExists(port, a.Sys); ok == true {
							a.resourcePrint(fileName, res)
						}
					}
				}
			}
		}
	}

	// Service
	if res, _, ok := gossConfig.Services.AppendSysResourceIfExists(key, a.Sys); ok == true {
		a.resourcePrint(fileName, res)
	}

	// user
	if res, _, ok := gossConfig.Users.AppendSysResourceIfExists(key, a.Sys); ok == true {
		a.resourcePrint(fileName, res)
	}

	return nil
}

func (a *Add) resourcePrint(fileName string, res resource.ResourceRead) {
	resMap := map[string]resource.ResourceRead{res.ID(): res}

	oj, _ := marshal(resMap)
	typ := reflect.TypeOf(res)
	typs := strings.Split(typ.String(), ".")[1]

	fmt.Fprintf(a.Writer, "Adding %s to '%s':\n\n%s\n\n", typs, fileName, string(oj))
}