Add with options (#39)

api/client.go

@@ -304,10 +304,11 @@ func (cli *Client) Download(url *url.URL, output string) (int64, error) {
 		return 0, err
 	}
 
-	resp, err := cli.Do(req)
+	resp, err := cli.httpClient.Do(req)
 	if err != nil {
 		return 0, err
 	}
+	defer resp.Body.Close() //nolint:errcheck
 
 	if resp.StatusCode == http.StatusOK {
 		w, err := os.Create(output)

cmd/flags/common.go

@@ -8,7 +8,7 @@ import (
 )
 
 func AddForceFlag(cmd *cobra.Command) {
-	cmd.Flags().BoolP("force", "f", false, "Force overwrite an existing file.")
+	cmd.Flags().BoolP("force", "f", false, "Force overwrite an existing file")
 }
 
 func AddAllFlag(cmd *cobra.Command) {

cmd/pro/file.go

@@ -45,8 +45,14 @@ var fileCmd = &cobra.Command{
 		query.Add("password", password)
 		url.RawQuery = query.Encode()
 
-		options := utils.NewDownloadOptions(client, url, filename, force)
-		err = utils.Download(options)
+		opts := utils.NewDownloadOptions(
+			utils.WithDownloadClient(client),
+			utils.WithDownloadURL(url),
+			utils.WithDownloadOutput(filename),
+			utils.WithDownloadForce(force),
+		)
+
+		err = utils.DownloadWithSpinner(opts)
 		if err != nil {
 			return err
 		}

cmd/pro/livescan/dom.go

@@ -49,8 +49,13 @@ var domCmd = &cobra.Command{
 		}
 
 		url := api.URL("/api/v1/livescan/%s/dom/%s", scannerId, scanId)
-		options := utils.NewDownloadOptions(client, url, output, force)
-		err = utils.Download(options)
+		opts := utils.NewDownloadOptions(
+			utils.WithDownloadClient(client),
+			utils.WithDownloadURL(url),
+			utils.WithDownloadOutput(output),
+			utils.WithDownloadForce(force),
+		)
+		err = utils.DownloadWithSpinner(opts)
 		if err != nil {
 			return err
 		}

cmd/scan/dom.go

@@ -4,7 +4,6 @@ import (
 	"fmt"
 
 	"github.com/spf13/cobra"
-	api "github.com/urlscan/urlscan-cli/api"
 	"github.com/urlscan/urlscan-cli/cmd/flags"
 
 	"github.com/urlscan/urlscan-cli/pkg/utils"
@@ -41,12 +40,16 @@ var domCmd = &cobra.Command{
 			return err
 		}
 
-		url := api.URL("%s", fmt.Sprintf("/dom/%s/", uuid))
 		if output == "" {
 			output = fmt.Sprintf("%s.html", uuid)
 		}
-		options := utils.NewDownloadOptions(client, url, output, force)
-		err = utils.Download(options)
+		opts := utils.NewDownloadOptions(
+			utils.WithDownloadClient(client),
+			utils.WithDownloadDOM(uuid),
+			utils.WithDownloadOutput(output),
+			utils.WithDownloadForce(force),
+		)
+		err = utils.DownloadWithSpinner(opts)
 		if err != nil {
 			return err
 		}

cmd/scan/screenshot.go

@@ -4,7 +4,6 @@ import (
 	"fmt"
 
 	"github.com/spf13/cobra"
-	api "github.com/urlscan/urlscan-cli/api"
 	"github.com/urlscan/urlscan-cli/cmd/flags"
 	"github.com/urlscan/urlscan-cli/pkg/utils"
 )
@@ -40,13 +39,16 @@ var screenshotCmd = &cobra.Command{
 			return err
 		}
 
-		url := api.URL("%s", fmt.Sprintf("/screenshots/%s.png", uuid))
 		if output == "" {
 			output = fmt.Sprintf("%s.png", uuid)
 		}
-
-		options := utils.NewDownloadOptions(client, url, output, force)
-		err = utils.Download(options)
+		opts := utils.NewDownloadOptions(
+			utils.WithDownloadClient(client),
+			utils.WithDownloadScreenshot(uuid),
+			utils.WithDownloadOutput(output),
+			utils.WithDownloadForce(force),
+		)
+		err = utils.DownloadWithSpinner(opts)
 		if err != nil {
 			return err
 		}

cmd/scan/submit.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 
 	"github.com/spf13/cobra"
+	"github.com/urlscan/urlscan-cli/cmd/flags"
 	"github.com/urlscan/urlscan-cli/pkg/utils"
 )
 
@@ -27,6 +28,16 @@ var submitCmd = &cobra.Command{
 		wait, _ := cmd.Flags().GetBool("wait")
 		maxWait, _ := cmd.Flags().GetInt("max-wait")
 
+		download, _ := cmd.Flags().GetBool("download")
+		screenshot, _ := cmd.Flags().GetBool("screenshot")
+		screenshot = screenshot || download
+		dom, _ := cmd.Flags().GetBool("dom")
+		dom = dom || download
+		force, _ := cmd.Flags().GetBool("force")
+
+		// override wait if any of with flag is set
+		wait = wait || screenshot || dom
+
 		reader := utils.StringReaderFromCmdArgs(args)
 		url, err := reader.ReadString()
 		if err != nil {
@@ -56,12 +67,45 @@ var submitCmd = &cobra.Command{
 
 		fmt.Print(waitResult.PrettyJson())
 
+		if screenshot {
+			downloadOpts := utils.NewDownloadOptions(
+				utils.WithDownloadClient(client),
+				utils.WithDownloadScreenshot(scanResult.UUID),
+				utils.WithDownloadOutput(fmt.Sprintf("%s.png", scanResult.UUID)),
+				utils.WithDownloadForce(force),
+				utils.WithDownloadSilent(true),
+			)
+			err := utils.DownloadWithSpinner(downloadOpts)
+			if err != nil {
+				return err
+			}
+		}
+
+		if dom {
+			downloadOpts := utils.NewDownloadOptions(
+				utils.WithDownloadClient(client),
+				utils.WithDownloadDOM(scanResult.UUID),
+				utils.WithDownloadOutput(fmt.Sprintf("%s.html", scanResult.UUID)),
+				utils.WithDownloadForce(force),
+				utils.WithDownloadSilent(true),
+			)
+			err := utils.DownloadWithSpinner(downloadOpts)
+			if err != nil {
+				return err
+			}
+		}
+
 		return nil
 	},
 }
 
 func init() {
 	addScanFlags(submitCmd)
+	flags.AddForceFlag(submitCmd)
+
+	submitCmd.Flags().Bool("screenshot", false, "Download only the screenshot (overrides wait)")
+	submitCmd.Flags().Bool("dom", false, "Download only the DOM contents (overrides wait)")
+	submitCmd.Flags().Bool("download", false, "Download screenshot and DOM contents (overrides wait/dom/screenshot)")
 
 	RootCmd.AddCommand(submitCmd)
 }

docs/urlscan_pro_live-scan_dom.md

@@ -16,7 +16,7 @@ urlscan pro live-scan dom [flags]
 ### Options
 
 ```
-  -f, --force               Force overwrite an existing file.
+  -f, --force               Force overwrite an existing file
   -h, --help                help for dom
   -o, --output string       Output file name (default <uuid>.html)
   -s, --scanner-id string   ID of the scanner (required)

docs/urlscan_scan_dom.md

@@ -16,7 +16,7 @@ urlscan scan dom <uuid> [flags]
 ### Options
 
 ```
-  -f, --force           Force overwrite an existing file.
+  -f, --force           Force overwrite an existing file
   -h, --help            help for dom
   -o, --output string   Output file name (default <uuid>.html)
 ```

docs/urlscan_scan_screenshot.md

@@ -16,7 +16,7 @@ urlscan scan screenshot <uuid> [flags]
 ### Options
 
 ```
-  -f, --force           Force overwrite an existing file.
+  -f, --force           Force overwrite an existing file
   -h, --help            help for screenshot
   -o, --output string   Output file name (default <uuid>.png)
 ```

docs/urlscan_scan_submit.md

@@ -18,10 +18,14 @@ urlscan scan submit <url> [flags]
 ```
   -c, --country string          Specify which country the scan should be performed from (2-Letter ISO-3166-1 alpha-2 country
   -a, --customagent string      Override User-Agent for this scan
+      --dom                     Download only the DOM contents (overrides wait)
+      --download                Download screenshot and DOM contents (overrides wait/dom/screenshot)
+  -f, --force                   Force overwrite an existing file
   -h, --help                    help for submit
   -m, --max-wait int            Maximum wait time per scan in seconds (default 60)
   -o, --overrideSafety string   If set to any value, this will disable reclassification of URLs with potential PII in them
   -r, --referer string          Override HTTP referer for this scan
+      --screenshot              Download only the screenshot (overrides wait)
   -t, --tags stringArray        User-defined tags to annotate this scan
   -v, --visibility string       One of public, unlisted, private
   -w, --wait                    Wait for the scan(s) to finish

pkg/utils/client.go

@@ -38,43 +38,91 @@ func NewAPIClient() (*APIClient, error) {
 }
 
 type DownloadOptions struct {
-	Client *APIClient
-	URL    *url.URL
-	Output string
-	Force  bool
+	client *APIClient
+	url    *url.URL
+	output string
+	force  bool
+	silent bool
 }
 
-func NewDownloadOptions(client *APIClient, url *url.URL, output string, force bool) DownloadOptions {
-	return DownloadOptions{
-		Client: client,
-		URL:    url,
-		Output: output,
-		Force:  force,
+type DownloadOption = func(*DownloadOptions)
+
+func WithDownloadClient(client *APIClient) DownloadOption {
+	return func(opts *DownloadOptions) {
+		opts.client = client
+	}
+}
+
+func WithDownloadURL(url *url.URL) DownloadOption {
+	return func(opts *DownloadOptions) {
+		opts.url = url
+	}
+}
+
+func WithDownloadOutput(output string) DownloadOption {
+	return func(opts *DownloadOptions) {
+		opts.output = output
+	}
+}
+
+func WithDownloadForce(force bool) DownloadOption {
+	return func(opts *DownloadOptions) {
+		opts.force = force
 	}
 }
 
-func Download(opt DownloadOptions) error {
-	// init and start the spinner
+func WithDownloadDOM(uuid string) DownloadOption {
+	return func(opts *DownloadOptions) {
+		opts.url = api.URL("%s", fmt.Sprintf("/dom/%s/", uuid))
+	}
+}
+
+func WithDownloadScreenshot(uuid string) DownloadOption {
+	return func(opts *DownloadOptions) {
+		opts.url = api.URL("%s", fmt.Sprintf("/screenshots/%s.png", uuid))
+	}
+}
+
+func WithDownloadSilent(silent bool) DownloadOption {
+	return func(opts *DownloadOptions) {
+		opts.silent = silent
+	}
+}
+
+func NewDownloadOptions(opts ...DownloadOption) *DownloadOptions {
+	downloadOpts := &DownloadOptions{}
+	for _, o := range opts {
+		o(downloadOpts)
+	}
+	return downloadOpts
+}
+
+func DownloadWithSpinner(opts *DownloadOptions) error {
 	s := spinner.New(spinner.CharSets[14], 100*time.Millisecond)
 	s.Start()
 
-	if !opt.Force {
-		// check if the file already exists
-		if _, err := os.Stat(opt.Output); err == nil {
+	if !opts.force {
+		if fileExists(opts.output) {
 			s.Stop()
-			return fmt.Errorf("%s already exists", opt.Output)
+			return fmt.Errorf("%s already exists", opts.output)
 		}
 	}
 
-	_, err := opt.Client.Download(opt.URL, opt.Output)
+	_, err := opts.client.Download(opts.url, opts.output)
 	if err != nil {
 		return err
 	}
 
 	// stop the spinner
 	s.Stop()
 
-	fmt.Printf("Downloaded: %s from %s\n", opt.Output, opt.URL.String())
+	msg := fmt.Sprintf("Downloaded: %s from %s\n", opts.output, opts.url.String())
+	if opts.silent {
+		// output it to stderr to make the rest of stdout clean (for piping with jq, etc.)
+		fmt.Fprint(os.Stderr, msg)
+	} else {
+		fmt.Print(msg)
+	}
 
 	return nil
 }