diff --git a/v2/internal/runner/cloud.go b/v2/internal/runner/cloud.go index 734660485..5d8c46ef2 100644 --- a/v2/internal/runner/cloud.go +++ b/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) diff --git a/v2/internal/runner/nucleicloud/cloud.go b/v2/internal/runner/nucleicloud/cloud.go index e06728c1d..1b5befa1d 100644 --- a/v2/internal/runner/nucleicloud/cloud.go +++ b/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") } diff --git a/v2/internal/runner/nucleicloud/types.go b/v2/internal/runner/nucleicloud/types.go index 1b5fcf73d..1643983eb 100644 --- a/v2/internal/runner/nucleicloud/types.go +++ b/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"` diff --git a/v2/internal/runner/runner.go b/v2/internal/runner/runner.go index b5e1cf786..4b6ef1618 100644 --- a/v2/internal/runner/runner.go +++ b/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 { diff --git a/v2/pkg/catalog/loader/loader.go b/v2/pkg/catalog/loader/loader.go index a8eae5798..739ed7970 100644 --- a/v2/pkg/catalog/loader/loader.go +++ b/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) } } diff --git a/v2/pkg/core/inputs/hybrid/hmap.go b/v2/pkg/core/inputs/hybrid/hmap.go index 0b9cd3298..8532af090 100644 --- a/v2/pkg/core/inputs/hybrid/hmap.go +++ b/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, ","))