diff --git a/cmd/flags/common.go b/cmd/flags/common.go index 3278611..9a20551 100644 --- a/cmd/flags/common.go +++ b/cmd/flags/common.go @@ -77,3 +77,7 @@ func AddDirectoryPrefixFlag(cmd *cobra.Command) { func AddRefangFlag(cmd *cobra.Command) { cmd.Flags().Bool("refang", false, "Refang an input (convert '[.]' back to '.' and so on)") } + +func AddDatasourceFlag(cmd *cobra.Command) { + cmd.Flags().StringP("datasource", "D", "scans", "Datasources to search: scans (urlscan.io), hostnames, incidents, notifications, certificates (urlscan Pro)") +} diff --git a/cmd/search/count.go b/cmd/search/count.go index 62bc77e..c7110f2 100644 --- a/cmd/search/count.go +++ b/cmd/search/count.go @@ -6,6 +6,7 @@ import ( "github.com/spf13/cobra" "github.com/urlscan/urlscan-cli/api" + "github.com/urlscan/urlscan-cli/cmd/flags" "github.com/urlscan/urlscan-cli/pkg/utils" ) @@ -24,6 +25,8 @@ var countCmd = &cobra.Command{ return cmd.Usage() } + datasource, _ := cmd.Flags().GetString("datasource") + reader := utils.StringReaderFromCmdArgs(args) q, err := reader.ReadString() if err != nil { @@ -35,7 +38,7 @@ var countCmd = &cobra.Command{ return err } - it, err := client.Search(q, api.IteratorSize(0)) + it, err := client.Search(q, api.IteratorSize(0), api.IteratorDatasource(datasource)) if err != nil { return err } @@ -53,5 +56,6 @@ var countCmd = &cobra.Command{ } func init() { + flags.AddDatasourceFlag(countCmd) RootCmd.AddCommand(countCmd) } diff --git a/cmd/search/root.go b/cmd/search/root.go index 0ef5815..5c9b560 100644 --- a/cmd/search/root.go +++ b/cmd/search/root.go @@ -95,7 +95,7 @@ func init() { flags.AddSizeFlag(RootCmd, 100) // non-pro user's max size is 100 flags.AddLimitFlag(RootCmd) flags.AddAllFlag(RootCmd) + flags.AddDatasourceFlag(RootCmd) RootCmd.Flags().String("search-after", "", "For retrieving the next batch of results, value of the sort attribute of the last (oldest) result you received (comma-separated)") - RootCmd.Flags().StringP("datasource", "D", "scans", "Datasources to search: scans (urlscan.io), hostnames, incidents, notifications, certificates (urlscan Pro)") RootCmd.Flags().StringP("collapse", "c", "", "Field to collapse results on") } diff --git a/docs/urlscan_search_count.md b/docs/urlscan_search_count.md index 85b1279..7ec2efb 100644 --- a/docs/urlscan_search_count.md +++ b/docs/urlscan_search_count.md @@ -16,7 +16,8 @@ urlscan search count [flags] ### Options ``` - -h, --help help for count + -D, --datasource string Datasources to search: scans (urlscan.io), hostnames, incidents, notifications, certificates (urlscan Pro) (default "scans") + -h, --help help for count ``` ### SEE ALSO