makeshift/cmd/upload.go

package cmd

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

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

var (
	inputFormat format.DataFormat = format.JSON
)

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

  # upload an archive (extracted and saved on server)
  makeshift upload -d @setup.tar.gz -t archive

  # upload a new profile with a specific path (used to set remote location)
  makeshift upload profile -d @kubernetes.json -p nodes/kubernetes.json
  makeshift upload profile -d @slurm.json -@compute.json -p nodes
`,
	Short: "Upload files and directories",
	PersistentPreRun: func(cmd *cobra.Command, args []string) {
		setenv(cmd, "host", "MAKESHIFT_HOST")
		setenv(cmd, "path", "MAKESHIFT_PATH")
	},
	Run: func(cmd *cobra.Command, args []string) {
		var (
			dataArgs, _ = cmd.Flags().GetStringArray("data")

			inputData        = processFiles(dataArgs)
			useDirectoryPath = len(inputData) > 1
		)
		for path, contents := range inputData {
			log.Info().Str("path", path).Int("size", len(contents)).Send()
			if useDirectoryPath {

			} else {

			}
		}
	},
}

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

  # upload a new profile with a specific path (used for lookup)
  makeshift upload profile -d @kubernetes.json -n k8s
`,
	Args:  cobra.NoArgs,
	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 = append(handleArgs(args), processDataArgs(dataArgs)...)
			host, _     = cmd.Flags().GetString("host")
			dataArgs, _ = cmd.Flags().GetStringArray("data")
			profiles    = processProfiles(dataArgs)

			c     = client.New(host)
			res   *http.Response
			query string
			body  []byte
			err   error
		)

		for _, profile := range profiles {
			if profile == nil {
				continue
			}

			body, err = json.Marshal(profile)
			if err != nil {
				log.Error().Err(err).Msg("failed to marshal profile")
				continue
			}

			// send data to server
			query = fmt.Sprintf("/profiles/%s", profile.ID)
			res, body, err = c.MakeRequest(client.HTTPEnvelope{
				Path:   query,
				Method: http.MethodPost,
				Body:   body,
			})
			if err != nil {
				log.Error().Err(err).
					Str("host", host).
					Str("query", query).
					Msg("failed to make request")
				os.Exit(1)
			}
			if res.StatusCode != http.StatusOK {
				log.Error().
					Any("status", map[string]any{
						"code":    res.StatusCode,
						"message": res.Status,
					}).
					Str("host", host).
					Msg("response returned bad status")
				os.Exit(1)
			}
		}
	},
}

var uploadPluginsCmd = &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)
		// temp := append(handleArgs(args), processDataArgs(dataArgs)...)
		var (
			inputData   []*makeshift.Profile
			dataArgs, _ = cmd.PersistentFlags().GetStringArray("data")
		)
		temp := processProfiles(dataArgs)
		for _, data := range temp {
			if data != nil {
				inputData = append(inputData, data)
			}
		}
	},
}

func init() {
	uploadCmd.PersistentFlags().String("host", "http://localhost:5050", "Set the makeshift remote host (can be set with MAKESHIFT_HOST)")
	uploadCmd.PersistentFlags().StringArrayP("data", "d", []string{}, "Set the data to send to specified host (prepend @ for files)")
	uploadCmd.Flags().StringP("path", "p", ".", "Set the path to list files (can be set with MAKESHIFT_PATH)")

	uploadProfilesCmd.Flags().VarP(&inputFormat, "format", "F", "Set the input format for profile")

	uploadCmd.AddCommand(uploadProfilesCmd, uploadPluginsCmd)
	rootCmd.AddCommand(uploadCmd)
}

func processFiles(args []string) map[string][]byte {
	// load data either from file or directly from args
	var collection = make(map[string][]byte, len(args))
	for _, 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
					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
				}

				// add loaded data to collection of all data
				collection[path] = contents
			} else {
				log.Warn().Msg("only files can be uploaded")

				continue
			}
		}
	}
	return collection
}

// processProfiles 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 processProfiles(args []string) []*makeshift.Profile {
	// load data either from file or directly from args
	var collection = make([]*makeshift.Profile, 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     *makeshift.Profile
					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 = parseProfile(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  *makeshift.Profile
					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 []*makeshift.Profile{data}
			}
		}
	}
	return collection
}

func parseProfile(contents []byte, dataFormat format.DataFormat) (*makeshift.Profile, error) {
	var (
		data *makeshift.Profile
		err  error
	)

	// convert/validate JSON input format
	err = format.Unmarshal(contents, &data, dataFormat)
	if err != nil {
		return nil, fmt.Errorf("failed to unmarshal profile: %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
}