makeshift/cmd/upload.go

package cmd

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

	"git.towk2.me/towk/makeshift/internal/format"
	makeshift "git.towk2.me/towk/makeshift/pkg"
	"git.towk2.me/towk/makeshift/pkg/client"
	"git.towk2.me/towk/makeshift/pkg/service"
	"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 multiple files with a specific path (used to set remote location)
  makeshift upload -d @kubernetes.json -p nodes/kubernetes.json
  makeshift upload -d @slurm.json -d @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 (
			host, _     = cmd.Flags().GetString("host")
			path, _     = cmd.Flags().GetString("path")
			dataArgs, _ = cmd.Flags().GetStringArray("data")

			inputData        = processFiles(dataArgs)
			useDirectoryPath = len(inputData) > 1
			c                = client.New(host)
			res              *http.Response
			query            string
			err              error
		)
		for inputPath, contents := range inputData {
			log.Info().Str("path", path).Int("size", len(contents)).Send()
			if useDirectoryPath {
				query = path + "/" + filepath.Clean(inputPath)
			} else {
				// use flag value if supplied
				if cmd.Flags().Changed("path") {
					query = path
				} else {
					query = inputPath
				}
			}

			query = fmt.Sprintf("/upload/%s", query)
			res, _, err = c.MakeRequest(client.HTTPEnvelope{
				Path:   query,
				Method: http.MethodPost,
				Body:   contents,
			})
			handleResponseError(res, host, query, err)
		}
	},
}

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

  # upload a new profile with a specific path
  makeshift upload profile -d @kubernetes.json
  makeshift upload profile -d '{"id": "custom", "data": {}}' kubernetes.json
`,
	Short: "Upload a new profile",
	Run: func(cmd *cobra.Command, args []string) {
		var (
			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
		)

		// load files from args
		for i, path := range args {
			body, err = os.ReadFile(path)
			if err != nil {
				log.Error().Err(err).
					Int("index", i).
					Str("path", path).
					Msg("failed to read profile file")
				continue
			}
			var profile *makeshift.Profile
			err = json.Unmarshal(body, &profile)
			if err != nil {
				log.Error().Err(err).
					Int("index", i).
					Str("path", path).
					Msg("failed to unmarshal profile")
			}
			profiles = append(profiles, profile)
		}

		// send each loaded profile to server
		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
			}

			query = fmt.Sprintf("/profiles/%s", profile.ID)
			res, body, err = c.MakeRequest(client.HTTPEnvelope{
				Path:   query,
				Method: http.MethodPost,
				Body:   body,
			})
			handleResponseError(res, host, query, err)
		}
	},
}

var uploadPluginsCmd = &cobra.Command{
	Use: "plugin [plugin_name]",
	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
`,
	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 (
			host, _     = cmd.Flags().GetString("host")
			dataArgs, _ = cmd.Flags().GetStringArray("data")

			plugins = processFiles(dataArgs)
			c       = client.New(host)
			res     *http.Response
			query   string
			body    []byte
			plugin  makeshift.Plugin
			err     error
		)

		// load files from args
		for i, path := range args {
			body, err = os.ReadFile(path)
			if err != nil {
				log.Error().Err(err).
					Int("index", i).
					Str("path", path).
					Msg("failed to read plugin file")
				continue
			}

			plugins[path] = body
		}

		for path, contents := range plugins {
			plugin, err = service.LoadPluginFromFile(path)
			if err != nil {
				log.Error().Err(err).
					Str("path", path).
					Msg("failed to load plugin from file")
			}

			query = fmt.Sprintf("/plugins/%s", plugin.Name())
			res, _, err = c.MakeRequest(client.HTTPEnvelope{
				Path:   query,
				Method: http.MethodPost,
				Body:   contents,
			})
			handleResponseError(res, host, query, err)
		}
	},
}

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 (add @ before the path)")

				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
}