makeshift/cmd/upload.go

package cmd

import (
	"bufio"
	"encoding/json"
	"fmt"
	"os"
	"strings"

	"git.towk2.me/towk/makeshift/internal/format"
	"github.com/rs/zerolog/log"
	"github.com/spf13/cobra"
)

var (
	inputFormat format.DataFormat = format.JSON
	dataArgs    []string
)

var uploadCmd = &cobra.Command{
	Use: "upload",
	Example: `
  # upload a single file
  makeshift upload -d @compute-base.yaml -t file

  # upload a single file with contents without specify type
  makeshift upload -d '{"name": "John Smith", "email": "john.smith@example.com"}'
  
  # upload a directory
  makeshift upload -d @setup/  -t directory

  # upload an archive (extracted and saved on server)
  makeshift upload -d @setup.tar.gz -t archive
`,
	Run: func(cmd *cobra.Command, args []string) {
		// make one request be host positional argument (restricted to 1 for now)
		var inputData []map[string]any
		temp := append(handleArgs(args), processDataArgs(dataArgs)...)
		for _, data := range temp {
			if data != nil {
				inputData = append(inputData, data)
			}
		}
	},
}

var uploadProfileCmd = &cobra.Command{
	Use: "profile",
	Example: `
  # upload a new profile
  makeshift upload profile -d @compute.json 

  # upload a new profile with a specific name (used for lookups)
  makeshift upload profile -d @kubernetes.json -n k8s
`,
	Args:  cobra.ExactArgs(1),
	Short: "Upload a new profile",
	Run: func(cmd *cobra.Command, args []string) {

		// make one request be host positional argument (restricted to 1 for now)
		var inputData []map[string]any
		temp := append(handleArgs(args), processDataArgs(dataArgs)...)
		for _, data := range temp {
			if data != nil {
				inputData = append(inputData, data)
			}
		}

	},
}

var uploadPluginCmd = &cobra.Command{
	Use: "plugin",
	Example: `
  # upload a new plugin
  makeshift upload plugin -d @slurm.so

  # upload a new plugin with a specific name (used for lookups)
  makeshift upload plugin -d @cobbler.so -n merge
`,
	Args:  cobra.ExactArgs(1),
	Short: "Upload a new plugin",
	Run: func(cmd *cobra.Command, args []string) {
		// make one request be host positional argument (restricted to 1 for now)
		var inputData []map[string]any
		temp := append(handleArgs(args), processDataArgs(dataArgs)...)
		for _, data := range temp {
			if data != nil {
				inputData = append(inputData, data)
			}
		}
	},
}

func init() {
	uploadProfileCmd.Flags().VarP(&inputFormat, "format", "F", "Set the input format for profile")

	uploadCmd.AddCommand(uploadProfileCmd, uploadPluginCmd)
	rootCmd.AddCommand(uploadCmd)
}

// processDataArgs takes a slice of strings that check for the @ symbol and loads
// the contents from the file specified in place (which replaces the path).
//
// NOTE: The purpose is to make the input arguments uniform for our request. This
// function is meant to handle data passed with the `-d/--data` flag and positional
// args from the CLI.
func processDataArgs(args []string) []map[string]any {
	// JSON representation
	type (
		JSONObject = map[string]any
		JSONArray  = []JSONObject
	)

	// load data either from file or directly from args
	var collection = make(JSONArray, len(args))
	for i, arg := range args {
		// if arg is empty string, then skip and continue
		if len(arg) > 0 {
			// determine if we're reading from file to load contents
			if strings.HasPrefix(arg, "@") {
				var (
					path     string = strings.TrimLeft(arg, "@")
					contents []byte
					data     JSONArray
					err      error
				)
				contents, err = os.ReadFile(path)
				if err != nil {
					log.Error().Err(err).Str("path", path).Msg("failed to read file")
					continue
				}

				// skip empty files
				if len(contents) == 0 {
					log.Warn().Str("path", path).Msg("file is empty")
					continue
				}

				// convert/validate input data
				data, err = parseInput(contents, format.DataFormatFromFileExt(path, inputFormat))
				if err != nil {
					log.Error().Err(err).Str("path", path).Msg("failed to validate input from file")
				}

				// add loaded data to collection of all data
				collection = append(collection, data...)
			} else {
				// input should be a valid JSON
				var (
					data  JSONArray
					input = []byte(arg)
					err   error
				)
				if !json.Valid(input) {
					log.Error().Msgf("argument %d not a valid JSON", i)
					continue
				}
				err = json.Unmarshal(input, &data)
				if err != nil {
					log.Error().Err(err).Msgf("failed to unmarshal input for argument %d", i)
				}
				return data
			}
		}
	}
	return collection
}

func handleArgs(args []string) []map[string]any {
	// JSON representation
	type (
		JSONObject = map[string]any
		JSONArray  = []JSONObject
	)
	// no file to load, so we just use the joined args (since each one is a new line)
	// and then stop
	var (
		collection JSONArray
		data       []byte
		err        error
	)

	if len(dataArgs) > 0 {
		return nil
	}
	data, err = ReadStdin()
	if err != nil {
		log.Error().Err(err).Msg("failed to read from standard input")
		return nil
	}
	if len(data) == 0 {
		log.Warn().Msg("no data found from standard input")
		return nil
	}
	fmt.Println(string(data))
	collection, err = parseInput([]byte(data), inputFormat)
	if err != nil {
		log.Error().Err(err).Msg("failed to validate input from arg")
	}
	return collection
}

func parseInput(contents []byte, dataFormat format.DataFormat) ([]map[string]any, error) {
	var (
		data []map[string]any
		err  error
	)

	// convert/validate JSON input format
	err = format.Unmarshal(contents, &data, dataFormat)
	if err != nil {
		return nil, fmt.Errorf("failed to unmarshal data: %v", err)
	}
	return data, nil
}

// ReadStdin reads all of standard input and returns the bytes. If an error
// occurs during scanning, it is returned.
func ReadStdin() ([]byte, error) {
	var b []byte
	input := bufio.NewScanner(os.Stdin)
	for input.Scan() {
		b = append(b, input.Bytes()...)
		b = append(b, byte('\n'))
		if len(b) == 0 {
			break
		}
	}
	if err := input.Err(); err != nil {
		return b, fmt.Errorf("failed to read stdin: %w", err)
	}
	return b, nil
}