Gitlab Custom Templates (#3570)

* Configuration options for GitLab template pulls

* GitLab client creation

* GitLab hooks and property renames

* Fix filesystem writing and update environment variables

* Fix type error in formatted error message

* Migrate directory config to new nucleiconfig file

* refactor + add custom templates to tm

* typo fix + only show installed ct with -tv

* add default gitlab url if not given

* fix template valid failure

---------

Co-authored-by: Tarun Koyalwar <tarun@projectdiscovery.io>

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)
 }
 

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
 		}
 

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")

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() {

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

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)
+	// updates all necessary paths
-	DefaultConfig.CustomS3TemplatesDirectory = filepath.Join(DefaultConfig.TemplatesDirectory, CustomS3TemplatesDirName)
+	DefaultConfig.SetTemplatesDir(DefaultConfig.TemplatesDirectory)
 }

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

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) {
+func (customTemplate *customTemplateGithubRepo) Download(ctx context.Context) {
-	downloadPath := filepath.Join(location, CustomGithubTemplateDirectory)
+	clonePath := customTemplate.getLocalRepoClonePath(config.DefaultConfig.CustomGithubTemplatesDirectory)
-	clonePath := customTemplate.getLocalRepoClonePath(downloadPath)
 
 	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) {
+func (customTemplate *customTemplateGithubRepo) Update(ctx context.Context) {
-	downloadPath := filepath.Join(location, CustomGithubTemplateDirectory)
+	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

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 {
+	ctm, err := NewCustomTemplatesManager(options)
-		ct.Download(templatesDirectory, context.Background())
+	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")

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
+}

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) {
+func (bk *customTemplateS3Bucket) Download(ctx context.Context) {
-	downloadPath := filepath.Join(location, CustomS3TemplateDirectory, bk.bucketName)
+	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
+// Update downloads custom templates from s3 bucket
-func (bk *customTemplateS3Bucket) Update(location string, ctx context.Context) {
+func (bk *customTemplateS3Bucket) Update(ctx context.Context) {
-	bk.Download(location, ctx)
+	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 {

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"
-)
+	errorutil "github.com/projectdiscovery/utils/errors"
-
-const (
-	CustomGithubTemplateDirectory = "github"
-	CustomS3TemplateDirectory     = "s3"
-	CustomAzureTemplateDirectory  = "azure"
 )
 
 type Provider interface {
-	Download(location string, ctx context.Context)
+	Download(ctx context.Context)
-	Update(location string, ctx context.Context)
+	Update(ctx context.Context)
 }
 
-// ParseCustomTemplates function reads the options.GithubTemplateRepo list,
+// CustomTemplatesManager is a manager for custom templates
-// Checks the given repos are valid or not and stores them into runner.CustomTemplates
+type CustomTemplatesManager struct {
-func ParseCustomTemplates(options *types.Options) []Provider {
+	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)
-	}
-	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)
-	}
-	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
+	// Add github providers
-		azTemplateContainer := &customTemplateAzureBlob{
+	githubProviders, err := NewGithubProviders(options)
-			azureBlobClient: azClient,
+	if err != nil {
-			containerName:   options.AzureContainerName,
+		return nil, errorutil.NewWithErr(err).Msgf("could not create github providers for custom templates")
+	}
+	for _, v := range githubProviders {
+		ctm.providers = append(ctm.providers, v)
 	}
 
-		// Add the Azure Blob Storage container object to the list of custom templates
+	// Add Aws S3 providers
-		customTemplates = append(customTemplates, azTemplateContainer)
+	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
 }

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