diff --git a/README.md b/README.md index 7f8004d..06980fd 100644 --- a/README.md +++ b/README.md @@ -227,6 +227,47 @@ trickest library search subdomain takeover [](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 [support@trickest.com](mailto:support@trickest.com). diff --git a/cmd/files/files.go b/cmd/files/files.go new file mode 100644 index 0000000..29272ab --- /dev/null +++ b/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 +} diff --git a/cmd/files/filesCreate.go b/cmd/files/filesCreate.go new file mode 100644 index 0000000..20c6be8 --- /dev/null +++ b/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 +} diff --git a/cmd/files/filesDelete.go b/cmd/files/filesDelete.go new file mode 100644 index 0000000..1274d55 --- /dev/null +++ b/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 +} diff --git a/cmd/files/filesGet.go b/cmd/files/filesGet.go new file mode 100644 index 0000000..5726a66 --- /dev/null +++ b/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 +} diff --git a/cmd/root.go b/cmd/root.go index e8dd784..1f6807f 100644 --- a/cmd/root.go +++ b/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) } diff --git a/types/files.go b/types/files.go new file mode 100644 index 0000000..3735cf6 --- /dev/null +++ b/types/files.go @@ -0,0 +1,25 @@ +package types + +import ( + "time" +) + +type Files struct { + Next string `json:"next"` + Previous string `json:"previous"` + Page int `json:"page"` + Last int `json:"last"` + Count int `json:"count"` + Results []File `json:"results"` +} + +type File struct { + ID string `json:"id"` + Name string `json:"name"` + Vault string `json:"vault"` + TweID string `json:"twe_id"` + ArtifactID string `json:"artifact_id"` + Size int `json:"size"` + PrettySize string `json:"pretty_size"` + ModifiedDate time.Time `json:"modified_date"` +} diff --git a/util/util.go b/util/util.go index c3b1ce6..cd774b4 100644 --- a/util/util.go +++ b/util/util.go @@ -3,15 +3,18 @@ package util import ( "encoding/json" "fmt" + "io" "log" "math" "net/http" "net/url" "os" + "path" "strconv" "strings" "time" + "github.com/schollz/progressbar/v3" "github.com/trickest/trickest-cli/client/request" "github.com/trickest/trickest-cli/types" @@ -457,3 +460,45 @@ func FormatDuration(duration time.Duration) string { return str } + +func DownloadFile(url, outputDir, fileName string) error { + err := os.MkdirAll(outputDir, 0755) + if err != nil { + return fmt.Errorf("couldn't create output directory (%s): %w", outputDir, err) + } + + filePath := path.Join(outputDir, fileName) + outputFile, err := os.Create(filePath) + if err != nil { + return fmt.Errorf("couldn't create output file (%s): %w", filePath, err) + } + defer outputFile.Close() + + response, err := http.Get(url) + if err != nil { + return fmt.Errorf("couldn't get URL (%s): %w", url, err) + } + defer response.Body.Close() + + if response.StatusCode != http.StatusOK { + return fmt.Errorf("unexpected HTTP status code: %s %d", url, response.StatusCode) + } + + if response.ContentLength > 0 { + bar := progressbar.NewOptions64( + response.ContentLength, + progressbar.OptionSetDescription(fmt.Sprintf("Downloading %s... ", fileName)), + progressbar.OptionSetWidth(30), + progressbar.OptionShowBytes(true), + progressbar.OptionShowCount(), + progressbar.OptionOnCompletion(func() { fmt.Println() }), + ) + _, err = io.Copy(io.MultiWriter(outputFile, bar), response.Body) + } else { + _, err = io.Copy(outputFile, response.Body) + } + if err != nil { + return fmt.Errorf("couldn't save file content to %s: %w", filePath, err) + } + return nil +}