Misc fixes to cloud target logic + use int for IDs

v2/internal/runner/cloud.go

@@ -3,6 +3,7 @@ package runner
 import (
 	"io/fs"
 	"path/filepath"
+	"strconv"
 	"strings"
 	"time"
 
@@ -69,7 +70,7 @@ func (r *Runner) listDatasources() error {
 		return err
 	}
 	for _, source := range datasources {
-		gologger.Silent().Msgf("[%s] [%s] [%s] [%s] %s", source.Updatedat.Format(DDMMYYYYhhmmss), source.ID, source.Type, source.Repo, source.Path)
+		gologger.Silent().Msgf("[%s] [%d] [%s] [%s] %s", source.Updatedat.Format(DDMMYYYYhhmmss), source.ID, source.Type, source.Repo, source.Path)
 	}
 	return err
 }
@@ -97,7 +98,9 @@ func (r *Runner) listTemplates() error {
 }
 
 func (r *Runner) removeDatasource(datasource string) error {
-	err := r.cloudClient.RemoveDatasource(datasource)
+	ID, _ := strconv.ParseInt(datasource, 10, 64)
+
+	err := r.cloudClient.RemoveDatasource(ID)
 	if err != nil {
 		gologger.Error().Msgf("Error in deleting datasource %s: %s", datasource, err)
 	} else {
@@ -170,52 +173,47 @@ func (r *Runner) removeTemplate(item string) error {
 }
 
 // initializeCloudDataSources initializes cloud data sources
-func (r *Runner) initializeCloudDataSources() ([]string, error) {
-	var ids []string
-
+func (r *Runner) initializeCloudDataSources() error {
 	if r.options.AwsBucketName != "" {
 		token := strings.Join([]string{r.options.AwsAccessKey, r.options.AwsSecretKey, r.options.AwsRegion}, ":")
-		if ID, err := r.processDataSourceItem(r.options.AwsBucketName, token, "s3"); err != nil {
-			return nil, err
-		} else {
-			ids = append(ids, ID)
+		if _, err := r.processDataSourceItem(r.options.AwsBucketName, token, "s3"); err != nil {
+			return err
 		}
 	}
 	for _, repo := range r.options.GithubTemplateRepo {
-		if ID, err := r.processDataSourceItem(repo, r.options.GithubToken, "github"); err != nil {
-			return nil, err
-		} else {
-			ids = append(ids, ID)
+		if _, err := r.processDataSourceItem(repo, r.options.GithubToken, "github"); err != nil {
+			return err
 		}
 	}
-	return ids, nil
+	return nil
 }
 
-func (r *Runner) processDataSourceItem(repo, token, Type string) (string, error) {
+func (r *Runner) processDataSourceItem(repo, token, Type string) (int64, error) {
 	var secret string
 	ID, err := r.cloudClient.StatusDataSource(nucleicloud.StatusDataSourceRequest{Repo: repo, Token: token})
 	if err != nil {
 		if !strings.Contains(err.Error(), "no rows in result set") {
-			return "", errors.Wrap(err, "could not get data source status")
+			return 0, errors.Wrap(err, "could not get data source status")
 		}
 
 		gologger.Info().Msgf("Adding new data source + syncing: %s\n", repo)
-		ID, secret, err = r.cloudClient.AddDataSource(nucleicloud.AddDataSourceRequest{Type: Type, Repo: repo, Token: token})
+		resp, err := r.cloudClient.AddDataSource(nucleicloud.AddDataSourceRequest{Type: Type, Repo: repo, Token: token})
 		if err != nil {
-			return "", errors.Wrap(err, "could not add data source")
+			return 0, errors.Wrap(err, "could not add data source")
 		}
-		if err = r.cloudClient.SyncDataSource(ID); err != nil {
-			return "", errors.Wrap(err, "could not sync data source")
+		ID = resp.ID
+		if err = r.cloudClient.SyncDataSource(resp.ID); err != nil {
+			return 0, errors.Wrap(err, "could not sync data source")
 		}
 		if secret != "" {
-			gologger.Info().Msgf("Webhook URL for added source: %s/datasources/%s/webhook", r.options.CloudURL, ID)
+			gologger.Info().Msgf("Webhook URL for added source: %s/datasources/%s/webhook", r.options.CloudURL, resp.Hash)
 			gologger.Info().Msgf("Secret for webhook: %s", secret)
 		}
 	}
 	if r.options.UpdateTemplates {
 		gologger.Info().Msgf("Syncing data source: %s (%s)\n", repo, ID)
 		if err = r.cloudClient.SyncDataSource(ID); err != nil {
-			return "", errors.Wrap(err, "could not sync data source")
+			return 0, errors.Wrap(err, "could not sync data source")
 		}
 	}
 	gologger.Info().Msgf("Got connected data source: %s\n", ID)

v2/internal/runner/nucleicloud/cloud.go

@@ -159,59 +159,57 @@ func (c *Client) DeleteScan(id string) (DeleteScanResults, error) {
 }
 
 // StatusDataSource returns the status for a data source
-func (c *Client) StatusDataSource(statusRequest StatusDataSourceRequest) (string, error) {
+func (c *Client) StatusDataSource(statusRequest StatusDataSourceRequest) (int64, error) {
 	var buf bytes.Buffer
 	if err := jsoniter.NewEncoder(&buf).Encode(statusRequest); err != nil {
-		return "", errors.Wrap(err, "could not encode request")
+		return 0, errors.Wrap(err, "could not encode request")
 	}
 	httpReq, err := retryablehttp.NewRequest(http.MethodPost, fmt.Sprintf("%s/datasources/status", c.baseURL), bytes.NewReader(buf.Bytes()))
 	if err != nil {
-		return "", errors.Wrap(err, "could not make request")
+		return 0, errors.Wrap(err, "could not make request")
 	}
 
 	resp, err := c.sendRequest(httpReq)
 	if err != nil {
-		return "", errors.Wrap(err, "could not do request")
+		return 0, errors.Wrap(err, "could not do request")
 	}
 	defer resp.Body.Close()
 
 	var data map[string]interface{}
 	if err := jsoniter.NewDecoder(resp.Body).Decode(&data); err != nil {
-		return "", errors.Wrap(err, "could not decode resp")
+		return 0, errors.Wrap(err, "could not decode resp")
 	}
-	id := data["id"].(string)
+	id := data["id"].(int64)
 	return id, nil
 }
 
 // AddDataSource adds a new data source
-func (c *Client) AddDataSource(req AddDataSourceRequest) (string, string, error) {
+func (c *Client) AddDataSource(req AddDataSourceRequest) (*AddDataSourceResponse, error) {
 	var buf bytes.Buffer
 	if err := jsoniter.NewEncoder(&buf).Encode(req); err != nil {
-		return "", "", errors.Wrap(err, "could not encode request")
+		return nil, errors.Wrap(err, "could not encode request")
 	}
 	httpReq, err := retryablehttp.NewRequest(http.MethodPost, fmt.Sprintf("%s/datasources", c.baseURL), bytes.NewReader(buf.Bytes()))
 	if err != nil {
-		return "", "", errors.Wrap(err, "could not make request")
+		return nil, errors.Wrap(err, "could not make request")
 	}
 	resp, err := c.sendRequest(httpReq)
 	if err != nil {
-		return "", "", errors.Wrap(err, "could not do request")
+		return nil, errors.Wrap(err, "could not do request")
 	}
 	defer resp.Body.Close()
 
-	var data map[string]interface{}
+	var data AddDataSourceResponse
 	if err := jsoniter.NewDecoder(resp.Body).Decode(&data); err != nil {
-		return "", "", errors.Wrap(err, "could not decode resp")
+		return nil, errors.Wrap(err, "could not decode resp")
 	}
-	id := data["id"].(string)
-	secret, _ := data["secret"].(string)
-	return id, secret, nil
+	return &data, nil
 }
 
 // SyncDataSource syncs contents for a data source. The call blocks until
 // update is completed.
-func (c *Client) SyncDataSource(ID string) error {
-	httpReq, err := retryablehttp.NewRequest(http.MethodGet, fmt.Sprintf("%s/datasources/%s/sync", c.baseURL, ID), nil)
+func (c *Client) SyncDataSource(ID int64) error {
+	httpReq, err := retryablehttp.NewRequest(http.MethodGet, fmt.Sprintf("%s/datasources/%d/sync", c.baseURL, ID), nil)
 	if err != nil {
 		return errors.Wrap(err, "could not make request")
 	}
@@ -317,8 +315,8 @@ func (c *Client) ListTemplates(query string) ([]GetTemplatesResponse, error) {
 	return items, nil
 }
 
-func (c *Client) RemoveDatasource(datasource string) error {
-	httpReq, err := retryablehttp.NewRequest(http.MethodDelete, fmt.Sprintf("%s/datasources/%s", c.baseURL, datasource), nil)
+func (c *Client) RemoveDatasource(datasource int64) error {
+	httpReq, err := retryablehttp.NewRequest(http.MethodDelete, fmt.Sprintf("%s/datasources/%d", c.baseURL, datasource), nil)
 	if err != nil {
 		return errors.Wrap(err, "could not make request")
 	}

v2/internal/runner/nucleicloud/types.go

@@ -38,6 +38,13 @@ type GetScanRequest struct {
 	Matches    int64     `json:"matches"`
 }
 
+// AddDataSourceResponse is a add data source response item.
+type AddDataSourceResponse struct {
+	ID     int64  `json:"id"`
+	Hash   string `json:"hash"`
+	Secret string `json:"secret,omitempty"`
+}
+
 type GetResultsResponseItem struct {
 	ID  int64  `json:"id"`
 	Raw string `json:"raw"`
@@ -65,14 +72,14 @@ type AddDataSourceRequest struct {
 // ExistsDataSourceItemRequest is a request to identify whether a data
 // source item exists.
 type ExistsDataSourceItemRequest struct {
-	ID       string `json:"-"`
 	Type     string `json:"type"`
 	Contents string `json:"contents"`
 }
 
 // GetDataSourceResponse is response for a get data source request
 type GetDataSourceResponse struct {
-	ID        string    `json:"id"`
+	ID        int64     `json:"id"`
+	Hash      string    `json:"hash"`
 	Type      string    `json:"type"`
 	Path      string    `json:"path"`
 	Repo      string    `json:"repo"`

v2/internal/runner/runner.go

@@ -77,6 +77,7 @@ type Runner struct {
 	pprofServer       *http.Server
 	customTemplates   []customtemplates.Provider
 	cloudClient       *nucleicloud.Client
+	cloudTargets      []string
 }
 
 const pprofServerAddress = "127.0.0.1:8086"
@@ -179,7 +180,16 @@ func New(options *types.Options) (*Runner, error) {
 	}
 
 	// Initialize the input source
-	hmapInput, err := hybrid.New(options)
+	hmapInput, err := hybrid.New(&hybrid.Options{
+		Options: options,
+		NotFoundCallback: func(target string) bool {
+			if err := runner.cloudClient.ExistsDataSourceItem(nucleicloud.ExistsDataSourceItemRequest{Contents: target, Type: "targets"}); err == nil {
+				runner.cloudTargets = append(runner.cloudTargets, target)
+				return true
+			}
+			return false
+		},
+	})
 	if err != nil {
 		return nil, errors.Wrap(err, "could not create input provider")
 	}
@@ -403,33 +413,20 @@ func (r *Runner) RunEnumeration() error {
 	}
 
 	var cloudTemplates []string
-	var cloudTargets []string
 	// Initialize cloud data stores if specified
 	if r.options.Cloud {
-		ids, err := r.initializeCloudDataSources()
-		if err != nil {
-			return err
+		if err := r.initializeCloudDataSources(); err != nil {
+			return errors.Wrap(err, "could not init cloud data sources")
 		}
 
 		// hook template loading
-		store.NotFoundCallback = func(template string) {
-			for _, id := range ids {
-				if err := r.cloudClient.ExistsDataSourceItem(nucleicloud.ExistsDataSourceItemRequest{ID: id, Type: "templates", Contents: template}); err == nil {
-					cloudTemplates = append(cloudTemplates, template)
-					break
-				}
+		store.NotFoundCallback = func(template string) bool {
+			if err := r.cloudClient.ExistsDataSourceItem(nucleicloud.ExistsDataSourceItemRequest{Type: "templates", Contents: template}); err == nil {
+				cloudTemplates = append(cloudTemplates, template)
+				return true
 			}
+			return false
 		}
-		// identify cloud targets
-		r.hmapInputProvider.Scan(func(value *contextargs.MetaInput) bool {
-			for _, id := range ids {
-				if err := r.cloudClient.ExistsDataSourceItem(nucleicloud.ExistsDataSourceItemRequest{ID: id, Contents: value.Input, Type: "targets"}); err == nil {
-					cloudTargets = append(cloudTargets, value.Input)
-					break
-				}
-			}
-			return true
-		})
 	}
 	if r.options.Validate {
 		if err := store.ValidateTemplates(); err != nil {
@@ -496,7 +493,7 @@ func (r *Runner) RunEnumeration() error {
 			err = r.removeTemplate(r.options.RemoveTemplate)
 		} else {
 			gologger.Info().Msgf("Running scan on cloud with URL %s", r.options.CloudURL)
-			results, err = r.runCloudEnumeration(store, cloudTemplates, cloudTargets, r.options.NoStore, r.options.OutputLimit)
+			results, err = r.runCloudEnumeration(store, cloudTemplates, r.cloudTargets, r.options.NoStore, r.options.OutputLimit)
 			enumeration = true
 		}
 	} else {

v2/pkg/catalog/loader/loader.go

@@ -61,7 +61,7 @@ type Store struct {
 
 	// NotFoundCallback is called for each not found template
 	// This overrides error handling for not found templatesss
-	NotFoundCallback func(template string)
+	NotFoundCallback func(template string) bool
 }
 
 // NewConfig returns a new loader config
@@ -394,9 +394,7 @@ func workflowContainsProtocol(workflow []*workflows.WorkflowTemplate) bool {
 
 func (s *Store) logErroredTemplates(erred map[string]error) {
 	for template, err := range erred {
-		if s.NotFoundCallback != nil {
-			s.NotFoundCallback(template)
-		} else {
+		if s.NotFoundCallback == nil || !s.NotFoundCallback(template) {
 			gologger.Error().Msgf("Could not find template '%s': %s", template, err)
 		}
 	}

v2/pkg/core/inputs/hybrid/hmap.go

@@ -36,9 +36,20 @@ type Input struct {
 	hostMapStream *filekv.FileDB
 }
 
+// Options is a wrapper around types.Options structure
+type Options struct {
+	// Options contains options for hmap provider
+	Options *types.Options
+	// NotFoundCallback is called for each not found target
+	// This overrides error handling for not found target
+	NotFoundCallback func(template string) bool
+}
+
 // New creates a new hmap backed nuclei Input Provider
 // and initializes it based on the passed options Model.
-func New(options *types.Options) (*Input, error) {
+func New(opts *Options) (*Input, error) {
+	options := opts.Options
+
 	hm, err := hybrid.New(hybrid.DefaultDiskOptions)
 	if err != nil {
 		return nil, errors.Wrap(err, "could not create temporary input file")
@@ -65,7 +76,7 @@ func New(options *types.Options) (*Input, error) {
 		}
 		input.hostMapStream = fkv
 	}
-	if initErr := input.initializeInputSources(options); initErr != nil {
+	if initErr := input.initializeInputSources(opts); initErr != nil {
 		return nil, initErr
 	}
 	if input.dupeCount > 0 {
@@ -83,7 +94,9 @@ func (i *Input) Close() {
 }
 
 // initializeInputSources initializes the input sources for hmap input
-func (i *Input) initializeInputSources(options *types.Options) error {
+func (i *Input) initializeInputSources(opts *Options) error {
+	options := opts.Options
+
 	// Handle targets flags
 	for _, target := range options.Targets {
 		switch {
@@ -105,11 +118,15 @@ func (i *Input) initializeInputSources(options *types.Options) error {
 	if options.TargetsFilePath != "" {
 		input, inputErr := os.Open(options.TargetsFilePath)
 		if inputErr != nil {
-			return errors.Wrap(inputErr, "could not open targets file")
+			// Handle cloud based input here.
+			if opts.NotFoundCallback == nil || !opts.NotFoundCallback(options.TargetsFilePath) {
+				return errors.Wrap(inputErr, "could not open targets file")
+			}
+		}
+		if input != nil {
+			i.scanInputFromReader(input)
+			input.Close()
 		}
-		defer input.Close()
-
-		i.scanInputFromReader(input)
 	}
 	if options.Uncover && options.UncoverQuery != nil {
 		gologger.Info().Msgf("Running uncover query against: %s", strings.Join(options.UncoverEngine, ","))