From c6de2ca4060323e13baf112bf1eed74202c38d89 Mon Sep 17 00:00:00 2001 From: Sajad Parra Date: Mon, 24 Jan 2022 16:48:12 +0530 Subject: [PATCH] rename few vars, extract duplicate functions to utils --- v2/pkg/catalog/loader/remote_loader.go | 56 +++++++++++--------------- v2/pkg/parsers/parser.go | 30 +------------- v2/pkg/templates/compile.go | 29 +------------ v2/pkg/utils/utils.go | 39 ++++++++++++++++++ 4 files changed, 65 insertions(+), 89 deletions(-) diff --git a/v2/pkg/catalog/loader/remote_loader.go b/v2/pkg/catalog/loader/remote_loader.go index f9b92331..2aa334fc 100644 --- a/v2/pkg/catalog/loader/remote_loader.go +++ b/v2/pkg/catalog/loader/remote_loader.go @@ -8,6 +8,7 @@ import ( "strings" "github.com/pkg/errors" + "github.com/projectdiscovery/nuclei/v2/pkg/utils" ) type ContentType string @@ -17,38 +18,38 @@ const ( Workflow ContentType = "Workflow" ) -type RemoteContentError struct { +type RemoteContent struct { Content []string Type ContentType Error error } func getRemoteTemplatesAndWorkflows(templateURLs, workflowURLs, remoteTemplateDomainList []string) ([]string, []string, error) { - remoteContentErrorChannel := make(chan RemoteContentError) + remoteContentChannel := make(chan RemoteContent) for _, templateURL := range templateURLs { - go getRemoteContent(templateURL, remoteTemplateDomainList, remoteContentErrorChannel, Template) + go getRemoteContent(templateURL, remoteTemplateDomainList, remoteContentChannel, Template) } for _, workflowURL := range workflowURLs { - go getRemoteContent(workflowURL, remoteTemplateDomainList, remoteContentErrorChannel, Workflow) + go getRemoteContent(workflowURL, remoteTemplateDomainList, remoteContentChannel, Workflow) } var remoteTemplateList []string var remoteWorkFlowList []string var err error for i := 0; i < (len(templateURLs) + len(workflowURLs)); i++ { - remoteContentError := <-remoteContentErrorChannel - if remoteContentError.Error != nil { + remoteContent := <-remoteContentChannel + if remoteContent.Error != nil { if err != nil { - err = errors.New(remoteContentError.Error.Error() + ": " + err.Error()) + err = errors.New(remoteContent.Error.Error() + ": " + err.Error()) } else { - err = remoteContentError.Error + err = remoteContent.Error } } else { - if remoteContentError.Type == Template { - remoteTemplateList = append(remoteTemplateList, remoteContentError.Content...) - } else if remoteContentError.Type == Workflow { - remoteWorkFlowList = append(remoteWorkFlowList, remoteContentError.Content...) + if remoteContent.Type == Template { + remoteTemplateList = append(remoteTemplateList, remoteContent.Content...) + } else if remoteContent.Type == Workflow { + remoteWorkFlowList = append(remoteWorkFlowList, remoteContent.Content...) } } } @@ -56,22 +57,22 @@ func getRemoteTemplatesAndWorkflows(templateURLs, workflowURLs, remoteTemplateDo return remoteTemplateList, remoteWorkFlowList, err } -func getRemoteContent(URL string, remoteTemplateDomainList []string, w chan<- RemoteContentError, contentType ContentType) { +func getRemoteContent(URL string, remoteTemplateDomainList []string, remoteContentChannel chan<- RemoteContent, contentType ContentType) { if strings.HasPrefix(URL, "http") && (strings.HasSuffix(URL, ".yaml") || strings.HasSuffix(URL, ".yml")) { - parsed, err := url.Parse(URL) + parsedURL, err := url.Parse(URL) if err != nil { - w <- RemoteContentError{ + remoteContentChannel <- RemoteContent{ Error: err, } return } - if !stringSliceContains(remoteTemplateDomainList, parsed.Host) { - w <- RemoteContentError{ - Error: errors.Errorf("Remote template URL host (%s) is not present in the `remote-template-domain` list in nuclei config", parsed.Host), + if !utils.StringSliceContains(remoteTemplateDomainList, parsedURL.Host) { + remoteContentChannel <- RemoteContent{ + Error: errors.Errorf("Remote template URL host (%s) is not present in the `remote-template-domain` list in nuclei config", parsedURL.Host), } return } - w <- RemoteContentError{ + remoteContentChannel <- RemoteContent{ Content: []string{URL}, Type: contentType, } @@ -79,14 +80,14 @@ func getRemoteContent(URL string, remoteTemplateDomainList []string, w chan<- Re } response, err := http.Get(URL) if err != nil { - w <- RemoteContentError{ + remoteContentChannel <- RemoteContent{ Error: err, } return } defer response.Body.Close() if response.StatusCode < 200 || response.StatusCode > 299 { - w <- RemoteContentError{ + remoteContentChannel <- RemoteContent{ Error: fmt.Errorf("get \"%s\": unexpect status %d", URL, response.StatusCode), } return @@ -103,23 +104,14 @@ func getRemoteContent(URL string, remoteTemplateDomainList []string, w chan<- Re } if err := scanner.Err(); err != nil { - w <- RemoteContentError{ + remoteContentChannel <- RemoteContent{ Error: errors.Wrap(err, "get \"%s\""), } return } - w <- RemoteContentError{ + remoteContentChannel <- RemoteContent{ Content: templateList, Type: contentType, } } - -func stringSliceContains(slice []string, item string) bool { - for _, i := range slice { - if strings.EqualFold(i, item) { - return true - } - } - return false -} diff --git a/v2/pkg/parsers/parser.go b/v2/pkg/parsers/parser.go index c060f679..39bc405e 100644 --- a/v2/pkg/parsers/parser.go +++ b/v2/pkg/parsers/parser.go @@ -2,9 +2,6 @@ package parsers import ( "fmt" - "io/ioutil" - "net/http" - "os" "regexp" "strings" @@ -130,7 +127,7 @@ func ParseTemplate(templatePath string) (*templates.Template, error) { if value, err := parsedTemplatesCache.Has(templatePath); value != nil { return value.(*templates.Template), err } - data, err := readFromTemplatePath(templatePath) + data, err := utils.ReadFromPathOrURL(templatePath) if err != nil { return nil, err } @@ -152,28 +149,3 @@ func ParseTemplate(templatePath string) (*templates.Template, error) { parsedTemplatesCache.Store(templatePath, template, nil) return template, nil } - -func readFromTemplatePath(templatePath string) (data []byte, err error) { - if utils.IsURL(templatePath) { - resp, err := http.Get(templatePath) - if err != nil { - return nil, err - } - defer resp.Body.Close() - data, err = ioutil.ReadAll(resp.Body) - if err != nil { - return nil, err - } - } else { - f, err := os.Open(templatePath) - if err != nil { - return nil, err - } - defer f.Close() - data, err = ioutil.ReadAll(f) - if err != nil { - return nil, err - } - } - return -} diff --git a/v2/pkg/templates/compile.go b/v2/pkg/templates/compile.go index f605bd02..24e334ed 100644 --- a/v2/pkg/templates/compile.go +++ b/v2/pkg/templates/compile.go @@ -2,9 +2,6 @@ package templates import ( "fmt" - "io/ioutil" - "net/http" - "os" "reflect" "strings" @@ -39,7 +36,7 @@ func Parse(filePath string, preprocessor Preprocessor, options protocols.Execute template := &Template{} - data, err := readFromTemplatePath(filePath) + data, err := utils.ReadFromPathOrURL(filePath) if err != nil { return nil, err } @@ -100,30 +97,6 @@ func Parse(filePath string, preprocessor Preprocessor, options protocols.Execute parsedTemplatesCache.Store(filePath, template, err) return template, nil } -func readFromTemplatePath(templatePath string) (data []byte, err error) { - if utils.IsURL(templatePath) { - resp, err := http.Get(templatePath) - if err != nil { - return nil, err - } - defer resp.Body.Close() - data, err = ioutil.ReadAll(resp.Body) - if err != nil { - return nil, err - } - } else { - f, err := os.Open(templatePath) - if err != nil { - return nil, err - } - defer f.Close() - data, err = ioutil.ReadAll(f) - if err != nil { - return nil, err - } - } - return -} // parseSelfContainedRequests parses the self contained template requests. func (template *Template) parseSelfContainedRequests() { diff --git a/v2/pkg/utils/utils.go b/v2/pkg/utils/utils.go index 57b09280..e7a1413f 100644 --- a/v2/pkg/utils/utils.go +++ b/v2/pkg/utils/utils.go @@ -2,7 +2,10 @@ package utils import ( "errors" + "io/ioutil" + "net/http" "net/url" + "os" "strings" "github.com/projectdiscovery/fileutil" @@ -53,3 +56,39 @@ func IsURL(input string) bool { return true } + +// ReadFromPathOrURL reads and returns the contents of a file or url. +func ReadFromPathOrURL(templatePath string) (data []byte, err error) { + if IsURL(templatePath) { + resp, err := http.Get(templatePath) + if err != nil { + return nil, err + } + defer resp.Body.Close() + data, err = ioutil.ReadAll(resp.Body) + if err != nil { + return nil, err + } + } else { + f, err := os.Open(templatePath) + if err != nil { + return nil, err + } + defer f.Close() + data, err = ioutil.ReadAll(f) + if err != nil { + return nil, err + } + } + return +} + +// StringSliceContains checks if a string slice contains a string. +func StringSliceContains(slice []string, item string) bool { + for _, i := range slice { + if strings.EqualFold(i, item) { + return true + } + } + return false +}