Merge pull request #97 from trickest/files-cmd

Add `files` sub-command for interacting with the Trickest file storage

README.md

@@ -227,6 +227,47 @@ trickest library search subdomain takeover
 
 [<img src="./banner.png" />](https://trickest.io/auth/register)
 
+
+## Files command
+Interact with the Trickest file storage
+
+#### Get files
+Use the **get** command with the **--file** flag to retrieve one or more files
+
+```
+trickest files get --file my_file.txt --output-dir out
+```
+
+| Flag                 | Type   | Default  | Description                                                         |
+|----------------------|--------|----------|---------------------------------------------------------------------|
+| --file               | string  | /       | File or files (comma-separated)                                     |
+| --output-dir         | string  | /       | Path to directory which should be used to store files (default ".") |
+| --partial-name-match | boolean | /       | Get all files with a partial name match                             |
+
+#### Create files
+Use the **create** command with the **--file** flag to upload one or more files
+
+```
+trickest files create --file targets.txt
+```
+
+| Flag                 | Type   | Default  | Description                                                         |
+|----------------------|--------|----------|---------------------------------------------------------------------|
+| --file               | string  | /       | File or files (comma-separated)                                     |
+
+
+#### Delete files
+Use the **delete** command with the **--file** flag to delete one or more files
+
+```
+trickest files delete --file delete_me.txt
+```
+
+| Flag                 | Type   | Default  | Description                                                         |
+|----------------------|--------|----------|---------------------------------------------------------------------|
+| --file               | string  | /       | File or files (comma-separated)                                     |
+
+
 ## Report Bugs / Feedback
 We look forward to any feedback you want to share with us or if you're stuck with a problem you can contact us at [[email protected]](mailto:[email protected]).
 

cmd/files/files.go

@@ -0,0 +1,55 @@
+package files
+
+import (
+	"encoding/json"
+	"fmt"
+	"net/http"
+
+	"github.com/spf13/cobra"
+	"github.com/trickest/trickest-cli/client/request"
+	"github.com/trickest/trickest-cli/types"
+	"github.com/trickest/trickest-cli/util"
+)
+
+var (
+	Files string
+)
+
+// filesCmd represents the files command
+var FilesCmd = &cobra.Command{
+	Use:   "files",
+	Short: "Manage files in the Trickest file storage",
+	Long:  ``,
+	Run: func(cmd *cobra.Command, args []string) {
+		cmd.Help()
+	},
+}
+
+func init() {
+	FilesCmd.PersistentFlags().StringVar(&Files, "file", "", "File or files (comma-separated)")
+	FilesCmd.MarkPersistentFlagRequired("file")
+
+	FilesCmd.SetHelpFunc(func(command *cobra.Command, strings []string) {
+		_ = FilesCmd.Flags().MarkHidden("workflow")
+		_ = FilesCmd.Flags().MarkHidden("project")
+		_ = FilesCmd.Flags().MarkHidden("space")
+		_ = FilesCmd.Flags().MarkHidden("url")
+
+		command.Root().HelpFunc()(command, strings)
+	})
+}
+
+func getMetadata(searchQuery string) ([]types.File, error) {
+	resp := request.Trickest.Get().DoF("file/?search=%s&vault=%s", searchQuery, util.GetVault())
+	if resp == nil || resp.Status() != http.StatusOK {
+		return nil, fmt.Errorf("unexpected response status code: %d", resp.Status())
+	}
+	var metadata types.Files
+
+	err := json.Unmarshal(resp.Body(), &metadata)
+	if err != nil {
+		return nil, fmt.Errorf("couldn't unmarshal file IDs response: %s", err)
+	}
+
+	return metadata.Results, nil
+}

cmd/files/filesCreate.go

@@ -0,0 +1,99 @@
+package files
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"mime/multipart"
+	"net/http"
+	"os"
+	"path/filepath"
+	"strings"
+
+	"github.com/schollz/progressbar/v3"
+	"github.com/spf13/cobra"
+	"github.com/trickest/trickest-cli/util"
+)
+
+// filesCreateCmd represents the filesCreate command
+var filesCreateCmd = &cobra.Command{
+	Use:   "create",
+	Short: "Create files on the Trickest file storage",
+	Long: "Create files on the Trickest file storage.\n" +
+		"Note: If a file with the same name already exists, it will be overwritten.",
+	Run: func(cmd *cobra.Command, args []string) {
+		filePaths := strings.Split(Files, ",")
+		for _, filePath := range filePaths {
+			err := createFile(filePath)
+			if err != nil {
+				fmt.Printf("Error: %s\n", err)
+			} else {
+				fmt.Printf("Uploaded %s successfully\n", filePath)
+			}
+		}
+	},
+}
+
+func init() {
+	FilesCmd.AddCommand(filesCreateCmd)
+}
+
+func createFile(filePath string) error {
+	file, err := os.Open(filePath)
+	if err != nil {
+		return fmt.Errorf("couldn't open %s: %s", filePath, err)
+	}
+	defer file.Close()
+
+	fileName := filepath.Base(file.Name())
+
+	body := &bytes.Buffer{}
+	writer := multipart.NewWriter(body)
+	defer writer.Close()
+
+	part, err := writer.CreateFormFile("thumb", fileName)
+	if err != nil {
+		return fmt.Errorf("couldn't create form file for %s: %s", filePath, err)
+	}
+
+	fileInfo, _ := file.Stat()
+	bar := progressbar.NewOptions64(
+		fileInfo.Size(),
+		progressbar.OptionSetDescription(fmt.Sprintf("Creating %s...", fileName)),
+		progressbar.OptionSetWidth(30),
+		progressbar.OptionShowBytes(true),
+		progressbar.OptionShowCount(),
+		progressbar.OptionOnCompletion(func() { fmt.Println() }),
+	)
+
+	_, err = io.Copy(io.MultiWriter(part, bar), file)
+	if err != nil {
+		return fmt.Errorf("couldn't process %s: %s", filePath, err)
+	}
+
+	_, err = part.Write([]byte("\n--" + writer.Boundary() + "--"))
+	if err != nil {
+		return fmt.Errorf("couldn't upload %s: %s", filePath, err)
+	}
+
+	client := &http.Client{}
+	req, err := http.NewRequest("POST", util.Cfg.BaseUrl+"v1/file/", body)
+	if err != nil {
+		return fmt.Errorf("couldn't create request for %s: %s", filePath, err)
+	}
+
+	req.Header.Add("Authorization", "Token "+util.GetToken())
+	req.Header.Add("Content-Type", writer.FormDataContentType())
+
+	var resp *http.Response
+	resp, err = client.Do(req)
+	if err != nil {
+		return fmt.Errorf("couldn't upload %s: %s", filePath, err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusCreated {
+		return fmt.Errorf("unexpected status code while uploading %s: %s", filePath, resp.Status)
+	}
+	return nil
+}

cmd/files/filesDelete.go

@@ -0,0 +1,69 @@
+package files
+
+import (
+	"fmt"
+	"net/http"
+	"strings"
+
+	"github.com/spf13/cobra"
+	"github.com/trickest/trickest-cli/client/request"
+)
+
+// filesDeleteCmd represents the filesDelete command
+var filesDeleteCmd = &cobra.Command{
+	Use:   "delete",
+	Short: "Delete files from the Trickest file storage",
+	Long:  ``,
+	Run: func(cmd *cobra.Command, args []string) {
+		fileNames := strings.Split(Files, ",")
+		for _, fileName := range fileNames {
+			err := deleteFile(fileName)
+			if err != nil {
+				fmt.Printf("Error: %s\n", err)
+			} else {
+				fmt.Printf("Deleted %s successfully\n", fileName)
+			}
+		}
+	},
+}
+
+func init() {
+	FilesCmd.AddCommand(filesDeleteCmd)
+}
+
+func deleteFile(fileName string) error {
+	metadata, err := getMetadata(fileName)
+	if err != nil {
+		return fmt.Errorf("couldn't search for %s: %s", fileName, err)
+	}
+
+	if len(metadata) == 0 {
+		return fmt.Errorf("couldn't find any matches for %s", fileName)
+	}
+
+	matchFound := false
+	for _, fileMetadata := range metadata {
+		if fileMetadata.Name == fileName {
+			matchFound = true
+			err := deleteFileByID(fileMetadata.ID)
+			if err != nil {
+				return fmt.Errorf("couldn't delete %s: %s", fileMetadata.Name, err)
+			}
+		}
+	}
+
+	if !matchFound {
+		return fmt.Errorf("couldn't find any matches for %s", fileName)
+	}
+
+	return nil
+}
+
+func deleteFileByID(fileID string) error {
+	resp := request.Trickest.Delete().DoF("file/%s/", fileID)
+	if resp == nil || resp.Status() != http.StatusNoContent {
+		return fmt.Errorf("unexpected response status code: %d", resp.Status())
+	}
+
+	return nil
+}

cmd/files/filesGet.go

@@ -0,0 +1,90 @@
+package files
+
+import (
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"strings"
+
+	"github.com/spf13/cobra"
+	"github.com/trickest/trickest-cli/client/request"
+	"github.com/trickest/trickest-cli/util"
+)
+
+var (
+	outputDir        string
+	partialNameMatch bool
+)
+
+// filesGetCmd represents the filesGet command
+var filesGetCmd = &cobra.Command{
+	Use:   "get",
+	Short: "Get files from the Trickest file storage",
+	Long:  ``,
+	Run: func(cmd *cobra.Command, args []string) {
+		fileNames := strings.Split(Files, ",")
+		for _, fileName := range fileNames {
+			err := getFile(fileName, outputDir, partialNameMatch)
+			if err != nil {
+				fmt.Printf("Error: %s\n", err)
+			} else {
+				fmt.Printf("Retrieved matches for %s successfully\n", fileName)
+			}
+		}
+	},
+}
+
+func init() {
+	FilesCmd.AddCommand(filesGetCmd)
+
+	filesGetCmd.Flags().StringVar(&outputDir, "output-dir", ".", "Path to directory which should be used to store files")
+
+	filesGetCmd.Flags().BoolVar(&partialNameMatch, "partial-name-match", false, "Get all files with a partial name match")
+}
+
+func getFile(fileName string, outputDir string, partialNameMatch bool) error {
+	metadata, err := getMetadata(fileName)
+	if err != nil {
+		return fmt.Errorf("couldn't search for %s: %s", fileName, err)
+	}
+
+	if len(metadata) == 0 {
+		return fmt.Errorf("couldn't find any matches for %s", fileName)
+	}
+
+	matchFound := false
+	for _, fileMetadata := range metadata {
+		if partialNameMatch || fileMetadata.Name == fileName {
+			matchFound = true
+			signedURL, err := getSignedURLs(fileMetadata.ID)
+			if err != nil {
+				return fmt.Errorf("couldn't get a signed URL for %s: %s", fileMetadata.Name, err)
+			}
+
+			err = util.DownloadFile(signedURL, outputDir, fileMetadata.Name)
+			if err != nil {
+				return fmt.Errorf("couldn't download %s: %s", fileMetadata.Name, err)
+			}
+		}
+	}
+
+	if !matchFound {
+		return fmt.Errorf("couldn't find any matches for %s", fileName)
+	}
+	return nil
+}
+
+func getSignedURLs(fileID string) (string, error) {
+	resp := request.Trickest.Get().DoF("file/%s/signed_url/", fileID)
+	if resp == nil || resp.Status() != http.StatusOK {
+		return "", fmt.Errorf("unexpected response status code: %d", resp.Status())
+	}
+	var signedURL string
+
+	err := json.Unmarshal(resp.Body(), &signedURL)
+	if err != nil {
+		return "", fmt.Errorf("couldn't unmarshal signedURL response: %s", err)
+	}
+
+	return signedURL, nil
+}

cmd/root.go

@@ -6,6 +6,7 @@ import (
 	"github.com/trickest/trickest-cli/cmd/create"
 	"github.com/trickest/trickest-cli/cmd/delete"
 	"github.com/trickest/trickest-cli/cmd/execute"
+	"github.com/trickest/trickest-cli/cmd/files"
 	"github.com/trickest/trickest-cli/cmd/get"
 	"github.com/trickest/trickest-cli/cmd/library"
 	"github.com/trickest/trickest-cli/cmd/list"
@@ -50,6 +51,7 @@ func init() {
 	RootCmd.AddCommand(output.OutputCmd)
 	RootCmd.AddCommand(execute.ExecuteCmd)
 	RootCmd.AddCommand(get.GetCmd)
+	RootCmd.AddCommand(files.FilesCmd)
 	// RootCmd.AddCommand(export.ExportCmd)
 }