diff --git a/v2/cmd/nuclei/main.go b/v2/cmd/nuclei/main.go index 5c3e4913e..20b86dbe6 100644 --- a/v2/cmd/nuclei/main.go +++ b/v2/cmd/nuclei/main.go @@ -417,12 +417,19 @@ func printVersion() { func printTemplateVersion() { cfg := config.DefaultConfig gologger.Info().Msgf("Public nuclei-templates version: %s (%s)\n", cfg.TemplateVersion, cfg.TemplatesDirectory) - if cfg.CustomS3TemplatesDirectory != "" { + + if fileutil.FolderExists(cfg.CustomS3TemplatesDirectory) { gologger.Info().Msgf("Custom S3 templates location: %s\n", cfg.CustomS3TemplatesDirectory) } - if cfg.CustomGithubTemplatesDirectory != "" { + if fileutil.FolderExists(cfg.CustomGithubTemplatesDirectory) { gologger.Info().Msgf("Custom Github templates location: %s ", cfg.CustomGithubTemplatesDirectory) } + if fileutil.FolderExists(cfg.CustomGitLabTemplatesDirectory) { + gologger.Info().Msgf("Custom Gitlab templates location: %s ", cfg.CustomGitLabTemplatesDirectory) + } + if fileutil.FolderExists(cfg.CustomAzureTemplatesDirectory) { + gologger.Info().Msgf("Custom Azure templates location: %s ", cfg.CustomAzureTemplatesDirectory) + } os.Exit(0) } diff --git a/v2/internal/installer/template.go b/v2/internal/installer/template.go index 8c0f671a0..9a0c84923 100644 --- a/v2/internal/installer/template.go +++ b/v2/internal/installer/template.go @@ -2,6 +2,7 @@ package installer import ( "bytes" + "context" "crypto/md5" "fmt" "io" @@ -14,6 +15,7 @@ import ( "github.com/olekukonko/tablewriter" "github.com/projectdiscovery/gologger" "github.com/projectdiscovery/nuclei/v2/pkg/catalog/config" + "github.com/projectdiscovery/nuclei/v2/pkg/external/customtemplates" errorutil "github.com/projectdiscovery/utils/errors" fileutil "github.com/projectdiscovery/utils/file" stringsutil "github.com/projectdiscovery/utils/strings" @@ -54,7 +56,9 @@ func (t *templateUpdateResults) String() string { // TemplateManager is a manager for templates. // It downloads / updates / installs templates. -type TemplateManager struct{} +type TemplateManager struct { + CustomTemplates *customtemplates.CustomTemplatesManager // optional if given tries to download custom templates +} // FreshInstallIfNotExists installs templates if they are not already installed // if templates directory already exists, it does nothing @@ -63,7 +67,13 @@ func (t *TemplateManager) FreshInstallIfNotExists() error { return nil } gologger.Info().Msgf("nuclei-templates are not installed, installing...") - return t.installTemplatesAt(config.DefaultConfig.TemplatesDirectory) + if err := t.installTemplatesAt(config.DefaultConfig.TemplatesDirectory); err != nil { + return errorutil.NewWithErr(err).Msgf("failed to install templates at %s", config.DefaultConfig.TemplatesDirectory) + } + if t.CustomTemplates != nil { + t.CustomTemplates.Download(context.TODO()) + } + return nil } // UpdateIfOutdated updates templates if they are outdated @@ -310,7 +320,7 @@ func (t *TemplateManager) calculateChecksumMap(dir string) (map[string]string, e return err } // skip checksums of custom templates i.e github and s3 - if stringsutil.HasPrefixAny(path, config.DefaultConfig.CustomGithubTemplatesDirectory, config.DefaultConfig.CustomS3TemplatesDirectory) { + if stringsutil.HasPrefixAny(path, config.DefaultConfig.GetAllCustomTemplateDirs()...) { return nil } diff --git a/v2/internal/runner/options.go b/v2/internal/runner/options.go index 7a89f7de3..851f7edf8 100644 --- a/v2/internal/runner/options.go +++ b/v2/internal/runner/options.go @@ -5,6 +5,7 @@ import ( "fmt" "os" "path/filepath" + "strconv" "strings" "github.com/pkg/errors" @@ -129,7 +130,7 @@ func validateOptions(options *types.Options) error { } validateCertificatePaths([]string{options.ClientCertFile, options.ClientKeyFile, options.ClientCAFile}) } - // Verify aws secrets are passed if s3 template bucket passed + // Verify AWS secrets are passed if a S3 template bucket is passed if options.AwsBucketName != "" && options.UpdateTemplates { missing := validateMissingS3Options(options) if missing != nil { @@ -145,6 +146,14 @@ func validateOptions(options *types.Options) error { } } + // Verify that all GitLab options are provided if the GitLab server or token is provided + if options.GitLabToken != "" && options.UpdateTemplates { + missing := validateMissingGitLabOptions(options) + if missing != nil { + return fmt.Errorf("gitlab server details are missing. Please provide %s", strings.Join(missing, ",")) + } + } + // verify that a valid ip version type was selected (4, 6) if len(options.IPVersion) == 0 { // add ipv4 as default @@ -186,6 +195,8 @@ func validateCloudOptions(options *types.Options) error { missing = validateMissingS3Options(options) case "github": missing = validateMissingGithubOptions(options) + case "gitlab": + missing = validateMissingGitLabOptions(options) case "azure": missing = validateMissingAzureOptions(options) } @@ -244,6 +255,18 @@ func validateMissingGithubOptions(options *types.Options) []string { return missing } +func validateMissingGitLabOptions(options *types.Options) []string { + var missing []string + if options.GitLabToken == "" { + missing = append(missing, "GITLAB_TOKEN") + } + if len(options.GitLabTemplateRepositoryIDs) == 0 { + missing = append(missing, "GITLAB_REPOSITORY_IDS") + } + + return missing +} + // configureOutput configures the output logging levels to be displayed on the screen func configureOutput(options *types.Options) { // If the user desires verbose output, show verbose output @@ -333,6 +356,29 @@ func readEnvInputVars(options *types.Options) { if repolist != "" { options.GithubTemplateRepo = append(options.GithubTemplateRepo, stringsutil.SplitAny(repolist, ",")...) } + + // GitLab options for downloading templates from a repository + options.GitLabServerURL = os.Getenv("GITLAB_SERVER_URL") + if options.GitLabServerURL == "" { + options.GitLabServerURL = "https://gitlab.com" + } + options.GitLabToken = os.Getenv("GITLAB_TOKEN") + repolist = os.Getenv("GITLAB_REPOSITORY_IDS") + // Convert the comma separated list of repository IDs to a list of integers + if repolist != "" { + for _, repoID := range stringsutil.SplitAny(repolist, ",") { + // Attempt to convert the repo ID to an integer + repoIDInt, err := strconv.Atoi(repoID) + if err != nil { + gologger.Warning().Msgf("Invalid GitLab template repository ID: %s", repoID) + continue + } + + // Add the int repository ID to the list + options.GitLabTemplateRepositoryIDs = append(options.GitLabTemplateRepositoryIDs, repoIDInt) + } + } + // AWS options for downloading templates from an S3 bucket options.AwsAccessKey = os.Getenv("AWS_ACCESS_KEY") options.AwsSecretKey = os.Getenv("AWS_SECRET_KEY") diff --git a/v2/internal/runner/runner.go b/v2/internal/runner/runner.go index 54039e909..18980d301 100644 --- a/v2/internal/runner/runner.go +++ b/v2/internal/runner/runner.go @@ -73,7 +73,6 @@ type Runner struct { hostErrors hosterrorscache.CacheInterface resumeCfg *types.ResumeCfg pprofServer *http.Server - customTemplates []customtemplates.Provider cloudClient *nucleicloud.Client cloudTargets []string } @@ -102,8 +101,16 @@ func New(options *types.Options) (*Runner, error) { gologger.Error().Msgf("nuclei version check failed got: %s\n", err) } } + + // check for custom template updates and update if available + ctm, err := customtemplates.NewCustomTemplatesManager(options) + if err != nil { + gologger.Error().Label("custom-templates").Msgf("Failed to create custom templates manager: %s\n", err) + } + // Check for template updates and update if available - tm := &installer.TemplateManager{} + // if custom templates manager is not nil, we will install custom templates if there is fresh installation + tm := &installer.TemplateManager{CustomTemplates: ctm} if err := tm.FreshInstallIfNotExists(); err != nil { gologger.Warning().Msgf("failed to install nuclei templates: %s\n", err) } @@ -116,6 +123,18 @@ func New(options *types.Options) (*Runner, error) { gologger.Warning().Msgf("failed to update nuclei ignore file: %s\n", err) } } + + if options.UpdateTemplates { + // we automatically check for updates unless explicitly disabled + // this print statement is only to inform the user that there are no updates + if !config.DefaultConfig.NeedsTemplateUpdate() { + gologger.Info().Msgf("No new updates found for nuclei templates") + } + // manually trigger update of custom templates + if ctm != nil { + ctm.Update(context.TODO()) + } + } } if options.Validate { @@ -125,8 +144,6 @@ func New(options *types.Options) (*Runner, error) { // TODO: refactor to pass options reference globally without cycles parsers.NoStrictSyntax = options.NoStrictSyntax yaml.StrictSyntax = !options.NoStrictSyntax - // parse the runner.options.GithubTemplateRepo and store the valid repos in runner.customTemplateRepos - runner.customTemplates = customtemplates.ParseCustomTemplates(runner.options) if options.Headless { if engine.MustDisableSandbox() { diff --git a/v2/pkg/catalog/config/constants.go b/v2/pkg/catalog/config/constants.go index 86cdf4e71..9e4a06e29 100644 --- a/v2/pkg/catalog/config/constants.go +++ b/v2/pkg/catalog/config/constants.go @@ -9,8 +9,6 @@ import ( const ( TemplateConfigFileName = ".templates-config.json" NucleiTemplatesDirName = "nuclei-templates" - CustomS3TemplatesDirName = "s3" - CustomGithubTemplatesDirName = "github" OfficialNucleiTeamplatesRepoName = "nuclei-templates" NucleiIgnoreFileName = ".nuclei-ignore" NucleiTemplatesCheckSumFileName = ".checksum" @@ -19,6 +17,12 @@ const ( ReportingConfigFilename = "reporting-config.yaml" // Version is the current version of nuclei Version = `v2.9.2-dev` + + // Directory Names of custom templates + CustomS3TemplatesDirName = "s3" + CustomGithubTemplatesDirName = "github" + CustomAzureTemplatesDirName = "azure" + CustomGitLabTemplatesDirName = "gitlab" ) // IsOutdatedVersion compares two versions and returns true diff --git a/v2/pkg/catalog/config/nucleiconfig.go b/v2/pkg/catalog/config/nucleiconfig.go index ad5f0ea17..59e0646c6 100644 --- a/v2/pkg/catalog/config/nucleiconfig.go +++ b/v2/pkg/catalog/config/nucleiconfig.go @@ -21,8 +21,12 @@ var DefaultConfig *Config type Config struct { TemplatesDirectory string `json:"nuclei-templates-directory,omitempty"` + // customtemplates exists in templates directory with the name of custom-templates provider + // below custom paths are absolute paths to respecitive custom-templates directories CustomS3TemplatesDirectory string `json:"custom-s3-templates-directory"` CustomGithubTemplatesDirectory string `json:"custom-github-templates-directory"` + CustomGitLabTemplatesDirectory string `json:"custom-gitlab-templates-directory"` + CustomAzureTemplatesDirectory string `json:"custom-azure-templates-directory"` TemplateVersion string `json:"nuclei-templates-version,omitempty"` NucleiIgnoreHash string `json:"nuclei-ignore-hash,omitempty"` @@ -104,6 +108,11 @@ func (c *Config) GetConfigDir() string { return c.configDir } +// GetAllCustomTemplateDirs returns all custom template directories +func (c *Config) GetAllCustomTemplateDirs() []string { + return []string{c.CustomS3TemplatesDirectory, c.CustomGithubTemplatesDirectory, c.CustomGitLabTemplatesDirectory, c.CustomAzureTemplatesDirectory} +} + // GetReportingConfigFilePath returns the nuclei reporting config file path func (c *Config) GetReportingConfigFilePath() string { return filepath.Join(c.configDir, ReportingConfigFilename) @@ -175,7 +184,9 @@ func (c *Config) SetTemplatesDir(dirPath string) { c.TemplatesDirectory = dirPath // Update the custom templates directory c.CustomGithubTemplatesDirectory = filepath.Join(dirPath, CustomGithubTemplatesDirName) - c.CustomS3TemplatesDirectory = filepath.Join(dirPath, CustomGithubTemplatesDirName) + c.CustomS3TemplatesDirectory = filepath.Join(dirPath, CustomS3TemplatesDirName) + c.CustomGitLabTemplatesDirectory = filepath.Join(dirPath, CustomGitLabTemplatesDirName) + c.CustomAzureTemplatesDirectory = filepath.Join(dirPath, CustomAzureTemplatesDirName) } // SetTemplatesVersion sets the new nuclei templates version @@ -202,8 +213,6 @@ func (c *Config) ReadTemplatesConfig() error { return errorutil.NewWithErr(err).Msgf("could not unmarshal nuclei config file at %s", c.getTemplatesConfigFilePath()) } // apply config - c.CustomGithubTemplatesDirectory = cfg.CustomGithubTemplatesDirectory - c.CustomS3TemplatesDirectory = cfg.CustomS3TemplatesDirectory c.TemplatesDirectory = cfg.TemplatesDirectory c.TemplateVersion = cfg.TemplateVersion c.NucleiIgnoreHash = cfg.NucleiIgnoreHash @@ -279,6 +288,11 @@ func init() { gologger.Error().Msgf("failed to write config file at %s got: %s", DefaultConfig.getTemplatesConfigFilePath(), err) } } + // Loads/updates paths of custom templates + // Note: custom templates paths should not be updated in config file + // and even if it is changed we don't follow it since it is not expected behavior + // If custom templates are in default locations only then they are loaded while running nuclei + DefaultConfig.SetTemplatesDir(DefaultConfig.TemplatesDirectory) } func getDefaultConfigDir() string { @@ -297,6 +311,6 @@ func getDefaultConfigDir() string { // Add Default Config adds default when .templates-config.json file is not present func applyDefaultConfig() { DefaultConfig.TemplatesDirectory = filepath.Join(DefaultConfig.homeDir, NucleiTemplatesDirName) - DefaultConfig.CustomGithubTemplatesDirectory = filepath.Join(DefaultConfig.TemplatesDirectory, CustomGithubTemplatesDirName) - DefaultConfig.CustomS3TemplatesDirectory = filepath.Join(DefaultConfig.TemplatesDirectory, CustomS3TemplatesDirName) + // updates all necessary paths + DefaultConfig.SetTemplatesDir(DefaultConfig.TemplatesDirectory) } diff --git a/v2/pkg/external/customtemplates/azure_blob.go b/v2/pkg/external/customtemplates/azure_blob.go index 76a6ea914..4e405caf5 100644 --- a/v2/pkg/external/customtemplates/azure_blob.go +++ b/v2/pkg/external/customtemplates/azure_blob.go @@ -3,19 +3,47 @@ package customtemplates import ( "bytes" "context" - "github.com/Azure/azure-sdk-for-go/sdk/azidentity" - "github.com/Azure/azure-sdk-for-go/sdk/storage/azblob" - "github.com/projectdiscovery/gologger" "os" "path/filepath" "strings" + + "github.com/Azure/azure-sdk-for-go/sdk/azidentity" + "github.com/Azure/azure-sdk-for-go/sdk/storage/azblob" + "github.com/projectdiscovery/gologger" + "github.com/projectdiscovery/nuclei/v2/pkg/catalog/config" + "github.com/projectdiscovery/nuclei/v2/pkg/types" + errorutil "github.com/projectdiscovery/utils/errors" ) +var _ Provider = &customTemplateAzureBlob{} + type customTemplateAzureBlob struct { azureBlobClient *azblob.Client containerName string } +// NewAzureProviders creates a new Azure Blob Storage provider for downloading custom templates +func NewAzureProviders(options *types.Options) ([]*customTemplateAzureBlob, error) { + providers := []*customTemplateAzureBlob{} + if options.AzureContainerName != "" { + // Establish a connection to Azure and build a client object with which to download templates from Azure Blob Storage + azClient, err := getAzureBlobClient(options.AzureTenantID, options.AzureClientID, options.AzureClientSecret, options.AzureServiceURL) + if err != nil { + return nil, errorutil.NewWithErr(err).Msgf("Error establishing Azure Blob client for %s", options.AzureContainerName) + } + + // Create a new Azure Blob Storage container object + azTemplateContainer := &customTemplateAzureBlob{ + azureBlobClient: azClient, + containerName: options.AzureContainerName, + } + + // Add the Azure Blob Storage container object to the list of custom templates + providers = append(providers, azTemplateContainer) + } + return providers, nil +} + func getAzureBlobClient(tenantID string, clientID string, clientSecret string, serviceURL string) (*azblob.Client, error) { // Create an Azure credential using the provided credentials credentials, err := azidentity.NewClientSecretCredential(tenantID, clientID, clientSecret, nil) @@ -34,12 +62,12 @@ func getAzureBlobClient(tenantID string, clientID string, clientSecret string, s return client, nil } -func (bk *customTemplateAzureBlob) Download(location string, ctx context.Context) { +func (bk *customTemplateAzureBlob) Download(ctx context.Context) { // Set an incrementer for the number of templates downloaded var templatesDownloaded = 0 // Define the local path to which the templates will be downloaded - downloadPath := filepath.Join(location, CustomAzureTemplateDirectory, bk.containerName) + downloadPath := filepath.Join(config.DefaultConfig.CustomAzureTemplatesDirectory, bk.containerName) // Get the list of all templates from the container pager := bk.azureBlobClient.NewListBlobsFlatPager(bk.containerName, &azblob.ListBlobsFlatOptions{ @@ -78,9 +106,9 @@ func (bk *customTemplateAzureBlob) Download(location string, ctx context.Context // Update updates the templates from the Azure Blob Storage container to the local filesystem. This is effectively a // wrapper of the Download function which downloads of all templates from the container and doesn't manage a // differential update. -func (bk *customTemplateAzureBlob) Update(location string, ctx context.Context) { +func (bk *customTemplateAzureBlob) Update(ctx context.Context) { // Treat the update as a download of all templates from the container - bk.Download(location, ctx) + bk.Download(ctx) } // downloadTemplate downloads a template from the Azure Blob Storage container to the local filesystem with the provided diff --git a/v2/pkg/external/customtemplates/github.go b/v2/pkg/external/customtemplates/github.go index 03168e583..e0cd75d9c 100644 --- a/v2/pkg/external/customtemplates/github.go +++ b/v2/pkg/external/customtemplates/github.go @@ -10,11 +10,15 @@ import ( "github.com/google/go-github/github" "github.com/pkg/errors" "github.com/projectdiscovery/gologger" + "github.com/projectdiscovery/nuclei/v2/pkg/catalog/config" + "github.com/projectdiscovery/nuclei/v2/pkg/types" fileutil "github.com/projectdiscovery/utils/file" "golang.org/x/oauth2" "gopkg.in/src-d/go-git.v4/plumbing/transport/http" ) +var _ Provider = &customTemplateGithubRepo{} + type customTemplateGithubRepo struct { owner string reponame string @@ -23,9 +27,8 @@ type customTemplateGithubRepo struct { } // This function download the custom github template repository -func (customTemplate *customTemplateGithubRepo) Download(location string, ctx context.Context) { - downloadPath := filepath.Join(location, CustomGithubTemplateDirectory) - clonePath := customTemplate.getLocalRepoClonePath(downloadPath) +func (customTemplate *customTemplateGithubRepo) Download(ctx context.Context) { + clonePath := customTemplate.getLocalRepoClonePath(config.DefaultConfig.CustomGithubTemplatesDirectory) if !fileutil.FolderExists(clonePath) { err := customTemplate.cloneRepo(clonePath, customTemplate.githubToken) @@ -38,13 +41,13 @@ func (customTemplate *customTemplateGithubRepo) Download(location string, ctx co } } -func (customTemplate *customTemplateGithubRepo) Update(location string, ctx context.Context) { - downloadPath := filepath.Join(location, CustomGithubTemplateDirectory) +func (customTemplate *customTemplateGithubRepo) Update(ctx context.Context) { + downloadPath := config.DefaultConfig.CustomGithubTemplatesDirectory clonePath := customTemplate.getLocalRepoClonePath(downloadPath) // If folder does not exits then clone/download the repo if !fileutil.FolderExists(clonePath) { - customTemplate.Download(location, ctx) + customTemplate.Download(ctx) return } err := customTemplate.pullChanges(clonePath, customTemplate.githubToken) @@ -55,6 +58,33 @@ func (customTemplate *customTemplateGithubRepo) Update(location string, ctx cont } } +// NewGithubProviders returns new instance of github providers for downloading custom templates +func NewGithubProviders(options *types.Options) ([]*customTemplateGithubRepo, error) { + providers := []*customTemplateGithubRepo{} + gitHubClient := getGHClientIncognito() + + for _, repoName := range options.GithubTemplateRepo { + owner, repo, err := getOwnerAndRepo(repoName) + if err != nil { + gologger.Error().Msgf("%s", err) + continue + } + githubRepo, err := getGithubRepo(gitHubClient, owner, repo, options.GithubToken) + if err != nil { + gologger.Error().Msgf("%s", err) + continue + } + customTemplateRepo := &customTemplateGithubRepo{ + owner: owner, + reponame: repo, + gitCloneURL: githubRepo.GetCloneURL(), + githubToken: options.GithubToken, + } + providers = append(providers, customTemplateRepo) + } + return providers, nil +} + // getOwnerAndRepo returns the owner, repo, err from the given string // eg. it takes input projectdiscovery/nuclei-templates and // returns owner=> projectdiscovery , repo => nuclei-templates diff --git a/v2/pkg/external/customtemplates/github_test.go b/v2/pkg/external/customtemplates/github_test.go index f52c2551f..60d469565 100644 --- a/v2/pkg/external/customtemplates/github_test.go +++ b/v2/pkg/external/customtemplates/github_test.go @@ -7,6 +7,7 @@ import ( "testing" "github.com/projectdiscovery/gologger" + "github.com/projectdiscovery/nuclei/v2/pkg/catalog/config" "github.com/projectdiscovery/nuclei/v2/pkg/testutils" "github.com/stretchr/testify/require" ) @@ -18,14 +19,16 @@ func TestDownloadCustomTemplatesFromGitHub(t *testing.T) { require.Nil(t, err, "could not create temp directory") defer os.RemoveAll(templatesDirectory) + config.DefaultConfig.SetTemplatesDir(templatesDirectory) + options := testutils.DefaultOptions options.GithubTemplateRepo = []string{"projectdiscovery/nuclei-templates", "ehsandeep/nuclei-templates"} options.GithubToken = os.Getenv("GITHUB_TOKEN") - customTemplates := ParseCustomTemplates(options) - for _, ct := range customTemplates { - ct.Download(templatesDirectory, context.Background()) - } + ctm, err := NewCustomTemplatesManager(options) + require.Nil(t, err, "could not create custom templates manager") + + ctm.Download(context.Background()) require.DirExists(t, filepath.Join(templatesDirectory, "github", "nuclei-templates"), "cloned directory does not exists") require.DirExists(t, filepath.Join(templatesDirectory, "github", "nuclei-templates-ehsandeep"), "cloned directory does not exists") diff --git a/v2/pkg/external/customtemplates/gitlab.go b/v2/pkg/external/customtemplates/gitlab.go new file mode 100644 index 000000000..ad502bcba --- /dev/null +++ b/v2/pkg/external/customtemplates/gitlab.go @@ -0,0 +1,144 @@ +package customtemplates + +import ( + "context" + "encoding/base64" + "os" + "path/filepath" + + "github.com/projectdiscovery/gologger" + "github.com/projectdiscovery/nuclei/v2/pkg/catalog/config" + "github.com/projectdiscovery/nuclei/v2/pkg/types" + errorutil "github.com/projectdiscovery/utils/errors" + "github.com/xanzy/go-gitlab" +) + +var _ Provider = &customTemplateGitLabRepo{} + +type customTemplateGitLabRepo struct { + gitLabClient *gitlab.Client + serverURL string + projectIDs []int +} + +// NewGitlabProviders returns a new list of GitLab providers for downloading custom templates +func NewGitlabProviders(options *types.Options) ([]*customTemplateGitLabRepo, error) { + providers := []*customTemplateGitLabRepo{} + if options.GitLabToken != "" { + // Establish a connection to GitLab and build a client object with which to download templates from GitLab + gitLabClient, err := getGitLabClient(options.GitLabServerURL, options.GitLabToken) + if err != nil { + return nil, errorutil.NewWithErr(err).Msgf("Error establishing GitLab client for %s %s", options.GitLabServerURL, err) + } + + // Create a new GitLab service client + gitLabContainer := &customTemplateGitLabRepo{ + gitLabClient: gitLabClient, + serverURL: options.GitLabServerURL, + projectIDs: options.GitLabTemplateRepositoryIDs, + } + + // Add the GitLab service client to the list of custom templates + providers = append(providers, gitLabContainer) + } + return providers, nil +} + +// Download downloads all .yaml files from a GitLab repository +func (bk *customTemplateGitLabRepo) Download(_ context.Context) { + + // Define the project and template count + var projectCount = 0 + var templateCount = 0 + + // Append the GitLab directory to the location + location := config.DefaultConfig.CustomGitLabTemplatesDirectory + + // Ensure the CustomGitLabTemplateDirectory directory exists or create it if it doesn't yet exist + err := os.MkdirAll(filepath.Dir(location), 0755) + if err != nil { + gologger.Error().Msgf("Error creating directory: %v", err) + return + } + + // Get the projects from the GitLab serverURL + for _, projectID := range bk.projectIDs { + + // Get the project information from the GitLab serverURL to get the default branch and the project name + project, _, err := bk.gitLabClient.Projects.GetProject(projectID, nil) + if err != nil { + gologger.Error().Msgf("error retrieving GitLab project: %s %s", project, err) + return + } + + // Add a subdirectory with the project ID as the subdirectory within the location + projectOutputPath := filepath.Join(location, project.Path) + + // Ensure the subdirectory exists or create it if it doesn't yet exist + err = os.MkdirAll(projectOutputPath, 0755) + if err != nil { + gologger.Error().Msgf("Error creating subdirectory: %v", err) + return + } + + // Get the directory listing for the files in the project + tree, _, err := bk.gitLabClient.Repositories.ListTree(projectID, &gitlab.ListTreeOptions{ + Ref: gitlab.String(project.DefaultBranch), + Recursive: gitlab.Bool(true), + }) + if err != nil { + gologger.Error().Msgf("error retrieving files from GitLab project: %s (%d) %s", project.Name, projectID, err) + } + + // Loop through the tree and download the files + for _, file := range tree { + // If the object is not a file or file extension is not .yaml, skip it + if file.Type == "blob" && filepath.Ext(file.Path) == ".yaml" { + gf := &gitlab.GetFileOptions{ + Ref: gitlab.String(project.DefaultBranch), + } + f, _, err := bk.gitLabClient.RepositoryFiles.GetFile(projectID, file.Path, gf) + if err != nil { + gologger.Error().Msgf("error retrieving GitLab project file: %d %s", projectID, err) + return + } + + // Decode the file content from base64 into bytes so that it can be written to the local filesystem + contents, err := base64.StdEncoding.DecodeString(f.Content) + if err != nil { + gologger.Error().Msgf("error decoding GitLab project (%s) file: %s %s", project.Name, f.FileName, err) + return + } + + // Write the downloaded template to the local filesystem at the location with the filename of the blob name + err = os.WriteFile(filepath.Join(projectOutputPath, f.FileName), contents, 0644) + if err != nil { + gologger.Error().Msgf("error writing GitLab project (%s) file: %s %s", project.Name, f.FileName, err) + return + } + + // Increment the number of templates downloaded + templateCount++ + } + } + + // Increment the number of projects downloaded + projectCount++ + gologger.Info().Msgf("GitLab project '%s' (%d) cloned successfully", project.Name, projectID) + } + + // Print the number of projects and templates downloaded + gologger.Info().Msgf("%d templates downloaded from %d GitLab project(s) to: %s", templateCount, projectCount, location) +} + +// Update is a wrapper around Download since it doesn't maintain a diff of the templates downloaded versus in the +// repository for simplicity. +func (bk *customTemplateGitLabRepo) Update(ctx context.Context) { + bk.Download(ctx) +} + +// getGitLabClient returns a GitLab client for the given serverURL and token +func getGitLabClient(server string, token string) (*gitlab.Client, error) { + client, err := gitlab.NewClient(token, gitlab.WithBaseURL(server)) + return client, err +} diff --git a/v2/pkg/external/customtemplates/s3.go b/v2/pkg/external/customtemplates/s3.go index 674fceb0d..0804590d2 100644 --- a/v2/pkg/external/customtemplates/s3.go +++ b/v2/pkg/external/customtemplates/s3.go @@ -4,6 +4,7 @@ import ( "context" "os" "path/filepath" + "strings" "github.com/aws/aws-sdk-go-v2/aws" "github.com/aws/aws-sdk-go-v2/config" @@ -11,9 +12,14 @@ import ( "github.com/aws/aws-sdk-go-v2/feature/s3/manager" "github.com/aws/aws-sdk-go-v2/service/s3" "github.com/projectdiscovery/gologger" + nucleiConfig "github.com/projectdiscovery/nuclei/v2/pkg/catalog/config" + "github.com/projectdiscovery/nuclei/v2/pkg/types" + errorutil "github.com/projectdiscovery/utils/errors" stringsutil "github.com/projectdiscovery/utils/strings" ) +var _ Provider = &customTemplateS3Bucket{} + type customTemplateS3Bucket struct { s3Client *s3.Client bucketName string @@ -22,8 +28,8 @@ type customTemplateS3Bucket struct { } // Download retrieves all custom templates from s3 bucket -func (bk *customTemplateS3Bucket) Download(location string, ctx context.Context) { - downloadPath := filepath.Join(location, CustomS3TemplateDirectory, bk.bucketName) +func (bk *customTemplateS3Bucket) Download(ctx context.Context) { + downloadPath := filepath.Join(nucleiConfig.DefaultConfig.CustomS3TemplatesDirectory, bk.bucketName) s3Manager := manager.NewDownloader(bk.s3Client) paginator := s3.NewListObjectsV2Paginator(bk.s3Client, &s3.ListObjectsV2Input{ @@ -47,9 +53,31 @@ func (bk *customTemplateS3Bucket) Download(location string, ctx context.Context) gologger.Info().Msgf("AWS bucket %s was cloned successfully at %s", bk.bucketName, downloadPath) } -// Update download custom templates from s3 bucket -func (bk *customTemplateS3Bucket) Update(location string, ctx context.Context) { - bk.Download(location, ctx) +// Update downloads custom templates from s3 bucket +func (bk *customTemplateS3Bucket) Update(ctx context.Context) { + bk.Download(ctx) +} + +// NewS3Providers returns a new instances of a s3 providers for downloading custom templates +func NewS3Providers(options *types.Options) ([]*customTemplateS3Bucket, error) { + providers := []*customTemplateS3Bucket{} + if options.AwsBucketName != "" { + s3c, err := getS3Client(context.TODO(), options.AwsAccessKey, options.AwsSecretKey, options.AwsRegion) + if err != nil { + return nil, errorutil.NewWithErr(err).Msgf("error downloading s3 bucket %s", options.AwsBucketName) + } + ctBucket := &customTemplateS3Bucket{ + bucketName: options.AwsBucketName, + s3Client: s3c, + } + if strings.Contains(options.AwsBucketName, "/") { + bPath := strings.SplitN(options.AwsBucketName, "/", 2) + ctBucket.bucketName = bPath[0] + ctBucket.prefix = bPath[1] + } + providers = append(providers, ctBucket) + } + return providers, nil } func downloadToFile(downloader *manager.Downloader, targetDirectory, bucket, key string) error { diff --git a/v2/pkg/external/customtemplates/templates_provider.go b/v2/pkg/external/customtemplates/templates_provider.go index ebf790c74..d016854ac 100644 --- a/v2/pkg/external/customtemplates/templates_provider.go +++ b/v2/pkg/external/customtemplates/templates_provider.go @@ -2,84 +2,79 @@ package customtemplates import ( "context" - "strings" - "github.com/projectdiscovery/gologger" "github.com/projectdiscovery/nuclei/v2/pkg/types" -) - -const ( - CustomGithubTemplateDirectory = "github" - CustomS3TemplateDirectory = "s3" - CustomAzureTemplateDirectory = "azure" + errorutil "github.com/projectdiscovery/utils/errors" ) type Provider interface { - Download(location string, ctx context.Context) - Update(location string, ctx context.Context) + Download(ctx context.Context) + Update(ctx context.Context) } -// ParseCustomTemplates function reads the options.GithubTemplateRepo list, -// Checks the given repos are valid or not and stores them into runner.CustomTemplates -func ParseCustomTemplates(options *types.Options) []Provider { +// CustomTemplatesManager is a manager for custom templates +type CustomTemplatesManager struct { + providers []Provider +} + +// Download downloads the custom templates +func (c *CustomTemplatesManager) Download(ctx context.Context) { + for _, provider := range c.providers { + provider.Download(ctx) + } +} + +// Update updates the custom templates +func (c *CustomTemplatesManager) Update(ctx context.Context) { + for _, provider := range c.providers { + provider.Update(ctx) + } +} + +// NewCustomTemplatesManager returns a new instance of a custom templates manager +func NewCustomTemplatesManager(options *types.Options) (*CustomTemplatesManager, error) { + ctm := &CustomTemplatesManager{providers: []Provider{}} + if options.Cloud { - return nil + // if cloud is enabled, custom templates are Nop + return ctm, nil } - var customTemplates []Provider - gitHubClient := getGHClientIncognito() - for _, repoName := range options.GithubTemplateRepo { - owner, repo, err := getOwnerAndRepo(repoName) - if err != nil { - gologger.Error().Msgf("%s", err) - continue - } - githubRepo, err := getGithubRepo(gitHubClient, owner, repo, options.GithubToken) - if err != nil { - gologger.Error().Msgf("%s", err) - continue - } - customTemplateRepo := &customTemplateGithubRepo{ - owner: owner, - reponame: repo, - gitCloneURL: githubRepo.GetCloneURL(), - githubToken: options.GithubToken, - } - customTemplates = append(customTemplates, customTemplateRepo) + // Add github providers + githubProviders, err := NewGithubProviders(options) + if err != nil { + return nil, errorutil.NewWithErr(err).Msgf("could not create github providers for custom templates") } - if options.AwsBucketName != "" { - s3c, err := getS3Client(context.TODO(), options.AwsAccessKey, options.AwsSecretKey, options.AwsRegion) - if err != nil { - gologger.Error().Msgf("error downloading s3 bucket %s %s", options.AwsBucketName, err) - return customTemplates - } - ctBucket := &customTemplateS3Bucket{ - bucketName: options.AwsBucketName, - s3Client: s3c, - } - if strings.Contains(options.AwsBucketName, "/") { - bPath := strings.SplitN(options.AwsBucketName, "/", 2) - ctBucket.bucketName = bPath[0] - ctBucket.prefix = bPath[1] - } - customTemplates = append(customTemplates, ctBucket) + for _, v := range githubProviders { + ctm.providers = append(ctm.providers, v) } - if options.AzureContainerName != "" { - // Establish a connection to Azure and build a client object with which to download templates from Azure Blob Storage - azClient, err := getAzureBlobClient(options.AzureTenantID, options.AzureClientID, options.AzureClientSecret, options.AzureServiceURL) - if err != nil { - gologger.Error().Msgf("Error establishing Azure Blob client for %s %s", options.AzureContainerName, err) - return customTemplates - } - // Create a new Azure Blob Storage container object - azTemplateContainer := &customTemplateAzureBlob{ - azureBlobClient: azClient, - containerName: options.AzureContainerName, - } - - // Add the Azure Blob Storage container object to the list of custom templates - customTemplates = append(customTemplates, azTemplateContainer) + // Add Aws S3 providers + s3Providers, err := NewS3Providers(options) + if err != nil { + return nil, errorutil.NewWithErr(err).Msgf("could not create s3 providers for custom templates") } - return customTemplates + for _, v := range s3Providers { + ctm.providers = append(ctm.providers, v) + } + + // Add Azure providers + azureProviders, err := NewAzureProviders(options) + if err != nil { + return nil, errorutil.NewWithErr(err).Msgf("could not create azure providers for custom templates") + } + for _, v := range azureProviders { + ctm.providers = append(ctm.providers, v) + } + + // Add GitLab providers + gitlabProviders, err := NewGitlabProviders(options) + if err != nil { + return nil, errorutil.NewWithErr(err).Msgf("could not create gitlab providers for custom templates") + } + for _, v := range gitlabProviders { + ctm.providers = append(ctm.providers, v) + } + + return ctm, nil } diff --git a/v2/pkg/types/types.go b/v2/pkg/types/types.go index 0cae03894..3a5b7fff4 100644 --- a/v2/pkg/types/types.go +++ b/v2/pkg/types/types.go @@ -326,17 +326,23 @@ type Options struct { ScanAllIPs bool // IPVersion to scan (4,6) IPVersion goflags.StringSlice - // Github token used to clone/pull from private repos for custom templates + // GitHub token used to clone/pull from private repos for custom templates GithubToken string - // GithubTemplateRepo is the list of custom public/private templates github repos + // GithubTemplateRepo is the list of custom public/private templates GitHub repos GithubTemplateRepo []string - // AWS access key for downloading templates from s3 bucket + // GitLabServerURL is the gitlab server to use for custom templates + GitLabServerURL string + // GitLabToken used to clone/pull from private repos for custom templates + GitLabToken string + // GitLabTemplateRepositoryIDs is the comma-separated list of custom gitlab repositories IDs + GitLabTemplateRepositoryIDs []int + // AWS access key for downloading templates from S3 bucket AwsAccessKey string - // AWS secret key for downloading templates from s3 bucket + // AWS secret key for downloading templates from S3 bucket AwsSecretKey string - // AWS bucket name for downloading templates from s3 bucket + // AWS bucket name for downloading templates from S3 bucket AwsBucketName string - // AWS Region name where aws s3 bucket is located + // AWS Region name where AWS S3 bucket is located AwsRegion string // AzureContainerName for downloading templates from Azure Blob Storage. Example: templates AzureContainerName string